Attention Optimizations¶

KV cache size¶

KV cache size  = S_prior x 2 (K and V) x n_layers × n_kv_heads × d_head × n_bytes

This entire 40GiB size cache

• needs to exist on HBM, limiting the maximum concurrency.
• needs to be loaded in each decode forward pass, limiting the forward pass rooflining.

The optimization directions are

• S_prior: Attend to fewer tokens → less compute, and sometimes less cache.
• S_prior: Replace the cache with a fixed-size state** → linear attention / SSMs.
• n_kv_heads and d_head: A smaller latent KV cache, store fewer K/V vectors per token.
• n_bytes quantization.

Sliding-window and interleaved local/global¶

Sliding window attention only keeps the most recent \(W\) KV cache. A pure-SWA model cannot do retrieval beyond w × n_layers in principle, and far worse in practice. Therefore, a naturally way is to interleave SWA layers with global attention layers.

Another note is that SWA is different from a max context length = W attention: RoPE will use the token's real position.

• Mistral 7B¹ — SWA with window 4096 in every layer. Cache is capped at the window size (rolling buffer), so cache is O(w), not O(L). Notably, later Mistral models dropped it.
• Gemma 2 / 3²³ — interleaved local:global, 1 global layer per 5 local in Gemma 3. A few global layers preserve long-range retrieval capability, while most layers keep an O(w) cache.
• gpt-oss⁴ same alternating local/global idea by 1:1 ratio.

Attention sinks¶

StreamingLLM⁵ found that if you naively slide the window and drop the first few tokens, perplexity explodes. Keeping just the first ~4 tokens ("attention sinks") plus a rolling window restores it. Explanation: softmax must sum to 1, so heads that want to attend to nothing dump probability mass onto the earliest tokens, which are visible to every query.

Modern models sometimes carry an explicit learned sink logit instead (gpt-oss⁴, others), which makes them more robust to windowing.

Query-aware retrieval¶

The cache stays whole but each decode step computes a cheap per-page summary statistic, takes the top-k pages by that statistic, and attends only over those. Note that it does not reduce HBM pressure on its own, but memory loading.

Early explorations are on the serving stack only: HiSparse⁶ — full KV in host memory, fixed-size GPU cache, hit-detection + LRU + prefetch fused inside the decode CUDA graph. Bit-identical outputs, up to 4.7× throughput. Merged into SGLang.

More recent models moved to native sparse attention. The model's own attention is shaped around the slice it will actually read, hence hold quality at far smaller k than inference-time chosen top-k.

DSA: DeepSeek Sparse Attention ⁷¶

DSA+MLA needs to work together: DSA is only affordable because MLA already collapsed K and V into a single vector per token.

