Sparse MoE FFN for DFlash-family drafters by sherlockwu · Pull Request #799 · sgl-project/SpecForge

sherlockwu · 2026-09-01T19:10:13Z

Stacked on #798 (retarget to main as the stack merges).

Adds a DeepSeek-V4-style sparse MoE FFN that any DFlash-family draft config can opt into, plus a ready-made dense-vs-MoE ablation pair for DeepSeek-V4-Flash.

What's here

specforge/modeling/draft/moe.py — self-contained MoE block:
- MoEGate: sqrtsoftplus top-k routing with the aux-loss-free balancing bias. The bias is an fp32 buffer updated by a sign controller from all-reduced expert loads (never by gradients) and survives module-wide dtype casts.
- GroupedExperts: routed experts as three stacked [E, out, in] parameters — grouped-GEMM-ready, and FSDP use_orig_params tracks 3 view tensors instead of 3*E. Loads both native stacked keys and the official per-expert naming (experts.{i}.w{1,2,3}.weight).
- SparseMoE: sorted-segment dispatch (one argsort → contiguous per-expert segments; a per-expert torch.where loop scales overhead with active experts, ~2× step time once the balancer spreads load), optional torch._grouped_mm path with no host syncs, one shared expert.
- Balance updates are deferred: forward N stashes routing counts, the model forward applies them before forward N+1, outside any activation-checkpoint region — a bias mutated between a forward and its recompute would route to different segment shapes and break checkpointing.
dflash.py: a draft JSON with n_routed_experts > 0 swaps the dense MLP for SparseMoE; the bare gate Parameter is initialized via _init_weights; moe_bias_update_rate comes from dflash_config.
backend.py: checkpoint files keep the official per-expert expert naming while modules keep stacked parameters — unstack/stack_grouped_expert_state_dict convert at the save/load boundary, so warm starts, resumes, and export bundlers are unaffected.
Ablation pair for DSV4-Flash: configs/deepseek-v4-flash-dspark-moe.json is deepseek-v4-flash-dspark.json plus the MoE fields (64 routed + 1 shared, top-6, moe_intermediate_size 2048 — activated width ≈ the dense 12288); examples/.../deepseek-v4-flash-dspark-moe.yaml mirrors the dense recipe so the draft config is the A/B diff.

Tests

