arXiv CS open access

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

2026 · arxiv_cs

arXiv CS · Papers · 2026

cloud distributed-computing parallel-computing

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

KV-Pipe: On the Relation Between KV Sharing and Pipeline Parallel Efficiency in LLMs Maryam Dialameh1,2 Hossein Rajabzadeh1,2 Walid Ahmed2 Weiwei Zhang2 1

Harish Krishnamoorthy Murali2 Hyock Ju Kwon1

School of Engineering, University of Waterloo, Waterloo, Canada 2

Ascend Team, Huawei Technologies, Toronto, Canada

arXiv:2608.15943v1 [cs.DC] 16 Aug 2026

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 KV-Pipe, a stageaware KV-sharing mechanism that turns KV reuse into a pipelinebalancing 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 9.2% higher training MFU and up to a 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.

Figure 1. Motivation. (a) Pipeline parallelism suffers from stage imbalance and bubble overhead. (b) KV sharing restructures crosslayer attention workloads, enabling better FLOP distribution across stages. (c) KV-Pipe leverages stage-aware KV sharing to reduce bubble time and improve MFU.

Large language models (LLMs) now achieve state-of-theart performance in complex reasoning tasks, mathematical problem solving, multi-step planning, and agentic behavior Zhao et al. (2024); Toroghi et al. (2024); Li et al. (2024); Schumann et al. (2024). As model sizes continue to scale Chowdhery et al. (2023); Hoffmann et al. (2022); Smith et al. (2022); Adler et al. (2024), both training and inference impose memory and compute demands far exceeding the capacity of a single accelerator.

and pipeline bubbles, which waste a non-trivial fraction of allocated accelerator time, especially at scale Huang et al. (2019); Narayanan et al. (2019, 2021). Recent systems work has pushed PP efficiency by (i) re-partitioning or rebalancing layers to address memory and compute skew (e.g., activation/memory balancing in BPipe and cost-modeldriven partition search in DawnPiper) Kim et al. (2023); Peng et al. (2025), (ii) mitigating bubbles via scheduling or opportunistic backfilling (e.g., PipeFill fills bubble time with other jobs) Arfeen et al. (2025), and (iii) handling prefill–decode heterogeneity in serving through phaseaware batching and throttling (e.g., chunked-prefill scheduling in Sarathi/Sarathi-Serve, token throttling in gLLM, temporally-disaggregated execution in TD-Pipe, and dynamic re-sharding in Seesaw) Agrawal et al. (2023, 2024); Guo et al. (2025); Zhang et al. (2025); Su et al. (2025). These approaches substantially improve utilization, yet they largely treat the model’s per-layer cost profile as fixed, leaving a key system-level degree of freedom underexplored: how the model’s internal KV layout interacts with pipeline partitioning.

Despite its attractive communication profile, pipeline parallelism (PP) is notoriously sensitive to stage imbalance

A common abstraction behind PP is that Transformer backbones comprise a stack of repeated blocks, suggesting a

Introduction

near-uniform compute distribution across depth. In modern LLMs, however, blocks are increasingly heterogeneous: models mix full attention with more efficient attention variants (e.g., sliding-window attention) Jiang et al. (2023), insert sparse Mixture-of-Experts (MoE) modules that change per-layer parameterization and token-wise compute Jiang et al. (2024a), and even interleave fundamentally different sequence modules in hybrid designs (e.g., Transformer– Mamba hybrids) Lieber et al. (2024); Blakeman et al. (2025). As a result, layers (and therefore pipeline stages) can differ materially in FLOPs, activation footprint, and KV-cache growth, making it difficult to maintain balanced stage times across training and inference.

Methodology

KV-Pipe improves pipeline-parallel efficiency by using cross-layer KV sharing Dialameh et al. (2025); Rajabzadeh (2026) as a stage-aware knob. Intuitively, converting a subset of full-attention layers into KV-sharing layers reduces (i) the FLOPs of those layers and (ii) the growth of the KV cache during autoregressive inference. KV-Pipe applies this conversion selectively to stages that dominate the pipeline step time, thereby reducing pipeline bubbles and improving overall utilization.

One natural response is to share parameters or cache structures across depth. Classical cross-layer parameter sharing (e.g., ALBERT) reduces redundancy and can homogenize memory footprint across layers Lan et al. (2019). More recently, cross-layer KV sharing has emerged as an effective way to reduce KV-cache memory and decoding overhead (e.g., MLKV shares KV structures across layers; HShare shares critical KV indices across layers/heads/queries; systematic studies unify and compare cross-layer KV-sharing variants) Zuhri et al. (2025); Wu et al. (2025a,b). Meanwhile, KV-cache memory itself is a first-order bottleneck in production-grade serving, motivating substantial systems effort in KV management (e.g., PagedAttention/vLLM) Kwon et al. (2023). However, existing KV-sharing methods primarily target memory compression or per-token decode speed, without explicitly optimizing for pipeline-stage balance.

2.1

Notation and Pipeline Partition

We consider a decoder-only Transformer with L attention layers. Pipeline parallelism with degree P partitions layers across P stages. Let stage indices be i ∈ {1, . . . , P }, and let Si denote the set of layers assigned to stage i. We allow non-uniform components, such as the embedding module (often on stage 1) and the LM head (often on stage P ). For clarity in our case study, we use LLaMA2-7B with L = 32 attention layers and P = 8 stages. Table 1 shows the baseline split (4 attention layers per stage) where the last stage is the bottleneck because it additionally includes the LM head.

2.2

This paper introduces KV-Pipe, a stage-aware KV-sharing framework that elevates KV reuse from a memorycompression trick to a pipeline-level control knob. The key novelty is to couple cross-layer KV sharing with PP partitioning: by converting carefully chosen attention layers (starting from the tail stage and then tracking the moving bottleneck), KV-Pipe reshapes the per-stage compute profile without changing the PP schedule, micro-batching, optimizer, or training dynamics. KV-Pipe diagnoses imbalance using a lightweight FLOPs Imbalance Ratio (FIR) and then applies bottleneck-targeted KV sharing to directly flatten stage times (FIR→ 1), eliminating bubbles at their source rather than merely hiding them with scheduling. This yields a double dividend: during training, KV-Pipe increases PP utilization by reducing stage skew and idle time; during inference, the same KV sharing reduces KV-cache growth and redundant KV projections, enabling larger effective batching and higher long-context throughput. In effect, KV-Pipe exposes KV layout as a new, orthogonal and composable system–architecture degree of freedom that can be applied on top of existing PP methods to improve both train-time MFU and inference speed.

Full-Attention vs. KV-Sharing Layers

Full attention. Given hidden states Xℓ ∈ RS×d at layer ℓ, full attention computes Qℓ = Xℓ WℓQ ,

Kℓ = Xℓ WℓK ,

Vℓ = Xℓ WℓV , (1)

and produces the attention output Qℓ K⊤ ℓ Oℓ = softmax √ + M Vℓ , dh

(2)

where M is the causal mask and dh is the per-head dimension. In autoregressive decoding, each layer appends (Kℓ , Vℓ ) into the KV cache. Cross-layer KV sharing. In a KV-sharing layer, keys/values are reused from a prior layer τ (ℓ) < ℓ: Kℓ ← Kτ (ℓ) ,

Vℓ ← Vτ (ℓ) .

(3)

The query projection and attention computation (Eq. 2) remain unchanged. This reduces redundant KV projections and reduces KV-cache growth, which is the primary memory bottleneck in long-context inference. 2

Emb Full attn KV sharing LM head Flops FWD time Memory

0 1 4 0 0 1.55E+13 59712.13 16.8425

1 0 4 0 0 1.55E+13 59237.7 15.5625

2 0 4 0 0 1.55E+13 59237.7 15.5625

3 0 4 0 0 1.55E+13 59237.7 15.5625

4 0 4 0 0 1.55E+13 59237.7 15.5625

5 0 4 0 0 1.55E+13 59237.7 15.5625

6 0 4 0 0 1.55E+13 59237.7 15.5625

7 0 4 0 1 1.763E+13 66796.0045 16.8425