DSA uses a top-k indexer to reduce per-head cache computations (but this won't reduce HBM pressure).

MLA⁸ reduces KV cache from 2 * n_kv_heads * d_head to a single d_latent + RoPE_coef = 512 + 64 = 576 (shapes as in DeepSeek-V3⁹). However, for each head in the SDPA compute, we need to expand the latent cache them back via a linear layer. Therefore, if we need to compute over the full context length, the computation is much larger.

Computation Path¶

# input: h = (1, 7168) 
# W_* are the linear layers

# Query path
c_Q = q_norm(W_qa(h)) # latent q (1, 1536)
q = W_qb(c_Q).view(n_q_heads, qk_head_dim)
q_nope, q_pe = torch.split(q, (qk_nope_dim, qk_rope_dim)) # Q nope + Q RoPE (1, 128, 128 + 64)
q_pe = RoPE(q_pe, freqs_cis)
q = torch.cat(q_nope, q_pe) # Q (1, 128, 192) 

# new latent KV
c_KV = W_kv(h) # latent KV (1, 512)
k_pe = RoPE(W_kr(h)) # K RoPE (1, 64)

# --- Cache size (S_prior, 576) ---
update_cache(torch.cat(c_KV, k_pe))
# ---------------------------------

# Indexer
q_i = W_qi(c_Q).view(n_heads, qk_head_dim)
q_i_nope, q_i_pe = torch.split(q_i, (indexer_nope_dim, indexer_rope_dim)) # Q nope + Q RoPE for indexer (1, 64, 64 + 64)
q_i_pe = RoPE(q_pe, freqs_cis)
q_i = torch.cat(q_i_nope, q_i_pe) # Q (1, 64, 128)

k_i = k_norm(W_ki(h)) # indexer k (1, 128)
k_i_nope, k_i_pe = torch.split(k_i, (indexer_nope_dim, indexer_rope_dim)) # K nope + K RoPE for indexer (1, 128, 64 + 64)
k_i_pe = RoPE(k_i_pe, freqs_cis)
k_i = torch.cat(k_i_nope, k_i_pe) # Q (1, 128)

# --- Cache indexer size (S_prior, 128) ---
update_cache(k_i)
# ---------------------------------

w = W_m(h) * n_heads ** -0.5 # mixing weight per token per head (1, 64)

# indexer score (1, 64, 1, 128) * (128k, 1, 128, 1) -> (128k, 64)
I = (w * torch.bmm(q_i[:, :, None, :], k_i[:, None, :, None]))[:, : , 0, 0]

indices = topk(I, k=2048) # pick up the top 2048 scores

# retrieve k and v for SDPA
c_prior, k_prior_rope = torch.split(c_cache[indices], (latent_dim, qk_rope_dim)) # (128k, 576) -> (2048, 512 + 64)
kv_prior = W_kv(c_prior).view(k, 2, n_heads, qk_nope_dim) # (2048, 2, 128, 128)
k_prior = torch.cat(kv_prior[:, 0, :, :], k_prior_rope) # (2048, 128, 192)
v_prior = kv_prior[:, 1, :, :] # (2048, 128, 128)

# Now we have the tensors required for SDPA
# q: (1, 128, 192)
# k_prior: (2048, 128, 192)
# v_prior: (2048, 128, 128)

KV Cache Memory Pressure¶

Per decode step at 128K context, batch 1, bf16 latent (61 layers):

Computation Considerations¶

For MLA, effectively we are trading the KV cache memory with a linear projection. Most considerably the W_kv(c_prior) which maps (S_prior, 512) -> (S_prior, 2, 128, 128). Compared with GQA (since in practice we will share KV heads), we are trading this with a memory loading of (S_prior, 2, 8, 128). On top of this, we also use DSA to reduce S_prior from 128k to 2k, with an overhead of indexer cache, some projections, and a topk operator (The topk can be expensive).

Some roofline analysis:

H200 SXM: 4.8 TB/s HBM3e, 989 TFLOP/s bf16 dense, 1979 TFLOP/s fp8 dense. Computational intensity = 206 FLOP/byte (bf16), 412 FLOP/byte (fp8).

Per layer, per decode step, batch 1, S_prior = 2048 selected tokens, 128 query heads, bf16 cache. "Bytes" counts cache traffic only — weights are handled separately below.

Note that since GQA is heavily memory bound, MLA's latent projections are almost given "for free", since computation engines are idle.

Another note is that MLA requires larger concurrency: the projection weights is ~33.6 MiB per layer, so the trade-off is (8.39MiB * BSZ) vs (2.36 MiB * BSZ + 33.6 MiB).

Caveats¶

• A 128k-element top-k per token per layer. This can be very expensive.
• DSA does not applied to prefill. Prefill is still \(O(L^2)\) on the score side.
• Quantization headroom is near zero on the latent.

Linear Attention¹⁰¶

Given \(Q, K, V \in \mathbb{R}^{L \times d}\), write softmax attention as

and view it as a special case of

for any non-negative similarity function \(\text{sim}(\cdot, \cdot)\) between queries and keys.

Having swapped \(\text{exp}(QK^T)\) for \(\phi(Q)\phi(K)^T\). — that is, the similarity metric is induced by applying a feature map \(\phi\) independently to \(Q\) and \(K\).

We can exploit the fact that matrix multiplication is associative to reorder operations:

Recurrent form¶

Linear attention only needs \(\sum_{i=1}^l \phi(k_i)^T\) and \(\sum_{i=1}^l \phi(k_i)^Tv_i\) which is independent of \(\phi(Q)\) so we can reuse the values computed through step \(l-1\).

Concretely, with \(Z_{l-1} = \sum_{i=1}^{l-1} \phi(k_i)^T\) and \(S_{l-1} = \sum_{i=1}^{l-1} \phi(k_i)^Tv_i\):

For \(o_{l+1}\), set \(S_l = S_{l-1} + \phi(k_l)^Tv_l\) and \(Z_l = Z_{l-1} + \phi(k_l)^T\) and repeat.

This is the recurrent view of linear attention: maintain a constant-size state \(S_l\) and normalizer \(Z_l\), and each decoding step costs \(O(1)\) time and memory. This is formally an RNN with matrix-valued state \(S_l \in \mathbb{R}^{d' \times d}\).

Some work¹²¹¹ drops the \(Z_l\) normalizer due to numerical instability and observes no problems empirically; taking \(\phi\) to be the identity also works, interestingly. That leaves the clean form

KDA: Kimi Delta Attention¹³¶

Kimi Linear¹³ is a 48B-total / 3B-active MoE at 1M context, built from Kimi Delta Attention (KDA) layers and full MLA layers at 3:1 — every fourth layer is full attention. Reported: 75% KV cache reduction, up to 6× decode throughput at 1M, and better quality than the full-attention baseline on their evals.

KDA replaces the state update with

Per head, with state \(S_t \in \mathbb{R}^{d_v \times d_k}\), key/query \(k_t, q_t \in \mathbb{R}^{d_k}\), value \(v_t \in \mathbb{R}^{d_v}\), channel-wise decay \(a_t \in (0,1)^{d_k}\), and write strength \(\beta_t \in (0,1)\). The \(\alpha_t\) and \(\beta_t\) are the "delta", computed from computed per token from the hidden state via linear layer \(W_\alpha,W_\beta\).

Full KDA Layer¶

The recurrence above is only the middle of the layer. It consumes \(q_t, k_t, v_t, a_t, \beta_t\), and every one of those is itself produced from the hidden state \(x_t \in \mathbb{R}^{d}\). The full path:

1. Project and locally mix with locally previous states Each of \(q, k, v\) gets a linear projection followed by a short convolution of window \(W = 4\). Per channel \(c\), with its own length-\(W\) kernel \(w^{(c)}\):

Likewise for \(k, v\).

# x: hidden state of token index t (1, H)
# q/k/v_conv_state: (3, W - 1, D)
q_proj = W_q(x) # (1, D)

# W_min: (W, C)
# q_proj and q_conv_state concat to (W, C)
q_mixed = conv1d(torch.cat(q_proj, q_conv_state), W_mix) # mixed and add along token dim (1, C)

q = SiLU(q_mixed).view(1, n_heads, d_head)

# same thing for k, v
k = ... 
v = ...

update_conv_state(q, k, v)

2. L2 Normalize keys and queries. \(q_t, k_t \leftarrow q_t/\|q_t\|_2,\; k_t/\|k_t\|_2\) per head, then \(q_t \leftarrow d_k^{-1/2} q_t\). The L2 normalization is what keeps \(\left(I - \beta_t k_t k_t^\top\right)\) a well-behaved (norm \(\le 1\)) projection, so the state can't blow up.

3. Compute the gates. The decay comes from a low-rank bottleneck — \(d \to d_h \to n_h d_k\) — squeezed through a Mamba-style parameterization:

4. Run the recurrence. The state update and \(o_t = S_t q_t\) from above, in one of the three execution modes.

5. Gated normalization and output. A second low-rank projection produces an output gate, applied multiplicatively after an RMSNorm on the head output:

Cache Size Memory Pressure¶

Per decode step at 1M context, batch 1, Kimi Linear's 27 layers (20 KDA + 7 MLA, n_h = 32, d_k = d_v = 128, MLA latent 512 + 64 = 576):

Chunk-wise Form¶

As we can see, this state update is very friendly for decode pass. All linear transformations. However, we can't parallelize tokens during prefill with this naive implementation.

Consider simple linear attention without delta terms, we still have the plain linear attention form \(S = K^T V \odot M\). So we can partition the sequence into chunks:

• across chunks: recurrent, \(L/C\) sequential steps instead of \(L\).
• within a chunk: quadratic attention, which is fine because \(C^2\) is small and it's one dense matmul.

Note that the causal mask \(M\) will be needed. The history term needs no mask — every token in \(S_{[c]}\) is strictly earlier than every token in the chunk, which is the whole trick. Causality only has to be enforced on the small block.

Delta decay gate¶

With a per-channel decay the state is no longer a plain sum, but the log gate makes the cumulative decay a cumsum.

# per chunk c, with S carried across chunks
g = g.cumsum(dim=0)                        # (C, d_k)  cumulative log decay
A = tril(beta[:, None] * ((K * exp(g)) @ (K / exp(g)).T), -1)
T = inv(I + A) @ diag(beta)                # (C, C)    forward substitution
W, U = T @ (K * exp(g)), T @ V             # pseudo-keys / pseudo-values
V_hat = U - W @ S                          # subtract what S already predicts

Aqk = tril(pairwise(Q, K, g))              # (C, C)    A_ij = q_i·k_j exp(g_i - g_j)
O = (Q * exp(g)) @ S + Aqk @ V_hat         # history + local
S = diag(exp(g[-1])) @ S + (K * exp(g[-1] - g)).T @ V_hat