perf(spec-decode): bound autoregressive draft prefill graph memory

[jackzampolin]

Purpose

Bound autoregressive draft-prefill CUDA-graph memory while retaining FULL replay for common decode request capacities. The draft-position-0 manager now captures a sparse power-of-two request ladder and leaves PIECEWISE available for non-uniform prefill shapes.

For MTP5 with max_num_seqs=32, the selected request capacities are 1, 2, 4, 8, 16, 32, corresponding to token capacities 6, 12, 24, 48, 96, 192. Candidate construction excludes every other FULL descriptor, so runtime dispatch cannot select a graph that was not captured.

Why this is separate

This PR is focused on the autoregressive speculator's graph-capture policy. It does not duplicate the GLM kernel work in #495, #496, or #498, nor the CKV/DCP changes in #488. PR #503 was an aggregate reproduction snapshot and was closed as a duplicate; this surviving PR contains only the generally useful graph-memory change.

Duplicate searches found no local-inference-lab PR implementing a sparse FULL ladder for this manager. Upstream vllm-project#43776 disables full EAGLE prefill graphs for live multimodal batches; it addresses multimodal correctness rather than bounded capture selection for autoregressive draft prefill.

Validation

Source checks:

• uvx ruff check on all touched Python files: passed.
• uvx ruff format on all touched Python files: passed.
• git diff --check: passed.
• Focused Linux pytest in the source-locked SM120 image: 2 passed, 24 deselected.

Runtime qualification used the reviewed sparse-FULL policy on 4x RTX PRO 6000 Blackwell at stock clocks with GLM-5.3-Flash-NVFP4, TP4, DCP4, MTP5, FP8 KV, B12X attention, Humming MoE, B12X linear, C32/B4096/524k, and FULL+PIECEWISE graphs.

• Capture evidence: target 27 PIECEWISE + 25 FULL; MTP prefill 27 PIECEWISE + 6 FULL; MTP decode 7 FULL.
• Boot: 10,835,285 usable KV tokens, 9.64 GiB graph pool, CKV full gather active, zero restarts.
• Estonia c30/high: five runs scored 29, 29, 29, 27, 30, averaging 28.8/30; 144/150 pass, zero errors, zero max-token truncations.
• Sustained prefill: 8K 8,666 tok/s, 64K 9,485 tok/s, 128K 9,393 tok/s.
• Sustained decode at 0/8K context: C1 105/112 tok/s, C8 550/541 tok/s, C32 1,166/1,183 tok/s; zero cell errors and zero server restarts.
• Compared with the prior qualified piecewise-prefill image, prefill changed by +1.2%/+0.2%/0.0%; C32 decode changed by -0.4%/+0.5%.

The reproducible community derivative is published at ghcr.io/jackzampolin/glm53-flash-nvfp4-jovian@sha256:8defdcb4a8b282ae9d5889e098cda289c09cdce67da006612df37fb858b0aa7b. Its release integration commit a66b522873907084f6240e8866a1d681c749b4db is the previously qualified production lineage plus commits patch-equivalent to this PR head.

Review disclosure

OpenAI Codex assisted with implementation, tests, runtime qualification, and PR preparation. This remains a draft until Jack or another maintainer reviews every changed line and can defend the behavior end-to-end.

Summary by CodeRabbit

• Performance Improvements

  • Improved autoregressive speculation CUDA graph handling by using sparse request-capacity captures.
  • Excluded unsupported request counts from uniform decode graph selection.
  • Preserved the configured maximum request capacity while reducing unnecessary graph captures.

• Tests

  • Added coverage for sparse prefill graph initialization and valid CUDA graph candidate filtering.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 4bef4094-94e9-4e0a-8968-a7184c4c86f9

📥 Commits

Reviewing files that changed from the base of the PR and between da4d7be and d303da3.

📒 Files selected for processing (3)

• tests/v1/worker/test_gpu_autoregressive_speculator.py
• vllm/v1/worker/gpu/cudagraph_utils.py
• vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

---

📝 Walkthrough

Walkthrough

The autoregressive speculator now selects sparse full-prefill CUDA graph capture sizes. CudaGraphManager accepts these sizes and filters uniform decode candidates. Tests cover initialization and candidate filtering.

Changes

Autoregressive CUDA graph capture

Layer / File(s)	Summary
Capture policy and manager filtering vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py, vllm/v1/worker/gpu/cudagraph_utils.py	The speculator computes power-of-two request capacities and passes them to SpeculatorCudaGraphManager. CudaGraphManager filters uniform decode descriptors to the configured capture sizes.
Capture behavior tests tests/v1/worker/test_gpu_autoregressive_speculator.py	Tests verify sparse full-prefill initialization, decode-only setup, and removal of uncaptured full-graph candidates.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: lukealonso, njhill, woosukkwon

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}