Table 1. Baseline – LLaMA2-7B, PP=8, sequence length 8K, MBS=1. The last stage is the bottleneck due to the LM head.

Decision variables. Let zℓ ∈ {0, 1} indicate whether layer ℓ is converted to KV sharing (zℓ = 1) or remains full attention (zℓ = 0). For each stage i, define the stage FLOPs after conversion as

Cost effect (used by KV-Pipe). KV sharing decreases perlayer compute and KV-cache bytes; we model this at the layer level by a reduction ∆Fℓ in FLOPs and a reduction ∆Bℓ in KV-cache bytes (per token) whenever layer ℓ is converted. In practice, ∆Fℓ is dominated by removing Kℓ , Vℓ projections, and ∆Bℓ is dominated by not storing a fresh (Kℓ , Vℓ ) for that layer.

Fstagei (z) =

Optimization target (intuitive). KV-Pipe aims to reduce the bottleneck stage relative to the average, i.e., bring FIR (Eq. 4) as close to 1 as possible, subject to a small budget of conversions: L X zℓ ≤ m. (7)

Let Fstagei denote the total FLOPs assigned to stage i. We define the FLOPs Imbalance Ratio (FIR) as maxi Fstagei FIR = , (4) Favg

ℓ=1

Rather than solving an expensive discrete optimization, KVPipe uses simple, interpretable allocation strategies that are easy to implement and robust in practice.

where Favg is the average FLOPs across stages. To account for architectural non-uniformity such as the LM head, we compute N X

(6)

where Fℓ (0) corresponds to full attention and Fℓ (1) corresponds to KV sharing.

To measure and compare FLOP balance across PP splits, we propose the following metric.

Fℓ (zℓ ) + δi,P FLM-Head ,

ℓ∈Si

2.3 Proposed Metric: FLOP Imbalance Ratio

1 X Favg = P i=1

! (5)

2.5 KV-Sharing Placement Strategies (with Figure)

Here, δi,P is the Kronecker delta (δi,P = 1 if i = P and 0 otherwise), FLM-Head is the FLOPs of the LM head, and layer Fstageik is the FLOPs of the k-th layer placed on stage i. (For uniform splits, N = L/P , while in general N can be interpreted as the number of layers on that stage.) A perfectly balanced pipeline yields FIR = 1, and values closer to 1 indicate better balance.

Figure 2 illustrates three placement strategies we study for converting layers into KV-sharing layers. Without loss of generality, these strategies serve as archetypes: (i) uniform smoothing (Uniform), (ii) symmetric redistribution (Symmetric Bipolar), and (iii) bottleneck targeting (ArchitectureBalanced). Many alternative policies either behave similarly to one of these archetypes or can be constructed as a combination of them (e.g., uniform allocation followed by a bottleneck-focused correction).

layer Fstageik

+ δi,P FLM-Head .

k=1

2.4 KV-Pipe: Balancing Pipeline Stages via KV Sharing 2.5.1

Uniform Split

Given a fixed PP partition {Si }P i=1 , KV-Pipe converts

Goal. a small number m of full-attention layers into KV-sharing layers to reduce the bottleneck stage FLOPs (and thus stage time), while simultaneously lowering KV-cache consumption.

Allocate KV-sharing layers approximately evenly across stages (starting from the last stage and moving left), reducing FLOPs broadly without explicitly prioritizing a single bottleneck. 3

PP split strategy Baseline Uniform Symmetric Bipolar Architecture-Balanced

FLOP Imbalance Ratio 1.119453925 1.103804831 1.068260764 1.00040391

Table 2. FLOP Imbalance Ratio (proposed) for LLaMA2-7B with 4 KV-sharing layers, PP=8, sequence length 8K, MBS=1.

ing pass), and it does not require running actual training or iterative online tuning. If we allow at most m conversions, each iteration performs constant-time updates to a single stage plus a scan over P stages to identify the current bottleneck. With typical pipeline sizes (P ≤ 16) and small m (a handful of layers), this cost is negligible compared to any real training or serving run, making KV-Pipe practical to deploy as a pre-processing step before launching large-scale jobs.

Figure 2. KV sharing strategies for LLaMA with 32 layers under PP.

2.5.2

Symmetric Bipolar Split

Allocate KV-sharing layers from two anchors: the last stage and the middle stage, moving left symmetrically. This is useful when imbalance is not localized only at the last stage. 2.5.3

2.7 Case Study Result via the Proposed Metric

Architecture-Balanced Split

Allocate KV-sharing layers preferentially to the current bottleneck stage (typically stage P due to the LM head), greedily reducing maxi Fstagei to drive the imbalance ratio toward 1.

2.6

Using the proposed FIR metric, Table 2 compares the baseline and the three strategies for LLaMA2-7B with m = 4 KV-sharing layers under PP=8. Architecture-Balanced split achieves the closest value to 1, indicating near-perfect FLOP balance.

Algorithm: KV-Pipe

Algorithm 1 describes the architecture-balanced KV-Pipe policy as a greedy, tail-first reconfiguration of attention layers. KV-Pipe always starts from the last pipeline stage (stage P ), which is typically the bottleneck due to architectural components such as the LM head. Within the selected stage, KV-Pipe converts layers to KV-sharing in a last-layer-first order, i.e., it picks the deepest remaining full-attention layer ℓ⋆ = max{ℓ ∈ Si⋆ : zℓ = 0} and turns it into a KVsharing layer. Each conversion decreases the FLOPs of that stage by ∆Fℓ⋆ = Fℓ⋆ (0) − Fℓ⋆ (1), and KV-Pipe recomputes the proposed imbalance metric FIR (Eqs. (4)–(5)). Once the tail stage becomes sufficiently close to the average (FstageP ≤ (1 + ϵ)Favg ), KV-Pipe automatically switches to the current bottleneck stage i⋆ = arg maxj Fstagej , ensuring that the procedure continues to drive the pipeline toward a balanced state (FIR → 1) even if the bottleneck shifts as sharing increases. The algorithm terminates when FIR ≤ 1 + ϵ, when a conversion budget is reached, or when no further convertible layers remain in the chosen stage. This behavior aligns with the architecture-balanced placement in Figure 2 and explains the observed FIR and MFU improvements (Table 3).

Experiments

3.1

Experimental Setup

Model and workload. We evaluate KV-Pipe on LLaMA27B with sequence length S = 8K. Unless otherwise stated, we report training efficiency under a fixed global batch size and fixed optimizer/precision settings. The results are primarily conducted on Huawei Ascend 910B and NVIDIA V100. Our main results are on NPUs; GPU results validate that KV-Pipe trends generalize across backends. Pipeline parallelism (PP). We purely investigate PP behavior. All experiments use pipeline parallelism with degree PP ∈ {2, 4, 8}. We keep other parallelism dimensions fixed (i.e., we do not set or vary TP/CP in the main experiments) to isolate PP effects. MFU definition. We report model FLOPs utilization (MFU) as Fmodel · Throughput | {z } | {z }

