add sigmoid router

Author: AngrySigma

d9d/module/block/moe/__init__.py

@@ -3,7 +3,7 @@
 from .grouped_experts import GroupedSwiGLU
 from .grouped_linear import GroupedLinear
 from .layer import MoELayer
-from .router import RoutingResult, TopKRouter
+from .router import RoutingResult, SigmoidGroupedTopKRouter, TopKRouter
 from .shared_expert import SharedExpertParameters, SharedSwiGLU
 
 __all__ = [
@@ -13,5 +13,6 @@
     "RoutingResult",
     "SharedExpertParameters",
     "SharedSwiGLU",
+    "SigmoidGroupedTopKRouter",
     "TopKRouter",
 ]

d9d/module/block/moe/layer.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 import torch
 from torch import nn
 from torch.distributed import ProcessGroup
@@ -9,7 +11,7 @@
     NoCommunicationHandler,
 )
 from .grouped_experts import GroupedSwiGLU
-from .router import RoutingResult, TopKRouter
+from .router import RoutingResult
 from .shared_expert import SharedExpertParameters, SharedSwiGLU
 
 
@@ -29,29 +31,23 @@ def __init__(
         hidden_dim: int,
         intermediate_dim_grouped: int,
         num_grouped_experts: int,
-        top_k: int,
-        router_renormalize_probabilities: bool,
+        router: nn.Module,
         shared_expert: SharedExpertParameters | None = None,
     ):
         """
-         Constructs the MoELayer.
+        Constructs the MoELayer.
 
         Args:
             hidden_dim: Hidden size.
             intermediate_dim_grouped: Intermediate dimension for the Expert FFNs.
             num_grouped_experts: Total number of experts.
-            top_k: Number of experts to route each token to.
-            router_renormalize_probabilities: Configures router probability normalization behavior.
+            router: Router module. Must accept ``(num_tokens, hidden_dim)`` and return a
+                ``RoutingResult``. Use ``TopKRouter`` or ``SigmoidGroupedTopKRouter``.
             shared_expert: Optional configuration for a shared expert.
         """
 
         super().__init__()
-        self.router = TopKRouter(
-            dim=hidden_dim,
-            num_experts=num_grouped_experts,
-            top_k=top_k,
-            renormalize_probabilities=router_renormalize_probabilities,
-        )
+        self.router = router
         self.grouped_experts = GroupedSwiGLU(
             hidden_dim=hidden_dim, intermediate_dim=intermediate_dim_grouped, num_experts=num_grouped_experts
         )
