The news. On August 12, 2026, researchers posted TokenStack, an architecture that targets the KV-cache bottleneck in LLM decode by splitting HBM stacks into dense capacity layers and PIM-enabled compute layers. It uses HBM4's logic-die substrate as a stack-local controller for cross-layer DMA, layered address translation, attention gather/broadcast coordination, and inline quantization during migration; its runtime keeps hot KV blocks near the PIM compute, moves colder KV state to the dense layers, and uses bounded replication so the scarce PIM capacity is not wasted on copies. Across production-derived trace experiments on four models, the paper reports 1.62× geometric-mean token throughput, 1.70× SLO-compliant serving capacity, and 30–47% lower per-token energy than AttAcc. Read the paper →
Picture a quarry. The rock has already been blasted and is sitting in the pit; the crushing plant is across the site. Every time the plant needs material, trucks drive out to the pit, load up, and haul the rock over — and the plant crushes it into something a fraction of the weight that arrived. Almost all of the day's fuel is spent moving rock that was going to be thrown away anyway. The obvious fix is not a faster plant and not smaller trucks. It is a crusher at the pit face, so that what leaves the pit is the output rather than the raw material.
That is the shape of the problem in LLM decode, and the reason it has this shape is that the KV cache is written once and then only read. Prefill fills it; every generated token after that walks the whole cache again to attend over the conversation so far. The math done on each byte is tiny — a multiply and an add — so decode attention has very low arithmetic intensity, and on a roofline plot it lands far to the left, under the sloped bandwidth roof rather than the flat compute roof. (The roofline framing here is the GPU & CUDA track's own teaching, not a claim from the TokenStack paper.)
Prefill vs Decode on the Roofline
Same GPU, fundamentally different bottlenecks
Sitting under the bandwidth roof has a specific consequence: adding matrix units does nothing, because the matrix units were never the thing that was busy. Nearly everything the field has built for this bottleneck therefore attacks the byte count instead. GQA shares one KV head across several query heads so there is less cache to begin with. KV quantization stores each entry in fewer bits. Eviction and sparse attention skip entries the model probably will not need. All three shrink the load. None of them change the fact that whatever survives still has to be trucked across the memory bus, once per token, forever.
TokenStack's move is to stop hauling. It treats an HBM stack as two different kinds of real estate rather than one uniform pool: some layers are ordinary dense DRAM, holding capacity and nothing else, and some layers are PIM-enabled — they carry arithmetic units next to the arrays and can perform the attention operation on blocks that live on them. Put a hot KV block on a PIM layer and it can be attended over on the layer it already occupies, so what leaves the stack is a partial attention result rather than the block that produced it. The crusher is at the pit face, and the trucks now carry crushed output.
That split creates a placement problem, which is where most of the design actually goes. PIM layers are the scarce resource — a stack cannot be all compute layers, or it stops being memory. So TokenStack's runtime keeps hot blocks near the PIM compute, migrates colder KV state down to the dense layers, and uses bounded replication to stop a popular block from being copied everywhere and eating the capacity it was supposed to accelerate. Coordinating that is the job of HBM4's logic die, which the paper uses as a stack-local controller: it drives cross-layer DMA, translates addresses across the layered map, orchestrates the gather and broadcast that attention needs, and quantizes blocks inline as they move between tiers. In the quarry, that is the foreman — deciding which benches get worked today and which get parked, and never letting the same pallet sit in five places at once.
Here is where it earns its keep, in bytes. Hold one request fixed at a 32,768-token context on a 70B-class model with grouped-query attention — 8 KV heads, 128-dimensional heads, 80 layers, fp16 — and price the cache. Each token costs 2 (one K entry and one V entry) × 8 heads × 128 dims × 2 bytes = 4 KiB per layer, so 320 KiB per token across 80 layers, and 32,768 tokens of context is 10 GiB of KV cache for that single request. Now the decode loop: every generated token re-reads all 10 GiB. On a part with roughly 4 TB/s of HBM bandwidth that is about 2.7 ms of pure memory traffic per token — a ceiling near 370 tokens per second from bandwidth alone, before a single FLOP, and before any other request in the batch. Halving the cache by quantizing it halves that to ~1.35 ms; it does not remove it. Attending over a block on the layer it already occupies takes that trip off the table for whatever the runtime manages to place there, and what the paper reports is 1.62× geometric-mean token throughput, 1.70× SLO-compliant serving capacity, and 30–47% lower per-token energy against AttAcc. (The 32K worked example is illustrative, built from the curriculum's own KV-cache arithmetic to show why the trip dominates; the paper publishes ratios against AttAcc on production-derived traces, not these byte figures.)
| Approach | What it changes | Bytes still crossing the bus per token | What it costs you |
|---|---|---|---|
| Grouped-query attention (GQA) | fewer KV heads to store | reduced roughly in proportion to the head-sharing ratio (~4–8×, varies by model configuration — illustrative) | fixed at architecture time; you cannot retune it for a deployment |
| KV quantization | fewer bits per stored entry | ~4–8× fewer at 2–4 bits, but still all moved | accuracy risk at low bit-widths, plus dequantization work |
| Eviction / sparse attention | fewer entries read at all | a selected subset only (ratio setup-dependent, illustrative) | a prediction about which entries matter, which can be wrong |
| Processing-in-memory (TokenStack) | where the attention math runs | partial results instead of KV blocks, for whatever sits on a PIM layer | new memory hardware, and a hot/cold placement problem to solve at runtime — reported 1.62× throughput vs AttAcc |
The reason this is worth noticing is that it is a different kind of answer than the field has been giving. Compression, eviction and prefill/decode disaggregation all take the memory system as given and rearrange the workload to fit it. TokenStack takes the workload as given and rearranges the memory system — which, in roofline terms, is raising the roof rather than sliding along it. That also sets the honest limits, and two of them are the paper's own: the evidence is production-derived trace experiments, not measurements on shipped silicon, and the comparison is against AttAcc rather than against a standard GPU. The third is this article's reading rather than the paper's — converting a bandwidth problem into a placement problem makes the result depend on how well the placement policy generalises past the traces it was tuned on.
Goes deeper in: GPU & CUDA → Roofline Model → Changing the Roof
Related explainers
- KVarN — Hadamard rotation for a 2-bit KV cache — the compression answer to the same bottleneck: keep hauling, but make each load lighter.
- FlashMemory — lookahead sparse attention — the skip-it answer: predict which entries the token will actually need and read only those.
- AMD ATOM — prefill/decode disaggregation — the scheduling answer: put the memory-bound phase on different hardware from the compute-bound one.