Importantly, KV-Pipe is a fast offline procedure: it requires only the PP partition and a lightweight per-layer FLOPs estimate (e.g., derived analytically or from a one-time profil-

MFU =

FLOPs/iter

iter/sec

Fpeak | {z }

device peak FLOPs/sec

(8)

Algorithm 1 KV-Pipe (Architecture-Balanced): Tail-first KV sharing until FIR → 1 L Require: PP partition {Si }P i=1 with ordered layers in each stage; full-attn and KV-share FLOPs {Fℓ (0), Fℓ (1)}ℓ=1 ; LM-head FLOPs FLM-Head ; tolerance ϵ > 0; max conversions m (optional, default m = L). Ensure: Conversion mask z ∈ {0, 1}L ; updated stage FLOPs {Fstagei }P i=1 ; achieved FIR. 1: z ← 0. 2: Compute baseline stage FLOPs: P Fstagei ← ℓ∈Si Fℓ (0) + δi,P FLM-Head , ∀i ∈ {1, . . . , P }. 3: Compute FIR using Eqs. (4)–(5). 4: t ← 0. 5: i ← P ▷ starts from the last stage 6: while FIR > 1 + ϵ and t < m do 7: i⋆ ← i ▷ current stage to lighten (initially the tail) 8: if Fstagei⋆ ≤ (1 + ϵ) · Favg then 9: i⋆ ← arg maxj∈{1,...,P } Fstagej ▷ tail is no longer bottleneck; pick the new one 10: end if 11: Choose the last convertible layer in stage i⋆ : ℓ⋆ ← max{ℓ ∈ Si⋆ : zℓ = 0}. 12: if ℓ⋆ is undefined then 13: break ▷ no convertible layers remain in the selected stage 14: end if 15: Convert to KV sharing: zℓ⋆ ← 1. 16: ∆Fℓ⋆ ≜ Fℓ⋆ (0) − Fℓ⋆ (1) > 0. 17: Update stage FLOPs: Fstagei⋆ ← Fstagei⋆ − ∆Fℓ⋆ . 18: t ← t + 1. 19: Recompute Favg (Eq. (5)) and FIR (Eq. (4)). 20: end while 21: return z, {Fstagei }P i=1 , FIR.

where Fmodel is the model compute per iteration (FLOPs/iter), Throughput is the measured iteration throughput (iter/s), and Fpeak is the device peak compute (FLOPs/s), scaled by the number of devices used.

gains grow with pipeline depth: from 62.45→64.39 (PP=2) to 61.28→63.93 (PP=4), and 56.32→61.49 (PP=8), corresponding to relative MFU improvements of 3.10%, 4.32%, and 9.17%, respectively (Table 3). Iteration time decreases by 4.90%–9.80%, indicating consistent end-to-end throughput gains.

KV-Pipe strategies. We compare three KV-sharing placement strategies under PP: U NIFORM, S YMMETRIC B IPO LAR , and A RCHITECTURE -BALANCED . We also vary the KV-sharing budget (number of KV-sharing layers) for ablation.

3.2

Imbalance metric tracks utilization. The best KV-Pipe setting also moves FIR closer to 1 in all three configurations (e.g., 1.1194→1.0004 for S=8K, PP=8), matching the trend in MFU improvements and supporting our hypothesis that reducing stage imbalance is a key driver of PP utilization (Table 3).Figures 3 and 4 reveal a consistent non-monotonic MFU pattern as we increase the KV-sharing budget (shared-KV size). At small shared-KV sizes, converting a few full-attention layers to KV-sharing layers in the A RCHITECTURE -BALANCED placement selectively reduces the bottleneck stage compute, thereby reducing the stage-time skew. Equivalently, the Architecture-Balanced strategy reduces FIR and moves it closer to the ideal value of 1 (Table 3), indicating a more balanced pipeline. This initial balancing effect increases MFU sharply (e.g., the green curve in Figure 3), because less time is wasted in pipeline bubbles and idle periods. Importantly, the KV-sharing size that maximizes MFU is an intermediate point rather than the largest shared-KV size. Across our three NPU settings, the MFU-optimal shared-KV sizes are KV=8 for (S=4K, PP=2), KV=6 for (S=8K, PP=4), and KV=4 for (S=8K,

Results on Ascend 910B NPUs

We evaluate KV-Pipe on 8× Ascend 910B NPUs using MindSpeed-LM’s 1F1B pipeline schedule (no interleaving) with LLaMA2-7B. We test three training configurations: (1) S=4K, PP=2; (2) S=8K, PP=4; (3) S=8K, PP=8. For each configuration, we sweep the KV-sharing budget (“KV size”) under the three placement strategies and select the best KV-Pipe variant, which is consistently the ArchitectureBalanced placement. Table 3 summarizes the baseline 1F1B (no KV sharing) versus the best KV-Pipe (ArchitectureBalanced) at the MFU-optimal KV-sharing budget for each configuration; the full MFU sweeps and time/FLOPs sweeps are shown in Figures 3 and 4. MFU and time improvements. Across all settings, KVPipe increases MFU and reduces iteration time. The MFU 5

PP=8), as highlighted by stars in Figure 3 and summarized in Table 3. These optima occur where KV-Pipe has sufficiently reduced the bottleneck stage to equalize stage latencies, i.e., FIR is closest to 1, so the pipeline achieves near-maximal overlap.

We repeat the PP-focused evaluation on a single node with 8×NVIDIA V100 GPUs. We compare 1F1B Harlap et al. (1806), Seq1F1B Ao et al. (2025), and KV-Pipe under PP∈ {2, 4, 8} and context length S ∈ {4K, 8K, 16K}, on four LLMs: LLaMA2-7B, LLaMA2-13B, LLaMA3-8B, and Qwen2.5-14B. Table 4 reports PP-only performance. Each cell contains a triplet “PP=2 / 4 / 8”, i.e., the first number corresponds to PP=2, the second to PP=4, and the third to PP=8. Across all models and sequence lengths, Seq1F1B improves over 1F1B by reducing pipeline bubbles through sequential scheduling, which is most visible as PP increases (e.g., LLaMA2-7B at S=8K: MFU 39.0→47.0 and time 305→265 when moving from 1F1B to Seq1F1B at PP=8). KV-Pipe provides an additional gain beyond Seq1F1B by directly reducing stage-level compute skew via stage-aware KV sharing, leading to both higher MFU and lower iteration time across the board. The benefit grows with pipeline depth: KV-Pipe’s improvements are modest at PP=2 but become consistently largest at PP=8, where imbalance and idle time are amplified (e.g., LLaMA2-13B at S=16K and PP=8: MFU 39.5→52.0 and time 960→750 from 1F1B to KV-Pipe). Longer contexts further strengthen the trend because KV/attention costs become more dominant and heterogeneous across stages, increasing the headroom for KV-Pipe to rebalance the pipeline. Overall, the table shows that Seq1F1B primarily hides bubbles through better scheduling, while KV-Pipe removes bubbles by making stage workloads more uniform; the combined effect is most pronounced at high PP and long context lengths.

Why MFU drops beyond the optimum: over-correction induces new imbalance. When shared-KV size increases beyond the MFU-optimal point, KV-Pipe begins to overlighten the previously bottlenecked stages. While total model FLOPs continue to decrease monotonically (bottom rows of Figure 4), the per-stage compute distribution is no longer aligned: the earlier (or middle) stages can become the new bottleneck, and FIR increases again (moving away from 1). This renewed imbalance reintroduces pipeline idle time, and MFU consequently declines after the peak (Figure 3). This explains why “more KV sharing” is not always better for PP efficiency: the dominant factor is balance, not just fewer FLOPs. Time/FLOPs sweeps corroborate the MFU behavior. Figure 4 further supports this interpretation. FLOPs decrease approximately monotonically with KV size for all strategies, yet iteration time shows diminishing returns and can flatten (or even slightly worsen) once imbalance shifts to other stages. In practice, the shared-KV size that minimizes time does not necessarily coincide with the shared-KV size that maximizes MFU, because MFU is sensitive to bubble overhead caused by imbalance. Overall, these results validate the design principle of KV-Pipe: stage-aware KV placement should be tuned to minimize imbalance (FIR→1), which yields an optimal KV-sharing budget that maximizes PP efficiency.

3.3

Across all evaluated models, the benefit of Seq1F1B grows with P (bubble mitigation), while KV-Pipe yields an additional improvement by directly reducing stage-level skew; the combined effect is most pronounced at PP=8 and longer contexts where both bubble overhead and per-stage KV/compute heterogeneity are amplified.

Results on GPUs (8×V100)

Inference case study (MHA) and double benefit of KV-Pipe. To complement our training results, we run an inference case study to illustrate that KV-Pipe provides a second, orthogonal benefit beyond pipeline balancing. We consider LLaMA2-7B with standard MultiHead Attention (MHA) (i.e., each layer maintains its own K, V cache per head, without grouped-query sharing), and evaluate long-context decoding at sequence lengths S ∈ {8K, 16K, 32K, 64K, 128K}. Figure 5 reports the endto-end inference throughput (tokens/s) comparing a Transformer baseline (no KV sharing) against KV-Pipe under identical execution assumptions. We observe a consistent throughput improvement that becomes more pronounced as context grows: KV-Pipe achieves modest gains at shorter contexts (where cache pressure is limited), but delivers substantially larger speedups at long contexts where KV-cache memory and attention cost dominate. Importantly, this inference improvement is additive to KV-Pipe’s training-side

Figure 5. Inference throughput (LLaMA2-7B, MHA). Tokens/s vs. context length for the Transformer baseline and KV-Pipe with 50% KV sharing, i.e. shared-KV size. KV-Pipe improves throughput across all contexts, with larger gains at longer contexts where KV-cache overhead dominates.

(a) S=4K, PP=2, SKV=8

(b) S=8K, PP=4, SKV=6

Metric name

Metric change

Improvement

MFU Iteration time (s) FIR FWD time

62.45→64.39 104.14→99.27

3.10% 4.90%

1.0182→1.0101 0.1152→0.1094

0.80% 5.30%

Metric name

Metric change

Improvement

MFU Iteration time (s) FIR FWD time

61.28→63.93 241.96→224.95

4.32% 7.56%

1.051→1.018 0.1260→0.1189

3.24% 5.97%

Metric name

Metric change

MFU Iteration time (s) FIR FWD time

56.32→61.49 263.3→237.5

Improvement 9.17% 9.80%

1.1194→1.0004 11.89% 0.0347→0.0302 14.90%

Table 3. NPU best-result summary (baseline vs KV-Pipe). Each subtable compares baseline 1F1B (no KV sharing) against the best KVPipe configuration (Architecture-Balanced placement at the MFU-optimal KV-sharing budget) for the corresponding PP/sequence-length setting. For MFU, we report the relative improvement and also the approximate absolute MFU-point gain in parentheses. All subtables use LLaMA2-7B as the underlying LLM. Here, S denotes the sequence length, and SKV denotes the shared-KV size.

(a) S=4K, PP=2

(b) S=8K, PP=4

Figure 3. MFU sweeps on 8× Ascend 910B. MFU vs. KV-sharing budget for the three KV-Pipe placements. Stars indicate the best MFU point used in Table 3.

gains: during training, KV-Pipe increases MFU by reducing stage-level compute skew and pipeline bubbles; during inference, it reduces KV-cache footprint and redundant KV projection work, enabling higher effective throughput. This double benefit highlights KV-Pipe as a generic and native system-level knob: it simultaneously (i) balances pipeline stages to improve training efficiency under PP and (ii) accelerates inference via cross-layer KV sharing, making KV layout an actionable degree of freedom for optimizing both train-time MFU and inference-time speed.

monotonic in the shared-KV size. Increasing KV sharing always reduces total model FLOPs, yet MFU increases only up to an optimal point and then declines. KV-Pipe explains this behavior through the proposed FLOPs Imbalance Ratio (FIR): small-to-moderate KV sharing reduces the bottleneck stage cost and drives FIR toward 1, minimizing idle time and bubbles. Past the optimum, KV-Pipe can over-correct the original bottleneck and shift the critical path to earlier stages, increasing FIR again and reintroducing bubbles even though total FLOPs keep decreasing. This highlights an important systems insight: for pipeline efficiency, the dominant objective is not minimizing FLOPs globally, but minimizing stage imbalance (i.e., aligning per-stage times).

Discussion

Many PP techniques improve utilization primarily through scheduling (e.g., reducing bubbles via refined 1F1B variants, sequential schedules, interleaving, or micro-batch shaping). These methods can be highly effective, but they largely treat the per-stage compute profile as fixed and aim to hide bubbles. In contrast, KV-Pipe directly changes the per-stage workload by altering the effective cost of attention layers, thereby removing bubbles at the source. This distinction is important because stage imbalance often persists even with aggressive scheduling, especially in realistic models where the last stage contains extra modules (e.g., LM head) or where attention/MLP costs are non-uniform across depth.

KV-Pipe turns cross-layer KV sharing into a pipeline-level optimization knob: by selectively converting full-attention layers into KV-sharing layers in FLOP-heavy pipeline stages, it reduces stage skew and improves end-to-end pipeline efficiency. Our Ascend 910B results (Table 3, Figures 3–4) show consistent gains across PP settings. Most notably, the relative MFU improvement grows with pipeline depth (PP): the baseline becomes more sensitive to stage imbalance and bubble overhead at larger PP, increasing the headroom for KV-Pipe. This is reflected both in higher MFU and reduced iteration time, indicating that KV-Pipe improves true utilization rather than merely shifting work between phases.

KV-Pipe is intentionally designed to be orthogonal to pipeline scheduling and parallel runtime choices. Conceptually, KV-Pipe outputs a modified per-layer cost profile by

A central observation from Figures 3–4 is that MFU is non7

(a) S=4K, PP=2

(b) S=8K, PP=4

Figure 4. Iteration time and FLOPs sweeps on 8× Ascend 910B. Each panel reports iteration time (top) and total model FLOPs (bottom) as KV-sharing budget increases, for all placements. Model

S 1F1B

LLaMA2-7B

MFU (%) at PP=2 / 4 / 8 Seq1F1B KV-Pipe

Iter time (s) at PP=2 / 4 / 8 1F1B Seq1F1B KV-Pipe

4K 46.5/44.0/38.0 49.5/48.0/45.0 51.5/50.0/48.0 8K 48.0/46.0/39.0 50.0/49.5/47.0 52.0/52.0/50.0 16K 49.0/46.5/40.0 51.0/50.0/48.0 53.0/52.5/51.0

112/125/150 215/265/305 420/520/610

106/114/128 101/109/120 205/245/265 196/232/248 400/470/520 385/448/485

LLaMA2-13B 4K 47.0/44.5/38.5 49.0/48.5/46.0 51.0/50.5/49.0 8K 48.5/45.0/39.0 50.0/49.5/47.5 52.0/52.5/50.5 16K 49.5/45.5/39.5 51.0/50.0/48.5 53.0/53.0/52.0

195/220/260 365/430/505 705/820/960

186/204/228 176/192/212 350/398/445 334/372/410 675/740/820 640/690/750

LLaMA3-8B

4K 46.0/43.5/37.5 49.0/47.8/45.0 51.0/49.8/48.0 8K 47.5/45.0/38.5 49.5/49.0/46.8 51.5/51.8/49.8 16K 48.5/45.2/39.0 50.5/49.6/47.8 52.5/52.0/51.0

118/132/158 228/278/320 445/545/635

111/121/136 105/114/126 217/257/278 205/240/258 420/495/545 402/468/505

Qwen2.5-14B

4K 47.5/45.0/39.0 49.8/49.5/47.0 52.0/51.8/50.0 210/236/278 200/218/242 188/202/222 8K 49.0/45.5/39.5 50.5/50.2/48.0 52.5/53.0/51.0 392/465/540 375/428/472 356/398/430 16K 50.0/46.0/40.0 51.5/51.0/49.0 53.5/54.0/52.5 760/885/1030 725/800/880 680/735/800

Table 4. GPU results on 8×V100 (compact). Each entry reports PP=2/4/8 in a single cell. KV-Pipe consistently improves MFU and reduces iteration time, with larger gains at PP=8.

introducing KV-sharing layers; it does not require changing the pipeline scheduler, micro-batching policy, or communication pattern. Therefore, KV-Pipe can be applied on top of a broad class of PP methods as long as they preserve the same notion of stage-level execution and the primary bottleneck remains determined by the maximum stage time. The key requirement for composability is that the PP method’s efficiency is sensitive to the distribution of compute across stages (i.e., the bottleneck stage dominates). Under this condition, any method that schedules or partitions pipeline work can benefit from KV-Pipe because KV-Pipe explicitly pushes the system toward a flatter stage-time profile.

architectural lever (KV reuse) and do not require separate mechanisms for training and inference. As a result, KV-Pipe can be viewed as a native system–architecture co-design knob: it improves PP efficiency during training and improves throughput/latency during serving, using a single consistent strategy grounded in stage-aware compute balancing.

While KV-Pipe is motivated by training-time PP efficiency, it also inherits the classical advantage of KV-sharing for inference: reducing redundant KV projections and KV-cache growth. This yields a double benefit: (i) improved training MFU by balancing stage workloads, and (ii) improved inference throughput by lowering per-token work and memory pressure, enabling larger effective batch sizes and faster decode. Importantly, these benefits arise from the same

This paper identifies an underexplored lever: the interaction between KV layout and pipeline partitioning. We introduce KV-Pipe, a stage-aware cross-layer KV-sharing framework that uses a simple FLOPs Imbalance Ratio (FIR) to locate bottleneck stages and then applies targeted KV sharing to flatten stage workloads (FIR→ 1) while simultaneously reducing KV-cache growth.

Conclusion

On 8×Ascend 910B across PP∈ {2, 4, 8}, KV-Pipe consis8

tently improves end-to-end utilization, delivering 3.1–9.2% MFU gains and 4.9–9.8% iteration-time reductions over baseline 1F1B. We also observe a clear operating principle: MFU peaks at an intermediate shared-KV budget—where imbalance is minimized—and can decline when excessive sharing over-corrects the original bottleneck and shifts the critical stage elsewhere. Overall, KV-Pipe establishes KV sharing as a system-level control knob that is orthogonal to PP scheduling: it directly reduces stage skew for higher train-time MFU and, by lowering KV-cache pressure, also improves long-context inference throughput.

with pathways. Journal of Machine Learning Research, 24(240):1–113, 2023. Dialameh, M., Karim, R., Rajabzadeh, H., Awad, O. M., Chen, B., Kwon, H. J., Ahmed, W., and Liu, Y. Echollama: Efficient caching for high-performance llama training. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing: Industry Track, pp. 2252–2269, 2025. Fan, S., Rong, Y., Meng, C., Cao, Z., Wang, S., Zheng, Z., Wu, C., Long, G., Yang, J., Xia, L., Diao, L., Liu, X., and Lin, W. Dapple: A pipelined data parallel approach for training large models, 2020. URL https://arxiv. org/abs/2007.01045.

References Adler, B., Agarwal, N., Aithal, A., Anh, D. H., Bhattacharya, P., Brundyn, A., Casper, J., Catanzaro, B., Clay, S., Cohen, J., et al. Nemotron-4 340b technical report. arXiv preprint arXiv:2406.11704, 2024.

Guo, T., Zhang, X., Du, J., Chen, Z., Xiao, N., and Lu, Y. gllm: Global balanced pipeline parallelism systems for distributed llms serving with token throttling. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1725–1741, 2025.

Agrawal, A., Panwar, A., Mohan, J., Kwatra, N., Gulavani, B. S., and Ramjee, R. Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills. arXiv preprint arXiv:2308.16369, 2023.

Harlap, A., Narayanan, D., Phanishayee, A., Seshadri, V., Devanur, N., Ganger, G., and Gibbons, P. Pipedream: Fast and efficient pipeline parallel dnn training, 2018. URL https://arxiv. org/abs, 1806.

Agrawal, A., Kedia, N., Panwar, A., Mohan, J., Kwatra, N., Gulavani, B., Tumanov, A., and Ramjee, R. Taming {Throughput-Latency} tradeoff in {LLM} inference with {Sarathi-Serve}. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pp. 117–134, 2024.

Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.

Ao, S., Zhao, W., Han, X., Yang, C., Zhang, X., Liu, Z., Shi, C., and Sun, M. Seq1f1b: Efficient sequence-level pipeline parallelism for large language model training. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 8998–9008, 2025.

Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019. Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., Lavaud, L. R., Lachaux, M.A., Stock, P., Scao, T. L., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mistral 7b, 2023. URL https: //arxiv.org/abs/2310.06825.

Arfeen, D., Zhang, Z., Fu, X., Ganger, G., and Wang, Y. Pipefill: Using gpus during bubbles in pipeline-parallel llm training. Proceedings of Machine Learning and Systems, 7, 2025. Blakeman, A., Basant, A., Khattar, A., Renduchintala, A., Bercovich, A., Ficek, A., Bjorlin, A., Taghibakhshi, A., Deshmukh, A. S., Mahabaleshwarkar, A. S., et al. Nemotron-h: A family of accurate and efficient hybrid mamba-transformer models. arXiv preprint arXiv:2504.03624, 2025.

Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M.A., Stock, P., Subramanian, S., Yang, S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of experts, 2024a. URL https://arxiv.org/abs/2401.04088.

Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling 9

Jiang, C., Jia, Z., Zheng, S., Wang, Y., and Wu, C. Dynapipe: Optimizing multi-task training through dynamic pipelines. In Proceedings of the Nineteenth European Conference on Computer Systems, EuroSys ’24, pp. 542–559. ACM, April 2024b. doi: 10.1145/3627703.3629585. URL http://dx.doi.org/10.1145/3627703. 3629585.

Peng, X., Shi, X., Zhang, H., Zhao, Y., and Qian, X. Dawnpiper: A memory-scablable pipeline parallel training framework. arXiv preprint arXiv:2505.05856, 2025. Rajabzadeh, H. Efficient learning for large language models. 2026.

Kim, T., Kim, H., Yu, G.-I., and Chun, B.-G. Bpipe: Memory-balanced pipeline parallelism for training large language models. In International Conference on Machine Learning, pp. 16639–16653. PMLR, 2023.

Schumann, R., Zhu, W., Feng, W., Fu, T.-J., Riezler, S., and Wang, W. Y. Velma: Verbalization embodiment of llm agents for vision and language navigation in street view. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 18924–18933, 2024.

Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems principles, pp. 611–626, 2023.

Smith, S., Patwary, M., Norick, B., LeGresley, P., Rajbhandari, S., Casper, J., Liu, Z., Prabhumoye, S., Zerveas, G., Korthikanti, V., et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201.11990, 2022.

Lan, Z., Chen, M., Goodman, S., Gimpel, K., Sharma, P., and Soricut, R. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv:1909.11942, 2019.

Su, Q., Zhao, W., Li, X., Andoorveedu, M., Jiang, C., Zhu, Z., Song, K., Giannoula, C., and Pekhimenko, G. Seesaw: High-throughput llm inference via model re-sharding. arXiv preprint arXiv:2503.06433, 2025.

Li, Q., Cui, L., Zhao, X., Kong, L., and Bi, W. Gsm-plus: A comprehensive benchmark for evaluating the robustness of llms as mathematical problem solvers. arXiv preprint arXiv:2402.19255, 2024.

Toroghi, A., Guo, W., Pesaranghader, A., and Sanner, S. Verifiable, debuggable, and repairable commonsense logical reasoning via llm-based theory resolution. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 6634–6652, 2024.

Li, Z., Liu, Y., Zhang, W., Yuan, T., Chen, B., Song, C., and Zhang, D. Slimpipe: Memory-thrifty and efficient pipeline parallelism for long-context llm training, 2025. URL https://arxiv.org/abs/2504.14519.

Wu, H., Li, L., Huang, H., Yi, T., Zhang, J., Yu, M., and Yan, J. Hshare: Fast llm decoding by hierarchical key-value sharing. In The Thirteenth International Conference on Learning Representations, 2025a.

Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., ShalevShwartz, S., Abend, O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Ratner, N., Rozen, N., Shwartz, E., Zusman, M., and Shoham, Y. Jamba: A hybrid transformer-mamba language model, 2024. URL https://arxiv.org/ abs/2403.19887.

Wu, Y., Wu, H., and Tu, K. A systematic study of cross-layer kv sharing for efficient llm inference. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 2: Short Papers), pp. 396–403, 2025b.

Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. Pipedream: Generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM symposium on operating systems principles, pp. 1–15, 2019.

