Coerce hybrid-model speculative decode off full decode CUDAGraphs on MUSA

[yeahdongcn]

Coerce hybrid-model speculative decode off full decode CUDAGraphs on MUSA

Summary

Enabling speculative decoding (e.g. MTP, dFlash) on a hybrid GatedDeltaNet /
Mamba model on MUSA with the default cudagraph_mode=FULL_AND_PIECEWISE
produces silently wrong output. This makes the default configuration
correct: for hybrid-model speculative decode the verify runs piecewise, and the
optimizations that would force it back to a full graph are disabled.

The bug

On MUSA a full decode CUDAGraph corrupts the multi-query speculative verify step
of a hybrid target:

• The verify runs the target at query_len = num_speculative_tokens + 1. The
  paged multi-query verify attention miscomputes when captured inside the full
  model graph at that shape, so the decoded output silently diverges from the
  correct (eager / PIECEWISE) result and degrades into garbage after ~a dozen
  tokens.
• query_len = 1 decode under the same full graph is byte-correct; the defect
  is specific to the query_len > 1 verify.
• The low apparent "acceptance" of the corrupt path is a symptom.

The defect was localized to the MUSA full-cudagraph capture context, below
vLLM: every input the verify attention receives is byte-identical FULL vs
PIECEWISE, each op is capture-correct in isolation, and it is independent of
the cudagraph memory pool and the mamba cache layout — yet the attention output
diverges only inside the full captured graph. It is a torch_musa / mate
FlashAttention capture-context defect, not a vLLM or single-kernel bug.

The change

vllm_musa/platform.py::MUSAPlatformBase.check_and_update_config: for any
speculative decode (num_speculative_tokens > 0) on a hybrid model
(model_config.is_hybrid) whose cudagraph mode includes full decode graphs:

1. set cudagraph_mode = PIECEWISE so the verify runs eager; or NONE when
   compilation is not VLLM_COMPILE (piecewise cudagraphs require compilation);
2. clear the full-graph-only fusion flags pass_config.enable_sp,
   fuse_gemm_comms, and fuse_attn_quant.

Step 2 is required because those optimizations need full-graph compilation, so
CompilationConfig.set_splitting_ops_for_v1() (run later in
VllmConfig.__post_init__, after this hook) would otherwise re-promote
PIECEWISE back to FULL and reintroduce the corruption. Opt back into full
decode capture with VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL=1.

• Method-agnostic: fires for any hybrid speculative decode — an
  unknown/renamed draft method cannot silently slip into the corrupt path.
• Covers dFlash: the dFlash target verify hits the same corruption, so it
  is coerced too (the dFlash drafter graph is configured separately).
• Non-spec runs and non-hybrid models are unaffected. Config-only; no kernel or
  native change.

Guarantee

This avoids the full-graph corruption and preserves the piecewise result. It
does not claim bit-exactness of speculative decode versus non-spec decode.

Validation

Env: MTT S5000 ×8, TP8, driver 3.3.5; vLLM 0.24.1.dev0, torchada 0.1.69,
mate 0.2.3, torch / torch_musa 2.9.0 / musa40305. Test model: Qwen3.6-35B-A3B
(hybrid GatedDeltaNet + 256-expert MoE, Qwen3_5MoeForConditionalGeneration),
bf16, greedy, num_spec=1.

• Correctness (MTP): with the guard the default MTP output is
  byte-identical to explicit cudagraph_mode=PIECEWISE on both prompts
  (192 tokens each); the token-17 corruption is eliminated.
• Finding 2 (re-promotion): with sequence parallelism enabled, the final
  resolved cudagraph_mode is PIECEWISE (not promoted to FULL) and the
  runtime captures piecewise; without the fusion-clearing it resolves to FULL.
• Finding 3 (compilation disabled): with compilation mode NONE, the guard
  resolves to cudagraph_mode=NONE — no "PIECEWISE requires compilation"
  startup assertion.

Unit + regression tests (tests/test_patches.py::TestMUSAPlatformDefaults)

• Gate: hybrid+spec → PIECEWISE; method-agnostic (unknown method still coerced);
  dFlash target verify → PIECEWISE; non-hybrid keeps FULL; non-spec keeps FULL;
  VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL=1 keeps FULL.
• Fusion / lifecycle: hybrid-spec clears enable_sp/fuse_gemm_comms/
  fuse_attn_quant; compilation-disabled → NONE; and two real-CompilationConfig
set_splitting_ops_for_v1 cases — SP promotes FULL_AND_PIECEWISE → FULL
  (the hazard), and with the fusions cleared PIECEWISE survives.