@@ -77,11 +73,11 @@ def enable_distributed_communicator(self, group: ProcessGroup):
         Args:
             group: The PyTorch process group spanning the expert parallel ranks.
         """
-        # Lazy load the handler to prevent early DeepEP bindings/evaluation
         from .communications.deepep import DeepEpCommunicationHandler  # noqa: PLC0415
 
         communicator = DeepEpCommunicationHandler(num_experts=self._num_grouped_experts)
-        communicator.setup(group, self._hidden_dim, self.router.gate.weight.dtype)
+        gate = cast(nn.Linear, self.router.gate)
+        communicator.setup(group, self._hidden_dim, gate.weight.dtype)
         self._communicator = communicator
 
     @torch.no_grad()
@@ -133,7 +129,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
     def reset_parameters(self):
         """Resets module parameters."""
-        self.router.reset_parameters()
+        cast(ModuleLateInit, self.router).reset_parameters()
         self.grouped_experts.reset_parameters()
         if self.shared_expert is not None:
             self.shared_expert.reset_parameters()

d9d/module/block/moe/router.py

@@ -3,6 +3,7 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
+from torch.distributed.tensor import DTensor
 
 from d9d.module.base import ModuleLateInit
 
@@ -105,3 +106,118 @@ def reset_parameters(self) -> None:
             nn.init.zeros_(self.expert_bias)
 
         self.gate.reset_parameters()
+
+
+class SigmoidGroupedTopKRouter(nn.Module, ModuleLateInit):
+    """
+    Selects top-K experts using sigmoid-activated scores with hierarchical group selection.
+
+    Used by the DeepSeek-V3 family (DSv3, GLM-4.6, Kimi K2/K2.6, MiniMax-M2, Mistral Large/Small 4).
+
+    Selection procedure per token:
+      1. Gate logits → sigmoid → raw per-expert probabilities (float32).
+      2. Add ``e_score_correction_bias`` to get selection scores (bias is not reflected in
+         the final expert weights — it is selection-only for load-balancing purposes).
+      3. Divide experts into ``n_group`` equal groups; compute each group's score as the sum
+         of its top-2 expert selection scores.
+      4. Keep the ``topk_group`` highest-scoring groups; mask out all other experts.
+      5. From the surviving experts run a final top-``top_k`` selection.
+      6. Gather weights from the unbiased sigmoid scores (step 1), not the biased scores.
+      7. Optionally renormalise to sum=1, then multiply by ``routed_scaling_factor``.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_experts: int,
+        top_k: int,
+        n_group: int,
+        topk_group: int,
+        routed_scaling_factor: float,
+        norm_topk_prob: bool,
+    ) -> None:
+        """
+        Constructs the SigmoidGroupedTopKRouter.
+
+        Args:
+            dim: Input feature dimensionality.
+            num_experts: Total number of routed experts. Must be divisible by ``n_group``.
+            top_k: Number of experts to select per token.
+            n_group: Number of expert groups. Experts are partitioned into contiguous groups
+                of size ``num_experts // n_group``.
+            topk_group: Number of groups whose experts are candidates for final top-k.
+            routed_scaling_factor: Scalar multiplied onto the final expert weights.
+            norm_topk_prob: If True, renormalise selected weights to sum to 1 before scaling.
+        """
+        super().__init__()
+        if num_experts % n_group != 0:
+            raise ValueError(f"num_experts ({num_experts}) must be divisible by n_group ({n_group})")
+
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+        self.e_score_correction_bias = nn.Buffer(
+            torch.zeros(num_experts, dtype=torch.float32),
+            persistent=True,
+        )
+
+        self._num_experts = num_experts
+        self._top_k = top_k
+        self._n_group = n_group
+        self._topk_group = topk_group
+        self._experts_per_group = num_experts // n_group
+        self._routed_scaling_factor = routed_scaling_factor
+        self._norm_topk_prob = norm_topk_prob
+
+    def forward(self, hidden_states: torch.Tensor) -> RoutingResult:
+        """
+        Calculates routing decisions for the input tokens.
+
+        Args:
+            hidden_states: Input tokens. Shape: ``(num_tokens, dim)``.
+
+        Returns:
+            The routing result with selected expert indices and their (possibly scaled) weights.
+        """
+        num_tokens = hidden_states.shape[0]
+
+        scores = F.linear(hidden_states.float(), self.gate.weight.float()).sigmoid()
+
+        bias = self.e_score_correction_bias
+        if isinstance(bias, DTensor):
+            bias = bias.to_local()
+        scores_for_sel = scores + bias
+
+        # Group scoring: sum of top-2 expert scores within each group
+        grouped = scores_for_sel.view(num_tokens, self._n_group, self._experts_per_group)
+        group_scores = grouped.topk(min(2, self._experts_per_group), dim=-1)[0].sum(dim=-1)
+
+        # Select topk_group groups
+        group_idx = torch.topk(group_scores, k=self._topk_group, dim=-1, sorted=False)[1]
+
+        # Build expert-level boolean mask from selected groups
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1.0)
+        expert_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(-1, -1, self._experts_per_group)
+            .reshape(num_tokens, self._num_experts)
+            .bool()
+        )
+
+        # Final top-k within surviving experts
+        masked_scores = scores_for_sel.masked_fill(~expert_mask, float("-inf"))
+        selected_indices = torch.topk(masked_scores, k=self._top_k, dim=-1, sorted=False)[1]
+
+        # Weights from unbiased scores (bias must not affect the FFN computation)
+        selected_weights = scores.gather(dim=-1, index=selected_indices)
+
+        if self._norm_topk_prob:
+            selected_weights = selected_weights / (selected_weights.sum(dim=-1, keepdim=True) + 1e-20)
+
+        selected_weights = selected_weights * self._routed_scaling_factor
+
+        return RoutingResult(selected_expert_indices=selected_indices, selected_probabilities=selected_weights)
+
+    def reset_parameters(self) -> None:
+        """Resets module parameters."""
+        self.gate.reset_parameters()
+        nn.init.zeros_(self.e_score_correction_bias)

d9d/module/model/qwen3_moe/decoder_layer.py

@@ -3,7 +3,7 @@
 
 from d9d.module.base import ModuleLateInit
 from d9d.module.block.attention import GroupedQueryAttention
-from d9d.module.block.moe import MoELayer
+from d9d.module.block.moe import MoELayer, TopKRouter
 from d9d.module.block.normalization import RMSNorm
 from d9d.module.block.positional import RotaryEmbeddingStyle
 
@@ -42,8 +42,12 @@ def __init__(self, params: Qwen3MoELayerParameters):
             hidden_dim=params.hidden_size,
             num_grouped_experts=params.num_experts,
             intermediate_dim_grouped=params.intermediate_size,
-            top_k=params.experts_top_k,
-            router_renormalize_probabilities=True,
+            router=TopKRouter(
+                dim=params.hidden_size,
+                num_experts=params.num_experts,
+                top_k=params.experts_top_k,
+                renormalize_probabilities=True,
+            ),
         )
 
         self.input_layernorm = RMSNorm(params.hidden_size, eps=params.rms_norm_eps)

deps/0007-dep-sigmoid-grouped-moe-router.md

@@ -0,0 +1,157 @@
+---
+DEP: 0006
+Title: Sigmoid Grouped MoE Router
+Author: Illarion Iov
+Status: Implemented
+Type: Feature
+Created: 2026-05-18
+---
+
+# DEP-0006: Sigmoid Grouped MoE Router
+
+## Abstract
+
+This proposal adds `SigmoidGroupedTopKRouter` to `d9d/module/block/moe/router.py` — a router that
+combines sigmoid-activated expert scores with a two-level hierarchical selection procedure
+(group selection → per-group top-k). It also makes `MoELayer` router-agnostic by accepting any
+router instance via constructor injection instead of constructing `TopKRouter` internally.
+
+The new router is the routing strategy used by DeepSeek-V3, GLM-4.6, Kimi K2 / K2.6,
+MiniMax-M2, and Mistral Large 3 / Small 4 — all of which are planned for d9d integration.
+
+## Motivation
+
+The existing `TopKRouter` uses softmax gating: token scores across all experts are normalised
+jointly before top-k selection. Softmax introduces a coupling between experts that is absent in
+sigmoid gating, and it does not allow per-token load-balancing corrections to be applied without
+re-normalising the whole distribution.
+
+DeepSeek-V3 and its derivative models address this with a three-part change:
+
+1. **Sigmoid instead of softmax.** Each expert is scored independently; there is no
+   inter-expert normalisation before selection.
+2. **Hierarchical group selection.** Experts are divided into `n_group` equal groups. A group
+   score (sum of the top-2 expert scores within the group) selects `topk_group` candidate groups
+   first; the final top-k experts are drawn only from those groups. This constrains which experts
+   a token can reach and simplifies load balancing.
+3. **Correction bias.** An `e_score_correction_bias` buffer (one scalar per expert, not updated
+   by gradient) shifts expert scores before group and expert selection. It is updated externally
+   by a load-balancing controller. The final expert weights fed into the FFN use the unbiased
+   sigmoid scores — the bias is selection-only.
+
+Without this router, none of the DeepSeek-V3 family checkpoints can be loaded correctly into d9d.
+
+## Design Proposal
+
+### New class: `SigmoidGroupedTopKRouter`
+
+```python
+class SigmoidGroupedTopKRouter(nn.Module, ModuleLateInit):
+    def __init__(
+        self,
+        dim: int,
+        num_experts: int,
+        top_k: int,
+        n_group: int,
+        topk_group: int,
+        routed_scaling_factor: float,
+        norm_topk_prob: bool,
+    ) -> None: ...
+
+    def forward(self, hidden_states: torch.Tensor) -> RoutingResult: ...
+    def reset_parameters(self) -> None: ...
+```
+
+**Forward pass** (tokens shape `(T, dim)`):
+
+```
+scores      = sigmoid(gate(hidden_states))                      # (T, E),  float32
+scores_sel  = scores + e_score_correction_bias                  # (T, E)   — for selection only
+group_score = scores_sel.view(T, G, E/G).topk(2, -1)[0].sum(-1) # (T, G)
+group_idx   = topk(group_score, topk_group)[1]                  # (T, topk_group)
+expert_mask = scatter group_idx → (T, E) boolean mask
+masked      = scores_sel.masked_fill(~expert_mask, -inf)
+indices     = topk(masked, top_k)[1]                            # (T, top_k)
+weights     = scores.gather(-1, indices)                        # unbiased weights
+if norm_topk_prob: weights /= weights.sum(-1, keepdim=True) + ε
+weights    *= routed_scaling_factor
+return RoutingResult(indices, weights)
+```
+
+The `e_score_correction_bias` is a `nn.Buffer(persistent=True)`, initialised to zero. It is not
+a gradient parameter; callers update it via an external load-balancing controller.
+
+### Change to `MoELayer`
+
+`MoELayer` is refactored to be fully router-agnostic. `router` becomes a required positional
+argument; `router_renormalize_probabilities` and `top_k` are removed — both were
+`TopKRouter`-specific concerns that had leaked into the wrong class.
+
+```python
+def __init__(
+    self,
+    hidden_dim: int,
+    intermediate_dim_grouped: int,
+    num_grouped_experts: int,
+    router: nn.Module,
+    shared_expert: SharedExpertParameters | None = None,
+):
+    self.router = router  # plain assignment, no branching
+```
+
+All existing call sites are updated to construct their router explicitly before passing it.
+`reset_parameters` calls `self.router.reset_parameters()` and requires no changes.
+
+### Exports
+
+`SigmoidGroupedTopKRouter` is added to `d9d/module/block/moe/__init__.py` `__all__`.
+
+## Usage
+
+```python
+from d9d.module.block.moe import MoELayer, SigmoidGroupedTopKRouter, TopKRouter
+from d9d.module.block.moe.shared_expert import SharedExpertParameters
+
+# DeepSeek-V3 style
+moe = MoELayer(
+    hidden_dim=7168,
+    num_grouped_experts=256,
+    intermediate_dim_grouped=2048,
+    router=SigmoidGroupedTopKRouter(
+        dim=7168,
+        num_experts=256,
+        top_k=8,
+        n_group=8,
+        topk_group=4,
+        routed_scaling_factor=2.5,
+        norm_topk_prob=True,
+    ),
+    shared_expert=SharedExpertParameters(intermediate_size=2048),
+)
+
+# Qwen3 style (unchanged semantics, explicit construction)
+moe = MoELayer(
+    hidden_dim=4096,
+    num_grouped_experts=64,
+    intermediate_dim_grouped=1024,
+    router=TopKRouter(dim=4096, num_experts=64, top_k=4, renormalize_probabilities=True),
+)
+```
+
+## Backward Compatibility
+
+**Breaking change** to `MoELayer`: `router_renormalize_probabilities` and `top_k` are removed;
+`router` is now required. All in-tree call sites are updated in the same commit. There are no
+external callers in this codebase.
+
+## Alternatives Considered
+
+**Subclassing `TopKRouter`**: the two routers share only the `gate` linear and
+`RoutingResult` return type. A base class would add overhead without benefit.
+
+**`routing_type: Literal["softmax", "sigmoid_grouped"]` enum on `MoELayer`**:
+pushes router hyperparameters into `MoELayer`, coupling it to every future routing variant.
+
+**Optional `router=None` with internal fallback**: keeps the dead
+`router_renormalize_probabilities` / `top_k` params alive and adds a conditional. The required
+argument is cleaner and forces callers to be explicit about what router they are using.