Zhang, H., Wei, T., Zheng, Z., Du, J., Chen, Z., and Lu, Y. Td-pipe: Temporally-disaggregated pipeline parallelism architecture for high-throughput llm inference. In Proceedings of the 54th International Conference on Parallel Processing, pp. 689–698, 2025.

Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., et al. Efficient large-scale language model training on gpu clusters using megatronlm. In Proceedings of the international conference for high performance computing, networking, storage and analysis, pp. 1–15, 2021.

Zhao, Y., Long, Y., Liu, H., Kamoi, R., Nan, L., Chen, L., Liu, Y., Tang, X., Zhang, R., and Cohan, A. Docmatheval: Evaluating math reasoning capabilities of llms in understanding long and specialized documents. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 16103–16120, 2024. 10

Zuhri, Z. M. K., Adilazuarda, M. F., Purwarianti, A., and Aji, A. F. Mlkv: Multi-layer key-value heads for memory efficient transformer decoding. In Findings of the Association for Computational Linguistics: NAACL 2025, pp. 5516–5525, 2025.

Emb Full attn KV sharing LM head

1 4 0 0

0 4 0 0

0 3 1 0

0 3 1 1

FLOPs FWD time Memory

1.55E+13 59712.13 16.8425

1.55E+13 59237.7 15.5625

1.49E+13 57301.7 15.5597

1.708E+13 64860.0028 16.8395703

Table 5. Uniform placement for LLaMA2-7B, PP=8, S = 8K, MBS=1. The SKV budget is spread across stages 4–7; stage 7 remains the bottleneck. 0

Emb Full attn KV sharing LM head

1 4 0 0

0 4 0 0

0 2 2 0

0 4 0 0

0 2 2 1

FLOPs FWD time Memory

1.55E+13 59712.13 16.8425

1.55E+13 59237.7 15.5625

1.44E+13 55365.7 15.55664

1.55E+13 59237.7 15.5625

1.653E+13 62924.0012 16.8366406

Table 6. Symmetric-Bipolar placement for LLaMA2-7B, PP=8, S = 8K, MBS=1. The SKV budget is split between stages 3 and 7; stage 7 still has the largest FLOPs and forward time.

Additional Analysis and Validation

This appendix provides detailed analyses and additional experiments that complement the main evaluation. We organize the material around the questions that arise naturally from KV-Pipe’s design: (i) how different KV-sharing placements reshape stage-level compute, (ii) whether the selected configurations preserve model quality, (iii) how sensitive the method is to placement, the FIR diagnostic, and the stopping tolerance, (iv) whether KV-Pipe remains useful on top of stronger pipeline schedules and prior partitioning methods, and (v) whether the inference-side benefit persists under grouped-query attention (GQA). We conclude with practical guidance and the scope of the claims supported by the current experiments.

