34. GQA-Aware KV-Cache Memory Layout

Problem

Grouped-query attention's real payoff isn't the forward-pass math (see problem 28) — it's what happens to the KV-cache during incremental decoding. A cache built for ordinary multi-head attention stores one key/value per query head; a GQA-aware cache stores only num_kv_heads worth (far fewer, e.g. 8x fewer for Mistral 7B's 32 query heads / 8 kv heads), and only expands up to the full num_query_heads width at attention time, inside the matmul, never in the stored cache itself. Getting this distinction backwards — storing the already-expanded, query-head-width tensors — silently throws away the entire memory savings GQA is supposed to provide during serving.

Implement a single incremental-decode step: given the existing compact (kv-head-width) cache and this step's new compact key/value, append the new step to the cache (still at kv-head width) and separately produce the expanded (query-head-width) key/value for this step's attention computation.

Source: src/34_gqa_kv_cache_layout.py

def gqa_kv_cache_step(
    cache_k: torch.Tensor | None, cache_v: torch.Tensor | None,
    new_k: torch.Tensor, new_v: torch.Tensor, num_query_heads: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ...

>>> new_k = new_v = torch.randn(1, 2, 1, 4)
>>> uk, uv, ek, ev = gqa_kv_cache_step(None, None, new_k, new_v, num_query_heads=6)
>>> uk.shape, ek.shape
(torch.Size([1, 2, 1, 4]), torch.Size([1, 6, 1, 4]))

Step-by-Step Approach

1. Handle the first step specially: if cache_k is None, the "updated cache" is just new_k, new_v — nothing to concatenate onto yet.
2. On later steps, append along the sequence axis at kv-head width: torch.cat([cache_k, new_k], dim=2) — the cache tensor's head dimension must stay num_kv_heads throughout, never grow.
3. Compute group = num_query_heads // num_kv_heads from the updated cache's own head dimension.
4. Produce the expanded view only now, only for this step's return value: updated_k.repeat_interleave(group, dim=1) — this is a fresh tensor computed on demand, not something stored back into the cache.
5. Return all four tensors: the compact updated_k/v (what the caller should keep as the persistent cache for the next step) and the expanded expanded_k/v (what the caller should actually pass into this step's attention computation).
6. Sanity-check memory: updated_k.numel() must equal batch * num_kv_heads * cached_len * head_dim at every step — if it ever matches the num_query_heads-width formula instead, the expansion leaked into the stored cache.

The key insight is a discipline, not a new algorithm: expansion via repeat_interleave is cheap to redo every step, so the cache should never pay to store what it can regenerate on demand — the compact cache plus a query-head count is a complete, minimal representation.

Reference solution

import torch


def gqa_kv_cache_step(
    cache_k, cache_v, new_k, new_v, num_query_heads,
):
    # append this step onto the EXISTING compact cache -- storage always
    # stays at kv-head width, that's the entire memory-saving property
    if cache_k is None:
        updated_k, updated_v = new_k, new_v
    else:
        updated_k = torch.cat([cache_k, new_k], dim=2)
        updated_v = torch.cat([cache_v, new_v], dim=2)

    # expand to query-head width only now, only for this step's attention
    # computation -- never write the expanded tensor back into the cache
    num_kv_heads = updated_k.shape[1]
    group = num_query_heads // num_kv_heads
    expanded_k = updated_k.repeat_interleave(group, dim=1)
    expanded_v = updated_v.repeat_interleave(group, dim=1)

    return updated_k, updated_v, expanded_k, expanded_v

Key Functions & Tricks

• torch.cat([cache_k, new_k], dim=2) — the ordinary cache-append operation, unchanged from a non-GQA cache; GQA only changes what width the head dimension stays at.
• tensor.repeat_interleave(group, dim=1) — recomputed fresh every step from the current cache, never persisted; this is what keeps the cache's memory footprint at kv-head width regardless of how many decode steps have run.
• updated_k.shape[1] — reading num_kv_heads back off the cache tensor itself keeps the function stateless with respect to that value, rather than needing it passed in separately.
• tensor.numel() — used in the test suite as a direct, quantitative check that the cache's storage cost tracks num_kv_heads, not num_query_heads — the concrete memory-savings claim GQA caching makes.

How to Recognize This Pattern