[Bugfix][Spec Decode] Honour the draft's moe_backend on Model Runner V2

[stecasta]

Purpose

The draft model is loaded with the target's VllmConfig, so --speculative-config '{"moe_backend": ...}' is ignored on Model Runner V2. When the target is quantized and the draft is not, which is the normal shape of an MTP head, a quantized-only backend rejects the draft and the engine fails to start:

ValueError: moe_backend='flashinfer_b12x' is not supported for unquantized MoE.
Expected one of ['triton', 'batched_triton', 'flashinfer_trtllm', 'flashinfer_cutlass', 'aiter'].

V1 applies this override in LLMBaseProposer._create_draft_vllm_config. V2 never grew the equivalent: load_eagle_model overrides kv_cache_dtype only, and vllm/v1/worker/gpu/ contains no reference to moe_backend at all. The same command line therefore serves under V1 and fails under V2, which became user-visible when MRV2 was made the default in #53183.

Fix

Apply the override alongside the existing kv_cache_dtype one in load_eagle_model. It is scoped to the draft; the target's config is left untouched. MTP reaches this path via mtp/speculator.py.

Validation

Qwen3.6-35B-A3B-NVFP4 with MTP-3 on an RTX PRO 6000 Blackwell Max-Q, --moe-backend flashinfer_b12x with the draft on flashinfer_cutlass. Three arms on one node with identical flags, vLLM v0.28.1rc1.dev199+g7c5dc571c:

arm	runner	fix	outcome
control	V2	no	fails with the ValueError above
treatment	V2	yes	serves, Application startup complete
regression	V1	yes	serves, no V1 change

pytest tests/v1/spec_decode/test_draft_moe_backend_override.py -v   # 3 passed
ruff check <changed files>                                          # All checks passed

The tests patch get_model to capture the config the draft would be built with, and assert that the override reaches the draft's kernel config, that its absence still inherits the target, and that the target's own config is not mutated. Reverting the fix fails them with the bug's signature:

assert used.kernel_config.moe_backend == "flashinfer_cutlass"
AssertionError: assert 'flashinfer_b12x' == 'flashinfer_cutlass'

Why this is not a duplicate

I ran the checks in AGENTS.md. Nothing open touches the V2 draft moe_backend path. The nearest related work is a cluster of V2 draft-config isolation fixes (#54731, #54716), which address different fields.

Model evaluation

Not applicable. This changes which MoE kernel the draft is built with, and only when the user explicitly asks for one. With no moe_backend in the speculative config the behaviour is unchanged, which the second test pins. Affected configurations currently fail to start, so there is no baseline to regress against. Under real rejection sampling the target verifies every drafted token, so the draft's kernel choice cannot alter output.

AI assistance

AI assistance was used to author this change. I have reviewed every changed line, ran the tests and linters above myself, and can defend the design.

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

[stecasta]

Hi @njhill, could you take a look? Currently we can't specify the MoE backend for the drafter head using runner V2

[njhill]

Thanks @stecasta could you just trim down the new comment, it is kind of wordy. We are trying to keep comments concise without superfluous detail.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: 853a0978-d158-4602-93c2-640080fd4ea6

📥 Commits

Reviewing files that changed from the base of the PR and between 425d723 and 0e6fc1f.

📒 Files selected for processing (2)

• tests/v1/spec_decode/test_draft_moe_backend_override.py
• vllm/v1/worker/gpu/spec_decode/eagle/utils.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/v1/spec_decode/test_draft_moe_backend_override.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Bug Fixes
  • Speculative draft models now honor an explicitly configured MoE backend.
  • Draft models inherit the target model’s backend when no override is provided.
  • Loading a draft model no longer alters the target model’s configuration.
  • Improved compatibility when using speculative decoding with different MoE backend settings between target and draft models.
  • Improved embedding sharing for draft models using pipeline parallelism, including clearer handling when embeddings are unavailable on a stage.

Walkthrough

load_eagle_model now applies speculative MoE backend overrides to the draft configuration and delegates target embedding sharing to a helper that handles pipeline-parallel and missing-embedding cases. Tests cover backend inheritance, overrides, and target configuration immutability.

Changes

EAGLE draft configuration and embedding sharing

Layer / File(s)	Summary
Draft backend override and regression coverage vllm/v1/worker/gpu/spec_decode/eagle/utils.py, tests/v1/spec_decode/test_draft_moe_backend_override.py	load_eagle_model copies the configuration with the speculative moe_backend when provided. Tests verify explicit overrides, target backend inheritance, and target configuration immutability.
Pipeline-parallel embedding sharing vllm/v1/worker/gpu/spec_decode/eagle/utils.py	maybe_share_target_embed handles PPMissingLayer, LoRA-wrapped embeddings, pipeline-parallel stages, and missing required target embeddings. load_eagle_model uses the helper for embedding sharing.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 0e6fc

This change lets speculative EAGLE drafts use their configured MoE backend without changing the target model configuration, resolving startup failures for differing target and draft backend requirements. No concrete current-head merge-blocking risk remains.

Suggested reviewers: njhill, mgoin, yongqinwang-cmd

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 37.50% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description clearly explains the Model Runner V2 bug, the moe_backend fix, its scope, and the validation performed.
Title check	✅ Passed	The title clearly and concisely identifies the bug fix and the affected speculative decoding behavior in Model Runner V2.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

---

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

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87322 for commit 9223d6b73bcd.

[github-actions]

✅ @stecasta, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87333 for commit 425d723c46fd.

[njhill]

@stecasta I think the CI failures are related

[chiptoe-svg]

Independent confirmation on a second architecture and a different backend pair — this PR's diagnosis matches what we hit today.

Control (fails): eugr/spark-vllm:nightly-20260905 = vLLM 0.28.1rc1.dev441+g2902ca17e, DGX Spark GB10 (SM121), nvidia/Qwen3.6-35B-A3B-NVFP4, --moe-backend marlin --speculative-config '{"method":"mtp","num_speculative_tokens":5,"moe_backend":"triton"}' — i.e. the NVIDIA DGX Spark playbook recipe for this model. Boot log shows [gpu_worker.py:436] Using V2 Model Runner, then the drafter dies at construction:

ValueError: moe_backend='marlin' is not supported for unquantized MoE. Expected one of ['triton', 'batched_triton', 'flashinfer_trtllm', 'flashinfer_cutlass', 'aiter'].
  qwen3_5_mtp.py:118 → qwen3_5.py:167 → qwen3_next.py:215 → fused_moe/layer.py:373 FusedMoEFactory
  → routed_experts.py:205 → unquantized_fused_moe_method.py:54 → oracle/unquantized.py:290 → :177 map_unquantized_backend

Working reference (V1): stock vllm/vllm-openai:v0.28.0, identical command line, boots and serves; we soaked it for 30 min at k=3 and 10 min at k=5 with no issues. So the same flags serve on V1 and fail on V2, exactly as described above — with marlin/triton rather than flashinfer_b12x/flashinfer_cutlass, so the bug isn't specific to the B12X path.

Practical note for anyone landing here from the playbook recipe: --moe-backend auto boots on the V2 build (the SM12x NVFP4 oracle still picks MARLIN for the target; the draft autoselects FlashInfer CUTLASS), but that's a workaround, not the intended config — the recipe pins marlin deliberately.

Haven't run the treatment on this box yet; happy to if a second-hardware treatment result would help merge.

[stecasta]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87521 for commit 0e6fc1f3ce01.

[stecasta]

@njhill CI is green now, thanks!

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87714 for commit 99d3178d0c83.

[stecasta]

/ci retry

[github-actions]

✅ Queued 4 failed job(s) for retry in Buildkite CI #87714.