A.1

Stage-Level Effect of KV-Sharing Placement

We first make the mechanism of KV-Pipe explicit at the stage level. The main paper shows that, for LLaMA2-7B with PP=8 and S = 8K, the last stage is initially the bottleneck because it contains both Transformer blocks and the LM head. Tables 5–7 show how the same SKV budget is distributed by the three placement strategies considered in the paper. The important distinction is where the compute reduction is applied. Uniform placement spreads the SKV budget across several stages; Symmetric Bipolar concentrates it at two locations; Architecture-Balanced concentrates the budget on the current critical stage. As a result, Architecture-Balanced brings the tail-stage FLOPs close to the other stages and nearly equalizes the stage-time profile. This stage-level view explains why reducing total FLOPs alone is not sufficient: reductions that do not lie on the critical path can leave the pipeline bottleneck largely unchanged.

A.2

Quality–Efficiency Trade-off and Final Layer Assignments

Cross-layer KV sharing permanently changes the model architecture, so system efficiency must be interpreted together with model quality. We therefore evaluate validation perplexity and downstream accuracy for the exact configurations selected by KV-Pipe and compare them with full attention and a fixed Echo-style 25% tail-sharing baseline. Table 8 also reports the final SKV assignments. For LLaMA2-7B, KV-Pipe selects layers 25–32 at PP=2, 27–32 at PP=4, and 29–32 at PP=8. Hence, in every quality-tested configuration, the earliest converted layer remains in the second half 12

Emb Full attn KV sharing LM head

1 4 0 0

0 4 0 0

0 0 4 1

FLOPs FWD time Memory

1.55E+13 59712.13 16.8425

1.55E+13 59237.7 15.5625

1.543E+13 59051.9978 16.8307812

Table 7. Architecture-Balanced placement for LLaMA2-7B, PP=8, S = 8K, MBS=1. All four SKV conversions are applied to the original tail bottleneck. The resulting stage FLOPs are nearly equal, after which stage 0 becomes marginally critical in measured forward time.

Table 8. Quality–efficiency Pareto for LLaMA2-7B. Efficiency changes are reported relative to the corresponding full-attention 1F1B baseline. SKV denotes the number of attention layers converted to shared-KV layers; layer indices are 1-indexed. Downstream Avg. is the average over MMLU, HellaSwag, ARC-Challenge, and TruthfulQA. Setting Method

PP Seq. SKV Converted layers Earliest Late-half? Val. PPL ↓ Downstream Avg. ↑ MFU / Iter. time change

S=4K

Full-attn baseline Echo-style 25% sharing KV-Pipe, optimal

2 2 2

4K 4K 4K

0 8 8

– 25–32 25–32

– 25 25

– Yes Yes

5.54 5.51 5.51

52.2 51.9 51.9

0.00% / 0.00% +2.85% / -4.35% +3.10% / -4.90%

S=8K

Full-attn baseline Echo-style 25% sharing KV-Pipe, optimal

4 4 4

8K 8K 8K

0 8 6

– 25–32 27–32

– 25 27

– Yes Yes

5.67 5.50 5.50

52.5 52.0 52.0

0.00% / 0.00% +3.70% / -6.35% +4.32% / -7.56%

S=8K

Full-attn baseline Echo-style 25% sharing KV-Pipe, optimal

8 8 8

8K 8K 8K

0 8 4

– 25–32 29–32

– 25 29

– Yes Yes

5.47 5.53 5.49

52.2 52.0 52.2

0.00% / 0.00% +6.80% / -8.10% +9.17% / -9.80%

of the 32-layer network. We stress that this is an empirical observation for the evaluated configurations, not a theoretical guarantee that the bottleneck can never move to an earlier stage in an arbitrary model or PP layout. Fixed Echo-style sharing versus adaptive budget selection. The Echo-style baseline always converts the last 25% of LLaMA2-7B, i.e., layers 25–32. KV-Pipe instead chooses how much sharing is required from the pipeline imbalance. At PP=8 and S = 8K, KV-Pipe uses four SKV layers rather than eight, obtains validation PPL 5.49 versus 5.53 and downstream average 52.2 versus 52.0, and reduces iteration time by 9.80% rather than 8.10%. The practical advantage in this setting is therefore a better quality–sharing-budget–efficiency operating point, rather than a claim that more KV sharing is universally worse. At PP=2, the fixed 25% rule and KV-Pipe select the same layers (25–32), and their system results are correspondingly close. The difference becomes more meaningful at higher PP, where KV-Pipe stops earlier once the stage imbalance has been sufficiently reduced. Matched-budget interpretation. For the evaluated LLaMA2 settings, Architecture-Balanced placement remains taillocalized. If a fixed tail-sharing heuristic is given the same budget selected by KV-Pipe, it therefore selects the same layer sets: 25–32 for SKV=8, 27–32 for SKV=6, and 29–32 for SKV=4. These experiments support a more precise interpretation of KV-Pipe: in these LLaMA2 cases, its main benefit is to determine the SKV budget from the measured pipeline imbalance instead of fixing a 25% or 50% sharing ratio a priori. The bottleneck-retargeting mechanism is designed for the more general case in which the critical stage moves away from the tail.

A.3

Optional Safe-Layer Guardrail

Although the tail-first policy stays in late layers in all quality-tested configurations above, KV-Pipe optimizes a systems objective and does not imply that early-layer conversion is universally quality-neutral. To make the quality constraint explicit, KV-Pipe can be run with a user-controlled admissible layer set. 13

Table 9. Sensitivity and diagnostic ablations. MFU and iteration-time changes are relative to the full-attention 1F1B baseline at PP=8 and S = 8K. Val. PPL ↓

