  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. First, 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.
