KV-Pipe: On the Relation Between KV Sharing and Pipeline Parallel Efficiency in LLMs

2026-08-16

AI Summary

This study addresses load imbalance and bubble efficiency bottlenecks in large model pipeline parallelism by proposing KV-Pipe, a mechanism that uniquely leverages cross-layer KV sharing as an architectural degree of freedom. Through stage-aware KV sharing, tail-priority layer transformation, and an offline FLOPs balancing algorithm, the method normalizes computational load ratios to jointly optimize training and inference performance without online tuning. Experimental results demonstrate that KV-Pipe improves training Model FLOPs Utilization (MFU) by up to 9.2% and reduces iteration time by 9.8%. Furthermore, it significantly enhances long-context decoding throughput, effectively mitigating pipeline efficiency limitations.

Abstract

Pipeline parallelism (PP) is widely used to scale large language model (LLM) training, but its efficiency is often limited by stage imbalance and pipeline bubbles. Meanwhile, cross-layer KV sharing has primarily been studied as a mechanism for reducing KV-cache costs during inference, without examining how KV reuse reshapes pipeline workloads. We present \textbf{KV-Pipe}, a stage-aware KV-sharing mechanism that turns KV reuse into a pipeline-balancing control knob. KV-Pipe starts from the tail stage, converts selected attention layers to cross-layer KV sharing in a tail-first order, and iteratively retargets the current bottleneck to drive the FLOPs Imbalance Ratio (FIR) toward $1$. The procedure is performed offline and requires only a pipeline partition and per-layer FLOPs estimates, introducing negligible runtime overhead and requiring no online tuning. Across multiple pipeline-parallel configurations, KV-Pipe consistently improves utilization and throughput, achieving up to \textbf{9.2\%} higher training MFU and up to a \textbf{9.8\%} reduction in iteration time, with larger gains at higher pipeline-parallel degrees where stage imbalance is amplified. Furthermore, the same KV-sharing mechanism provides an inference-side benefit by reducing KV-cache growth and redundant KV projection work, resulting in higher decoding throughput for long-context workloads. These results identify KV layout as a system--architecture degree of freedom for jointly improving pipeline-parallel training efficiency and long-context inference.

Problem

• Pipeline Parallelism
• Stage Imbalance
• Pipeline Bubbles
• KV Sharing
• Large Language Models

Innovation

• KV-Pipe
• Pipeline Parallelism
• Cross-layer KV Sharing
• Stage-aware Load Balancing
• Training-Inference Co-optimization