Whether an LLM decoding kernel is compute-bound or memory-bound changes while the system runs — batches drain as requests finish and speculation length gets retuned — so the fixed GPU-versus-PIM kernel mappings of prior heterogeneous accelerators are wrong much of the time; PAPI predicts the flip with a two-integer multiply (batch size times speculation length approximates arithmetic intensity), reschedules fully-connected kernels between GPU tensor cores and a compute-rich FC-PIM (4 processing units per DRAM bank, affordable only because batching lets one DRAM row activation feed many computations), keeps the permanently memory-bound attention kernel on capacity-lean disaggregated Attn-PIM, and beats static hybrids by 1.8–1.9× and a PIM-only design by 11.1×.
The argument in one glance
-
Problem
Decoding kernels straddle the roofline — and their position moves
A transformer decoder runs two kernel species: fully-connected (FC) kernels (QKV generation, projection, feed-forward — big weight-matrix multiplies) and the attention kernel (each request's query against its own stored KV cache). How much computation each performs per byte fetched — its arithmetic intensity — depends on two parallelism levels the serving system changes constantly: how many requests are batched together (request-level parallelism, RLP) and how many speculated tokens are verified at once (token-level parallelism, TLP). At batch 4, FC kernels are memory-bound; at batch 32 and above they are compute-bound. Attention never flips: batching gives it zero data reuse, so it stays memory-bound at every setting.
-
What people did
Heterogeneous accelerators: GPU for compute, processing-in-memory for bandwidth
A wave of designs pairs computation-centric processors with processing-in-memory (PIM) — compute cores placed inside the memory stack, where bandwidth is enormous. AttAcc puts every attention kernel on PIM and every FC kernel on GPUs. IANUS does the reverse split with an NPU. SpecPIM co-schedules both across NPU and PIM using a genetic algorithm plus Monte Carlo tree search. All show large wins over GPU-only serving — at the operating point they were tuned for.
-
Why it fell short
Static mappings are wrong whenever the operating point drifts — which is always
Real serving never sits still: latency objectives cap a 854-request-capable machine at batch 22, KV-cache capacity allows 282 requests at length 128 but only 18 at length 2048, requests in a batch finish at different times so RLP decays during execution, and speculation length is retuned for utilization. Measured consequence, FC kernel latency: at batch 4 PIM beats an A100 GPU, at batch 16+ the GPU wins decisively — so AttAcc's FC-on-GPU is wrong at low load, and a PIM-only system is catastrophically wrong at high load. SpecPIM's scheduler needs 50 genetic-algorithm rounds and 10,000 tree searches — offline only. And one PIM flavor cannot serve both memory-bound species: attention needs about 4.5× more memory bandwidth per unit of computation than a memory-bound FC kernel.
-
This paper's idea
Predict the flip with one multiply; give each kernel species its own PIM
PAPI (from ETH Zurich's SAFARI group, the Chinese Academy of Sciences, and the University of Toronto) makes the mapping dynamic and the PIM hybrid. A hardware scheduler estimates FC arithmetic intensity as simply RLP × TLP — mathematically justified because the hidden dimension is huge — counts end-of-sequence tokens after each iteration to track the draining batch, and moves FC kernels between GPU tensor cores and FC-PIM whenever the estimate crosses a calibrated threshold. FC-PIM packs 4 processing units per DRAM bank — affordable within HBM's power budget only because batching lets one DRAM row activation feed many computations. Attention lives permanently on Attn-PIM: 1 unit per 2 banks, disaggregated over PCIe/CXL so KV capacity can grow. Result: 1.8×/1.9× over the static A100+AttAcc / A100+HBM-PIM hybrids, 11.1× over PIM-only, with 3.4× better energy efficiency.
Why this problem is worth solving
Serving a large language model is a tug of war between two
kinds of starvation.
how many floating-point
operations a kernel performs per byte it fetches from memory.
The roofline model turns this one number into a verdict: below
the machine’s balance point the kernel is memory-bound (the
compute units idle, waiting on bytes), above it the kernel is
compute-bound (the memory system idles, waiting on math). An
A100 GPU with 312 TFLOPS and 1.9 TB/s flips verdicts around
160 operations per byte
the stored Key and Value matrices from all
earlier tokens of a request, kept so attention does not
recompute them; its size grows with every generated token, and
each request has its own private copy
a parallelism trick where a
small, cheap draft model proposes the next several tokens and
the large model verifies them all in one pass — turning serial
token generation into parallel verification of 2 to 10 tokens
per iteration
The prior generation of processing-in-memory accelerators answered it once, at design time. The cost of that frozen answer is what PAPI’s motivation quantifies: real serving sweeps the operating point continuously (a 30 ms latency objective forces batch 22 on hardware that could hold 854 requests; a batch that starts at 64 decays request by request as short answers finish), and each sweep drags the FC kernels across the roofline boundary. A system that guessed “compute-bound” runs its FC kernels on a GPU that is bandwidth-starved half the time; a system that guessed “memory-bound” is 11× too slow when traffic surges. The machine is only ever right by coincidence.
What people did before, and where it broke
- AttAcc (covered elsewhere on this site). Attention
always on HBM-based
units, FC always on GPUs. Right at high parallelism, wrong at low — where PIM would beat the A100 on FC kernels too.
placing compute units inside the memory device — here, floating-point units next to DRAM banks in an HBM stack — so kernels run at internal-bank bandwidth instead of squeezing through the external memory bus. The catch: DRAM dies have hard power and area budgets, so PIM compute is necessarily modest
- IANUS (covered elsewhere on this site). FC always on PIM, attention on an NPU — tuned for single-request or small-batch serving, and upside-down once requests batch up and FC turns compute-bound.
- SpecPIM. Co-schedules both kernel types across NPU and PIM for speculative decoding — but its scheduler runs 50 rounds of a genetic algorithm plus 10,000 Monte Carlo tree searches, viable offline for one fixed batch size and speculation length, impossible to re-run every time the batch drains by one request.
- One-size-fits-all PIM units in all of them. Even among memory-bound kernels the needs diverge: at batch 4 and speculation length 8, FC intensity is 31.7 operations per byte and attention’s is 7.0 — attention wants about 4.5× more bandwidth per unit of compute. A single PIM configuration must shortchange one species or the other.
How PAPI works
Three device classes, one two-integer predictor
-
Heterogeneous platform new
A host CPU; a high-performance processor holding GPU tensor cores (the PUs), FC-PIM memory devices, and a hardware scheduler, linked by NVLink (FC weights are bulky); and a pool of Attn-PIM devices disaggregated over PCIe or CXL — attention traffic is just small query vectors, and CXL scales to thousands of devices, letting KV-cache capacity grow with sequence lengths independently of the processor.
-
Attention placement (static, and correctly so)
The attention kernel is memory-bound at every batch size and speculation length — batching gives it no reuse because every request owns a private KV cache. So it always runs on Attn-PIM. No decision needed.
-
FC boundedness predictor new
FC arithmetic intensity ≈ RLP × TLP. This is exact algebra, not a heuristic: intensity = 2·RLP·TLP·h² / (2·(2·RLP·TLP·h + h²)) and with hidden dimension h in the thousands (12,288 for GPT-3 175B) the h² term dominates the denominator, cancelling to RLP × TLP. Verified against measured intensity on GPT-3 66B: near-exact, with harmless overestimation only at RLP=128 — already deep in compute-bound territory, so the verdict is unchanged.
-
Runtime scheduling loop new
After each decoding iteration: gather all requests' output tokens, count end-of-sequence tokens to detect finished requests (RLP drop, Attn-PIM resources freed); read TLP from a register the host CPU updates when speculation length is retuned; compute RLP × TLP; compare against a threshold α calibrated offline by running FC on both device types across parallelism levels. Above α → FC runs on the GPU tensor cores (FC-PIM devices then serve as plain memory holding the weights). Below α → FC runs inside FC-PIM.
-
FC-PIM: performance-flavored PIM (4 units per bank) new
A 1-unit-per-bank PIM spends 96.7% of its energy on DRAM row activation when there is no data reuse — more compute units would blow HBM's 116 W budget. But batched FC reuses each activated row for many computations; at reuse level 4 and above, the DRAM-access share collapses (33.1% at reuse 64), freeing power for 4 floating-point units per bank. Area is the other wall: each unit is 0.1025 mm² against a 121 mm² die, so FC-PIM keeps 96 banks instead of 128 — 12 GB per device instead of 16, capacity traded for throughput.
-
Attn-PIM: capacity-flavored PIM (1 unit per 2 banks) new
Attention has no reuse to exploit, so even 1 unit per bank exceeds the power budget; 1 per 2 banks fits. Full 16 GB capacity retained — the KV cache is what needs the space. Data laid out per the AttAcc scheme: attention heads spread across devices, K matrix split column-wise at pseudo-channel/bank-group level and row-wise at bank level, V the transpose.
The idea
The roofline position of a kernel is a runtime variable, so read it at runtime. PAPI’s enabling observation is that the expensive-looking question “is this FC kernel compute-bound right now?” reduces, after algebra, to multiplying two small integers the serving system already knows. Batch size and speculation length are not measurements — they are control inputs, sitting in registers. The hidden dimension being huge (thousands) makes the weight matrix dominate the byte count and makes every unit of parallelism translate one-for-one into intensity, so RLP × TLP against a calibrated threshold is the whole predictor. Where SpecPIM needed a genetic algorithm and ten thousand tree searches to place kernels for one operating point, PAPI re-decides every iteration for the cost of counting end-of-sequence tokens.
Data reuse is a power budget in disguise. The reason prior PIM units were compute-anemic is not conservatism — a DRAM die that spends 96.7% of its energy activating rows has no headroom for more arithmetic units. PAPI notices that the same parallelism that makes FC kernels compute-hungry also hands PIM the cure: a batch of 64 reuses each activated row 64 times, collapsing the activation energy share and freeing the power budget for 4 units per bank. The design closes a loop: high parallelism demands more PIM compute, and high parallelism is precisely what makes more PIM compute affordable. Attention, with no reuse ever, gets the opposite design — half a unit per bank, full capacity — because pretending one PIM flavor fits both species means overpaying in power for attention or underdelivering compute for FC.
Statically place what is statically true; schedule only what actually varies. PAPI is not “make everything dynamic.” Attention is memory-bound at every reachable operating point, so it is pinned to Attn-PIM forever — the scheduling machinery spends itself exclusively on the one kernel species that genuinely flips. This keeps the scheduler trivial (one multiply, one comparison) and the architecture honest: each hardware unit is shaped by the worst case of the kernels it will actually host, not by an average over kernels it might.
A batch’s life under PAPI
LLaMA-65B serving, watched through the scheduler's eyes
RLP = requests decoding in parallel (batch size). TLP = tokens verified per iteration (speculation length). α = memory-boundedness threshold, calibrated offline. PU = GPU tensor cores. Numbers from the paper's simulator evaluation.
Step 1 of 6
A full batch arrives — FC goes to the GPU
A batch of 64 requests starts decoding with speculation length 4. The scheduler computes RLP × TLP = 256, far above α: FC kernels are compute-bound, so they run on the GPU tensor cores at full 312-TFLOPS throughput. The FC-PIM devices act as ordinary memory holding the weights. Attention, as always, runs on the Attn-PIM pool.
a PIM-only system here is the disaster case — 11.1× slower end-to-end compute-bound work belongs on compute-centric silicon
Why it works
- Batch size and speculation length are control inputs the system already knows → arithmetic intensity ≈ RLP × TLP (exact when the hidden dimension is large) → boundedness prediction costs one multiply and one comparison, cheap enough to re-run every decoding iteration — where prior schedulers needed offline search.
- The weight matrix is shared across all parallel tokens but each KV cache is private → FC intensity scales with parallelism while attention intensity stays flat → only FC needs dynamic placement, and attention can be pinned, keeping the mechanism minimal.
- Row-activation energy dominates PIM power only when data is not reused → batching converts one activation into tens of computations → the power budget that limited prior PIM to one unit per bank stretches to four, giving FC-PIM real compute — the same parallelism that creates the demand funds the supply.
- Attention traffic is small query vectors, not bulky weights → Attn-PIM can hang off commodity PCIe/CXL links → KV-cache capacity scales by adding devices, disaggregated from the processor, instead of being capped by GPU memory.
- Each device class is shaped by the kernels it actually hosts (GPU for compute-bound FC, 4-unit FC-PIM for memory-bound FC, 2-bank Attn-PIM for reuse-less attention) → no device is provisioned for a workload it never sees → 1.8–1.9× over static hybrids and 3.4× energy efficiency.
What it costs
- Bespoke silicon everywhere. Two custom HBM-PIM die variants, GPUs whose memory is partly those custom dies (60 GB instead of 80 GB per A100-class device), a hardware scheduler, and CXL/PCIe-attached PIM pools — this is a clean-slate system design, not a drop-in.
- Capacity traded for compute. FC-PIM dies sacrifice 25% of their banks (12 GB vs 16 GB) for the added units, forcing six GPUs where the model alone would need fewer.
- Attention runs slower than it could. Attn-PIM’s 1-unit-per-2-banks is 1.7× slower on attention than AttAcc’s denser design — an explicit area-and-power economy the end-to-end numbers absorb.
- Calibration and simulation. The threshold α comes from offline profiling per deployment; all results are from a Ramulator2-based simulator (with baselines reimplemented in the same framework), not silicon.
- Communication is the next wall. Even in the PIM-only analysis, 28.2% of decoding time is interconnect traffic — the paper itself points to better networks as future work.
The evaluation
Simulator-based, but disciplined about realism: real request traces (Dolly creative-writing and general-qa, with naturally varying lengths and therefore naturally draining batches), three model scales (LLaMA-65B, GPT-3 66B, GPT-3 175B), equal hardware budgets across all systems (90 HBM devices each), and sensitivity sweeps over batch size 4–128 and speculation length 1–8.
Four ways to build an LLM decoding machine
| Property | A100+AttAcc / A100+HBM-PIM (static heterogeneous) | AttAcc-only (PIM-only) | PAPI (dynamic + hybrid PIM) |
|---|---|---|---|
| Kernel-to-device mapping | |||
| PIM hardware | Same single flavor for all kernels | ||
| Scheduling cost | None | ||
| End-to-end performance (Dolly creative-writing) | |||
| Energy efficiency | Reference — FC always burns GPU power, even when memory-bound | Frugal but slow |
Check yourself
Check your understanding
1. A batch starts at 64 requests (speculation length 4) and drains to 8 as short answers finish. Predict what happens on the static A100+AttAcc system versus on PAPI.
2. PAPI estimates FC arithmetic intensity as just RLP × TLP — two integers multiplied. Why is this algebraically honest rather than a rough heuristic, and why does its one known error not matter?
3. Prior HBM-PIM designs put at most 1 compute unit per bank (or per 2 banks) — power-limited. PAPI's FC-PIM packs 4 per bank inside the same 116 W budget. What changed?
4. PAPI dynamically schedules FC kernels but pins attention to Attn-PIM permanently. Why is attention exempt from the roofline flip that governs everything else?