How Local and Global Attention Share a Memory Budget

Why People Wanted It

A model’s weights are not its entire inference-memory bill. During a long conversation, attention also accumulates keys and values: representations of earlier tokens that later tokens can consult. The device-memory target here is 32 GB, so context storage matters alongside the dense model’s weights.

The problem is not simply storing fewer tokens. If every layer discards distant history, it also loses direct access to those discarded entries. So the real question is: which layers should retain the full history, and how much should each retained token cost?

The Core Move

Give every layer its own cache, retain only a recent window in local layers, retain the full prefix in periodic global layers, and share stored keys and values across groups of query heads.

Let $C_{\ell,t}$ denote layer $\ell$’s cache after processing token $t$, and let $(k_{\ell,t},v_{\ell,t})$ be that token’s new key-value entry. The cache update is:

$$C_{\ell,t}=\begin{cases} \operatorname{tail}_{W}\!\left(C_{\ell,t-1}\mathbin{\Vert}(k_{\ell,t},v_{\ell,t})\right), & \ell\text{ local},\\ C_{\ell,t-1}\mathbin{\Vert}(k_{\ell,t},v_{\ell,t}), & \ell\text{ global}. \end{cases}$$

Here, $\Vert$ means append, and $\operatorname{tail}_{W}$ keeps the most recent $W$ entries. The configured window is $W=2048$, with a repeating local, local, local, global pattern. For this illustration, the window includes the current token.

The important boundary is per layer, across tokens. A cache is not a list of entries collected from successive layers.

Quiz

After a conversation exceeds the local window, how should the attention caches evolve?

Once local windows fill, the entire model’s cache stops growing.

All layers append entries to one shared cache, which is trimmed whenever a local layer runs.

Each local layer keeps its recent window, while each global layer continues accumulating its own history.

$ check

Why Each Layer Needs Its Own History

Here is the part that gets technical: cache ownership determines which representations a query is allowed to compare, not merely where tensors are stored.

At a given layer, a learned projection turns the current hidden representation into a query $q$, a key $k$, and a value $v$. Keys are lookup descriptors; values are the vectors returned by the lookup. A deeper layer receives a transformed representation and makes its own keys and values. Consequently, an entry from one layer cannot simply stand in for the corresponding entry at another layer.

Number tokens starting at one. The permitted token positions at time $t$ are:

$$\mathcal{I}_{\ell,t}=\begin{cases} \{\max(1,t-W+1),\ldots,t\}, & \ell\text{ local},\\ \{1,\ldots,t\}, & \ell\text{ global}. \end{cases}$$

Both sets exclude the future. Local attention additionally excludes positions older than its window. With a cache containing exactly these positions, no extra causal mask is needed for a single-token decoding step.

For one query head $h$, the ordinary attention lookup over that cache is:

$$a_{\ell,t,h,s}=\operatorname{softmax}_{s\in\mathcal{I}_{\ell,t}} \left(\frac{q_{\ell,t,h}^{\top}k_{\ell,s,g(h)}}{\sqrt{d}}\right), \qquad o_{\ell,t,h}=\sum_{s\in\mathcal{I}_{\ell,t}}a_{\ell,t,h,s}v_{\ell,s,g(h)}.$$

The query scores each available key, softmax normalizes those scores over retained positions, and the result weights the values. Here $d$ is the head dimension; $g(h)$ selects the key-value head shared by query head $h$.

A global layer can directly consult distant entries in its own history. It does not reopen the discarded entries of a local layer. Its output instead feeds subsequent layers through the hidden representation. This provides a route for long-range information, not a guarantee that every old detail survives or is retrieved.

That explains which tokens remain accessible. What determines the storage cost of each token?

Why Shared Heads Reduce Storage but Do Not Stop Growth

Now the version for readers who want more of the machinery: grouped-query attention changes the head dimension of the cache, while windowing changes its sequence dimension.

