MoE FFN for DFlash-family drafts: the module by yushengsu-thu · Pull Request #1033 · sgl-project/SpecForge

yushengsu-thu · 2026-10-28T06:10:25Z

Summary

The MoE FFN for DFlash-family drafts, split out of #1003: MoELayer, grouped_mm experts, TopKRouter, the qwen3_5_moe family, build_ffn (one line in the DFlash decoder layer), and the Qwen3-8B MoE DSpark config and recipe: 5 layers, every FFN an MoE layer of 8 experts x 1536 with top-1 routing, 6 epochs, use_liger_kernel on (RMSNorms only). No balancing: the router is stateless, the trainer registers no hook.

Interface the stacked PRs build on: the Router contract below. #1034 adds the expert bias to Router and the kept-slot counts to MoELayer; #1035 adds the load-balance loss to MoELayer.

Stacked: #1034 (expert bias), #1035 (load-balance loss). One-branch reference: #1031.

Router

Router is a score -> select -> routing_weights template. forward is fixed and ends with return scores, indices, weights; a family overrides only the step it changes.

step	in -> out	who overrides it	contract
`score`	`x` [T, H] -> fp32 `scores` [T, E]	the family; `TopKRouter`: linear gate `weight` [E, H] (checkpoint key `gate.weight`), softmax, logits in the input dtype	probabilities in [0, 1]; the balancing in #1034 and #1035 assumes it
`select`	`scores` [T, E] -> `indices` [T, K]	a different pick, e.g. DeepSeek-V3's group-limited top-k	reads only what it is given (#1034 hands it the bias-shifted scores)
`routing_weights`	`scores`, `indices` -> `weights` [T, K]	rarely	gather from the unshifted scores, renormalize when `norm_topk_prob` (DeepSeek's 1e-20), times `route_scale`

DeepSeek-V4 as in transformers main needs no new class: TopKRouter(..., norm_topk_prob=True, score_fn=torch.sigmoid, route_scale=config.routed_scaling_factor) plus #1034's expert bias. DeepSeek-V3's group-limited routing is a select override only.

Tests: test_moe_block.py, test_moe_draft_launch.py (2 GPUs), test_dflash_host_syncs.py (MoE case).

The MoE layer, experts, top-k router, the qwen3_5_moe family and the build_ffn switch, with the Qwen3-8B MoE DSpark config and recipe. No load balancing: the router is stateless and the trainer is untouched.

Router fixes what every family's router provides, forward(x) -> scores, indices, weights, and is the point where other families (DeepSeek) add their own routing; TopKRouter keeps the algorithm. Load balancing builds on this contract in a later PR.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

maocheng23 · 2026-10-28T23:47:28Z

+        self.activation = activation
+        self.reset_parameters(initializer_range)
+
+    def reset_parameters(self, std: float) -> None:


Nit, for the semantic alignment, let's add on the reset on the shared experts parts as well

Done: SharedExpert.reset_parameters, called from MoELayer.reset_parameters together with the shared gate. 90967b2

maocheng23 · 2026-10-28T23:48:30Z

+        self.shared_expert_gate = shared_expert_gate
+
+    def reset_parameters(self, std: float) -> None:
+        # Only the router and experts hold bare Parameters that HF's init skips.


fixed by the commit: 90967b2

maocheng23 · 2026-10-28T23:53:06Z

+    """``dense(config)``, or the MoE layer of the family the ``moe_preset`` names."""
+
+    moe_preset = getattr(config, "moe_preset", None)
+    if moe_preset is None:


nit, let's read dense first.

Also, we may need to consider the split of MoE and dense over layers.

Done: within an MoE draft, ffn_layer_types marks the dense layers; build_ffn reads it per layer. 7aab24c

Done in the same commit: 7aab24c

maocheng23

There are some extra code not belongs to this PR, please move.

maocheng23 · 2026-10-28T23:53:54Z

+    """``dense(config)``, or the MoE layer of the family the ``moe_preset`` names."""
+
+    moe_preset = getattr(config, "moe_preset", None)
+    if moe_preset is None:


Also, we may need to consider the split of MoE and dense over layers.

maocheng23 · 2026-10-28T23:56:47Z

+  "max_window_layers": 2,
+  "model_type": "qwen3",
+  "moe_intermediate_size": 1536,
+  "moe_preset": "qwen3_5_moe",


Split the layer wise MoE/dense selection?

{ "num_hidden_layers": 4, "ffn_layer_types": ["dense", "moe", "dense", "moe"], "moe_preset": "qwen3_5_moe", "num_experts": 8, "num_experts_per_tok": 2 }

Done: the config now is able to assign the "dense" or "moe"

maocheng23 · 2026-10-28T23:57:45Z

+from transformers import PretrainedConfig
+
+from specforge.modeling.draft.moe.experts import RoutedExperts, SharedExpert
+from specforge.modeling.draft.moe.router import TopKRouter


Use Router instead of TopKRouter, as we need to support more Routers later

Done: MoELayer takes any Router. 03ca258

maocheng23 · 2026-10-29T00:01:09Z

+    return logits.softmax(dim=-1)
+
+
+class Router(nn.Module):


detail this implementation please.
And for the fwd, the return should be "return scores, indices, weights"

Done: Router implements score -> select -> routing_weights and forward returns scores, indices, weights; TopKRouter supplies the score. https://github.com/yushengsu-thu/SpecForge/blob/dev-moe-module/specforge/modeling/draft/layers/ffn/moe/router.py

maocheng23 · 2026-10-29T00:04:40Z

        target_hidden: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: bool = False,
+        moe_token_mask: Optional[torch.Tensor] = None,


Why do we need this?

it's an argument that will be used in balance loss. Let me move it to next the PR

Moved to #1034, where the balancing counts with it.

maocheng23 · 2026-10-29T00:05:10Z

    return result


+class StepHook(Protocol):


This is not used in this PR?

Moved to #1034, where the balancing registers its hook interface. (I just thought puting the hook interface here for #1034 and #1035 to easilier build on.)

maocheng23 · 2026-10-29T00:19:21Z

And for the dir, please follow the setup of #1025 please, as we need to extract is as a common component.

DSpark accepts use_liger_kernel since sgl-project#963. Liger replaces the draft's RMSNorms; the MoE FFN keeps its grouped_mm experts, which Liger does not cover yet.

SharedExpert.reset_parameters, and MoELayer.reset_parameters covers the shared expert and its gate too, so resetting the layer means the whole layer. HF's init already reached these Linears; the values are the same.

Within an MoE draft (moe_preset), ffn_layer_types marks the layers that keep the dense MLP (intermediate_size); the others take the family's MoE layer (moe_intermediate_size per expert). No list means every layer is MoE, as before. The DSpark MoE config trains ["dense", "moe"].

…he score Router implements the shared steps, top-k selection and the combine weights (renormalized, scaled), and forward always returns scores, indices, weights. A family overrides the steps it needs: TopKRouter adds the linear gate and its softmax score; a DeepSeek-style router would override score (fp32 sigmoid) and select (group-limited). MoELayer takes any Router. TopKRouter's numerics are unchanged.

The step that gathers, renormalizes and scales the chosen experts' scores returns the weights that combine their outputs; routing_weights says so (combine did not).

Follows the draft-layer layout of sgl-project#1025: the MoE FFN is a common FFN component next to the dense one, in specforge/modeling/draft/layers/ffn/moe. A pure move plus import paths; layers/__init__.py and layers/ffn/__init__.py stay empty so sgl-project#1025's versions merge in cleanly.

yushengsu-thu · 2026-10-29T23:15:22Z

And for the dir, please follow the setup of #1025

Done: moved to modeling/draft/layers/ffn/moe next to #1025's dense FFN; when #1026 lands, build_ffn joins layers/factory.py. 20702d2

Sigmoid top-k sums can underflow to 0 and turn the weights into NaN; HF DeepSeek-V3/V4 and torchtitan's route_norm_epsilon guard it with 1e-20. For softmax the sum is at least K / E, so the addition is a bitwise no-op in fp32: Qwen parity and the e2e losses are unchanged.

The shipped draft config now matches the 8 x 1536 experiment: 5 layers, 32 query / 8 KV heads, every FFN an MoE layer (8 experts, top-1), and the recipe trains for 6 epochs with Liger on. ffn_layer_types stays explicit so the per-layer choice is visible; the dense + MoE mix is still covered by the tests.

yushengsu-thu added 2 commits October 27, 2026 22:52

moe: MoE FFN module for DFlash-family drafts

75bff71

The MoE layer, experts, top-k router, the qwen3_5_moe family and the build_ffn switch, with the Qwen3-8B MoE DSpark config and recipe. No load balancing: the router is stateless and the trainer is untouched.

Copilot AI balanced review requested due to automatic review settings October 28, 2026 06:10

yushengsu-thu requested a review from FrankLeeeee as a code owner October 28, 2026 06:10

yushengsu-thu mentioned this pull request Oct 28, 2026

MoE drafts: the expert bias (stacked on #1033) #1034

Open

yushengsu-thu requested review from FlamingoPg, shuaills and sleepcoo as code owners October 28, 2026 06:10

Copilot AI reviewed Oct 28, 2026

This was referenced Oct 28, 2026

MoE FFN for DFlash-family drafts, staged: module, selection bias, aux loss (final) #1031

Open

MoE FFN support for DFlash-family drafts #1003

Open

maocheng23 reviewed Oct 28, 2026

View reviewed changes

maocheng23 reviewed Oct 29, 2026

View reviewed changes

yushengsu-thu added 5 commits October 29, 2026 15:33

configs: turn on Liger kernels in the MoE DSpark recipe

e90bb15

DSpark accepts use_liger_kernel since sgl-project#963. Liger replaces the draft's RMSNorms; the MoE FFN keeps its grouped_mm experts, which Liger does not cover yet.

moe: reset the shared expert with the rest of the MoE layer

b61e1d0

SharedExpert.reset_parameters, and MoELayer.reset_parameters covers the shared expert and its gate too, so resetting the layer means the whole layer. HF's init already reached these Linears; the values are the same.

moe: name the Router's last step routing_weights

6b4ada5

The step that gathers, renormalizes and scales the chosen experts' scores returns the weights that combine their outputs; routing_weights says so (combine did not).

yushengsu-thu force-pushed the dev-moe-module branch from 195f982 to 6b4ada5 Compare October 29, 2026 22:36

yushengsu-thu added 3 commits October 29, 2026 19:25

moe: Router's default score returns its input

0da0ada

This was referenced Oct 31, 2026

MoE drafts: the DeepSeek family (stacked on #1035) #1055

Closed

MoE drafts: the DeepSeek-V4 family (stacked on #1035) #1056

Open

MoE drafts: DeepSeek-V4's own DSpark draft architecture (stacked on #1056) #1057

Draft

This was referenced Nov 4, 2026

MoE drafts: DeepSeek-V4.1's own DSpark draft (stacked on #1035) #1073

Open

[Feat] dflash: Serve MoE FFN drafts for DFlash and DSpark sgl-project/sglang#48763

Open

Conversation

yushengsu-thu commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Router

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maocheng23 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maocheng23 commented Oct 29, 2026

Uh oh!

yushengsu-thu commented Oct 29, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yushengsu-thu commented Oct 28, 2026 •

edited

Loading

yushengsu-thu Oct 29, 2026 •

edited

Loading

yushengsu-thu Oct 29, 2026 •

edited

Loading

yushengsu-thu Oct 29, 2026 •

edited

Loading

yushengsu-thu Oct 29, 2026 •

edited

Loading

yushengsu-thu Oct 29, 2026 •

edited

Loading