Downstream Avg. ↑

MFU / Iter. time change

17, 22, 27, 32 17, 18, 31, 32 29–32

5.50 5.50 5.49

52.1 52.1 52.2

+6.40% / -7.00% +7.10% / -7.60% +9.17% / -9.80%

5 4 3

28–32 29–32 30–32

5.50 5.49 5.48

52.1 52.2 52.2

+9.05% / -9.65% +9.17% / -9.80% +8.60% / -9.20%

4 4

29–32 28–31

5.49 5.49

52.2 52.2

+9.17% / -9.80% +8.90% / -9.45%

Ablation

Setting

SKV

Uniform placement Symmetric bipolar Architecture-Balanced

PP=8, S=8K PP=8, S=8K PP=8, S=8K

4 4 4

ϵ = 0.02 ϵ = 0.05 ϵ = 0.10

PP=8, S=8K PP=8, S=8K PP=8, S=8K

FIR-guided selection Max-stage-time guided

PP=8, S=8K PP=8, S=8K

Converted layers

Let α ∈ [0, 1) denote the minimum normalized depth at which KV sharing is allowed. We define Cα = {ℓ ∈ {1, . . . , L} : ℓ ≥ ⌈αL⌉} .

(9)

The layer-selection step is then restricted to ℓ⋆ = max {ℓ ∈ Si⋆ ∩ Cα : zℓ = 0} .

(10)

For example, α = 0.5 restricts conversions to the second half of the network. If no admissible layer remains before the target FIR tolerance is reached, KV-Pipe terminates and reports the residual imbalance rather than converting an earlier layer. This guard does not alter any result in Table 8, because all selected layers already lie in the second half of LLaMA2-7B.

A.4

Sensitivity to Placement, FIR, and the Stopping Tolerance

We next isolate three design choices: the placement policy, the stopping tolerance ϵ, and the imbalance signal. All rows in Table 9 use LLaMA2-7B with PP=8 and S = 8K under the same 1F1B setting as the NPU experiments. Placement policy. At the same SKV=4 budget, Architecture-Balanced placement gives the largest efficiency improvement among the three representative strategies. Uniform and Symmetric Bipolar placement spend part of the SKV budget on stages that are not always on the critical path, whereas Architecture-Balanced concentrates the reduction on the bottleneck stage. Tolerance sensitivity. The method is relatively stable over the tested range. Changing ϵ from 0.02 to 0.10 changes the selected budget from five to three SKV layers, while the MFU improvement remains between 8.60% and 9.17% and validation PPL remains between 5.48 and 5.50. Thus, the observed optimum does not require fine-grained tuning of ϵ for this workload. FIR versus measured stage-time guidance. FIR is intended as a lightweight analytic imbalance proxy, not as a claim of universal superiority over timing-based diagnostics. In this ablation, FIR-guided selection obtains +9.17% MFU and -9.80% iteration time, compared with +8.90% and -9.45% for max-stage-time guidance. This supports FIR as a useful low-cost offline signal for the evaluated setup, while direct timing may remain preferable when accurate stage profiles are already available. The current experiment should therefore be read as a diagnostic ablation rather than a universal predictive-power guarantee.

A.5

Composability with a Stronger Pipeline Schedule

KV-Pipe changes the per-stage workload and does not require a particular pipeline scheduler. To test whether its benefit remains after schedule-level bubble reduction, we combine KV-Pipe with Seq1F1B. Relative to Seq1F1B alone, adding KV-Pipe increases MFU by approximately 6.4% on LLaMA2-7B and 7.2% on LLaMA213B, while reducing iteration time by approximately 6.4% and 8.5%, respectively. This provides direct evidence for the composability claim: schedule-level bubble mitigation and stage-workload balancing address different sources of inefficiency. 14

Table 10. Composability with Seq1F1B. KV-Pipe provides an additional gain after applying a stronger pipeline schedule. Model / Setting

Method

MFU

Iter. time

LLaMA2-7B, PP=8, 8K

1F1B Seq1F1B Seq1F1B + KV-Pipe

39.0% 47.0% 50.0%

305s 265s 248s

LLaMA2-13B, PP=8, 16K

1F1B Seq1F1B Seq1F1B + KV-Pipe

39.5% 48.5% 52.0%

960s 820s 750s

Table 11. Training-time comparison with prior pipeline partitioning methods. Experiments use the DawnPiper/vPipe evaluation setting with 8 NVIDIA A100 40GB GPUs and 8 pipeline stages. Throughput is averaged over the feasible batch-size sweep. Speedup is normalized to vPipe-AS for each model. Model

Method

Pipeline mode

Max batch

Avg. speed (samples/s)

Speedup vs. vPipe-AS

Stage imbalance maxi ti /avgi ti

GPT-2 770M

GPipe vPipe-S vPipe-AS DawnPiper-AS KV-Pipe

Sync. Sync. Async. Async. Async. + SKV rebalance

7 7 16 20 20

16.0 16.2 28.7 33.0 34.5

0.56× 0.56× 1.00× 1.15× 1.20×

1.06 1.05 1.05 1.03 1.02

T5 780M

GPipe vPipe-S vPipe-AS DawnPiper-AS KV-Pipe

Sync. Sync. Async. Async. Async. + SKV rebalance

80 80 180 220 220

56.0 62.0 116.0 155.0 163.0

0.48× 0.53× 1.00× 1.34× 1.41×

1.42 1.31 1.25 1.12 1.07

Hardware coverage. The Ascend 910B experiments use MindSpeed-LM’s standard 1F1B schedule. Seq1F1B requires schedule-level runtime support that is available in our CUDA/GPU stack but is not available in our current Ascend 910B implementation. We therefore evaluate Seq1F1B composability on GPUs rather than implying identical scheduler coverage across software stacks. This mismatch is an implementation limitation of the current evaluation, not a requirement of KV-Pipe itself.

A.6

Comparison with Prior Pipeline-Parallel Methods

Prior PP systems attack inefficiency through several different levers: scheduling, partition search, dynamic micro-batching, memory balancing, and bubble filling. KV-Pipe is different in that it changes the effective cost of selected attention layers. We report both a quantitative comparison under the vPipe/DawnPiper benchmark setting and a high-level taxonomy of representative systems. A.6.1

Quantitative comparison with vPipe and DawnPiper

Compared directly with DawnPiper-AS, KV-Pipe improves average throughput by approximately 4.5% on GPT-2 and 5.2% on T5, while reducing the measured stage-imbalance ratio from 1.03 to 1.02 and from 1.12 to 1.07, respectively. Relative to vPipe-AS, the corresponding speedups are 1.20× and 1.41×. These results should be interpreted as evidence that changing the per-layer cost profile can provide additional headroom beyond partition search, not as a claim that KV-Pipe subsumes general PP partitioning methods. A.6.2

Relationship to other PP optimization levers

Table 12 summarizes the role of representative PP methods at a high level. This table is intended as a taxonomy rather than a performance ranking. Methods that primarily change the schedule or search for a partition can in principle be combined with KV-Pipe because KV-Pipe operates on the model’s per-layer cost profile. 15

Table 12. High-level positioning of KV-Pipe relative to representative PP optimizations. The table summarizes the primary optimization lever rather than claiming that the categories are mutually exclusive. Method

Primary optimization lever

Relationship to KV-Pipe

KV-Pipe

Stage-aware cross-layer KV sharing that changes the effective per-layer and per-stage compute profile. Micro-batched synchronous pipeline scheduling.

Directly targets stage skew; can be layered with scheduling or partition-search methods.

GPipe Huang et al. (2019)

PipeDream Narayanan et al. (2019) DAPPLE Fan et al. (2020)

Asynchronous pipeline execution and forward/backward overlap. Joint planning of data and pipeline parallel execution.

DynaPipe Jiang et al. (2024b)

Dynamic micro-batching and runtime adaptation.

BPipe Kim et al. (2023)

Memory-oriented pipeline balancing and activation movement.

SlimPipe Li et al. (2025)