Each layer has 32 query heads, 2 key-value heads, and head dimension 128. In a conventional contiguous grouping, $g(h)=\lfloor h/16\rfloor$ for zero-indexed query heads. Each group of queries reads the same keys and values, but different queries can assign different attention weights to them.

For one sequence, a cache tensor for keys has shape $[n_{\ell,t},2,128]$, and values have the same shape. Here $n_{\ell,t}$ is the number of retained positions. Old queries need not be stored for future attention lookups.

A direct storage accounting, excluding allocations and metadata, is:

$$M_{\mathrm{KV}}(t)=2bH_{\mathrm{KV}}d \left[L_{\mathrm{local}}\min(t,W)+L_{\mathrm{global}}t\right].$$

The leading factor of two counts keys and values; $b$ is bytes per stored scalar; $H_{\mathrm{KV}}$ is the key-value head count; and the two layer counts separate local from global storage. This is tensor-size accounting, not a measured device-memory result.

Once the window fills, the local contribution stops growing. The global contribution still grows with sequence length. Sharing key-value heads reduces the cost of both contributions, but does not change that distinction.

The remaining question is not whether the cache becomes constant-size—it does not—but whether its reduced growth fits the intended session length.

Minimal Causal Code

This plain PyTorch loop isolates the cache policy and attention lookup during single-token decoding. Each layer owns a separate cache, and eviction happens across token positions. It uses the configured head counts and treats the local window as including the current token.

The layer hooks supply projected queries, keys, and values and finish the surrounding block. Position handling, attention gating, normalization, and feed-forward computation are outside this cache example; this is not a complete model implementation. Repeating key-value heads below is temporary lookup work, not extra persistent cache storage.

import torch


def cached_lookup(q, k, v, cache, is_local, window=2048):
    # One token: q [32, 128], k and v [2, 128].
    new_k = k.unsqueeze(0)  # [1, 2, 128]
    new_v = v.unsqueeze(0)

    if cache is None:
        keys, values = new_k, new_v
    else:
        old_k, old_v = cache
        keys = torch.cat([old_k, new_k], dim=0)
        values = torch.cat([old_v, new_v], dim=0)

    if is_local:
        # Copy so an evicted prefix is not retained by a tensor view.
        keys = keys[-window:].clone()
        values = values[-window:].clone()

    # The persistent cache still has only two KV heads.
    next_cache = (keys, values)
    lookup_k = keys.repeat_interleave(16, dim=1)
    lookup_v = values.repeat_interleave(16, dim=1)

    scores = torch.einsum("hd,shd->hs", q, lookup_k)
    weights = torch.softmax(scores / (q.shape[-1] ** 0.5), dim=-1)
    heads = torch.einsum("hs,shd->hd", weights, lookup_v)
    return heads, next_cache


# token_states supplies one new input embedding at a time.
# Each layer hook projects the current representation and finishes its block.
caches = [None for _ in layers]

with torch.no_grad():
    for position, hidden in enumerate(token_states):
        for layer_index, layer in enumerate(layers):
            q, k, v = layer.project_qkv(hidden, position)
            is_local = (layer_index % 4 != 3)
            heads, caches[layer_index] = cached_lookup(
                q, k, v, caches[layer_index], is_local
            )
            hidden = layer.finish_block(hidden, heads)

The outer loop advances time. The inner loop advances depth. Only the cache belonging to the current layer is read and updated. Because each lookup contains only current and earlier tokens, the loop is causal.

What You Gain, and What You Lose

Local layers stop accumulating cache entries after their windows fill. Grouped-query attention stores fewer key-value heads than query heads. Periodic global layers retain direct access to their full prefixes, providing a long-range path through the stack.

The costs are equally concrete. A local layer cannot directly attend to an evicted entry. Shared key-value representations are not independent representations for every query head. Global caches still grow with the conversation, so this is reduced cache growth, not a fixed-size total memory.

The design therefore spends memory selectively: recent detail is available in every local window, while full-history lookup is concentrated in the global layers.