Large Language Model or LLM inference has two phases, the prompt (or prefill) phase to output the first token and the extension (or decoding) phase to the generate subsequent tokens. In this work, we propose an efficient parallelization scheme, KV-Runahead to accelerate the prompt phase. The key observation is that the extension phase generates tokens faster than the prompt phase because of key-value cache (KV-cache). Hence, KV-Runahead parallelizes the prompt phase by orchestrating multiple processes to populate the KV-cache and minimizes the time-to-first-token (TTFT). Dual-purposing the KV-cache scheme has two main benefits. Fist, since KV-cache is designed to leverage the causal attention map, we minimize computation and computation automatically. Second, since it already exists for the extension phase, KV-Runahead is easy to implement. We further propose context-level load-balancing to handle uneven KV-cache generation (due to the causal attention) and to optimize TTFT. Compared with an existing parallelization scheme such as tensor or sequential parallelization where keys and values are locally generated and exchanged via all-gather collectives, our experimental results demonstrate that KV-Runahead can offer over 1.4x and 1.6x speedups for Llama 7B and Falcon 7B respectively.

在这项工作中，我们提出了一种高效的并行化方案KV-Runahead，旨在加速大型语言模型的前置阶段。通过利用键值缓存（KV-cache），KV-Runahead通过协调多个进程填充KV-cache以最小化时间到达第一个标记（TTFT）。我们进一步提出上下文级负载平衡来处理不均匀的KV-cache生成，并优化TTFT。与现有的并行化方案相比，KV-Runahead在Llama 7B和Falcon 7B上分别提供了1.4倍和1.6倍的加速。

KV-Runahead: 并行键值缓存生成的可扩展因果LLM推断