[yushengsu-thu]

What

SGLang cannot load SpecForge's MoE-FFN drafters, so the Qwen3.8-27B DSpark MoE drafters trained by #811 / #812 / kan/moe-3-qwen38 (e.g. RadixArk/qwen38-dspark-moe-3ep-cont-step9916) could only be served from a private SGLang copy, and untuned. This PR adds a serving package SGLang loads through SGLANG_EXTERNAL_MODEL_PACKAGE, plus the tuned fused-MoE kernel configs for that drafter's expert shape.

• specforge/serving/sglang_models/dflash_moe.py: Qwen3MoeDSparkModel (and DFlashMoEDraftModel / DFlash2MoEDraftModel) = SGLang's own Qwen3DSparkModel / DFlashDraftModel / DFlash2DraftModel with each layer's MLP replaced by SGLang's own MoE block (TopK + FusedMoE + a sigmoid-gated shared expert, the same modules as Qwen2MoeSparseMoeBlock). Checkpoints load through the modules' own weight loaders (FusedMoE.make_expert_params_mapping for the per-expert tensors), so expert-width padding, runner-specific layouts (Triton / flashinfer TRT-LLM), quantised checkpoints and TP sharding are handled by SGLang. A strict check refuses partially matching exports. Router logits are computed in fp32 with the folded centering bias exactly as the trainer does, so a token selects the same experts as in training. Qwen (experts.{i}.gate_proj) and DeepSeek (experts.{i}.w1) checkpoint naming are both accepted.
• specforge/serving/sglang_models/moe_ffn.py: the same FFN in plain PyTorch (no SGLang import), used as the reference by the gate and the tests.
• patches/sglang/moe_configs/: tuned fused-MoE Triton configs E=512,N=512 for H200 and GB300 (SGLang ships none, so this shape runs with the generic tile config), the stand-in Qwen3MoeForCausalLM config and a wrapper for SGLang's tuner, README with usage and numbers.
• scripts/gates/check_dspark_moe_sglang_equivalence.py: served block vs plain reference at a small shape and at the real one (512 x 512 top-10, shared 2048): identical top-k, combine weights within 1e-6, bf16-level output drift, truncated export refused; --moe-runner-backend triton|flashinfer_trtllm. Compares the reference against the trainer's MoELayer too when that package is in the checkout (it is not on main yet).
• tests/test_serving/: CPU tests for the reference recipe, preset defaults and name mapping.

Supported recipe: the qwen3_5_moe preset (softmax, renormalised top-k, optional router bias, one sigmoid-gated shared expert). DeepSeek-V4-style recipes (sqrtsoftplus, selection bias, groups, SwiGLU clamp) raise a clear error.

Usage

pip install -e .            # SpecForge with this package
# tuned kernel config (optional, +0.5-2%): copy into SGLang's config dir
C=$(python -c "import sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config as m, os; print(os.path.dirname(m.__file__))")/configs/triton_3_7_1
cp patches/sglang/moe_configs/configs/triton_3_7_1/E=512,N=512,device_name=*.json "$C/"

SGLANG_EXTERNAL_MODEL_PACKAGE=specforge.serving.sglang_models python -m sglang.launch_server \
  --model-path Qwen/Qwen3.8-27B --speculative-algorithm DSPARK --speculative-dspark-block-size 7 \
  --speculative-draft-model-path RadixArk/qwen38-dspark-moe-3ep-cont-step9916 \
  --speculative-draft-model-quantization unquant --attention-backend triton --speculative-draft-attention-backend triton \
  --moe-runner-backend triton      # or flashinfer_trtllm on Blackwell

The export's config.json must say architectures: ["Qwen3MoeDSparkModel"] with the Qwen MoE keys (num_experts, num_experts_per_tok, moe_intermediate_size, shared_expert_intermediate_size, moe_router_bias), which SpecForge's normalizer on the MoE branches writes.

Validation

Gate (check_dspark_moe_sglang_equivalence.py), real shape, 7 / 56 / 448 tokens: PASS on GB300 with the Triton and the flashinfer TRT-LLM runner, and on H200 (Triton); routing identical to the reference, max |dy| 1.6e-2 to 3.1e-2 at mean |y| 0.5 (bf16 level).

End-to-end, bf16 target Qwen/Qwen3.8-27B, drafters bf16, specforge benchmark gsm8k / mt-bench (100 / 200 prompts, 512 new tokens), DSPARK block 7, triton attention, mamba cache 160, max 32 running requests. accept length / tok/s (speedup vs no draft):

GB300 (4 GPUs, one server per GPU):

dataset	c	none	official dense v1	MoE, default config	MoE, tuned config	MoE, TRT-LLM runner
gsm8k	1	101	5.16 / 339 (3.37x)	5.21 / 336 (3.34x)	5.21 / 340 (3.38x)	5.21 / 343 (3.41x)
gsm8k	8	665	5.15 / 1465 (2.20x)	5.20 / 1435 (2.16x)	5.20 / 1457 (2.19x)	5.20 / 1469 (2.21x)
gsm8k	32	1804	5.17 / 3695 (2.05x)	5.23 / 3473 (1.92x)	5.23 / 3487 (1.93x)	5.23 / 3521 (1.95x)
mt-bench	1	101	3.63 / 218 (2.17x)	3.72 / 217 (2.15x)	3.72 / 220 (2.18x)	3.72 / 221 (2.20x)
mt-bench	8	678	3.89 / 1095 (1.62x)	3.99 / 1065 (1.57x)	3.99 / 1084 (1.60x)	3.99 / 1097 (1.62x)
mt-bench	32	1885	4.02 / 2682 (1.42x)	4.11 / 2380 (1.26x)	4.11 / 2392 (1.27x)	4.11 / 2411 (1.28x)

(dense measured alone after the others; the MoE rows ran concurrently.)

Loader check, this package (SGLang's weight loaders) vs a stacked direct-copy loader on the same checkpoint, kernel and config, started together: GB300 gsm8k c=1/8/32 339 / 1367 / 3395 vs 330 / 1342 / 3381 tok/s, mt-bench 218 / 1039 / 2353 vs 215 / 1025 / 2336; H200 (1 GPU, mamba 48, c=1/8) gsm8k 243 / 1173 vs 243 / 1128, mt-bench 157 / 815 vs 157 / 841. Same speed within noise; the 39 GB export loads in 8 s (GB300) / 38 s (H200). Accept lengths identical.

Tuning the kernel config is worth 0.4-1.8% here (0.5-1.1% with an FP8 target, see the README): the drafter's MoE FFN is bandwidth-bound at this expert shape and is 9-21% of a speculative step.

Relation to other PRs

• Serves the exports of MoE FFN skeleton for DFlash-family drafters (contracts, config, seams) #811 / DeepSeek-V4 MoE FFN for DFlash-family drafters + DSV4-Flash ablation arm #812 and kan/moe-3-qwen38; no trainer change.
• Supersedes the SGLang patch of exp - feat(serving): SGLang MoE DSpark draft loader for Kan's MoE drafter (#811/#812), normalizer arch name and equivalence gate (stacked on #931) #940 for the Qwen MoE drafter (that patch ran the experts through torch._grouped_mm inside a patched sglang 0.5.18; this uses stock SGLang >= 0.5.19 and its FusedMoE).
• Includes the configs of feat/sglang-fused-moe-configs-qwen38-dspark-moe.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[yushengsu-thu]

Follow-up: if need to support - DeepSeek-V4 MoE drafters (#812)

sgl-project/SpecForge#941 serves the Qwen recipe only. The DSV4-Flash drafter (64 x 2048 experts, top-6, ungated shared expert, sqrtsoftplus scoring, noaux_tc selection bias, scaling 1.5, SwiGLU clamp 10) needs:

Routing: build TopK with scoring_func="sqrtsoftplus", correction_bias= the export's fp32 gate.bias (selection only, not added to logits), use_grouped_topk / num_expert_group / topk_group, routed_scaling_factor=1.5. SGLang already implements all of this.
Experts: FusedMoE(..., swiglu_limit=10.0); Triton runner only (the TRT-LLM runner has no sqrtsoftplus routing). Loading is unchanged.
Shared expert: ungated, with the same clamp (add a small swiglu_clamped activation; SiluAndMul has no clamp).
Model class: a DeepSeekMoESparseBlock sibling of QwenMoESparseBlock, and a draft class based on SGLang's DeepseekV4ForCausalLMDSpark (V4 attention/residual layout, not Qwen3's); extend check_loaded for gate.bias, no shared_expert_gate.
Export: the normalizer writes the new architecture name; the DeepSeek config keys are already there.
Validation: add the deepseek_v4 case to the gate at 64 x 2048 (the plain reference already supports it); tune E=64,N=2048 for the serving GPU.

About 150 lines, no SGLang changes.