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

2024-05-08 · Minsik Cho, Mohammad Rastegari, Devang Naik

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.

📄 PDF Abstract BibTeX arXiv:2405.05329

Code (0)

등록된 구현이 없습니다.

Tasks

Language ModellingLarge Language Model

Methods 이 논문이 사용한 방법론

LLaMA LLaMA is a collection of foundation language models ranging from 7B to 65B parameters. It is based on the transformer architecture with various improvements that were…

Parallel ensemble methods for causal direction inference

2020-06-05 · Yulai Zhang, Jiachen Wang, Gang Cen, Guiming Luo

Inferring the causal direction between two variables from their observation data is one of the most fundamental and challenging topics in data science. A causal direction inference algorithm maps the observation data int…

NVR: Vector Runahead on NPUs for Sparse Memory Access

2025-02-19 · Hui Wang, Zhengpeng Zhao, Jing Wang, Yushu Du 외

Deep Neural Networks are increasingly leveraging sparsity to reduce the scaling up of model parameter size. However, reducing wall-clock time through sparsity and pruning remains challenging due to irregular memory acces…

A Fast Kernel-based Conditional Independence test with Application to Causal Discovery

2025-05-16 · Oliver Schacht, Biwei Huang

Kernel-based conditional independence (KCI) testing is a powerful nonparametric method commonly employed in causal discovery tasks. Despite its flexibility and statistical reliability, cubic computational complexity limi…

Causal Discovery Causal Inference Gaussian Processes Mixture-of-Experts

Scalable Matrix-valued Kernel Learning for High-dimensional Nonlinear Multivariate Regression and Granger Causality

2014-08-09 · Vikas Sindhwani, Ha Quang Minh, Aurelie Lozano

We propose a general matrix-valued multiple kernel learning framework for high-dimensional nonlinear multivariate regression problems. This framework allows a broad class of mixed norm regularizers, including those that …

Causal Inference Generalization Bounds regression

ConfoundingSHAP: Quantifying confounding strength in causal inference

2026-05-11 · Marie Brockschmidt, Santo M. A. R. Thies, Maresa Schröder, Dennis Frauen 외 arxiv

In causal inference, confounders are variables that influence both treatment decisions and outcomes. However, unlike as in randomized clinical trials, the treatment assignment mechanism in observational studies is not kn…

Causal Inference