Review response (maintainer review on the prior revision)

1. dFlash carve-out was a correctness hole — removed; the target verify is
   now coerced for dFlash too.
2. Later normalization could undo the guard — fixed by clearing the
   full-graph-only fusions so set_splitting_ops_for_v1 cannot re-promote; a
   lifecycle test proves PIECEWISE survives.
3. Compilation-disabled full modes failed startup — now fall back to NONE.
4. Wording — narrowed to "avoids the corruption, preserves the piecewise
   result".
5. Tests — added the fusion / compilation-disabled / dFlash / promotion
   cases above.

Test plan

• Gate + lifecycle unit tests (fusion-clearing, compilation-disabled→NONE,
  dFlash target coerced, SP promotion + survival)
• MTP default byte-identical to explicit PIECEWISE (2 prompts)
• SP enabled → resolved cudagraph_mode PIECEWISE (not promoted)
• Compilation disabled → resolved cudagraph_mode NONE (no crash)
• Non-hybrid MLA (DeepSeek) control — gate does not fire
• [~] dFlash target coercion — unit test + a dense-dFlash negative control pass;
  the hybrid-dFlash E2E is deferred (the Qwen3.5-4B-DFlash drafter
  checkpoint is not staged on the test cluster)

[augmentcode]

🤖 Augment PR Summary

Summary: Avoids silently wrong outputs on MUSA when using MTP speculative decoding with hybrid (GatedDeltaNet/Mamba) models by preventing full-decode CUDAGraph capture.
Change: Adds a guarded config coercion in MUSAPlatformBase.check_and_update_config to force cudagraph_mode=PIECEWISE for hybrid+MTP unless VLLM_MUSA_HYBRID_MTP_ALLOW_FULL=1 is set.

_{🤖 Was this summary useful? React with 👍 or 👎}

[augmentcode]

Review completed. 2 suggestions posted.

Comment augment review to trigger a new review at any time.

[yeahdongcn]

I independently re-reviewed and reran exact head e28e84e565efec2b362295e5a89a306198b8687f. My verdict is request changes / do not merge as-is.

Remote verification used a fresh, non-privileged container on 8x S5000, driver 3.3.5-server, image sha256:2deb7a34b747..., torch 2.9.0 / MUSA 4.3.5, torchada 0.1.69, mate 0.2.3, and the pinned vLLM v0.24.0 tree. The tracked source was reset to the exact PR head. The only native source newer than the image was an unrelated FP8 per-token-group kernel; these runs used BF16 models.

Why PIECEWISE works but FULL_AND_PIECEWISE does not

FULL_AND_PIECEWISE is dispatch routing, not a nested “piecewise inside full” safety mode. A speculative target verification batch is uniform decode with query_len = 1 + num_speculative_tokens; the dispatcher prefers FULL before PIECEWISE (dispatcher). Therefore:

• pure PIECEWISE keeps attention/GDN splitting ops outside the captured fragments;
• FULL_AND_PIECEWISE routes the uniform multi-query verification step through the outer full-model graph.

The real TP8 Qwen3.6-35B-A3B (num_speculative_tokens=1, BS1, greedy seed 0, two 192-token outputs) reproduced this exactly:

Row	Effective verify mode	Result
explicit PIECEWISE	PIECEWISE	coherent reference
PR guarded default	PIECEWISE	token-for-token equal to the reference for both prompts
FULL_AND_PIECEWISE with the escape hatch	FULL at qlen=2	first divergence at generated tokens 17 and 14; then constant-token suffixes of 107 and 105 tokens
non-spec default	FULL at qlen=1	coherent; only late differences from the speculative reference at tokens 159 and 183

This is not a blanket “MUSA full graphs never work” issue. It is the uniform multi-token verify under the full-model capture context. As an additional control, isolated qlen=1/2 attention write+read graph replays matched eager with max|diff|=0 in four cases, so the evidence does not justify claiming a narrower single-kernel root cause yet.

Blocking findings

