[WIP][Qwen3.8-Flash-Next] Enable breakable CUDA graphs without torch.compile

[gcanlin]

Purpose

Related to #54688.

This draft removes the NVIDIA Qwen3.8-Flash-Next (Qwen4Exp) model-level torch.compile dependency and runs the model-specific operations through breakable CUDA graphs:

• remove support_torch_compile from the main model and MTP paths;
• opt Qwen4Exp into breakable CUDA graphs and mark the metadata-dependent QSA, PLE, and GDN transactions as eager capture boundaries;
• warm the reachable QSA, PLE, GDN, and FlashInfer BF16 MoE specializations before serving;
• fuse the PLE outer residual add and QSA sigmoid output gate into their existing kernels; and
• defer small-batch FlashInfer BF16 MoE finalization and the shared-expert gate, then fuse routed top-k reduction, the shared gate, and the shared/routed add into one Triton tail kernel.

DeepGEMM/Mega-MoE work is intentionally not included in this PR.

Why this is not duplicate work

The existing Qwen4Exp PRs listed in #54688 provide the underlying HC, QSA, and PLE kernels. This PR implements the remaining model integration needed to run the NVIDIA path without model-level Inductor compilation, plus the PLE/QSA epilogue fusions and BF16 MoE tail fusion. No other open PR found during the duplicate-work check covers this integration.

Performance

Initial end-to-end results on 4 x B200, BF16, TP=4, EP=4, using Qwen3.8-Flash-Next. The POC uses VLLM_USE_BREAKABLE_CUDAGRAPH=1; no FP8 or DeepGEMM path is involved.

Sustained 8K/1K serving

200 prompts, 8192 input tokens, 1024 output tokens, max concurrency 16, seed 17, with prefix caching disabled. Both variants used the same container, model, GPUs, server arguments, and per-variant warmup. The baseline was main at 848ab131bcdb5264bcff0d802f47f7d4adb0f548 with the default compiled path; the POC was 6c9fcc3fa78c727cdf2fa2a67eee46892eab4410 with breakable CUDA graphs. Both runs completed 200/200 requests successfully.

Metric	Compiled baseline	Breakable/custom-op POC	Change
Benchmark duration	131.53 s	129.80 s	-1.3%
Request throughput	1.521 req/s	1.541 req/s	+1.3%
Output token throughput	1,557.04 tok/s	1,577.78 tok/s	+1.3%
Total token throughput	14,013.38 tok/s	14,199.98 tok/s	+1.3%
Mean TTFT	1,067.44 ms	1,042.13 ms	-2.4%
Median TTFT	1,215.99 ms	1,041.77 ms	-14.3%
P99 TTFT	1,820.79 ms	1,701.86 ms	-6.5%
Mean TPOT	8.934 ms	8.836 ms	-1.1%
P99 TPOT	9.847 ms	9.811 ms	-0.4%

For this longer, sustained workload, the end-to-end improvement is much smaller than in the short preliminary workloads below: output throughput improves by 1.3%, mean TTFT by 2.4%, and mean TPOT by 1.1%. The median TTFT result should be confirmed with alternating repeated runs before being treated as stable.

Heavy prefill

16 prompts, 8192 input tokens, 32 output tokens, max concurrency 4:

Metric	Compiled baseline	Breakable/custom-op POC	Change
Total token throughput	34,580 tok/s	46,470 tok/s	+34.4%
Mean TTFT	621.00 ms	402.55 ms	-35.2%
Mean TPOT	10.59 ms	9.79 ms	-7.6%

Heavy decode

16 prompts, 128 input tokens, 512 output tokens, max concurrency 16, seed 11:

Metric	Compiled baseline	Breakable/custom-op POC	Change
Output token throughput	1,619.00 tok/s	1,913.76 tok/s	+18.2%
Mean TTFT	925.00 ms	189.18 ms	-79.5%
Mean TPOT	8.08 ms	8.00 ms	-1.0%

Across decode seeds 11-13, output throughput was 1,913.76-1,997.19 tok/s and mean TPOT was 7.69-8.02 ms.

These are preliminary POC measurements. More repetitions, accuracy evaluation, mixed-batch coverage, and per-change ablations are required before the PR is ready for review.

Test plan and results

• pytest -q tests/models/qwen4_exp/test_qsa_reference.py: 50 passed on the initial fusion slice.
• pytest -q tests/models/qwen4_exp/test_ple.py: 32 passed on the initial fusion slice.
• pre-commit run: passed on the initial fusion slice.
• pytest tests/models/qwen4_exp/test_hc_ops.py -k finalize_moe_with_shared: 6 passed, 4 deselected.
• git diff --check: passed for the current branch.
• Deterministic completion check: The capital of France is produced Paris. The capital of Germany is Berlin. The capital of.
• End-to-end vllm bench serve completed for all workloads reported above; the sustained 8K/1K comparison completed 200/200 requests for both variants.

Before marking ready, rerun the complete Qwen4Exp tests and pre-commit on the expanded branch, then add accuracy and mixed-batch results.

AI assistance

AI assistance was used for implementation, profiling orchestration, benchmark analysis, and drafting this PR. The human submitter will review every changed line and rerun the required checks before marking the PR ready for review.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR and related issue.
• The test plan and commands.
• Initial end-to-end performance results.
• Complete test/evaluation refresh before ready-for-review.
• Optional documentation update, if required.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

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.}

[gcanlin]

Superseded by #55272, which is smaller. And split one fusion to #55309