Fine-grained sequence slicing for pipeline execution. Filling pipeline bubbles with additional work/jobs.

PipeFill Arfeen et al. (2025)

vPipe

Pipeline partitioning with memory/compute considerations.

DawnPiper Peng et al. (2025)

Cost-model-driven pipeline partition search.

Unity / Galvatron / FlexPipe / Varuna

Automatic or hybrid search over pipeline/parallel execution strategies.

A.7

Primarily changes execution schedule; KVPipe can change the stage costs seen by the schedule. Primarily changes execution overlap; stagecost balancing is orthogonal. Planner/search can operate on a model whose layer costs have already been modified by KV sharing. Runtime adaptation and offline stage-cost reshaping address different sources of imbalance. Targets memory imbalance; KV-Pipe primarily targets compute skew while also reducing KV-related memory. Sequence-level scheduling can be combined with stage-level compute reshaping. Uses otherwise idle slots; KV-Pipe instead attempts to reduce the underlying stage skew that creates idle time. Partition search changes where work is placed; KV-Pipe changes the cost of selected attention layers. Complementary in objective: DawnPiper changes placement, while KV-Pipe can reshape the cost profile supplied to a partitioner. Potentially composable with KV-Pipe because their search space can operate on the modified stage-cost profile.

Attention FLOPs under MHA and GQA

The original inference case study uses multi-head attention (MHA), whereas many modern LLMs use grouped-query attention (GQA). We therefore derive the attention FLOPs for both cases and quantify what cross-layer KV sharing can still remove under GQA. We follow the standard matrix-multiplication FLOP rule and count a fused multiply-add (MAC) as two FLOPs.

A.7.1

Matrix-multiplication FLOPs

For matrices A ∈ Rm×k and B ∈ Rk×n , FLOPs(AB) = 2mkn.

A.7.2

(11)

Notation

Let sequence length be s, batch size be b, hidden size be h, number of query heads be a, and per-head dimension be d = h/a. The input activations to an attention layer are X ∈ R(bs)×h . 16

For GQA, let g be the number of KV heads and define r≜

a ≥ 1, g

(12)

so that the total KV projection dimension is gd = h/r. A.7.3

Q/K/V projections

MHA. MHA uses full-size projections for all three tensors: Q = XWQ , with W(·) ∈ Rh×h . Hence

K = XWK ,

V = XWV ,

2 2 FLOPsMHA QKV = 3 · 2(bs)h = 6bsh .

GQA. GQA keeps the query projection full-size but reduces the K/V dimension to h/r: 2 h 2 2 = 2bsh 1 + . = 2(bs)h + 2 · 2(bs)h FLOPsGQA QKV r r A.7.4

(13) (14)

(15)

Attention score and weighted sum

For each query head, QK⊤ costs 2s2 d FLOPs. Summed over the batch and all query heads, FLOPsscore = 2bas2 d = 2bs2 h.

(16)

The weighted sum Attn · V has the same cost: FLOPsctx = 2bs2 h.

(17)

Although GQA shares K/V across groups of query heads, the score and weighted-sum operations are still performed per query head in the standard formulation, so Eqs. (16)–(17) apply to both MHA and GQA. A.7.5

Output projection

The concatenated attention output is projected with WO ∈ Rh×h : FLOPsout = 2bsh2 . A.7.6

(18)

Total attention FLOPs per token

For MHA, 2 2 2 2 FLOPsMHA attn = 6bsh + 2bs h + 2bs h + 2bsh

= 8bsh2 + 4bs2 h,

(19)

and therefore MHA

FLOPs/tokenattn = 8h2 + 4sh.

(20)

For GQA, 2 1+ + 4bs2 h + 2bsh2 r 4bsh2 = 4bsh2 + + 4bs2 h, r

2 FLOPsGQA attn = 2bsh

(21)

Table 13. GQA inference validation. KV-Pipe improves decoding throughput on GQA-based models. Model

Variant

Throughput

Time/token

LLaMA3-8B, 32K

GQA baseline KV-Pipe, SKV=4

620 tok/s 665 tok/s

1.61 ms 1.50 ms

Qwen2.5-14B, 128K

GQA baseline KV-Pipe, SKV=6

170 tok/s 184 tok/s

5.88 ms 5.43 ms

and therefore GQA

FLOPs/tokenattn = 4h2 +

4h2 + 4sh. r

(22)

For a decoder-only Transformer with L layers, the forward attention FLOPs per token are L times the corresponding per-layer quantity. For training, a common approximation is forward plus backward ≈ 3× forward FLOPs.

A.7.7

What KV sharing removes under GQA

Cross-layer KV sharing removes the new K/V projections for a shared layer. Therefore, the forward projection FLOPs saved per shared layer and token are ∆FGQA =

4h2 . r

(23)

If each cached element occupies belem bytes, the per-token KV-cache reduction for a shared layer is h ∆BGQA = 2 belem . r

(24)

Thus, GQA reduces the amount of K/V redundancy available to remove by the grouping ratio r, but the savings remain non-zero.

A.8

Inference Validation under GQA

We complement the FLOP analysis with decoding measurements on two GQA models. Table 13 shows that the inference-side benefit of KV sharing persists even when K/V projections have already been reduced by grouped-query attention. KV-Pipe improves throughput by approximately 7.3% on LLaMA3-8B and 8.2% on Qwen2.5-14B. These results extend the inference case study beyond MHA and support the “double benefit” claim under a more deployment-relevant attention configuration, while the analysis above explains why the absolute K/V projection saving is smaller than under MHA.

A.9

Practical Operating Regime

The experiments identify a clear operating regime in which KV-Pipe is most useful. The benefit grows when stage imbalance is large enough to materially limit utilization. On Ascend 910B, the relative MFU improvement increases from 3.10% at PP=2 to 4.32% at PP=4 and 9.17% at PP=8, while iteration-time reduction increases from 4.90% to 7.56% and 9.80%, respectively. The GPU experiments in the main paper show the same qualitative trend across model families and context lengths. Accordingly, KV-Pipe is most attractive for moderate-to-high PP degrees, long-context workloads, and partitions with non-uniform components such as a heavy LM head or heterogeneous attention layers. At low PP degree or when the baseline stage profile is already well balanced, the available headroom is smaller and the architectural change may provide only a modest benefit. 18

Limitations and Future Work The quality results in this work should be interpreted within their measured scope. We validate perplexity and downstream accuracy for LLaMA2-7B; the LLaMA2-13B, LLaMA3-8B, and Qwen2.5-14B experiments provide system-level efficiency evidence but are not used to claim architecture-independent quality preservation. In particular, we do not assume that observations established for dense LLaMA-family models automatically transfer to MoE, hybrid Transformer–SSM architectures, or arbitrary training regimes. Likewise, the current study evaluates PP degrees up to 8. Evaluation on PP=16/32 and very large models remains future work. At larger scales the relative contribution of the LM head and the location of the critical stage may change, making the bottleneck-retargeting behavior more important. Finally, the greedy Architecture-Balanced procedure is designed as a lightweight offline heuristic; we do not claim global optimality over all possible layer-sharing assignments. Joint optimization of pipeline placement, KV-sharing placement, and an explicit quality constraint is an interesting direction for future work.

Impact Statement This paper presents work whose goal is to advance the field of Machine Learning by improving the computational efficiency and scalability of Large Language Models (LLMs). By introducing KV-Pipe, we provide a mechanism to optimize pipeline parallelism through architectural KV-sharing, which directly reduces the energy and hardware resources required to train and deploy frontier models. The primary societal consequence of this research is the potential to lower the environmental impact of large-scale AI development by increasing Model FLOPs Utilization (MFU) and reducing the total training time for a given model quality. Additionally, making high-performance training more efficient can lower the barrier to entry for institutions with limited computational budgets, fostering a more diverse and equitable research ecosystem. While the broader impacts of AI development include well-established ethical considerations regarding model bias, safety, and the use of generated content, our work focuses on the underlying systems optimization level and does not introduce new specific ethical risks beyond those generally associated with the advancement of the field.

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

Analyze & explore related research in Conceptio

Related documents