An LLM answers each request in two phases with opposite appetites — PREFILL reads the whole prompt at once and maxes out the GPU's compute while barely touching memory bandwidth, and DECODE emits one token at a time and maxes out memory bandwidth while using under 10% of compute; servers already mix prefill and decode requests into one batch to keep the GPU busy, but the ATTENTION step stubbornly resists, because today's attention kernels (FlashAttention and friends) run the prefill and decode parts one after the other, so each leaves idle exactly the resource the other is starving for; this paper (Microsoft Research) builds POD-Attention, the first kernel that runs prefill and decode attention CONCURRENTLY on the same GPU multiprocessor by fusing them into one kernel and steering the scheduling in software — each threadblock reads which multiprocessor it landed on and grabs a ticket telling it to do prefill or decode work — so compute-bound and bandwidth-bound work share each unit and saturate both at once, speeding up attention by up to 59% (mean 28%) and lifting end-to-end serving throughput 12 to 27%.
The argument in one glance
-
Problem
An LLM's two phases want opposite things from the GPU
Every LLM request is served in two phases with opposite resource appetites. PREFILL reads the entire input prompt at once and processes all its tokens in parallel — this is COMPUTE-bound, saturating the GPU's tensor cores while using only a few percent of memory bandwidth. DECODE then generates the answer one token at a time, and each step must read the cached keys and values of every prior token from memory — this is MEMORY-BANDWIDTH-bound, saturating bandwidth while using under 10% of the compute. Because the two phases stress opposite resources, a GPU serving many requests swings between compute-saturated and bandwidth-saturated, wasting whichever resource the current phase does not happen to need.
-
What people did
Mix prefill and decode requests into one hybrid batch
The standard fix is HYBRID BATCHING (chunked-prefill plus continuous batching, as in systems like Sarathi-Serve): put a chunk of a new request's prefill together with the ongoing decodes of other requests in a single batch. Now the model's weights are fetched from memory once and used to compute over both the prefill and the decode inputs at the same time, so the compute-heavy prefill work and the bandwidth-heavy decode work share the GPU and keep more of it busy. For the big linear layers, this works well.
-
Why it fell short
Attention still runs the two phases one after the other
Hybrid batching fixes the linear layers but not ATTENTION — the one operation where each token must look back at all prior tokens' keys and values in the KV-cache. Existing attention kernels (FlashAttention, FlashInfer) are specialized separately for the two phases and, even inside a hybrid batch, run prefill attention and decode attention SEQUENTIALLY, one kernel after the other. So the compute-bound prefill attention and the memory-bound decode attention never OVERLAP: while prefill attention runs, memory bandwidth sits nearly idle; while decode attention runs, the tensor cores sit nearly idle. And you cannot simply launch them as two concurrent kernels: CUDA streams do not guarantee the two land on the SAME multiprocessor at the same time (only 3 to 7% faster, unreliably), warp-level fusion (HFuse) hits a "straggler" problem where a whole threadblock must finish before the next starts (sometimes 13% SLOWER), and the synchronization barriers inside attention block finer-grained interleaving. The overlap that hybrid batching promises is left on the table exactly where it matters most.
-
This paper's idea
Fuse prefill and decode attention into one kernel that shares each SM
This paper (Microsoft Research — Kamath, Prabhu, Mohan, and colleagues) builds POD-Attention, the first GPU attention kernel that runs prefill and decode CONCURRENTLY on the same Streaming Multiprocessor (SM), so the compute-bound and bandwidth-bound work share each unit and saturate both compute AND memory bandwidth at once. The trick is to stop relying on the GPU's opaque scheduler and take control in software. POD fuses prefill and decode attention into a SINGLE kernel launched with (prefill + decode) threadblocks; each threadblock, at runtime, reads a hardware register to discover WHICH SM it was placed on, then atomically grabs a "ticket" from an SM-local counter that tells it whether to do prefill or decode work — guaranteeing a balanced, co-located mix on every SM that the hardware scheduler could not. Careful resource management makes the sharing real: decode uses tiny tiles to free the tensor cores for prefill, the kernel runs 2 or 4 threadblocks per SM depending on the batch mix, decode threadblocks are split into single-warp "virtual" blocks to balance shared memory, and the parallel splitting of prefill is capped to avoid bandwidth contention. The result is attention up to 59% faster (mean 28%), 12 to 27% higher end-to-end serving throughput, and far fewer decode stalls.
Why this problem is worth solving
LLM serving is one of the largest compute workloads in the world right
now, so squeezing more out of each GPU directly lowers the cost of every
generated token. The structure that creates the waste is the two-phase
nature of inference. A
an LLM
answers a request in two phases. PREFILL ingests the whole input prompt
at once: because all prompt tokens are processed in parallel, it is
COMPUTE-bound and keeps the GPU’s tensor cores busy. DECODE then produces
the output token by token, autoregressively — each new token depends on
the previous one, so only one token’s worth of work happens per step,
making decode MEMORY-BANDWIDTH-bound and leaving compute mostly idle. The
operation that ties the phases to memory is ATTENTION: for each query
token, attention computes a weighted combination of the keys and values
of ALL earlier tokens in the sequence, which are stored in the KV-CACHE.
During decode, every single step must stream the entire KV-cache from
memory to attend over the whole context, which is why decode is bound by
bandwidth. So prefill wants compute and barely touches bandwidth (a few
percent), while decode wants bandwidth and barely touches compute (under
10%) — two workloads whose bottlenecks are exact opposites
The reason “run them together” is hard is buried in how GPUs execute.
A
a GPU is
an array of Streaming Multiprocessors (SMs), each with its own tensor
cores, memory pipes, registers, and shared memory. Work is launched as a
KERNEL, which spawns many threadblocks (also called CTAs, cooperative
thread arrays); the hardware scheduler places threadblocks onto SMs as
they free up, and the programmer does not control which threadblock lands
on which SM. To make a compute-bound and a bandwidth-bound operation
actually share a resource, their threadblocks must run at the SAME TIME
on the SAME SM — but launching two separate kernels gives no such
guarantee: the scheduler may run them back-to-back, or place them on
different SMs, so they never overlap. Fixes that work at the boundary,
like CUDA streams (which do not guarantee co-location) or warp-level
fusion (which stalls when one fused piece finishes long before the
other), leave most of the benefit unrealized. Genuinely co-scheduling the
two requires reaching below the abstraction the GPU offers
What people did before, and where it broke
- Separate prefill and decode attention kernels (FlashAttention, FlashInfer). Each phase gets a kernel tuned for its own resource profile. Individually optimal, but in a hybrid batch they run sequentially, so compute-bound and bandwidth-bound attention never overlap and half the GPU idles at any moment.
- CUDA streams (kernel-parallel). Launch the two kernels on separate streams and hope the GPU overlaps them. Streams guarantee neither concurrency nor SM co-location, so the gain is only about 3 to 7% and unreliable.
- Warp-parallel fusion (HFuse). Fuse the two kernels at the warp level within a threadblock. It hits the straggler problem — a whole threadblock must finish before it is replaced, so when one fused piece lags, the next is delayed — giving about 11% median but up to 13% slower on prefill-heavy chunks.
- Intra-thread fusion. Interleave prefill and decode instructions within a thread. Attention’s CTA-level synchronization barriers (used to stage data into shared memory) block this, since instructions before a barrier cannot overlap those after it.
- The gap. No approach made compute-bound prefill attention and memory-bound decode attention run concurrently on the same SM with a guarantee. That is the gap POD-Attention closes.
How POD-Attention overlaps prefill and decode
Fuse into one kernel, self-schedule per SM, and manage resources so they truly share
-
The resource mismatch
Prefill attention is compute-bound (a few percent of memory bandwidth used) and decode attention is memory-bandwidth-bound (under 10% of compute used). Run sequentially, each leaves idle exactly the resource the other is starving for, so half the GPU is wasted at any moment.
-
Fuse prefill and decode attention into one kernel new
Instead of two separate attention kernels that the GPU runs back-to-back, POD launches a SINGLE kernel with (prefill + decode) threadblocks. FlashAttention's prefill and decode kernels are refactored into generic device functions that a wrapper kernel calls with remapped threadblock identities, so both live in one launch.
-
Self-assign roles per SM in software new
Because the hardware scheduler will not guarantee co-location, each threadblock takes control: a leader thread reads the SM-ID hardware register to learn which SM it landed on, then atomically increments an SM-local counter to get a "ticket" that decides whether this block runs prefill or decode. This guarantees a balanced, co-located mix on every SM.
-
Shrink decode tiles to free the tensor cores new
Prefill uses large 128-by-128 tiles to keep the tensor cores full. Decode is bandwidth-bound, so large tiles would waste compute; POD shrinks decode's query-dimension tile to 16, dropping decode's compute use to about 10% and freeing the tensor cores for the concurrent prefill while decode still saturates memory bandwidth.
-
Balance shared memory and threadblock count new
POD runs 2 threadblocks per SM for prefill-heavy batches (more shared memory each, for bigger tiles) or 4 for decode-heavy batches (finer interleaving). Since decode needs about a quarter of prefill's shared memory, decode blocks are split into single-warp "virtual" blocks with warp-level barriers, so shared memory is allocated evenly and not wasted.
-
Cap prefill splitting to avoid contention new
Splitting attention along the KV dimension (FlashDecoding-style) adds parallelism but too many splits make prefill and decode threadblocks fight over memory bandwidth. POD limits splits to fill at most two full waves of SMs, balancing parallelism against contention and avoiding wave quantization.
The idea
Overlap two workloads with opposite bottlenecks so each fills the other’s idle resource. The whole opportunity is that prefill attention and decode attention are complementary: one is starved for compute and awash in spare bandwidth, the other the reverse. Running them concurrently on the same SM lets the prefill use the tensor cores while the decode uses the memory pipes, so both are saturated at once rather than one at a time. This overlap-complementary-work-in-time discipline is exactly the one the site’s in-/near-memory-fusion page uses when it runs regular parallel work in the bitlines while irregular work runs beside the banks, and the accelerator-companion page uses when it computes non-matmul work on one tile while the matmul unit handles the next — pipeline two operations with different resource profiles so neither waits. The complementary bottleneck is not a nuisance here; it is the resource that makes the overlap pay.
When the hardware scheduler will not cooperate, schedule in software. The obstacle is that a GPU does not let you say “run this prefill block and this decode block together on SM 7.” Launching two kernels, or using streams, leaves placement to an opaque hardware scheduler that gives no co-location guarantee. POD’s answer is to reach below the abstraction: each threadblock reads its own SM-ID and atomically claims a role from an SM-local counter, so the software itself enforces the balanced, co-located mix the hardware would not. Taking control of a decision the default mechanism handles badly — rather than tuning around it — is the same move the site’s coherent-NIC page makes when it manages cache placement and signaling explicitly instead of trusting the protocol’s defaults, and the memory-timing-defense page makes when it shapes traffic itself rather than leaving it to the memory scheduler. The lever is control: fuse into one kernel so there is a single launch to steer, then steer it.
Make the sharing real by managing the resources, not just co-locating. Putting prefill and decode on the same SM is necessary but not sufficient; if decode grabbed large tiles it would steal the tensor cores it does not need, and if shared memory were split evenly between unequal-sized blocks it would be wasted. So POD tunes the resource split: tiny decode tiles surrender the tensor cores to prefill, virtual single-warp decode blocks right-size shared memory, and capped KV-splits keep decode from flooding the bandwidth prefill also uses. And it balances the block counts PROPORTIONALLY to the batch’s prefill-to-decode ratio rather than a fixed 50:50, worth up to 14%. This measure-what-each-side- actually-needs-and-divide-accordingly discipline — the same one behind the site’s heterogeneity-aware scheduler that sizes allocations to real demand — is what turns co-location into genuine simultaneous saturation.
Filling the idle resource
From half a GPU idle to both resources saturated
Overlapping prefill and decode attention in a hybrid batch. Numbers from the paper's A100 evaluation.
Step 1 of 5
The waste
Prefill attention saturates the tensor cores but leaves roughly 95% of memory bandwidth idle; decode attention saturates bandwidth but leaves about 90% of compute idle. Run one after the other, as existing kernels do even in a hybrid batch, half the GPU sits unused at any instant.
attention is where the phases touch memory: decode must stream the whole KV-cache every single step the linear layers already benefit from hybrid batching; attention is the operation left behind
Why it works
- Prefill attention is compute-bound and decode attention is bandwidth-bound → running them concurrently lets each use the resource the other leaves idle → both compute and memory bandwidth are saturated at once instead of one at a time.
- The hardware threadblock scheduler gives no co-location guarantee → each block reads its SM-ID and atomically claims a prefill-or-decode role from an SM-local counter → software enforces the balanced, co-located mix that separate kernels or streams cannot.
- Decode is bandwidth-bound and needs no large tiles → shrinking decode tiles drops its compute use to about 10% → the tensor cores are freed for the concurrent prefill while decode still saturates bandwidth.
- Prefill and decode need different amounts of shared memory and parallelism → choosing 2 or 4 blocks per SM, splitting decode into single-warp virtual blocks, and capping KV-splits → shared memory is used evenly and neither side floods the bandwidth the other needs.
- Most of the runtime’s wasted resource was in attention, not the linear layers → overlapping prefill and decode attention specifically → a 59% attention speedup translates into 12 to 27% higher end-to-end throughput and far fewer decode stalls.
What it costs (trade-offs and scope)
- The win needs genuinely hybrid batches. Peak speedup comes at prefill-to-decode token ratios around 12 to 18, where most batches contain both kinds of work; a batch that is nearly all prefill or all decode has little to overlap, so the gains shrink.
- It targets attention, so the end-to-end effect is bounded. The linear layers already benefit from hybrid batching; POD fixes the one operation that did not, so its overall impact is limited by how much of the runtime attention occupies.
- The SM-aware scheduling is a GPU-specific workaround. Reading the SM-ID register and using atomic tickets circumvents the opaque hardware scheduler, but it is tuned and evaluated on A100 and built on FlashAttention refactored into device functions — a different GPU or attention backend needs re-engineering.
- The resource knobs must be chosen per batch. The number of blocks per SM, the decode tile size, the split cap, and proportional versus even allocation all matter; the paper sets them by heuristic and sensitivity analysis, and a wrong choice (for example even instead of proportional) costs up to 14%.
- It improves the hybrid-batching regime, not every metric absolutely. POD dramatically cuts time-to-first-token versus Sarathi (which stalls new prefills behind decodes) but can still trail plain vLLM on that one metric; its value is a better balance of throughput and latency, not a universal win on every axis.
The evaluation
FlashAttention refactored into fused device functions and integrated into Sarathi-Serve on vLLM, evaluated on Yi-6B, Llama-2-7B, and Llama-3-8B across one to two NVIDIA A100 80GB GPUs, over more than 1000 hybrid batches (contexts 4K to 20K, prefill chunks 512 to 2K), measuring attention-kernel speedup, energy, end-to-end throughput, and online latency against FlashAttention (serial), CUDA streams, HFuse, vLLM, and Sarathi.
Running prefill and decode attention together
| Property | Separate kernels (FlashAttention) | CUDA streams / warp fusion | POD-Attention (this work) |
|---|---|---|---|
| Prefill and decode overlap | |||
| SM co-location | |||
| Resource saturation | |||
| Stragglers and overhead | |||
| Attention speedup |
Check yourself
Check your understanding
1. LLM serving already uses hybrid batching to share the GPU, yet POD-Attention argues attention still wastes half the GPU. Predict why, and what makes attention special.
2. POD-Attention has each threadblock read its own SM-ID and atomically grab a role, rather than launching separate prefill and decode kernels. Predict why it takes over scheduling in software.
3. After co-locating prefill and decode on the same SM, POD still shrinks decode's tile size to 16 and caps how much prefill is split. Predict why co-location alone is not enough.
4. POD-Attention allocates threadblocks proportionally to the batch's prefill-to-decode ratio instead of a fixed 50:50 split, and its peak gains occur only at certain ratios. Predict why the mix matters.
Sources
- POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference (arXiv:2410.18038) — ASPLOS 2025 (full text read via ar5iv).
- Microsoft Research PDF — official paper.
- ACM DL: POD-Attention (DOI 10.1145/3676641.3715996) — published version.