1. The dFlash carve-out is a real correctness hole. The PR excludes dFlash because it manages a draft graph (PR code), but the target forward happens before the later draft proposal (target, draft). I staged the compatible hybrid pair Qwen/Qwen3.5-4B + z-lab/Qwen3.5-4B-DFlash and ran TP1 with 16 draft tokens, exact capture size 17, max_num_seqs=1, prompts longer than 17 tokens, sequential BS1 requests, and VLLM_MUSA_DRAFT_FULL_WRAP unset:

   • PIECEWISE: runtime logged PIECEWISE for every qlen=17 verify; coherent; both repeats token-identical.
   • VLLM_MUSA_DFLASH_FULL_WRAP=0 + FULL_AND_PIECEWISE: runtime logged FULL for every qlen=17 verify, isolating the target full graph without the special draft wrapper. Both prompts first diverged at token 5 and collapsed into deterministic constant-token suffixes of 50 and 40 tokens.
   • The PR-default dFlash full-wrap path resolved to pure FULL and reproduced the corruption.
   • Eager dFlash and non-spec qlen=1 FULL controls remained coherent.

   Please remove the exclusion, or separate target and drafter graph configuration so the target verify is PIECEWISE while the dFlash drafter can retain its required full graph.

2. Later config normalization can undo the correctness guard. vLLM calls the platform hook before set_splitting_ops_for_v1 (call order). Real EngineArgs -> VllmConfig construction on both Qwen3.6-35B-A3B and Qwen3.6-27B confirmed:

   • sequence parallel / fused GEMM comms promotes the PR's PIECEWISE result back to FULL (logic);
   • attention-quant fusion also promotes it to FULL (logic).

   Enforce the safety invariant after splitting-op normalization, or fall back to NONE when PIECEWISE is incompatible.

3. Compilation-disabled full modes now fail startup. With compilation mode NONE, both explicit FULL and FULL_DECODE_ONLY are changed to PIECEWISE after the earlier compatibility handling, then fail the final assertion that PIECEWISE requires compilation. The safe result in this case should be CUDAGraphMode.NONE.

4. The “correct + lossless” wording is too broad. The fresh MTP run proved guarded-default == explicit PIECEWISE, but speculative PIECEWISE was not byte-equal to non-spec at tokens 159/183. In the dFlash run, PIECEWISE was coherent and matched eager exactly on one prompt, while the other had a benign wording divergence at token 6. Please state the demonstrated guarantee narrowly: it avoids the early full-graph corruption and preserves the PIECEWISE result.

The exact-head platform test class passed (26 passed, 1 xfailed), but its direct-hook tests do not exercise the final config lifecycle above. Please add final EngineArgs/VllmConfig tests for the re-promotion and compilation-disabled cases, plus a hybrid+dFlash correctness regression.

[yeahdongcn]

Thanks for the thorough review — all four findings are addressed in 0e6baddff, with the requested tests. Re-validated E2E on Qwen3.6-35B-A3B (TP8, MTP, capture on).

1. dFlash carve-out removed. The target verify is coerced for dFlash too (it hits the same corruption; the drafter graph is configured separately, earlier in the hook). Covered by a unit test and a dense-dFlash negative control.

2. Later normalization no longer undoes the guard. Confirmed the ordering: check_and_update_config (vllm/config/vllm.py:1350) runs before set_splitting_ops_for_v1 (:1365), which promotes PIECEWISE → FULL for sequence-parallel / fused-GEMM-comms / attn-quant fusion (they require full-graph compilation). Rather than re-assert after, the guard now also clears pass_config.{enable_sp, fuse_gemm_comms, fuse_attn_quant} for the hybrid-spec case, so there is nothing left to promote (an SP+piecewise config would itself be invalid). Verified: create_engine_config with enable_sp=true resolves cudagraph_mode=PIECEWISE (guard-bypassed → FULL_AND_PIECEWISE), and the E2E worker captures PIECEWISE. Two real-CompilationConfig tests lock this in (SP promotes FULL_AND_PIECEWISE → FULL; with the fusions cleared, PIECEWISE survives).

3. Compilation-disabled full modes no longer crash. The guard falls back to NONE when compilation_config.mode != VLLM_COMPILE. Verified: mode NONE resolves cudagraph_mode=NONE with no "PIECEWISE requires compilation" assertion.

4. Wording narrowed to "avoids the full-graph corruption and preserves the piecewise result"; dropped the "correct + lossless" claim.

Regression: default + MTP is byte-identical to explicit PIECEWISE on both prompts; non-spec keeps FULL_AND_PIECEWISE. Unit suite: 30 passed, 1 pre-existing xfail.

The hybrid-dFlash E2E is deferred — the Qwen/Qwen3.5-4B + z-lab/Qwen3.5-4B-DFlash pair isn't staged on our cluster (the only staged dFlash drafter targets a dense model); the config logic is covered by the unit test + the dense-dFlash negative control. Happy to run the full E2E if you can point me at a staged hybrid + dFlash pair.

[yeahdongcn]

