KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation

本文提出了一种高效的并行化方案KV-Runahead，以加速大型语言模型（LLM）的预测阶段。LLM预测有两个阶段，即预测阶段和扩展阶段，其中预测阶段生成第一个令牌，扩展阶段生成后续令牌。KV-Runahead方案通过协调多个进程填充KV-cache来并行化预测阶段，从而最小化时间到第一个令牌（TTFT）。双重使用KV-cache方案有两个主要好处。首先，由于KV-cache是设计用于利用因果关注映射，因此我们可以最小化计算和计算自动化。其次，由于它已经存在于扩展阶段中，因此KV-Runahead易于实现。我们进一步提出上下文级负载平衡来处理不均匀的KV-cache生成（由于因果关注）并优化TTFT。与现有的并行化方案（如张量或顺序并行化）相比，其中键和值是通过all-gather集合本地生成和交换的，我们的实验结果表明，KV-Runahead可以为Llama 7B和Falcon 7B分别提供超过1.4倍和1.6倍的加速。

加速大型语言模型的预测阶段（prompt phase）

通过KV-Runahead并行化预测阶段，利用KV-cache加速模型生成第一个token的时间

KV-Runahead方案利用KV-cache，避免了重复计算，且易于实现。同时，提出了上下文级别的负载均衡来优化TTFT。实验结果表明，相比于现有的tensor或sequential parallelization方案，KV-Runahead可以在Llama 7B和Falcon 7B上分别提供超过1.4x和1.6x的加速。

近期的相关研究包括GShard、Megatron、Turing-NLG等。