[Bugfix] Remove nested torch.compile in GDN rearrange_mixed_qkv causing CUDA graph capture failure

[tdoublep]

Summary

• Remove @torch.compile(fullgraph=True) from rearrange_mixed_qkv which triggers Triton autotuning (torch.cuda.synchronize()) during CUDA graph capture
• The method is already compiled by the outer AOT compilation pass, making the nested decorator redundant and harmful

Reproduction

vllm serve Qwen/Qwen3.5-35B-A3B   --language-model-only   --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}'

Fails with:

torch.AcceleratorError: CUDA error: operation not permitted when stream is capturing

Related

• [Bugfix] Fix Triton stream capture error on A100 in GDN attention with MTP speculative decoding #39483 — fixes a similar issue in causal_conv1d Triton kernels via warmup path
• fix mtp launch error in vllm-0.17.1-rc, about cuda graph during memory profile #36634 — same class of bug (GDN + MTP spec decode triggering JIT during capture)

Test plan

• Verify Qwen/Qwen3.5-35B-A3B with MTP spec decode starts successfully on GB200

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request removes the @torch.compile(fullgraph=True) decorator from the rearrange_mixed_qkv method in vllm/model_executor/layers/mamba/gdn_linear_attn.py. I have no feedback to provide as there were no review comments to evaluate.

[ZJY0516]

Doesn't this only run on rocm?

[tjtanaa]

@tdoublep a question. In the PR that introduced this seems to be able to run Qwen3.5 on B200 https://buildkite.com/vllm/ci/builds/65043/canvas?jid=019e05df-d38b-4432-b813-5f66b42e419a&tab=output

If I understand correctly, the rearrange_mixed_qkv is invoked within the scope of custom op gdn_attention_core which is opaque to the torch compile.

[tdoublep]

@ZJY0516 @tjtanaa I am seeing latest main failing on GB200 without this change.

[tpopp]

Sorry about this and I give my spiritual approval. When this started, there were some noticeable additional fusions with torch.compile (with an older vLLM though and before some other GDN changes), so there might be some lost perf, but of course this should be done to fix breakages.

[ZJY0516]

reproduced using

vllm serve Qwen/Qwen3.5-35B-A3B   --language-model-only   --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}'

And it only happens with spec decoding

[ZJY0516]

Sorry about this and I give my spiritual approval. When this started, there were some noticeable additional fusions with torch.compile (with an older vLLM though and before some other GDN changes), so there might be some lost perf, but of course this should be done to fix breakages.

I think the best way is to change the kernel to accept non-contiguous input

[tjtanaa]

reproduced using

vllm serve Qwen/Qwen3.5-35B-A3B   --language-model-only   --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}'

And it only happens with spec decoding

That's why the CI also didn't capture this issue.

[tjtanaa]

@ZJY0516 @tjtanaa I am seeing latest main failing on GB200 without this change.

@tdoublep may I know which test group is that?

[tdoublep]

@tjtanaa No test group, I was just deploying model with MTP - similar to the above example. Surprised it is not caught by tests though.

[tjtanaa]

@tjtanaa No test group, I was just deploying model with MTP - similar to the above example. Surprised it is not caught by tests though.

There is no Qwen 3..5 test on CI it seems. There is only one test group lm-eval Qwen3.5 (B200) that I had to triggered manually.

[ZJY0516]

@tjtanaa related failure https://buildkite.com/vllm/ci/builds/65133/canvas?sid=019e06ae-7db5-4cdb-b6ad-1d49ae6cb583&tab=output

[SoluMilken]

Should we rebase this on latest main and rerun the relevant Buildkite jobs? Thanks.

[SoluMilken]

Could a maintainer please rerun the two failing Buildkite jobs?

• fusion-e2e-tp2-b200 appears to be infra: pytest never started, and the B200 k8s pod stayed Pending for 15m with imagecheck-0 incomplete.
• amd-multi-modal-models-standard-2-qwen3-plus-gemm is not reproducible for me locally because I do not have AMD/MI300 access; the log ends with termination/cancellation rather than a clear assertion failure.

The PR change is limited to removing one nested torch.compile(fullgraph=True), so these look worth rerunning first.

Thanks.