I re-reviewed and rebuilt exact head 0e6baddffbb98317e1e7e928b2fcce5be774abd9 after the reply. Verification used a fresh non-privileged container on 8x S5000, driver 3.3.5-server, torch 2.9.0 / MUSA 4.3.5, torchada 0.1.69, mate 0.2.3, and the vLLM v0.24.0 source tree with this exact PR applied. The dFlash exclusion, compilation-disabled FULL/FULL_DECODE_ONLY startup failure, and over-broad "lossless" wording are fixed on the paths they now cover. The default hybrid+spec path also resolves to PIECEWISE with the three promotion flags cleared.

I still recommend fixing one lifecycle hole before merge.

An incoming PIECEWISE mode bypasses the fusion cleanup

The cleanup is inside if _cg is None or _cg.has_full_cudagraphs() (platform.py L533-L546). An explicitly requested PIECEWISE mode therefore skips the entire block and leaves enable_sp / fuse_gemm_comms / fuse_attn_quant enabled. Later, vLLM calls set_splitting_ops_for_v1() and promotes that mode to FULL (lifecycle, SP/GEMM, attn-quant).

I ran the complete EngineArgs -> VllmConfig lifecycle on both Qwen3.6-35B-A3B and Qwen3.6-27B; both produced the same client-side result:

Requested configuration	Completed config	Final flags
combined/default + attn-quant	PIECEWISE	cleared
initial PIECEWISE + attn-quant	FULL	attn-quant remains on
initial PIECEWISE + SP/GEMM	FULL	SP/GEMM remain on

The real TP8 run exposed a second fallback rather than wrong output: the frontend config remained FULL with attn-quant enabled, but each spawned worker reconstructed the config, ran the platform hook a second time, changed it to PIECEWISE, then vLLM saw the already-empty splitting_ops and downgraded it again to NONE. Runtime qLen=2 dispatch was therefore NONE, not FULL, and the output remained coherent. In contrast, the normal guarded default dispatched qLen=2 as PIECEWISE and was byte-equal to the explicit-PIECEWISE reference on both 128-token prompts.

So this exact spawned-worker topology is safe, but it reveals a non-idempotent, construction-count-dependent guard: the first VllmConfig finalizes to FULL, while worker reconstruction falls back to NONE and loses PCG. Reconstruction is not a guaranteed safety boundary: vLLM supports an in-process client that passes the same finalized config directly into EngineCore (selection, InprocClient); at world size 1 it selects UniProcExecutor, which passes that same config object directly to the worker (backend, worker init). The safety invariant should hold after the first lifecycle rather than depending on a later reconstruction.

The new tests do not cover this: the fusion-cleanup test starts from FULL_AND_PIECEWISE, while the separate preservation test starts from PIECEWISE only after manually setting all three flags false (tests L573-L596, L686-L701). Please enforce the invariant for an already-PIECEWISE input too and add full-lifecycle regressions for PIECEWISE + each promotion flag, asserting both the canonical and worker-effective modes.

dFlash follow-up

The target verify is now routed safely to PIECEWISE, but the reply/test comment that the drafter remains "configured separately" is not what the code does. The earlier dFlash block writes FULL and capture sizes to the shared compilation_config (platform.py L470-L497); the new generic guard then overwrites that same mode with PIECEWISE. The draft dispatcher is constructed from the same config and adds its FULL keys only when the supplied mode still has full graphs (dispatcher patch).

The exact-head hybrid+dFlash E2E (Qwen/Qwen3.5-4B + z-lab/Qwen3.5-4B-DFlash, 16 draft tokens, BS1, capture size 17) resolved to PIECEWISE; every qLen=17 runtime dispatch was PIECEWISE, both 96-token prompts were deterministic across repeats, and both were byte-equal to the eager dFlash control. That verifies the target fix. However, startup captured only PIECEWISE graphs and draft-FULL eligibility was false. This disables the claimed dFlash full-draft path rather than preserving it separately. If that performance tradeoff is intentional, please document it and correct the test comment; otherwise it needs a genuinely separate draft graph mode/configuration.

Two smaller follow-ups:

• breakable CUDAGraph supports PIECEWISE with compilation mode NONE (upstream handling), but the new mode != VLLM_COMPILE branch forces it to NONE, unnecessarily disabling PCG;
• under Inductor graph partitioning, upstream does not perform the incompatible SP/fusion promotion (upstream gates), but an incoming combined mode still has all flags unconditionally cleared. Also, the compilation-disabled warning says the flags were disabled although that branch does not clear them.

Exact-head platform tests passed (30 passed, 1 xfailed).