PolyGraph CUDA graphs go after a hidden cost in AI serving. Every time a speculative decoder changes how many tokens it guesses, the GPU’s pre-recorded work plan stops fitting. A new arXiv paper shows that one recorded graph plus a 28-byte note, written before each step, can stand in for 64 graphs without slowing the decode step.
Key takeaways
- PolyGraph, by Junchao Yang, appeared in arXiv’s listings on 20 October 2026. It records one maximum-size CUDA graph and tells it the real workload size before each replay.
- In a prototype on one NVIDIA H800, it replaced 64 shape-specific graphs with one and cut pre-serving build time from 3,915 ms to 93 ms.
- Median step latency matched the fully pre-captured baseline: 17.7 ms against 17.8 ms.
- The win over padding grows with waste. P99 latency fell 18.7% when padding doubled the work and 33.4% when padding was 700%.
- It is a research prototype. It has not been integrated into vLLM or any other production engine.
What problem do PolyGraph CUDA graphs solve?
Start with CUDA Graphs. Normally the CPU launches GPU kernels one at a time, and each launch costs time. NVIDIA’s documentation puts that overhead at roughly 20 to 200 microseconds per operation in deep learning work. A CUDA Graph records the whole sequence once and then replays it as a single submission.
The catch is that a graph replays exactly what it recorded, including the launch sizes. It works like a player-piano roll: the playing is flawless, but you only get the one song.
Now add speculative decoding. A small draft model guesses the next few tokens, and the big model checks all of them in one pass. Adaptive systems change the batch size (B, the number of requests) and the draft length (d, the number of guessed tokens) from one step to the next. The model therefore has to verify a different number of tokens each step, M = B × d.
The paper says serving systems end up with three awkward choices. They can record a graph for every possible shape, record new graphs while serving live traffic, or pad every step up to the largest draft and throw away the extra work.
How does the 28-byte plan buffer work?
PolyGraph records a single graph at the maximum size, B_max × d_max. It also sets aside a small piece of GPU memory called the plan buffer. The buffer has eight 32-bit integer slots, and seven of them are used.
Those slots describe the real shape of the step: the token rows (M), the hidden and MLP element counts, the context length, the token-head pairs, the draft length and the KV-cache slot base. The KV cache is the model’s memory of earlier tokens. Before each replay, the runtime copies these 28 bytes to the GPU with one asynchronous copy. On the H800 that took about 3.6 microseconds, against a step time of roughly 18 ms.
The custom Triton kernels then read the plan rather than a size fixed at capture time. Each tile of work runs a one-line check: if the tile starts beyond the active row count, it returns at once. As the paper puts it, capture fixes the buffer’s address, not its values. The piano roll stays the same, and a note on the music stand tells it where to stop.
The numbers: PolyGraph vs three common workarounds
The test used model dimensions from Qwen2.5-3B, batch sizes from 1 to 16 and draft lengths from 1 to 8. Together they produce 64 distinct values of M. Results are averaged over three seeds on a synthetic trace, with times in milliseconds.
| Approach | P50 | P99 | Graphs | Pre-build (ms) | Padding waste |
|---|---|---|---|---|---|
| AllShapes (record every shape) | 17.8 | 18.4 | 64 | 3,915 | 0% |
| OnDemand (record while serving) | 17.8 | 62.0 | 39 | 0 | 0% |
| MaxPad (always use d = 8) | 18.4 | 18.4 | 16 | 911 | 31% |
| PolyGraph | 17.7 | 18.4 | 1 | 93 | 0% |
Look at OnDemand’s P99 of 62.0 ms. That is what happens when a new shape has to be recorded in the middle of live traffic: the unlucky request waits. MaxPad avoids that stall but spends 31% of its token work on positions it discards.
Fixed-draft control runs show where the gain comes from. At d = 8, which needs no padding, PolyGraph and MaxPad were within 0.3%. At d = 4, PolyGraph’s P99 was 15.0 ms against 18.4 ms. At d = 1 it was 12.2 ms. The author states it plainly: the gain is conditional, and it appears only when there was padding to remove.
Why does this matter for AI serving?
Graph sprawl is a real headache for engine builders. A vLLM pull request shows capture sizes clashing: speculative decoding needs multiples of the speculative token count plus one, while sequence parallelism needs multiples of the tensor-parallel size. In a 14 August 2026 post on adaptive verification, the vLLM team wrote that variable-length verification needs special variable-length decode graphs, and that cost rises in steps because a batch slightly above a graph size pays for the full padded graph.
PolyGraph offers a cleaner contract: record once, then describe each step in a few bytes. Faster startup means faster autoscaling, and fewer graphs means less memory spent on recordings. My view, as opinion: the idea is almost boringly simple, and that is why it could spread. It also fits the wider push to make GPU software less tied to fixed shapes, a theme in our look at Amazon TorchNeuron and Nvidia’s CUDA moat.
What are the catches?
The paper is candid about its limits, and they are significant:
- Prototype only: it replays custom Triton kernels and leaves out a production scheduler, paged KV-cache management and the draft model.
- Every kernel must cooperate: each shape-dependent kernel has to read the plan and guard its tiles. Closed libraries such as cuBLAS would need new interfaces or separate graph boundaries.
- Narrow evidence: one H800, model-shaped kernels without trained weights, and synthetic traffic.
- Build time is broad: the 93 ms figure measures construction time overall and can include compilation, so it is not pure CUDA capture time.
What happens next
The author lists three next steps: plug the plan-buffer contract into a real adaptive speculative-decoding engine, test it with trained models and production traces, and study a middle path with a few size buckets (K greater than 1). Buckets would add some graphs but cut the number of idle tiles that run only their guard check.
The real test is adoption. If you see vLLM or SGLang pull requests mentioning plan buffers or shape-guarded Triton kernels, the idea has moved from paper to practice.
FAQ
What is PolyGraph in simple terms?
It is a technique that records one maximum-size CUDA graph and passes it a 28-byte description of each step’s real workload, so kernels can skip work that isn’t needed.
Does PolyGraph make LLMs faster?
Only where padding was wasting work. It matched a fully pre-captured setup at the median and beat max-padding by up to 33.4% at P99 in the paper’s prototype tests.
Can I use PolyGraph in vLLM?
Not yet. The paper says it has not been integrated into vLLM or any other production engine, and it needs plan-aware custom kernels.
If you run your own inference stack, check how many graphs it captures at startup and how much of each step goes to padding. That number tells you whether PolyGraph is worth watching. For another way to cut serving costs, read our explainer on Vercel AI Gateway session affinity.
1 comment