exp - feat(serving): SGLang MoE DSpark draft loader for Kan's MoE drafter (#811/#812), normalizer arch name and equivalence gate (stacked on #931)

[yushengsu-thu]

Serving-side companion to Kan's MoE drafter PRs (#811 skeleton + #812 DeepSeek-V4 MoE FFN, @sherlockwu). Kan's PRs add the MoE FFN on the training side and ship the DeepSeek-V4-Flash dense-vs-MoE DSpark ablation arm, but SGLang has no draft class that can load the MoE export, so the MoE arm could not be served or benchmarked. This PR adds the missing loader (as a patch, pending SGLang upstream), the export normalizer wiring, and an equivalence gate against Kan's MoELayer, so the #812 arm can be measured end to end. The serving numbers below are the first real acceptance-length results for Kan's MoE arm.

Stacked on #931 (moe-stack-rebased = main + Kan's #811/#812). This branch is #931 plus one commit; review that commit only (git diff fd0745e3..HEAD). Once #811/#812/#931 land this PR shrinks to that commit.

Problem

SGLang's stock Qwen3DSparkModel has a dense MLP. Loading a DSpark export whose FFN is a DeepSeek-style MoE (moe_preset: deepseek_v4, the MoE arm of Kan's #812) silently drops every expert, gate and shared-expert weight: the server starts, CUDA graphs capture, but the drafter is random (acceptance length 1.006 on gsm8k after a full two-epoch run). Nothing on the serving side could load the MoE arm, so it could only be measured training-side.

What this adds

• patches/sglang/v0.5.18/dspark-moe-draft.patch (same convention as spec-capture.patch; apply with git apply on top of it):
  • new sglang/srt/models/dspark_moe.py with Qwen3MoEDSparkModel: DSparkDraftModel with each layer's dense MLP replaced by DraftMoEFFN (gate + stacked routed experts + ungated shared expert). The routing mirrors Kan's specforge/modeling/draft/moe (DeepSeek-V4 MoE FFN for DFlash-family drafters + DSV4-Flash ablation arm #812) 1:1 in fp32: sqrtsoftplus scores, aux-loss-free selection bias (gate.bias), renormalized combine weights × routed_scaling_factor, SwiGLU clamp, optional group-limited selection. Experts are loaded from the official per-expert naming into stacked [E, out, in] weights and run through torch._grouped_mm (loop fallback). Draft is replicated (TP1 draft worker), no EP.
  • strict loading in dflash.py / dspark.py: checkpoint weights that map to no parameter, or backbone / markov-head parameters the checkpoint leaves untouched, now raise ValueError instead of serving uninitialised. Mislabelling the MoE export as the dense class fails at load (985 checkpoint weight(s) do not map to any parameter).
• scripts/gates/normalize_dflash_export.py: writes architectures: ["Qwen3MoEDSparkModel"] when the export has n_routed_experts > 0; refuses MoE DFlash/DFlash2 exports (no MoE-capable serving class exists for them yet). Two tests added in tests/test_scripts/test_gate_orchestration.py.
• scripts/gates/check_dspark_moe_sglang_equivalence.py: builds SpecForge's MoELayer, converts it through to_checkpoint_state_dict, loads it into DraftMoEFFN via the serving stacking logic, and compares routing (identical top-k and weights) and outputs for the small config and the real 64-expert / top-6 / 2048 sizes on both dispatch paths; also checks that a truncated expert set is rejected.
• Recipe doc: "Serving the MoE arm" paragraph in docs/recipes/deepseek-v4-flash-dspark-disaggregated.md.

Validation

• Equivalence gate: PASS on B200 (bit-exact outputs, identical routing) for small and DSV4-preset sizes, grouped_mm and loop.

• End-to-end with Kan's DeepSeek-V4 MoE FFN for DFlash-family drafters + DSV4-Flash ablation arm #812 DSV4-Flash MoE arm (SGLang v0.5.18 + spec-capture + this patch, --speculative-algorithm DSPARK, greedy, 512 tokens): strict load lists 64 experts / top-6 / shared / sqrtsoftplus / bias / ×1.5 / clamp 10 with nothing skipped; CUDA graph capture bs 1–8 fine; generation correct.

  MoE checkpoint	gsm8k accept (c=1)	mt-bench accept (c=1)	gsm8k accept (c=8)	tok/s (c=1)
  step 128	1.118	1.095 (c=4)	1.117	86
  step 1024	1.30	1.34	1.31	94
  dense step 1024 (reference, stock class)	1.38	1.40	–	117

  Acceptance grows with training in step with the dense arm (previously stuck at 1.006 with the dense class), and matches the training-side per-position accuracy estimate within ~0.08.

• Local: python -m pytest tests/test_scripts/test_gate_orchestration.py tests/test_runtime/test_package_architecture.py → 30 passed; black / isort clean on the changed Python files.

Follow-ups

• Send Qwen3MoEDSparkModel and the strict loader checks to SGLang upstream (the patch targets v0.5.18; upstream dspark.py / dflash.py have moved, so it needs a rebase), then drop the patch here and keep only the normalizer change.
• MoE-capable DFlash / DFlash2 serving classes (the exp - feat(moe): qwen3 preset and the Qwen3.8-27B DFlash2 MoE ablation arm (stacked on #811/#812) #930 Qwen3.8-27B DFlash2 MoE arm is training-side only until then).
• specforge benchmark cannot render DeepSeek-V4 prompts (no Jinja template); the numbers above used a script around the checkpoint's Python encoder.

🤖 Generated with Claude Code