17 GB of Attention Buffers Nobody Read

A sparse-attention MoE model was allocating dense-attention scratch it never touched. A per-allocation trace found it, and prefill jumped to 327–399 tok/s.

Summary

Our engine prefilled a 320B mixture-of-experts model in 512-token chunks, because 2,048-token chunks ran the GPU out of memory. Its activation scratch was 7.07 GB at 512 tokens per chunk and 26.36 GB at 2,048. We first cut the one chunk-sized buffer we knew about, a 4.8 GB attention copy, and 2,048 still didn't fit. Then we added a trace that names every allocation of 32 MB or more as it happens. It showed that 4.3 GB at 512 and 17 GB at 2,048 were split-attention partials that only dense and MLA attention layers read. This model has no such layers. Its attention is sparse (DSA) over a paged latent cache, but the allocation was keyed on "has KV-cache layers" instead of on the layer types. We gated those partials, and the eviction-score buffer next to them, on the layer types actually present. That cut scratch to 2.7 GB and 8.8 GB, the 2,048-token chunk fit under the default working-set limit, and prefill went from 327 to 399 tok/s, +22%. Nothing had ever read those buffers, so no computation changed. The whole speedup comes from the bigger chunk they had been blocking.

Why the chunk size mattered on this model

At 512-token chunks, the grouped expert matmul took 46% of prefill time. The model routes each token to 8 of its 288 experts, so a 512-token chunk gives each expert about 14 tokens. The expert kernel works in tiles of 32 token slots, so most tiles ran half empty.

We tried the obvious kernel fix first, a 16-slot tile, and it lost: 326.6 / 327.5 → 323.3 / 323.1 tok/s. The 32-slot tile already skips its empty 8-slot blocks, and a smaller tile pays more in per-tile dequantization than it saves. The real lever was more tokens per chunk. At 2,048 tokens, each expert sees about 57.

The other large MoE model on the engine shows what that lever is worth. Its prefill went from 360.7 tok/s at 512-token chunks to 407.2 at 1,024 and 435.0 at 2,048. This model was stuck at 512.

The first fix was a guess

We already knew about one large buffer that grew with the chunk: the sparse attention kernel's gathered copy of the key rows each token selects. It was 1.2 GB at 512-token chunks and 4.8 GB at 2,048. We rewrote the kernel to stage those rows itself, straight from the cache. That removed the copy and made prefill 3.3% faster on its own. With the copy gone and a scalar scratch buffer shrunk, 2,048-token chunks should have fit.

They didn't. 512-token chunks ran at 304.4 tok/s, and 2,048 failed with a GPU out-of-memory error.

The kernel rewrite was a real improvement. As a memory fix, though, it was a guess: we cut the buffer we knew about and assumed it was the one that mattered.

Name every allocation

The engine already kept a running total of activation scratch. We added two reporting levels behind an environment variable:

• level 1 prints the total at load, for the configured chunk size and number of sequence slots;
• level 2 prints every allocation of 32 MB or more, by name and size, as it happens.

Level 1 put a number on the problem: 7.07 GB at 512 tokens and 26.36 GB at 2,048, so 19 GB of growth for a chunk four times larger. Level 2 showed which buffers accounted for that growth:

That one family of buffers was 61–65% of all scratch, and this model never read any of it.

What the partials are, and why this model got them

Our dense attention prefill kernel runs over the paged KV cache in page-sized splits. Each split writes into three partial buffers, and a merge pass combines them. The buffers are sized for the worst case: chunk rows (rounded up to whole 8-row tiles) × heads × maximum pages, and for the output partial, times the head dimension as well. They grow linearly with the chunk.

The allocation condition was cache_layers > 0, and it came from an earlier fix. On DeepSeek-V4, every layer is sparse attention, and the same partials had been, in the words of the code comment, "pure address space: 66 KB per token of max_context, which is 69 GB at 1M context." None of that model's layers count as cache layers, so the "has cache layers" check correctly stubbed the partials out there.

This model broke that shortcut. Its sparse attention keeps a paged latent cache, so it does have cache layers, but none of its layers ever reaches the dense attention phase. The condition held, and the engine allocated buffers sized for dense attention that nothing would ever read.

The fix replaces the shortcut with the condition the buffers actually depend on:

const bool dense_attn = std::any_of(layers.begin(), layers.end(), [](const auto& l) {
  return l.type == LayerType::kFullAttention || l.type == LayerType::kMLA;
});
const size_t pages = cache_layers > 0 && dense_attn ? max_pages : 0;

The eviction-score buffer, [chunk, heads, max_seq] in fp32, got the same condition. That buffer has its own history. Six days earlier, a research pass had proposed gating it on "eviction enabled". That change broke a test model's tokens outright, because the dense attention kernel read the scores back as its own cache. The fix then was to move each split's scores into threadgroup memory, leaving eviction as the buffer's only reader. It is now allocated only when eviction is on and a dense or MLA layer exists.

The whole change was 13 lines added and 2 removed.

The result

• Scratch: 2.7 GB at 512-token chunks, 8.8 GB at 2,048.
• Fit: 185 GB of pinned weights plus 8.8 GB of scratch comes to 94% of the default 206 GB GPU working-set limit. It fits, and the engine's startup check warns that it's close (the post on that limit covers why it exists and how we raised it).
• Speed: on a 9,000-token prompt over two interleaved rounds, 327.0 → 399.8 / 399.1 tok/s (+22%) at 2,048-token chunks, with 57 token slots per expert instead of 14.
• Shipped: the model's recipe now uses 2,048-token chunks, its Neural Engine bank moved to the same size, and the converter now derives 2,048-token chunks for every model family.

Why nobody noticed

• No correctness signal. A buffer nothing reads can't change the output, so every parity test passed.
• No speed signal. Unused memory costs nothing until it blocks something. At 512-token chunks it was 4.3 GB inside a working set of nearly 200 GB.
• Totals don't point anywhere. "7.07 GB of scratch" is a number you nod at. A list of sizes by name is what shows you where to look.
• The shortcut held for every model it was written for. On the dense models, "has cache layers" meant "has dense attention". On the all-sparse model, no cache layers meant no partials were needed. This was the first model where the two came apart.

Seventeen gigabytes of buffers nobody reads is the kind of waste the setup audit finds before it becomes your budget.

Key takeaways

1. Gate a buffer on the code that reads it, not on a proxy that happens to line up. "Has KV-cache layers" stood in for "has dense attention" until a sparse model with a paged latent cache showed up.
2. Trace allocations by name, with sizes, as they happen. The scratch total told us there was a problem; only the per-allocation trace told us where.
3. Cutting the buffer you already know about is a guess. Ours freed 4.8 GB, and the chunk still didn't fit.
4. Read the comments above the allocation. These partials had already been stubbed once for another sparse-attention model, using a condition that only covered that model.
5. On MoE prefill, chunk size decides how full the expert kernels run. Going from 14 to 57 slots per expert was worth +22%, and a smaller tile lost.
6. Parity tests can't see unused memory. It only shows up as the next thing you can't do.