[DFlash2] Add configurable MoE draft FFN

src/speculators/models/dflash2/config.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from speculators import SpeculatorModelConfig
 from speculators.models.dflash.config import DFlashSpeculatorConfig
@@ -43,3 +43,51 @@ class DFlash2SpeculatorConfig(DFlashSpeculatorConfig):
         ge=1,
         description="Number of unary candidates reranked during inference.",
     )
+    draft_ffn_type: Literal["dense", "moe"] = Field(
+        default="dense",
+        description="Feed-forward implementation used by every draft layer.",
+    )
+    num_experts: int = Field(
+        default=256,
+        ge=1,
+        description="Number of routed experts when draft_ffn_type='moe'.",
+    )
+    num_experts_per_tok: int = Field(
+        default=8,
+        ge=1,
+        description="Number of routed experts selected per token.",
+    )
+    moe_intermediate_size: int = Field(
+        default=512,
+        ge=1,
+        description="Intermediate width of each routed expert.",
+    )
+    shared_expert_intermediate_size: int = Field(
+        default=512,
+        ge=1,
+        description="Intermediate width of the always-on shared expert.",
+    )
+    moe_experts_implementation: Literal[
+        "grouped_mm", "batched_mm", "deepgemm", "sonicmoe", "reference"
+    ] = Field(
+        default="grouped_mm",
+        description=(
+            "Kernel used for the routed-expert GEMMs. This is a closed set on "
+            "purpose: Transformers silently falls back to the per-expert Python "
+            "loop when it does not recognize the name, and that loop costs about "
+            "68x a grouped GEMM at 256 experts / hidden 2048 / intermediate 512 "
+            "(measured on H200, ~505ms vs ~7.4ms per layer forward+backward). "
+            "'reference' names that loop explicitly for debugging; 'batched_mm' "
+            "materializes a dense per-expert activation and will exhaust device "
+            "memory at production expert counts."
+        ),
+    )
+
+    @model_validator(mode="after")
+    def validate_moe_routing(self) -> "DFlash2SpeculatorConfig":
+        if self.draft_ffn_type == "moe" and self.num_experts_per_tok > self.num_experts:
+            raise ValueError(
+                "num_experts_per_tok cannot exceed num_experts: "
+                f"{self.num_experts_per_tok} > {self.num_experts}."
+            )
+        return self

src/speculators/models/dflash2/core.py

@@ -31,7 +31,7 @@ class DFlash2DraftModel(DFlashDraftModel):
     """DFlash with local convolution and bilinear candidate reranking."""
 
     config_class: ClassVar[type[DFlash2SpeculatorConfig]] = DFlash2SpeculatorConfig  # type: ignore[misc,assignment]
-    _no_split_modules = ["Qwen3DFlash2DecoderLayer"]
+    _no_split_modules = ["Qwen3DFlash2DecoderLayer", "Qwen35DFlashMoeBlock"]
 
     def __init__(self, config: DFlash2SpeculatorConfig) -> None:
         target_vocab_size = config.transformer_layer_config.vocab_size
@@ -48,13 +48,14 @@ def __init__(self, config: DFlash2SpeculatorConfig) -> None:
             )
         super().__init__(config=config)
 
+        initializer_range = getattr(
+            config.transformer_layer_config, "initializer_range", 0.02
+        )
         for layer_ in self.layers:
             assert isinstance(layer_, Qwen3DFlash2DecoderLayer)  # noqa: S101
             layer_.reset_convolutions()
+            layer_.reset_moe_parameters(initializer_range)
 
