Speculative inference speeds up large language models by pairing small draft models (which write ahead, one memory-bound token at a time) with the big target model (which verifies whole batches of drafts compute-intensively) — two workloads with opposite hardware appetites that today share identical hardware; SpecPIM gives each model its own custom device, choosing per model how much of the memory to convert into in-memory compute (processing-in-memory dies that trade capacity for bandwidth), maps every operator to host or PIM per model, overlaps draft generation with verification by letting drafters run ahead optimistically, and searches this whole architecture-plus-dataflow space automatically — beating speculative inference on GPUs by 1.52× and on existing fixed-recipe PIM accelerators by 2.02×, with 6.67×/2.68× better energy efficiency.
The argument in one glance
-
Problem
Decoding one token at a time is the tax on every LLM answer
Large language model inference has two phases: prefill (process the whole prompt at once — parallel, efficient) and decoding (generate one token per step, each depending on the last). The dependency makes decoding serial and dominant: for GPT-3 175B producing just 128 output tokens, decoding is over 88% of total latency. Each decoding step re-reads enormous weight matrices to produce a single token — the definition of memory-bound work — and cloud providers serve this workload continuously.
-
What people did
Speculate with small models, verify with the big one
Speculative inference breaks the serial chain: small draft language models (DLMs) generate several candidate tokens cheaply, and the target language model (TLM) verifies the whole batch in ONE pass, accepting the longest correct prefix and re-sampling where a draft goes wrong — with a probability-ratio rule that provably yields output equivalent to the big model decoding alone. Multiple drafters' sequences merge into a token tree verified together under a masked "tree attention." On GPUs this is worth 1.17–2.06×.
-
Why it fell short
Two opposite workloads, one kind of hardware
The trick manufactures heterogeneity. A drafter still decodes one token at a time — its matrix operations are vector products with arithmetic intensity near 1 multiply-accumulate per byte, pure memory-bound. The verifier processes 2–10 tree tokens per pass, multiplying its intensity by that factor — decisively more compute-hungry. The system bottleneck SHIFTS between the two as drafter size and verification length change. GPUs handle the memory-bound side poorly at the small batch sizes online serving demands. And the emerging processing-in-memory accelerators — Samsung's HBM-PIM, SK-Hynix's AiM, AttAcc — were all designed for a SINGLE model: every device gets the same fixed PIM allotment, and one fixed rule decides which operators run in memory. Applied to speculative inference, the standard offload-all-FC-layers recipe is SLOWER than a plain GPU — the verifier's compute-intensive layers strangle in memory, and overflowing the halved PIM capacity adds 38.8% of data-shuffling overhead.
-
This paper's idea
Give each model the hardware its physics wants — and search the space
SpecPIM (Peking University) treats the algorithm's heterogeneity and the hardware's heterogeneity as one co-design problem. Each model gets dedicated devices, and PER MODEL the design space chooses: memory type (HBM — more external bandwidth, better for host compute; or GDDR — more in-memory compute) and how many memory dies to convert into PIM dies, each conversion halving that die's capacity to buy internal bandwidth. Per model, a dataflow space maps every operator group to host or PIM with explicit space-time schedules (exploiting, for example, PaLM's parallel attention-and-FFN structure to run both sides concurrently), and a speculative pipeline lets drafters decode OPTIMISTICALLY — overlapping with verification, keeping the work when drafts pass. A genetic-algorithm-plus-Monte-Carlo-tree-search explorer finds the best point per scenario. Results: 1.52× geomean speedup over speculative inference on GPUs and 2.02× over fixed-recipe PIM accelerators, with 6.67× and 2.68× better energy efficiency — and the search consistently rediscovers a rule: give drafters most of the PIM, give the verifier most of the capacity.
Why this problem is worth solving
Every chatbot answer, code completion, and translation is
paid for one token at a time.
the now-standard acceleration
for autoregressive decoding — small draft language models
(DLMs) write several tokens ahead; the large target language
model (TLM) checks the whole draft in one parallel pass. A
per-token probability-ratio test decides acceptance: if the
target model’s confidence in a drafted token, relative to the
drafter’s, falls below a random threshold, the token is
rejected and re-sampled — a rule that makes the final output
mathematically equivalent to the target model decoding alone.
Several drafters can run at once, their candidate sequences
merged into a token tree that a masked attention verifies in
a single pass without recomputing shared prefixes
The diagnostic tool is
operations performed per byte
of memory traffic — the quantity that decides whether a
processor’s compute units or its memory system limits speed.
A drafter generating one token multiplies a vector by huge
weight matrices: every weight byte is used once, intensity
near 1, and the memory bus is the wall. The verifier pushes
its 2–10 tree tokens through the same matrices together, so
each loaded weight serves several tokens and intensity rises
by that factor — enough to shift the wall toward
compute
What people did before, and where it broke
- Speculative inference on GPUs. The algorithmic win is real (1.45× geomean in this paper’s setup) but the drafters’ vector-matrix operations leave the GPU’s compute idle while its memory bus saturates — the classic small-batch LLM problem, now a permanent component of every iteration.
- DRAM processing-in-memory for LLMs. Placing processing engines beside memory banks gives roughly 8× the internal bandwidth of the external bus and far cheaper data movement — a natural home for memory-bound operators. Samsung’s HBM-PIM pairs PIM stacks with an MI100 GPU and offloads all fully-connected layers (built for single-model GPT-J 6B); SK-Hynix’s AiM card offloads fully-connected AND attention operators (GPT-3 6.7B/20B); AttAcc keeps the KV cache in PIM and offloads attention. All three share two fixed choices: every device carries the same PIM allotment, and one rule decides forever which operators live in memory.
- Where fixed recipes break. PIM engines cost capacity — Samsung’s and SK-Hynix’s designs both HALVE a module’s memory to fit the compute — so more PIM means more devices needed to hold a big model. A verification model’s fully-connected layers are too compute-intensive for PIM engines; forcing them there (the FC-offload recipe) is slower than a plain GPU, compounded by 38.8% overhead shuttling overflowed tensors between normal and PIM memory. The attention-only recipe survives but leaves the drafters — the MOST PIM-friendly models — unaccelerated, managing only 1.09× over the GPU baseline. And fixed operator mappings serialize what model structure would let run in parallel: PaLM computes attention and feed-forward blocks independently, an opportunity no fixed offload exploits.
How SpecPIM works
Per-model hardware, per-model dataflow, and a search over both
-
Per-model device allocation new
Each model — every drafter and the verifier — gets its own dedicated devices connected by NVLink-class interconnect (the PaLM benchmark: four devices for the 62B verifier, one/one/two for the 1.3B/2.7B/8B drafters). Devices serving the same model are identical; devices serving different models may differ — that asymmetry is the whole point.
-
The architecture design space new
Per model, two knobs: memory technology and PIM dosage. HBM offers higher external bandwidth and cheaper access energy — better when the host processor does the work; GDDR-PIM offers ~1.7× the in-memory compute per gigabyte — better when the PIM engines do. Each chip's modules can be memory (full capacity) or PIM (half capacity, in-memory compute): 0–8 PIM dies per HBM stack, 0–4 PIM packages per GDDR chip. Capacity and internal bandwidth trade directly.
-
The speculative pipeline new
Vanilla speculative inference alternates: drafters wait for verification; the verifier waits for drafts — separate devices guarantee someone is always idle. SpecPIM lets drafters continue OPTIMISTICALLY from their own unverified tokens, overlapping this iteration's drafting with last iteration's verification. Drafts that pass keep their head start; a drafter that fails adopts the verifier's corrected tokens while the verifier pauses one round to let it resync. Worth 1.20× geomean — more when there are fewer drafters, since all-pass rounds are then more likely.
-
Operator mapping as a first-class space new
Within each model's layer, operators are partitioned into ordered groups; each group is assigned to host or PIM, and groups on DIFFERENT engines run concurrently when independent — recovering the parallelism fixed recipes discard (PaLM's parallel attention/FFN blocks overlap host and PIM work). Inside a group, a space-time mapping assigns operators to specific instances with explicit schedules; tiling strategies split fully-connected weights and attention heads across the many small PIM engines, none of which can hold an operator alone.
-
The DSE framework new
A genetic algorithm evolves dataflow candidates (generation lengths, pipeline stages, groupings, schedules); per model, a Monte Carlo tree search fills in architecture choices, group placements, and tiling factors; a capacity checker enforces feasibility (PIM memory is inaccessible to the host during PIM compute — overflowing designs are repaired by re-placing tensors or discarded); analytical latency/energy models built on profiled kernels score every candidate. Default objectives: latency inside the search, energy-delay product for final selection.
-
What the search keeps finding
Across benchmarks the optimizer converges on a legible policy: drafters get MOST of the PIM modules (50% or more), the verifier gets FEW (50% or less) but keeps capacity — allocation tracks each model's ratio of memory-bound operators. The paper distills four takeaways for future PIM hardware: more PIM compute capacity, more external bandwidth, and offload strategies that adapt to both the model and the memory's limits.
The idea
The algorithm made the workload heterogeneous; the hardware should follow. Speculative inference is two jobs wearing one trench coat: a memory-bound token-writer and a compute-bound batch-checker, coupled in a loop. Serving both from identical devices guarantees one of them runs on the wrong machine — the paper’s measurements show the FC-offload PIM recipe (correct for a lone small model) turning counterproductive on the verifier, while all-GPU squanders the drafters. SpecPIM’s central move is to let the RESOURCE ALLOCATION itself be per-model: how much of each device’s memory becomes compute is decided by which model lives there. The searched-out policy — PIM-rich devices for drafters, capacity-rich devices for the verifier — is exactly the arithmetic-intensity analysis made silicon.
Optimism converts dependency into overlap. With models on separate devices, the draft-verify dependency becomes idle hardware: drafters wait, then the verifier waits. The speculative pipeline notices that the common case — verification passes — makes waiting pointless: drafters continue from their own unverified output, and when verification confirms it (the frequent outcome, by construction — speculative inference only pays off when drafts usually pass), the head start is banked. The system thus speculates on its own speculation, with the same insurance structure: a cheap rollback (adopt corrections, skip one verification round) covering the uncommon failure. The measured 1.20× is bounded by exactly the acceptance rate that made the algorithm worthwhile in the first place.
Neither architecture nor dataflow is optimal alone — search their product. The title’s “co-exploration” is earned by the breakdown: starting from GPU speculative inference on PaLM, attention-offload PIM adds 1.22–1.29×, freeing the operator mapping adds 1.12×, the speculative pipeline adds 1.12–1.14×, and freeing the ARCHITECTURE adds another 1.12–1.16× — multiplicative layers, each unlocked by the previous one’s flexibility. Restricting the search to off-the-shelf devices (A100, Samsung HBM-PIM, SK-Hynix GDDR-PIM as-is) costs 1.22× of the win but still beats every fixed baseline — evidence that the exploration, not exotic hardware, carries the value.
One decoding iteration, heterogeneously
PaLM 62B with three drafters, on custom silicon each
Verifier: four capacity-rich devices (few PIM dies, layers pipelined). Drafters: 1.3B/2.7B/8B on PIM-rich devices. Numbers from the paper's largest benchmark.
Step 1 of 5
Drafters write ahead — in memory
Each drafter generates its next tokens autoregressively. Every step is vector-times-matrix at intensity near 1: on the drafters' devices, most memory modules have been converted to PIM, so the multiply happens beside the banks at roughly 8× the external bandwidth. Crucially, the drafters are NOT waiting — they continue optimistically from their own last drafts, while the verifier is still checking the previous round.
the search consistently gives drafters ≥50% PIM modules — the memory-bound half gets the bandwidth dataflow exploration alone is worth 1.29–1.41× over fixed FC/attention offload rules on drafters
Why it works
- Drafters and verifier sit at opposite ends of the arithmetic-intensity spectrum → per-model PIM dosage puts bandwidth where intensity is low and capacity plus host compute where it is high → both halves of the algorithm run on hardware matched to their physics.
- PIM engines pay for bandwidth with capacity (half the module) → treating the conversion ratio as a per-model variable, checked for feasibility, turns a fixed design flaw into a tunable trade → out-of-memory failures and overflow-shuffling overheads disappear from the search’s survivors.
- Draft acceptance is the common case by construction → optimistic continuation overlaps drafting with verification and rarely rolls back → idle time from inter-model dependency becomes throughput, bounded by the acceptance rate itself.
- Operator groups mapped to different engines can run concurrently when independent → model structures with internal parallelism (PaLM’s parallel blocks) overlap host and PIM work → utilization a fixed offload rule cannot reach, worth 1.12× on top of the best fixed rule.
- The joint space is too large and too coupled for manual design → genetic search over dataflows with per-model MCTS over architectures, scored by analytical models on profiled kernels → optimal points per scenario, and legible design rules extracted from where the optimizer converges.
What it costs
- Simulation, not silicon. Results come from an extended vendor PIM simulator with profiled GPU kernel latencies and event-driven multi-device modeling — carefully constructed, but no fabricated system; the architecture space assumes devices (custom PIM dosages) nobody currently ships, though the fixed-device variant partially answers this.
- Search cost per scenario. Fifty genetic generations of a thousand individuals, with ten thousand MCTS leaves per model — offline and amortizable, but a new model mix, batch size, or device budget means re-exploring.
- Benchmark artifacts. The 1.3B/2.7B LLaMA and PaLM drafters do not exist as released models — they are constructed from OPT-shaped specs and profiled with OPT weights — so the acceptance-rate statistics driving the pipeline’s gains are proxies.
- Scope. Decoding only (prefill stays on host), FP16/BF16, batch sizes 1–8, dedicated devices per model — the design targets online serving of one speculative-inference job, not multi-tenant clusters.
- Dependence on speculation quality. The pipeline’s overlap and the whole algorithm’s advantage shrink together when drafts are rejected often — a bad drafter-verifier pairing degrades every layer of the design at once.
The evaluation
Three model scales (OPT 13B, LLaMA 32.5B, PaLM 62B as verifiers, one to three drafters each) at batch sizes 1–8, against GPU autoregressive decoding, GPU speculative inference (DeepSpeed on A100s), and the two fixed PIM offloading recipes on identical PIM-enabled devices.
Three homes for speculative inference
| Property | GPUs (A100 + DeepSpeed) | Fixed-recipe PIM accelerators (FC- or attention-offload) | SpecPIM |
|---|---|---|---|
| Memory-bound drafter operators | |||
| Compute-intensive verifier | Well served — this is what GPUs are for | ||
| Draft-verify dependency | Serialized on shared hardware | Serialized across dedicated devices — someone is always idle | |
| End-to-end latency (geomean) | GPU-SI: 1.45× over plain decoding — the reference | 2.91× (FC-offload) and 1.39× (attention-offload) SLOWER than SpecPIM | |
| Energy efficiency (geomean) | Reference (GPU-SI: 2.07× over plain decoding) | FC-offload only 1.25× over GPU-SI, worse at large batch; attention-offload 4.95× |
Check yourself
Check your understanding
1. Samsung's FC-offload recipe accelerated single-model GPT-J on PIM. Applied to speculative inference, it is SLOWER than a plain GPU. Predict the two failure mechanisms.
2. The search converges on giving drafters at least half the PIM modules and the verifier at most half. Explain why from the arithmetic-intensity analysis — and what the verifier gets instead.
3. The speculative pipeline lets drafters generate from their own UNVERIFIED tokens instead of waiting. Why is this safe, when does it pay, and what bounds the gain?
4. The PaLM speedup breakdown reads: PIM attention offload 1.22–1.29×, then mapping exploration +1.12×, then speculative pipeline +1.12–1.14×, then architecture exploration +1.12–1.16×. What does this compositional structure argue, and what does the fixed-device experiment add?