PaperlayerView exact revision on arXiv

Source-linked AI summary

Grouped Value Attention: Efficient KV Caching via On-Demand Key Reconstruction

Vishesh Tripathi, Abhay Kumar, Ramsha Khan

arXiv:2609.13285v2cs.ARcs.LG

TL;DR

KV-cache memory and read traffic grow with sequence length, while GQA still stores both keys and values. GVA stores grouped values and reconstructs content keys, pairing that cache with a shared positional channel; its studied configurations reduce cache scalars by about 45–47% while reaching accuracy close to GQA.

  • Problem

    The KV cache grows with context length and can become a dominant capacity and bandwidth cost, while GQA still stores both keys and values at every step.

    Source: §1 · p. 1

  • Method

    GVA caches grouped values, reconstructs per-head content keys with learned linear maps absorbed into the query, and uses a small shared decoupled RoPE channel for position.

    Source: §1 · pp. 1–2

  • Results

    45–47% fewer cache scalars than matched GQA; the 16-dimensional positional variant reaches 44.35 average accuracy versus 44.36 for GQA and 43.88 for MLA.

    Source: §1 · p. 3

  • Takeaways & Limitations

    GVA retains distinct head-specific content keys while reducing the intended persistent cache relative to GQA.

    Source: §1 · pp. 2–3

    Evaluation covers one approximately 350M-parameter scale and one 30B-token data mix, and decode latency and throughput gains are not established.

    Source: §6 · p. 9

Abstract

from arXiv · show

The KV cache is a primary bottleneck for Transformer decoding: its memory footprint and cache-read traffic grow with sequence length. Grouped-query attention (GQA) reduces this cost by sharing key-value heads, but still stores both a key and a value at every step. We introduce Grouped Value Attention (GVA), which stores grouped values and reconstructs content keys with a learned linear map. At inference, the map can be absorbed into the query, eliminating the need to materialize content keys in the intended decode path. A small shared decoupled RoPE channel retains positional information through a separately cached positional key. For the configurations studied, this representation reduces persistent cache scalars by approximately 45-47% relative to matched GQA. At the 350M-parameter scale with 30B FineWeb-Edu tokens, the 16-dimensional positional variant reaches 44.18 average accuracy across five tasks, compared with 44.36 for GQA and 43.88 for MLA. These results demonstrate near-GQA benchmark accuracy with a more compact cache representation. To translate this compact representation into faster autoregressive inference, we have developed custom decoding kernels and are currently evaluating their end-to-end inference performance with an open-source release planned soon.

1 Introduction

The introduction frames KV-cache growth as a decoding cost and presents GVA, which stores grouped values and reconstructs per-head content keys. It pairs this design with a shared positional channel to reduce cache size while retaining benchmark accuracy close to GQA.

1 Introduction

  • The cache grows linearly with context length and can become a dominant capacity and bandwidth cost in long-context serving.GVA targets decode-time cache costs; it seeks no prefill advantage, and latency or throughput gains are not established in the reported experiments.Source: §1 · pp. 1–2
  • GQA shares key–value heads but still writes both streams, while MLA compresses them into a joint latent with an extra projection and a more involved decode path.GVA instead uses grouped values as its persistent content state.Source: §1 · p. 1
  • GVA stores G grouped value streams and reconstructs distinct content keys for H query heads with per-head linear maps absorbed into the query at inference.This keeps head-specific content keys while removing the content-key stream from the intended cache.Source: §1 · pp. 1–2
  • A small shared decoupled RoPE channel preserves positional information without restoring a full key cache.The shared positional key adds dr cached scalars per token rather than Gdr.Source: §1 · p. 2
  • 45–47% fewer cache scalars than matched GQA; the 16-dimensional positional variant reaches 44.35 average accuracy, versus 44.36 for GQA and 43.88 for MLA.The scores are averaged over three runs with different random seeds.Source: Figure 1 · p. 2, §1 · p. 3

2 Related Work

Related work describes attention-cache designs that trade cache size against head sharing or latent compression. MQA and GQA share key–value heads across query heads, while MLA caches a joint latent and shared positional slice.

2 Related Work

  • Multi-head attention stores separate keys and values for each head, with per-layer, per-sequence cache size expressed as 2THdh [1] [2].The notation passage introduces this storage count.Source: §2 · p. 3
  • MQA keeps H query heads but shares one key–value head, making its cache a factor of H smaller than multi-head attention, with reported quality and sometimes training-stability costs.The passage attributes these trade-offs to prior work [2] [3].Source: §2 · p. 3
  • GQA splits query heads into G groups that share key–value heads; prior experiments report quality close to multi-head attention and speed close to MQA for modest G [3].G=H and G=1 recover multi-head attention and MQA, respectively [3].Source: §2 · p. 3
  • MLA caches a low-rank joint key–value latent plus a small shared RoPE slice, and can absorb the key up-projection into the query at decode [4].This avoids materializing a full content-key tensor [4].Source: §2 · p. 3

3 Grouped Value Attention

Grouped Value Attention (GVA) stores grouped values and reconstructs head-specific content keys with learned linear maps. Its decode path absorbs those maps into queries, while a shared decoupled RoPE slice retains positional information and keeps the persistent cache roughly half the size of matched GQA for the configurations studied.

3.1 Value-only cache and Kh = Vg(h)Mh

  • Directly reusing values as keys halves the cache but does not recover the GQA training loss, so it is only a reference design.The authors report that one vector being used both for scoring and retrieval was inadequate in this initial approach.Source: §3.1 · p. 4

3 Grouped Value Attention

  • GVA caches grouped values and uses a learned per-head linear map to reconstruct distinct content keys, rather than writing a separate content-key stream.The values carry content to the output, while each map selects features used for scoring.Source: §3 · p. 3, §3.1 · p. 4

3.1 Value-only cache and Kh = Vg(h)Mh

  • GVA can reconstruct H distinct content-key streams from G cached value streams, unlike GQA’s shared key within each query-head group.The reconstructed keys remain linear transforms of grouped values, not unconstrained independent projections.Source: §3.1 · p. 4

3.2 Absorption at decode

  • Absorbing each fixed reconstruction map into the query lets content attention score directly against cached values without materializing keys.The transformed query is computed once per query token and head; the identity is exact.Source: §3.2 · p. 5

3.3 Decoupled RoPE

  • Decoupled RoPE preserves relative position with an unrotated content slice and a short rotated positional slice shared across heads.The shared positional key costs dr scalars per token, rather than Gdr.Source: §3.3 · p. 5, §3.2 · p. 5

3.4 Cache size

  • GVA’s persistent cache is roughly half GQA’s for the widths studied, with the shared positional stream adding only a few percent.The cache count describes persistent state, not measured peak serving memory.Source: §3.4 · p. 6

4 Experiments

The experiments train matched-width decoder-only Transformers at 350M parameters on FineWeb-Edu and compare GQA, MLA, and several GVA variants. Evaluation uses training loss and five zero-shot task accuracies, while end-to-end decoding performance remains under evaluation.

4 Experiments

  • Models are trained from scratch at 350M parameters on a 30B-token FineWeb-Edu sample, with matched query-head counts across GQA, MLA, and GVA.The runs share data order, token budget, optimizer, and context length unless a variant is specified.Source: §4 · p. 6

4.1 Variants

  • The GVA comparisons vary map initialization, query RMSNorm, and whether a shared decoupled RoPE positional slice is used.The reported variants also include shared KV as a reference; it halves cache size but does not recover the GQA loss and is not a proposed system.Source: §4.1 · pp. 6–7

4.2 Evaluation

  • Evaluation reports training loss and zero-shot accuracy on HellaSwag, WinoGrande, OpenBookQA, and both ARC splits, averaging three random-seed runs.Average accuracy is the unweighted mean of the five task-level means; small differences are not statistically significant without assessing seed variability.Source: §4.2 · p. 7
  • End-to-end inference measurements, including fused decoding throughput against MLA and GQA, are not reported because custom kernels remain under evaluation.An open-source release is planned soon.Source: §4.2 · p. 7

5 Results

The 16-dimensional decoupled-RoPE variant reaches 44.35 average accuracy, 0.01 points below GQA, while reducing intended cache scalars by about 47%. Scale-matched GVA tracks GQA in training loss, whereas shared KV halves the cache but does not recover the loss gap.

Downstream accuracy

  • Average accuracy is the unweighted mean of five task-level means, each computed across three runs with different random seeds.The tasks are HellaSwag, WinoGrande, OpenBookQA, and the Easy and Challenge splits of ARC.Source: §5 · p. 7
  • Scale-matched GVA with query RMSNorm reaches 44.41 average accuracy, close to GQA at 44.36; scale matching mainly improves early loss, not final averages consistently.The baseline without matched initialization reaches 43.91, while the default without query RMSNorm reaches 43.77.Source: §5 · p. 7

Training loss

  • After the initial transient, scale-matched GVA tracks GQA and MLA in training loss; decoupled-RoPE runs follow the same broad trajectory.The scale-matched trajectory is shown in Figure 2.Source: §5 · p. 7, Figure 2 · p. 8

Cache

  • About 47% (dr=16) and 45% (dr=24) fewer cache scalars are required than with GQA for the studied grouping.These are representation-level counts for the intended decode state, not measured serving-memory reductions.Source: §5 · p. 8

Downstream accuracy

  • 44.35 (dr=16) and 44.29 (dr=24) average accuracy; neither decoupled-RoPE row beats GQA, and width/content allocation effects remain unisolated.The table reports accuracy across five benchmarks averaged over three runs.Source: §5 · pp. 7–8, Table 2 · p. 8

Shared KV

  • Shared KV halves the GQA cache, but its training loss stays above GQA throughout the displayed run, motivating key reconstruction from a dedicated value.All later GVA runs retain a dedicated value and reconstruct the key.Source: §5 · p. 7, Figure 3 · p. 9

6 Limitations

The reported cache reduction describes the intended decode state, while kernel performance evaluation remains in progress. Comparisons are limited to one model scale and data mix, with several design dimensions left for future study.

6 Limitations

  • 45–47% cache reduction describes the intended decode state: grouped values plus a shared positional key.Omitting the positional slice gives an idealized 50% value-only reduction.Source: §6 · p. 9
  • Fused decode throughput, peak serving memory, and batch capacity measurements are not reported.Custom decoding kernels are being tested, with an open-source release planned.Source: §6 · p. 9
  • Comparisons cover one approximately 350M-parameter scale and one 30B-token FineWeb-Edu mix, without systematic sweeps of RoPE width, allocation, or longer contexts.Each benchmark configuration averages three random seeds.Source: §6 · p. 9

7 Conclusion

GVA stores grouped values and reconstructs keys with a linear map, targeting roughly half the persistent cache while keeping downstream quality in the same broad band as GQA. Its custom decoding kernels are still undergoing inference-performance evaluation.

7 Conclusion

  • Roughly half the persistent cache is intended relative to GQA, while downstream quality remains in the same broad band.The conclusion reports the comparison in Table 2.Source: §7 · p. 9
  • GVA reconstructs keys from grouped values with a linear map, while a small shared decoupled RoPE channel retains position compatibly with absorbing the map at decode.The conclusion contrasts this with shared key-value vectors, which halve cache size but do not match GQA in the observed run.Source: §7 · p. 9
  • End-to-end inference performance is still being tested, with broader sweeps of RoPE width, model scale, and random seeds left for further work.An open-source release is planned soon.Source: §7 · p. 9

A Query–Key Scale

Attention scale affects whether softmax weights are flat or concentrated. Because GVA key scale depends on its learned map, initialization matches reconstructed-key RMS to query RMS; positional and content scales are assessed separately.

A Query–Key Scale

  • With centered, weakly correlated query and key coordinates, too-small scales yield nearly equal logits and flat attention, while too-large scales make a few positions dominate.Figure 4 illustrates both regimes.Source: A Query–Key Scale · p. 11
  • Default map initialization left GVA keys far below query scale in initial runs, producing nearly uniform attention and spending early training recovering from the mismatch.In GVA, key scale also depends on the map because keys are reconstructed from values.Source: A Query–Key Scale · p. 11
  • The map is initialized so reconstructed keys start at approximately query RMS, using the query scale after any query normalization.The map remains learnable, and the variance-fixed row remains competitive without optional query RMSNorm; this initialization does not normalize cached values.Source: A Query–Key Scale · p. 11
  • Positional and content scales are initialized per slice: the shared positional-key projection is set to match the initial RMS of the positional query.Pooling both slices into one key variance hid the mismatch in early runs.Source: Decoupled RoPE · p. 11

Decoupled RoPE.

The figures illustrate how query and key scales shape attention heatmaps. Small scales produce flat weights, while large scales concentrate them, even when query and key scales match.

Decoupled RoPE.

  • Figure 4 contrasts matched query and key scales with keys that are too small or too large; small keys yield nearly uniform attention and large keys concentrate weights.Each panel uses its own color scale.Source: Figure 4 · p. 11
  • Small σ gives nearly uniform weights (≈0.1 per key), while large σ gives peaked, near one-hot rows, despite matched query and key scales.The columns use σQ = σK values 0.006, 1, 4, and 8 from left to right.Source: Decoupled RoPE · p. 11, Figure 5 · p. 12
  • At unit σ, the heatmaps show moderate variation between the nearly uniform low-scale and peaked high-scale patterns.Each panel uses its own color scale.Source: Figure 5 · p. 12

B Extra Evaluation Curves

The appendix plots per-task accuracy over training for seven configurations, retaining the original run names and horizontal-axis units.

B Extra Evaluation Curves

Loading 2609.13285v2…