Share keys and values without confusing query-head identity

For a regular grouping, let Hq be divisible by Hkv and assign consecutive groups of Hq/Hkv query heads to each KV head. Each query head keeps its own Q projection and output, but uses the K and V tensors selected by that mapping. MHA is the endpoint Hkv=Hq; MQA is Hkv=1. The exact group ordering belongs to the model and API contract, not merely to the shapes of two head axes.

FlashAttention 2.8.3 documents this consecutive grouping for its fewer-KV-head interface. For example, six query heads and two KV heads assign query heads zero through two to KV head zero, and three through five to KV head one. A kernel should reuse the shared data where possible rather than physically repeating it into independent buffers. An implementation that expands K and V can preserve arithmetic results while losing much of the storage or bandwidth advantage.

A logical KV cache contains both keys and values for each stored token, KV head, and head coordinate. With equal key/value width d and s bytes per element, its payload scales with Hkv rather than Hq. Multiply by batch size and layer count when counting a full model. Allocator rounding, page tables, unused capacity, quantization scales, and temporary expanded tensors are additional storage. State whether a number describes live token payload or total reserved memory.

Sharing KV heads does not necessarily reduce leading query-key and probability-value multiply work by the same factor: each query head still forms its own attention distribution and output. Fewer KV projections and better cache reuse can help, especially in incremental decoding where cache reads matter. Actual speedup depends on the implementation, batch, sequence lengths, and hardware. A cache payload ratio is a memory statement, not a measured latency ratio.

During training, a shared key or value tensor receives gradient contributions from every query head in its group. A reference can explicitly repeat K and V along the query-head axis and let differentiation sum back to the original tensors. A fused implementation should agree with that reference within declared numerical tolerances. Incorrect head mapping or failure to accumulate group contributions can pass a forward test with identical heads while failing on deliberately distinct inputs.

The original MQA and GQA papers study the quality and efficiency consequences of shared projections, including conversion or uptraining procedures. Those empirical outcomes do not prove that every existing MHA checkpoint can be changed at inference with identical outputs. Evaluate model quality separately from kernel equivalence. For a fixed GQA model, the fused and expanded implementations target the same operator; converting the model’s projection structure is another experiment with its own validation requirements.