-        initializer_range = getattr(
-            config.transformer_layer_config, "initializer_range", 0.02
-        )
         self.candidate_selector = CandidateSelector(
             vocab_size=target_vocab_size,
             hidden_size=config.transformer_layer_config.hidden_size,
@@ -73,6 +74,12 @@ def _make_decoder_layer(
             block_size=config.block_size,
             conv_kernel_size=config.conv_kernel_size,
             conv_group_size=config.conv_group_size,
+            draft_ffn_type=config.draft_ffn_type,
+            num_experts=config.num_experts,
+            num_experts_per_tok=config.num_experts_per_tok,
+            moe_intermediate_size=config.moe_intermediate_size,
+            shared_expert_intermediate_size=config.shared_expert_intermediate_size,
+            moe_experts_implementation=config.moe_experts_implementation,
         )
 
     @classmethod
@@ -90,6 +97,16 @@ def from_training_args(
             conv_group_size=kwargs.get("conv_group_size", 16),
             selector_rank=kwargs.get("selector_rank", 256),
             selector_top_k=kwargs.get("selector_top_k", 16),
+            draft_ffn_type=kwargs.get("draft_ffn_type", "dense"),
+            num_experts=kwargs.get("num_experts", 256),
+            num_experts_per_tok=kwargs.get("num_experts_per_tok", 8),
+            moe_intermediate_size=kwargs.get("moe_intermediate_size", 512),
+            shared_expert_intermediate_size=kwargs.get(
+                "shared_expert_intermediate_size", 512
+            ),
+            moe_experts_implementation=kwargs.get(
+                "moe_experts_implementation", "grouped_mm"
+            ),
         )
         model = cls(config=config)
         model.load_vocab_mappings(t2d, d2t)

src/speculators/models/dflash2/model_definitions.py

@@ -25,13 +25,20 @@
 https://github.com/z-lab/dflash/blob/07ebd93db9f472af339b644bb70221ad8428328a/dflash/model.py
 """
 
+from copy import copy
+
 import torch
 from torch import nn
 from transformers.cache_utils import Cache
 from transformers.models.qwen3.modeling_qwen3 import (
     FlashAttentionKwargs,
     Qwen3Config,
 )
+from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
+    Qwen3_5MoeExperts,
+    Qwen3_5MoeMLP,
+    Qwen3_5MoeTopKRouter,
+)
 from typing_extensions import Unpack
 
 from speculators.models.dflash.model_definitions import Qwen3DFlashDecoderLayer
@@ -40,10 +47,84 @@
     "CandidateSelector",
     "GroupedDynamicCausalConv",
     "Qwen3DFlash2DecoderLayer",
+    "Qwen35DFlashMoeBlock",
     "grouped_dynamic_conv",
 ]
 
 
+class Qwen35DFlashMoeBlock(nn.Module):
+    """Qwen3.5-compatible routed and shared experts for a DFlash2 layer.
+
+    Only the feed-forward is replaced; the attention, the local convolutions
+    and the draft-block geometry are the dense DFlash2 ones. Weight names match
+    the Qwen3.5 MoE contract so a verifier layer warm-starts tensor for tensor.
+    """
+
+    def __init__(
+        self,
+        config: Qwen3Config,
+        *,
+        num_experts: int,
+        num_experts_per_tok: int,
+        moe_intermediate_size: int,
+        shared_expert_intermediate_size: int,
+        experts_implementation: str = "grouped_mm",
+    ) -> None:
+        super().__init__()
+        # Copy rather than annotate the shared attention config: it is
+        # serialized as transformer_layer_config, and MoE architecture fields
+        # have their one source of truth in DFlash2SpeculatorConfig.
+        moe_config = copy(config)
+        moe_config.num_experts = num_experts
+        moe_config.num_experts_per_tok = num_experts_per_tok
+        moe_config.moe_intermediate_size = moe_intermediate_size
+        moe_config._experts_implementation = (  # noqa: SLF001
+            None if experts_implementation == "reference" else experts_implementation
+        )
+        self.gate = Qwen3_5MoeTopKRouter(moe_config)
+        self.experts = Qwen3_5MoeExperts(moe_config)
+        self.shared_expert = Qwen3_5MoeMLP(
+            moe_config, intermediate_size=shared_expert_intermediate_size
+        )
+        self.shared_expert_gate = nn.Linear(moe_config.hidden_size, 1, bias=False)
+
+    def reset_parameters(self, initializer_range: float) -> None:
+        """Initialize the raw expert tensors and never leave the router at zero.
+
+        ``Qwen3_5MoeTopKRouter`` allocates its weight with ``torch.zeros`` and
+        relies on ``Qwen3_5MoePreTrainedModel._init_weights`` to replace it.
+        A DFlash2 draft is not on that class's MRO, so without this every token
+        would see identical router logits at step 0 and collapse onto one fixed
+        expert subset.
+        """
+        nn.init.normal_(self.gate.weight, mean=0.0, std=initializer_range)
+        nn.init.normal_(self.experts.gate_up_proj, mean=0.0, std=initializer_range)
+        nn.init.normal_(self.experts.down_proj, mean=0.0, std=initializer_range)
+        for projection in (
+            self.shared_expert.gate_proj,
+            self.shared_expert.up_proj,
+            self.shared_expert.down_proj,
+            self.shared_expert_gate,
+        ):
+            nn.init.normal_(projection.weight, mean=0.0, std=initializer_range)
+
+    def _route(
+        self, flat_hidden: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        router_logits, routing_weights, selected_experts = self.gate(flat_hidden)
+        routed = self.experts(flat_hidden, selected_experts, routing_weights)
+        shared_gate = torch.sigmoid(self.shared_expert_gate(flat_hidden))
+        shared = shared_gate * self.shared_expert(flat_hidden)
+        return routed + shared, router_logits, selected_experts, shared_gate
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        original_shape = hidden_states.shape
+        output, _logits, _selected, _gate = self._route(
+            hidden_states.reshape(-1, original_shape[-1])
+        )
+        return output.reshape(original_shape)
+
+
 def grouped_dynamic_conv(
     hidden_states: torch.Tensor,
     delta_kernel: torch.Tensor,
@@ -183,8 +264,24 @@ def __init__(
         block_size: int,
         conv_kernel_size: int,
         conv_group_size: int,
+        draft_ffn_type: str = "dense",
+        num_experts: int = 256,
+        num_experts_per_tok: int = 8,
+        moe_intermediate_size: int = 512,
+        shared_expert_intermediate_size: int = 512,
+        moe_experts_implementation: str = "grouped_mm",
     ) -> None:
         super().__init__(config=config, layer_idx=layer_idx)
+        self.block_size = block_size
+        if draft_ffn_type == "moe":
+            self.mlp = Qwen35DFlashMoeBlock(
+                config,
+                num_experts=num_experts,
+                num_experts_per_tok=num_experts_per_tok,
+                moe_intermediate_size=moe_intermediate_size,
+                shared_expert_intermediate_size=shared_expert_intermediate_size,
+                experts_implementation=moe_experts_implementation,
+            )
         conv_kwargs = {
             "block_size": block_size,
             "kernel_size": conv_kernel_size,
@@ -200,6 +297,11 @@ def reset_convolutions(self) -> None:
         self.attention_conv.reset_parameters()
         self.mlp_conv.reset_parameters()
 
+    def reset_moe_parameters(self, initializer_range: float) -> None:
+        """Initialize MoE tensors; a no-op for a dense draft layer."""
+        if isinstance(self.mlp, Qwen35DFlashMoeBlock):
+            self.mlp.reset_parameters(initializer_range)
+
     def forward(
         self,
         target_hidden: torch.Tensor | None = None,

src/speculators/train/config/schema.py

@@ -484,6 +484,34 @@ class DFlash2Args(_Group):
         ge=0.0,
         description="DFlash2: weight of the candidate-selector K-way CE term.",
     )
+    draft_ffn_type: Literal["dense", "moe"] = Field(
+        default="dense",
+        description="DFlash2: dense MLP or routed MoE plus a shared expert. "
+        "Only the feed-forward changes; attention, convolutions and the draft "
+        "block layout stay the same.",
+    )
+    num_experts: int = Field(
+        default=256, ge=1, description="DFlash2 MoE: number of routed experts."
+    )
+    num_experts_per_tok: int = Field(
+        default=8, ge=1, description="DFlash2 MoE: routed experts selected per token."
+    )
+    moe_intermediate_size: int = Field(
+        default=512, ge=1, description="DFlash2 MoE: per-routed-expert width."
+    )
+    shared_expert_intermediate_size: int = Field(
+        default=512, ge=1, description="DFlash2 MoE: shared-expert width."
+    )
+    moe_experts_implementation: Literal[
+        "grouped_mm", "batched_mm", "deepgemm", "sonicmoe", "reference"
+    ] = Field(
+        default="grouped_mm",
+        description="DFlash2 MoE: routed-expert GEMM kernel. Restricted on "
+        "purpose -- an unrecognized name makes Transformers fall back to a "
+        "per-expert Python loop that costs ~68x a grouped GEMM at 256 experts. "
+        "'reference' selects that loop deliberately; 'batched_mm' exhausts "
+        "device memory at production expert counts.",
+    )
 
 
 class DSparkArgs(_Group):
@@ -734,6 +762,20 @@ def _validate_dpace(self) -> "TrainConfig":
                 )
         return self
 
+    @model_validator(mode="after")
+    def _validate_moe(self) -> "TrainConfig":
+        """Reject MoE geometry that would otherwise fail deep in a run."""
+        if self.dflash2.draft_ffn_type != "moe":
+            return self
+        if self.speculator_type != "dflash2":
+            raise ValueError("--draft-ffn-type=moe is currently supported by dflash2")
+        if self.dflash2.num_experts_per_tok > self.dflash2.num_experts:
+            raise ValueError(
+                "--num-experts-per-tok cannot exceed --num-experts: "
+                f"{self.dflash2.num_experts_per_tok} > {self.dflash2.num_experts}"
+            )
+        return self
+
     def flatten(self) -> dict[str, Any]:
         """The flat ``vars(args)``-shaped dict the ``SpeculatorModel`` classes
         consume via ``**kwargs``.