tests/test_modeling/test_moe.py: balance-update semantics (train stashes / eval doesn't / second apply is a no-op), fp32 bias through casts, stacked↔official state-dict round trip (including wrapper-level stack), grouped-GEMM vs sorted-loop parity in outputs/grads on CUDA, DFlash wiring (MoE vs dense selection, gate init, model-level round trip).

🤖 Generated with Claude Code

…lation pair - specforge/modeling/draft/moe.py: DeepSeek-V4-style MoE block — sqrtsoftplus top-k gate with the aux-loss-free balancing bias (fp32 buffer, sign-controller updates from all-reduced loads), one shared expert, routed experts as stacked [E, out, in] parameters with sorted segment dispatch and an opt-in grouped-GEMM path - dflash.py: draft configs with n_routed_experts > 0 swap the dense MLP for the MoE block; gate init through _init_weights; the deferred balance update runs in the model forward, outside checkpoint regions - backend.py: checkpoint FILES keep the official per-expert expert naming while modules expose FSDP-friendly stacked parameters (unstack/stack at the state-dict boundary) - configs/deepseek-v4-flash-dspark-moe.json + example recipe: the MoE arm of the drafter ablation vs the dense deepseek-v4-flash-dspark recipe (only the draft JSON's MoE fields differ) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Adds to the framework core the options the Qwen3.5/3.8 and DeepSeek-V4 recipes need, each a family-agnostic component a preset can select: - router.py: sigmoid and sqrtsoftplus scores, group-limited selection, a fixed fp32 router logit bias, and per-sample router centering (the mode the final kan/moe-3-qwen38 run used against from-scratch router collapse) with the export fold of its running mean into that bias. sqrtsoftplus clamps softplus at the smallest normal value before the sqrt: softplus underflows to 0 below a logit of about -104, where sqrt's infinite derivative made the router gradient NaN although the forward was finite (reported on sgl-project#800 for sgl-project#799's MoEGate). - balance.py: noaux_tc (DeepSeek's aux-loss-free selection bias, with the optional complementary sequence loss) and aux_loss (the Switch/Qwen auxiliary loss). - experts.py: the DeepSeek-V4 SwiGLU clamp and the shared expert's sigmoid gate (Qwen). - base.py: apply_keeping_fp32 keeps the fp32 routing buffers (selection bias, centering mean, router bias) exact through module-wide dtype casts; rounding them to bf16 changed which experts a warm-started draft picks. - checkpoint.py: per-layout serving config and serving class; exports fold router centering into the gate and write config.json for the folded draft (an export reloads as the draft it describes). to_hf applies both and the export normalizer names the MoE serving class. - model_loading.py: warm start reads MoE files through the layout without building a second model, and ranks load one at a time, so large drafts no longer exhaust host memory.

…chitecture - specforge/modeling/draft/dspark_v4.py: training-side port of the DSpark head bundled with DeepSeek-V4-Flash-0731 — three DeepSeek-V4 blocks (MLA-style low-rank attention, MoE via draft/moe.py, mHC hyper-connections) under the official mtp.* naming, so trained state dicts round-trip into the bundled layout without renaming - training.attention_backend=native: the draft model implements its own block/window attention (no wrapper-built masks); dspark providers and assembly accept the DSparkV4DraftModel architecture; the confidence head reads its own feature stream via pop_confidence_hidden - scripts: official-weight download, init-from-target (identity/average main_proj, --stage-layers for spread-capture drafts), export bundling into the official servable layout, numeric parity gate vs the reference implementation, step-time benchmark - recipes: official-arch two-node H200 + runbook, 0813-traffic 32K on H200 and B200 2-node stacks, with launch scripts (cherry picked from commit 67fea1a) [Yusheng Su: cherry-picked onto the MoE framework stack (sgl-project#811 rework) without its parent sgl-project#799, so dspark_v4.py's `from .moe import MoEGate, SparseMoE` no longer resolves until a follow-up commit moves the FFN onto MoELayer. examples/configs/README.md conflict: kept both paragraphs and dropped the sentence about sgl-project#799's `...-moe.yaml` ablation recipe.] [Yusheng Su: rebased onto main a580ec5, which gave _dspark_objective_chunk_terms a positional block_index and a keyword-only total_blocks; confidence_hidden now follows block_index, so checkpointed_chunk_reduce still slices it with the other block inputs.]

DSparkV4DraftModel now builds its FFN from the MoE framework (MoELayer with the deepseek_v4 preset) instead of sgl-project#799's SparseMoE: balance updates run through apply_pending_balance_updates, init through MoELayer.reset_parameters, and checkpoint files keep the official names through the deepseek layout (stacked experts and gate.balance.bias in the module). from_pretrained delegates to AutoDraftModel, which reads files through that layout. Scripts: - dspark_v4_official_weights / init / bundle use the family's FP4/FP8 (de)quantizers; init reads only the shards it needs, validates through the layout and takes a strided expert subset when the draft has fewer experts than the target; the bundle maps the root-owned heads back to the official mtp.<last>.* names and checks the drafter's expert counts against the official config. - the parity gate compares MoE routing per stage and position plus every stage's attention and MoE output. Exact mode (fp32 with --tiny or --no-fake-quant) requires no routing flip and rel-err < 1e-4 everywhere; the default serving-dtype mode (bf16) requires half of all routing decisions to agree and cos > 0.99, rel-err < 10% at positions whose routing agrees. It takes --snapshot / --reference-dir and fixes --tiny (the reference now receives official tensor names).

…oject#800 sgl-project#800 also brought tooling that training, exporting and serving the drafter do not use. This keeps the model, its integration, the H200 recipe, the weight loading/init/bundle scripts and the parity gate, and drops: - the B200 drafter config, the 0813-32k two-node B200 recipe and its launcher (the B200 config also carried sgl-project#799's moe_grouped_dispatch key); - the 0813-32k H200 recipe, an experiment recipe for a private dataset; - scripts/launch_dsv4_official_2node.sh, a hand-written launcher with fixed addresses; its capture-server command, Mooncake environment and the trainer's PYTORCH_CUDA_ALLOC_CONF move into the runbook; - the two benchmarks, scripts/bench_dspark_compute.py and scripts/gates/dspark_v4_step_bench.py (the runbook keeps the measured step time and peak memory); - four CLI-bypass allow-list entries for the removed launchers (two of them named scripts sgl-project#800 never added). The runbook no longer points at the approximation recipe's runbook, which this branch does not have, or at that run's training history.

Adds to the framework core the options the Qwen3.5/3.8 and DeepSeek-V4 recipes need, each a family-agnostic component a preset can select: - router.py: sigmoid and sqrtsoftplus scores, group-limited selection, a fixed fp32 router logit bias, and per-sample router centering (the mode the final kan/moe-3-qwen38 run used against from-scratch router collapse) with the export fold of its running mean into that bias. sqrtsoftplus clamps softplus at the smallest normal value before the sqrt: softplus underflows to 0 below a logit of about -104, where sqrt's infinite derivative made the router gradient NaN although the forward was finite (reported on sgl-project#800 for sgl-project#799's MoEGate). - balance.py: noaux_tc (DeepSeek's aux-loss-free selection bias, with the optional complementary sequence loss) and aux_loss (the Switch/Qwen auxiliary loss). - experts.py, shared.py: the DeepSeek-V4 SwiGLU clamp and the shared expert's sigmoid gate (Qwen). - base.py: apply_keeping_fp32 keeps the fp32 routing buffers (selection bias, centering mean, router bias) exact through module-wide dtype casts; rounding them to bf16 changed which experts a warm-started draft picks. - checkpoint.py: per-layout serving config and serving class; exports fold router centering into the gate and write config.json for the folded draft (an export reloads as the draft it describes). to_hf applies both and the export normalizer names the MoE serving class. - model_loading.py: warm start reads MoE files through the layout without building a second model, and ranks load one at a time, so large drafts no longer exhaust host memory.

…chitecture - specforge/modeling/draft/dspark_v4.py: training-side port of the DSpark head bundled with DeepSeek-V4-Flash-0731 — three DeepSeek-V4 blocks (MLA-style low-rank attention, MoE via draft/moe.py, mHC hyper-connections) under the official mtp.* naming, so trained state dicts round-trip into the bundled layout without renaming - training.attention_backend=native: the draft model implements its own block/window attention (no wrapper-built masks); dspark providers and assembly accept the DSparkV4DraftModel architecture; the confidence head reads its own feature stream via pop_confidence_hidden - scripts: official-weight download, init-from-target (identity/average main_proj, --stage-layers for spread-capture drafts), export bundling into the official servable layout, numeric parity gate vs the reference implementation, step-time benchmark - recipes: official-arch two-node H200 + runbook, 0813-traffic 32K on H200 and B200 2-node stacks, with launch scripts (cherry picked from commit 67fea1a) [Yusheng Su: cherry-picked onto the MoE framework stack (sgl-project#811 rework) without its parent sgl-project#799, so dspark_v4.py's `from .moe import MoEGate, SparseMoE` no longer resolves until a follow-up commit moves the FFN onto MoELayer. examples/configs/README.md conflict: kept both paragraphs and dropped the sentence about sgl-project#799's `...-moe.yaml` ablation recipe.] [Yusheng Su: rebased onto main a580ec5, which gave _dspark_objective_chunk_terms a positional block_index and a keyword-only total_blocks; confidence_hidden now follows block_index, so checkpointed_chunk_reduce still slices it with the other block inputs.]

DSparkV4DraftModel now builds its FFN from the MoE framework (MoELayer with the deepseek_v4 preset) instead of sgl-project#799's SparseMoE: balance updates run through apply_pending_balance_updates, init through MoELayer.reset_parameters, and checkpoint files keep the official names through the deepseek layout (stacked experts and gate.balance.bias in the module). from_pretrained delegates to AutoDraftModel, which reads files through that layout. Scripts: - dspark_v4_official_weights / init / bundle use the family's FP4/FP8 (de)quantizers; init reads only the shards it needs, validates through the layout and takes a strided expert subset when the draft has fewer experts than the target; the bundle maps the root-owned heads back to the official mtp.<last>.* names and checks the drafter's expert counts against the official config. - the parity gate compares MoE routing per stage and position plus every stage's attention and MoE output. Exact mode (fp32 with --tiny or --no-fake-quant) requires no routing flip and rel-err < 1e-4 everywhere; the default serving-dtype mode (bf16) requires half of all routing decisions to agree and cos > 0.99, rel-err < 10% at positions whose routing agrees. It takes --snapshot / --reference-dir and fixes --tiny (the reference now receives official tensor names).

…oject#800 sgl-project#800 also brought tooling that training, exporting and serving the drafter do not use. This keeps the model, its integration, the H200 recipe, the weight loading/init/bundle scripts and the parity gate, and drops: - the B200 drafter config, the 0813-32k two-node B200 recipe and its launcher (the B200 config also carried sgl-project#799's moe_grouped_dispatch key); - the 0813-32k H200 recipe, an experiment recipe for a private dataset; - scripts/launch_dsv4_official_2node.sh, a hand-written launcher with fixed addresses; its capture-server command, Mooncake environment and the trainer's PYTORCH_CUDA_ALLOC_CONF move into the runbook; - the two benchmarks, scripts/bench_dspark_compute.py and scripts/gates/dspark_v4_step_bench.py (the runbook keeps the measured step time and peak memory); - four CLI-bypass allow-list entries for the removed launchers (two of them named scripts sgl-project#800 never added). The runbook no longer points at the approximation recipe's runbook, which this branch does not have, or at that run's training history.

- router.py: sigmoid (DeepSeek-V3) and sqrtsoftplus (DeepSeek-V4) score functions. sqrtsoftplus clamps softplus at the smallest normal value before the sqrt: softplus underflows to 0 below a logit of about -104, where sqrt's infinite derivative made the router gradient NaN although the forward was finite (reported on sgl-project#800 for sgl-project#799's MoEGate). group_limited_mask keeps the topk_group groups with the highest top-2 score sums (DeepSeek's group scoring); TopKRouter applies it when topk_group < n_group. - config.py: n_group and topk_group (architecture keys); n_group must divide n_routed_experts and topk_group lies in [1, n_group]. - DESIGN.md: the score functions and the group-limited mask. - tests: the test family scores with sqrtsoftplus; the score functions, a finite sqrtsoftplus gradient for underflowing logits (alone and through a layer), and the group-limited mask.

…chitecture - specforge/modeling/draft/dspark_v4.py: training-side port of the DSpark head bundled with DeepSeek-V4-Flash-0731 — three DeepSeek-V4 blocks (MLA-style low-rank attention, MoE via draft/moe.py, mHC hyper-connections) under the official mtp.* naming, so trained state dicts round-trip into the bundled layout without renaming - training.attention_backend=native: the draft model implements its own block/window attention (no wrapper-built masks); dspark providers and assembly accept the DSparkV4DraftModel architecture; the confidence head reads its own feature stream via pop_confidence_hidden - scripts: official-weight download, init-from-target (identity/average main_proj, --stage-layers for spread-capture drafts), export bundling into the official servable layout, numeric parity gate vs the reference implementation, step-time benchmark - recipes: official-arch two-node H200 + runbook, 0813-traffic 32K on H200 and B200 2-node stacks, with launch scripts (cherry picked from commit 67fea1a) [Yusheng Su: cherry-picked onto the MoE framework stack (sgl-project#811 rework) without its parent sgl-project#799, so dspark_v4.py's `from .moe import MoEGate, SparseMoE` no longer resolves until a follow-up commit moves the FFN onto MoELayer. examples/configs/README.md conflict: kept both paragraphs and dropped the sentence about sgl-project#799's `...-moe.yaml` ablation recipe.] [Yusheng Su: rebased onto main a580ec5, which gave _dspark_objective_chunk_terms a positional block_index and a keyword-only total_blocks; confidence_hidden now follows block_index, so checkpointed_chunk_reduce still slices it with the other block inputs.]

DSparkV4DraftModel now builds its FFN from the MoE framework (MoELayer with the deepseek_v4 preset) instead of sgl-project#799's SparseMoE: balance updates run through apply_pending_balance_updates, init through MoELayer.reset_parameters, and checkpoint files keep the official names through the deepseek layout (stacked experts and gate.balance.bias in the module). from_pretrained delegates to AutoDraftModel, which reads files through that layout. Scripts: - dspark_v4_official_weights / init / bundle use the family's FP4/FP8 (de)quantizers; init reads only the shards it needs, validates through the layout and takes a strided expert subset when the draft has fewer experts than the target; the bundle maps the root-owned heads back to the official mtp.<last>.* names and checks the drafter's expert counts against the official config. - the parity gate compares MoE routing per stage and position plus every stage's attention and MoE output. Exact mode (fp32 with --tiny or --no-fake-quant) requires no routing flip and rel-err < 1e-4 everywhere; the default serving-dtype mode (bf16) requires half of all routing decisions to agree and cos > 0.99, rel-err < 10% at positions whose routing agrees. It takes --snapshot / --reference-dir and fixes --tiny (the reference now receives official tensor names).

…oject#800 sgl-project#800 also brought tooling that training, exporting and serving the drafter do not use. This keeps the model, its integration, the H200 recipe, the weight loading/init/bundle scripts and the parity gate, and drops: - the B200 drafter config, the 0813-32k two-node B200 recipe and its launcher (the B200 config also carried sgl-project#799's moe_grouped_dispatch key); - the 0813-32k H200 recipe, an experiment recipe for a private dataset; - scripts/launch_dsv4_official_2node.sh, a hand-written launcher with fixed addresses; its capture-server command, Mooncake environment and the trainer's PYTORCH_CUDA_ALLOC_CONF move into the runbook; - the two benchmarks, scripts/bench_dspark_compute.py and scripts/gates/dspark_v4_step_bench.py (the runbook keeps the measured step time and peak memory); - four CLI-bypass allow-list entries for the removed launchers (two of them named scripts sgl-project#800 never added). The runbook no longer points at the approximation recipe's runbook, which this branch does not have, or at that run's training history.

sherlockwu requested review from FlamingoPg, FrankLeeeee, shuaills and sleepcoo as code owners September 1, 2026 19:10

sherlockwu mentioned this pull request Sep 1, 2026

DSparkV4DraftModel: official DeepSeek-V4 DSpark drafter architecture #800

Closed

sherlockwu marked this pull request as draft September 1, 2026 19:33

sherlockwu force-pushed the kan/pr2-dsv4-dspark branch from 1b5031b to f74b2d9 Compare September 1, 2026 20:35

sherlockwu force-pushed the kan/pr3-moe-drafter branch 2 times, most recently from 5abf509 to f1ceac8 Compare September 1, 2026 21:45

sherlockwu force-pushed the kan/pr2-dsv4-dspark branch from f74b2d9 to b627abd Compare September 1, 2026 21:45

sherlockwu force-pushed the kan/pr3-moe-drafter branch from f1ceac8 to a60f1de Compare September 1, 2026 21:55

sherlockwu force-pushed the kan/pr2-dsv4-dspark branch 2 times, most recently from 09ff4cb to decc438 Compare September 2, 2026 00:24

sherlockwu force-pushed the kan/pr3-moe-drafter branch from a60f1de to 3c81322 Compare September 2, 2026 00:24

sherlockwu force-pushed the kan/pr2-dsv4-dspark branch from decc438 to 1b0fbee Compare September 2, 2026 00:44

sherlockwu force-pushed the kan/pr3-moe-drafter branch 2 times, most recently from f90c498 to c99dd7c Compare September 2, 2026 06:22

sherlockwu force-pushed the kan/pr2-dsv4-dspark branch from 1b0fbee to 84dd507 Compare September 2, 2026 06:22

sherlockwu force-pushed the kan/pr3-moe-drafter branch 2 times, most recently from dfbf40c to 8a985a8 Compare September 2, 2026 06:52

sherlockwu force-pushed the kan/pr2-dsv4-dspark branch from df01afd to 22cb116 Compare September 2, 2026 06:52

sherlockwu force-pushed the kan/pr3-moe-drafter branch from 8a985a8 to a4b3769 Compare September 2, 2026 17:43

sherlockwu force-pushed the kan/pr2-dsv4-dspark branch from 22cb116 to 3405412 Compare September 2, 2026 17:43

sherlockwu force-pushed the kan/pr3-moe-drafter branch 2 times, most recently from d3d1964 to e7e20ed Compare September 2, 2026 18:40

sherlockwu force-pushed the kan/pr3-moe-drafter branch from e7e20ed to f9a6a80 Compare September 2, 2026 21:14

This was referenced Sep 2, 2026

MoE FFN skeleton for DFlash-family drafters (contracts, config, seams) #811

Closed

DeepSeek-V4 MoE FFN for DFlash-family drafters + DSV4-Flash ablation arm #812

Closed

Base automatically changed from kan/pr2-dsv4-dspark to kan/pr1-infra September 8, 2026 08:58

Base automatically changed from kan/pr1-infra to main September 8, 2026 15:18

This was referenced Sep 16, 2026

feat(moe): partition the routed experts across an expert-parallel group #883

Open

feat(moe): partition the routed experts across an expert-parallel group (moe.py variant) #884

Closed

yushengsu-thu mentioned this pull request Oct 19, 2026

[Feature] MOE Draft model Training #979

Open

3 of 4 tasks

Navigation Menu

Repository navigation

Sparse MoE FFN for DFlash-family drafters - #799

Sparse MoE FFN for DFlash-family drafters#799
sherlockwu wants to merge 1 commit into
mainfrom
kan/pr3-moe-drafter

sherlockwu commented Sep 1, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

sherlockwu commented Sep 1, 2026

What's here

Tests

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant