[Bugfix][TurboQuant] Fix CUDA graph capture crash with spec-decode + chunked-prefill (#40807)

[oneraghavan]

Purpose

Fixes #40807.

TurboQuant's _prefill_attention continuation branch calls .tolist() on GPU tensors (query_start_loc and seq_lens), which is illegal during CUDA graph capture. This crashes engine initialization when TurboQuant KV cache (turboquant_k8v4, turboquant_4bit_nc, or turboquant_3bit_nc) is combined with --speculative-config method=mtp and --enable-chunked-prefill:

RuntimeError: Cannot copy between CPU and CUDA tensors during CUDA graph capture
unless the CPU tensor is pinned.

The crash occurs at turboquant_attn.py:570 in the continuation-prefill path, which is only entered when max_query_len != max_seq_len — exactly the condition created by spec-decode K+1 verify batches during CUDA graph capture warmup.

Root Cause

1. _cudagraph_support was set to UNIFORM_BATCH, telling vLLM's compilation framework that TurboQuant supports full CUDA graph capture for multi-token batches (including spec-decode K+1 verify batches).
2. During capture warmup, spec-decode constructs batches where max_query_len != max_seq_len, falling into the continuation branch that calls .tolist() on GPU tensors — illegal during graph capture.
3. build_for_cudagraph_capture did not guarantee CPU-resident copies of query_start_loc and seq_lens were populated, so the fallback .tolist() path could be reached.

Fix

Three-layer defense:

1. Downgrade _cudagraph_support from UNIFORM_BATCH to UNIFORM_SINGLE_TOKEN_DECODE. This tells the compilation framework that TurboQuant only supports full CUDA graph capture for single-token decode batches. When spec-decode is active (uniform_decode_query_len > 1), the framework automatically downgrades to PIECEWISE mode — K+1 verify batches run eager (correct), while 1-token decode retains CUDA graph (fast).

2. Populate CPU-resident copies in build_for_cudagraph_capture. If seq_lens_cpu or query_start_loc_cpu are None, they are now explicitly created. This ensures the prefill path always has CPU-resident data and never falls through to .tolist() on GPU tensors.

3. Add is_current_stream_capturing() guard in _prefill_attention continuation branch. Defense-in-depth: if the continuation branch is somehow reached during CUDA graph capture despite fix 1, it returns zeros instead of crashing. Safe because attention outputs during capture are only used for memory profiling, not graph content.

Relation to Other PRs

• PR [Bugfix][Spec-Decode] TurboQuant K+1 spec-verify routing (fixes #40880) #40914 (open) — addresses Layer 2 (degenerate output from cu_seqlens_k = cu_seqlens_q assumption). Complementary; this PR fixes the crash, [Bugfix][Spec-Decode] TurboQuant K+1 spec-verify routing (fixes #40880) #40914 fixes the correctness.
• PR [Bugfix][V1][TurboQuant] Fix workspace AssertionError on first decode (#42544) #42551 (draft) / PR [Bugfix][V1] Warm attention through dummy profiles #42215 / PR [TurboQuant] Share decode scratch workspace across layers #40798 — address Layer 3 (workspace allocation). Complementary; different failure mode.
• This PR fixes Layer 1 (the .tolist() crash) which had no open PR.

Related Issues

• [Bug]: TurboQuant KV + spec-decode + chunked-prefill crashes CUDA graph capture at query_start_loc.tolist() in continuation-prefill path (Qwen3-Next hybrid dense) #40807 — primary issue (this PR)
• [Bug]: MTP × TurboQuant × CUDA graph capture produces degenerate output on Qwen3-Next hybrid (not closed by v7.13 ngram fix tree) #40880 — MTP × TurboQuant × CUDA graph degenerate output (closed, workaround via Genesis P65)
• [Bug]: TurboQuant KV × any speculative decoding (MTP or ngram) produces degenerate token loops — confirmed across dense and hybrid attention #40831 — TurboQuant × any spec-decode produces degenerate loops (closed)
• [Bug]: TurboQuant decode attention: workspace allocation not sized during warmup → AssertionError on first decode after lock_workspace() #42544 — workspace allocation assertion (open, addressed by [Bugfix][V1][TurboQuant] Fix workspace AssertionError on first decode (#42544) #42551)
• [Bug]: Latest Nightly build with TurboQuant KV cache crashes on large chunked continuation prefill after workspace lock ( testing PR #39931 implementing TQ on Hybrid Attention Models e.g Qwen3.5-9B) #41726 — large chunked continuation prefill crash (open)
• [Tracking issue]: TurboQuant/HIGGS Attention follow-ups #40069 — TurboQuant/HIGGS tracking issue

Test Plan

python -m pytest tests/v1/attention/test_turboquant_cudagraph_safety.py -v

7 tests covering:

• _cudagraph_support is UNIFORM_SINGLE_TOKEN_DECODE
• Support level blocks spec-decode full capture
• CPU copies present/absent behavior
• .tolist() on CPU tensors is safe
• is_current_stream_capturing() guard returns zeros (requires CUDA)

Test Result

tests/...::TestCudagraphSupport::test_cudagraph_support_level PASSED
tests/...::TestCudagraphSupport::test_cudagraph_support_blocks_spec_decode_full_capture PASSED
tests/...::TestBuildForCudagraphCaptureCPUCopies::test_cpu_copies_present_when_builder_provides_them PASSED
tests/...::TestBuildForCudagraphCaptureCPUCopies::test_cpu_copies_absent_shows_the_gap PASSED
tests/...::TestBuildForCudagraphCaptureCPUCopies::test_tolist_on_cpu_tensor_is_safe PASSED
tests/...::TestBuildForCudagraphCaptureCPUCopies::test_tolist_on_cpu_seq_lens_is_safe PASSED
tests/...::TestPrefillCaptureGuard::test_capture_guard_returns_zeros SKIPPED (no CUDA)
6 passed, 1 skipped

[oneraghavan]

Hi @LucasWilkinson @MatthewBonanni — could you take a look when you get a chance? This fixes a CUDA graph capture crash in TurboQuant's _prefill_attention when used with spec-decode + chunked-prefill (GPU .tolist() call during capture). Fixes #40807. CI needs the ready label to run. Thanks!

[oneraghavan]

Friendly ping — this PR is ready for review. Let me know if there are any changes needed or blockers. Happy to address feedback. Thanks!

[oneraghavan]

@mgoin — Friendly ping for review. This fixes #40807 where TurboQuant KV + spec-decode + chunked-prefill crashes CUDA graph capture at query_start_loc.tolist() in the continuation-prefill path. Would appreciate a look when you get a chance. Thanks!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @oneraghavan.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork