Repository navigation
[Feat] dflash: Serve MoE FFN drafts for DFlash and DSpark - #48763
Open
yushengsu-thu wants to merge 28 commits into
Open
yushengsu-thu wants to merge 28 commits into
yushengsu-thu wants to merge 28 commits into
Conversation
Three msgspec structs under layers/draft: DraftDenseFFNSpec, DraftMoeFFNSpec and DraftMixFFNSpec (one FFN kind per layer). Config-level only, no torch.
yushengsu-thu
force-pushed
the
dflash-family-moe-ffn
branch
from
November 6, 2026 03:36
d40425d to
67a5251
Compare
DraftMixFFNSpec.from_config reads the HF MoE fields and ffn_layer_types; the dense and MoE sub-specs build themselves from the same config. The config helpers move to layers/draft/config_utils.py for the next components.
DraftMoeFFNSpec now carries scoring_func, group-limited routing, routed scaling, the SwiGLU clamp and both shared-expert layouts, so Qwen, Mixtral and DeepSeek-style drafts parse alike. The only refusals are what sglang's TopK and FusedMoE cannot run; moe_preset is not read.
parse_dflash_draft_config now fills DFlashDraftConfig.ffn, so every DFlash and DSpark draft config is parsed at server-args time. The spec reads the per-layer schedule keys of other families too (moe_layer_freq, first_k_dense_replace, mlp_only_layers, mlp_layer_types) and sizes experts by intermediate_size when moe_intermediate_size is absent, so published drafts whose config carries target MoE fields (MiniMax DSpark) still parse as dense.
Dense, Qwen-style, DeepSeek-style and gpt-oss-style configs, every per-layer schedule key, the kernel-capability refusals, inconsistent fields, and the DFlashDraftConfig wiring. CPU only, suite base-a-test-cpu.
Gate-up and down parallel linears with silu, DeepSeek's clamped SwiGLU or gpt-oss's swigluoai, chosen from the spec's activation fields. It is the dense layer kind and the shared expert of the MoE layer that follows.
The gate carries the trained fp32 expert_bias, TopK takes the spec's scoring, grouping, bias and scaling (torch path for biased softmax), FusedMoE takes the activation and clamp, and the shared expert is a DraftMLP with an optional sigmoid gate. Forward follows.
…e reduce Router logits go to TopK in fp32 as the draft was trained, FusedMoE runs the routed experts, the optional sigmoid-gated shared expert is added, and the block reduces once over TP through reduce_moe_output like the target MoE blocks. A rank with no tokens routes an empty top-k output.
A draft layer asks the factory for its FFN: the caller's dense module when the spec marks the layer dense, a DraftMoE otherwise.
DFlashDecoderLayer takes the model's parsed FFN spec and asks the factory for its FFN, with DFlashMLP as the dense module, so DFlash, DFlash2, Domino and DSpark drafts get a DraftMoE on the layers the spec marks MoE. Without a spec the layer builds the dense MLP exactly as before.
…ft config Native drafts keep mirroring their target: DFlashDraftModel passes no FFN spec to its layers. CustomizedDFlashDraftModel, CustomizedDFlash2DraftModel and CustomizedDSparkDraftModel set customized=True and build their FFN from the draft config, dense or MoE per layer. The architecture names live next to the config parser for the server-args hook.
…drafts A customized DFlash or DSpark draft with MoE layers runs FusedMoE on the dense draft path, which has no DP gather/scatter or idle-rank participation around the experts and no all-to-all dispatcher. The DFLASH and DSPARK handlers refuse --enable-dp-attention and a draft a2a backend other than none for such drafts at server-args time; native drafts are not inspected.
…hecks owns/load_weight accept the router, both expert-bias names, SpecForge's stacked experts, Qwen and DeepSeek per-expert names, the shared expert under either prefix and its gate, each checked against the spec's shapes; FusedMoE's own loader slices for TP and EP. check_loaded refuses a layer with any expert or part missing, so a mismatched export fails instead of serving uninitialized weights.
…ding
DFlashDraftModel.load_weights hands layers.{i}.mlp.* of an MoE layer to
its DraftMoE, refuses MoE tensors aimed at a dense layer, and leaves every
other name on the existing path. post_load_weights asks each DraftMoE for
check_loaded, so a draft never serves with a half-loaded FFN.
load_draft_ffn_weight and check_draft_ffns_loaded move into layers/draft/ffn.py, where the checkpoint naming lives; DFlashDraftModel keeps only the two calls, so another draft family joins with the same two lines.
yushengsu-thu
marked this pull request as ready for review
November 6, 2026 07:18
yushengsu-thu
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
ch-wan,
hnyls2002,
ispobock and
merrymercy
as code owners
November 6, 2026 07:18
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
… DraftMoE on CUDA TopK.forward_cuda resets torch_native, so the fused softmax router dropped the trained expert_bias; biased softmax now routes through fused_topk_torch_native as a custom_routing_function. routed_scaling_factor also goes to FusedMoE, where the Triton runner applies it. The shared expert runs before the in-place fused experts overwrite their input. swiglu_limit values other than 10, the only clamp the fused MoE kernel implements, are refused at construction.
…or DraftMoE test_draft_ffn.py keeps 15 of the 29 spec tests, folding the duplicate parse cases. test_draft_moe_loading.py (CPU) checks that only the Customized* drafts build DraftMoE, that SpecForge's stacked export, per-expert names and the model. prefix load the same experts as FusedMoE.weight_loader, and that partial or mismatched checkpoints are refused. test_draft_ffn_cuda.py (one GPU, base-b) runs DraftMoE against a torch reference of SpecForge's MoELayer: biased softmax with a gated shared expert, plain softmax, DeepSeek-V3-style grouped sigmoid routing, DeepSeek-V4-style sqrtsoftplus with the clamp, and CUDA graph replay.
One CPU file for the draft FFN: the spec's parsing, schedules and refusals (trimmed to 7 tests) plus the 4 cases that build Customized* drafts, load SpecForge's MoE exports and refuse partial or mismatched checkpoints.
main moved add_prefix out of sglang.srt.utils, which broke every import of the DFlash-family models on the merged tree.
main moved it out of sglang.srt.utils; the customized-draft check in _handle_dflash and _handle_dspark failed at every DFLASH/DSPARK launch.
This comment has been minimized.
This comment has been minimized.
parse_dflash_draft_config read the FFN fields of every draft, so native drafts with minimal configs (no intermediate_size) failed to resolve, and the hook's config read failed on the dummy draft paths unit tests use. DFlashDraftConfig.ffn is gone: the Customized* classes and the hook build the spec themselves, and the hook skips its check when the config is unreadable.
This comment has been minimized.
This comment has been minimized.
The biased softmax router routes in torch; its int64 ids were refused by flashinfer_trtllm's routed bf16 MoE. With int32 ids, as the fused routers emit, DraftMoE passes the reference check on the trtllm runner at the 8 x 1536 draft shape as well as on triton.
…e Customized classes Every DFlash-family draft now parses DraftMixFFNSpec and builds dense or MoE per layer from it; a config without expert fields parses as all dense with nothing to size, so existing checkpoints are untouched and SpecForge MoE exports serve under their own architecture name. The Customized* classes, the customized attribute and CUSTOMIZED_DRAFT_ARCHITECTURES go away; the hook refuses DP attention and a non-none a2a backend for any draft whose spec has MoE layers.
…e hook The hook read every draft's config and treated any expert fields as a DraftMoE draft, which broke drafts that bundle their own MoE implementation (ERNIE-5.5-Flash's DSpark draft under DP attention + deepep). DraftMoE now checks attn_dp_size and the a2a backend when it is built, so only drafts that really build it are refused; the hook no longer reads draft configs. speculative_hook.py is back to main.
The dense kind comes from the spec too: DraftDenseFFNSpec carries the SwiGLU fields and DraftMLP.from_spec builds it, so build_ffn needs no dense callback and DFlashDecoderLayer asks it for every layer's FFN. DFlashMLP stays as a DraftMLP built from the config; its weights and math are unchanged. A mixed dense-and-MoE draft loads each layer's own tensors (new test).
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
SpecForge #1033 -> #1035 train DFlash-family drafts (DFlash, DFlash2, DSpark) whose modules differ from the target, starting with an MoE FFN. SGLang's draft classes only mirror the target's layers, so this PR adds an opt-in path for such drafts, with the MoE FFN first.
Design
A draft need not mirror its target. Every DFlash-family draft reads each layer's FFN kind from its own config: dense, or MoE where it names experts (
ffn_layer_typesor the family's schedule keys). A config without expert fields is today's dense draft, so a SpecForge MoE export serves as is.Modifications
layers/draft/(new):config_utils.py;ffn_spec.py, whereDraftMixFFNSpec.from_configreads one FFN kind per layer and the MoE fields of Qwen, DeepSeek, Mixtral and gpt-oss configs, refusing only what sglang cannot run;ffn.py:DraftMLPandDraftMoE(TopK+FusedMoE; bias, groups, scale, clamp, shared expert, strict loader; biased softmax in torch, on triton or flashinfer_trtllm), picked per layer bybuild_ffn.models/dflash.py:DFlashDraftModelparses the spec; every layer's FFN comes frombuild_ffn;DFlashMLPis now aDraftMLPbuilt from the config; MoE tensors go toDraftMoE, checked after loading.DraftMoErefuses DP attention and a non-noneMoE a2a backend when built.Verification
test_draft_ffn.py(13 tests): spec parsing, schedules and refusals; a config with expert fields buildsDraftMoE, one without stays dense, mixed schedules load both kinds; stacked, per-expert andmodel.-prefixed exports load alike; partial or mismatched checkpoints are refused.test_draft_ffn_cuda.py(6):DraftMoEagainst a torch reference of SpecForge'sMoELayerfor biased softmax, grouped sigmoid, sqrtsoftplus with clamp, shared experts and CUDA graph replay. It caught three CUDA bugs, fixed here: dropped expert bias, unappliedrouted_scaling_factor, shared expert run after the in-place experts.Accuracy Tests
One B200, sglang main + this PR, Qwen3-8B bf16 TP1,
--speculative-algorithm DSPARK --speculative-dspark-block-size 8(draft block 8, verify 9), greedy, thinking off, 1024 new tokens, 8 concurrent requests, radix cache off. Accept length is sglang'savg_spec_accept_length: accepted tokens per verify step over the whole set.yushengsu/Qwen3-8B-5layer-MoE-8E1A-v2Under the same settings the native dense draft
jiapingW/Qwen3-8b-dspark-SpecForgemeasures 3.64 and 6.78.Future work (separate PRs)
DraftDecoderLayerinlayers/draft/, composing attention, FFN, norms and residual from the specs a config names; residual kindsgrandmhc(SpecForge'sdraft_residual) on sglang'sGatedResidualand mHC; attention kindsgqa,mla,deepseek_v41under a contract that keeps DSpark's fused context-KV path. Configs without those fields keep today's layers.dp_moe_sync/ idle-rank participation fromdraft_is_deepseek_v4()to "the draft declares MoE layers", add the DP gather/scatter around the experts, then drop the hook refusal.build_ffn(one call per layer plus the two loader calls).loading.pywith the component loading protocol (owns/load_weight/check_loaded) once a second component exists.Related
#47181 (earlier, wider attempt that also carried tuning artifacts; superseded by this PR), SpecForge #1016 (interim external-package serving), SpecForge #1073 (DeepSeek-V4.1's own draft, served through the official
mtp.*layout).Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ✅ Run #42319039871
Latest PR Test (Extra): ❌ Run #42319039819
Latest PR Test (AMD ROCm 10): ❌ Run #42319039835
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: