[Model] Kolibri 1 via the Transformers modeling backend (#60827)

Signed-off-by: Luca Motz <luca.motz@icloud.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

docs/models/supported_models.md

@@ -486,6 +486,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
 | `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | ✅︎ | ✅︎ |
+| `Kolibri1ForCausalLM` | Kolibri 1 | `Aleph-Alpha/Kolibri-1`, `Aleph-Alpha/Kolibri-1-BF16`, etc. | ✅︎ | ✅︎ |
 | `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | ✅︎ | ✅︎ |
 | `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ |
 | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -360,6 +360,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "KimiLinearForCausalLM": _HfExamplesInfo(
         "moonshotai/Kimi-Linear-48B-A3B-Instruct", trust_remote_code=True
     ),
+    "Kolibri1ForCausalLM": _HfExamplesInfo(
+        "Aleph-Alpha/Kolibri-1", min_transformers_version="5.20.0"
+    ),
     "LagunaForCausalLM": _HfExamplesInfo("poolside/Laguna-XS.2"),
     "Lfm2ForCausalLM": _HfExamplesInfo("LiquidAI/LFM2-1.2B"),
     "Lfm2MoeForCausalLM": _HfExamplesInfo(

tests/models/transformers/fusers/test_moe.py

@@ -144,6 +144,16 @@ class DisconnectedRouter(TopKRouter):
         return logits, value, index
 
 
+class ExtraInputRouter(TopKRouter):
+    """Router taking a block-owned correction bias as a second input (HYV3)."""
+
+    def forward(self, hidden_states, bias=0.0):
+        logits = F.linear(hidden_states, self.weight)
+        _, index = torch.topk(logits.sigmoid() + bias, self.top_k, dim=-1)
+        value = logits.sigmoid().gather(1, index)
+        return logits, value, index
+
+
 class MoEExperts(nn.Module):
     """Packed experts (3D weights); only its name (`experts`) matters here."""
 
@@ -166,17 +176,19 @@ class MoEBlock(nn.Module):
 
     def _shared(self, x, logits):
         """The term added to the experts' output (none for a plain block)."""
-        return 0
+        return None
 
     def forward(self, hidden_states):
         x = hidden_states.reshape(-1, hidden_states.shape[-1])
         logits, weights, index = self.gate(x)
-        out = self.experts(x, index, weights) + self._shared(x, logits)
+        out = self.experts(x, index, weights)
+        if (shared := self._shared(x, logits)) is not None:
+            out = out + shared
         return out.reshape(hidden_states.shape)
 
 
 class MoEBlockNoShared(MoEBlock):
-    """No shared-expert child but a gate-derived add -> trace skipped, still fuses."""
+    """No shared-expert child but a gate-derived add -> declined, not misread."""
 
     def _shared(self, x, logits):
         return logits.sum()
@@ -293,6 +305,115 @@ class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted):
         self.extra = BufferScale()
 
 
+class MoEBlockExpertsKeywords(MoEBlock):
+    """Passes the experts' routing by keyword -> still fuses."""
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        return self.experts(x, weights=weights, index=index)
+
+
+class ScaledExperts(MoEExperts):
+    def forward(self, hidden_states, index, weights, scale=1.0):
+        return hidden_states * scale
+
+
+class MoEBlockExpertsExtraArg(MoEBlock):
+    """Passes the experts an extra argument the fused forward would drop."""
+
+    def __init__(self):
+        super().__init__()
+        self.experts = ScaledExperts()
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        return self.experts(x, index, weights, scale=2.0)
+
+
+class MoEBlockSharedUnflattened(MoEBlockSharedNoGate):
+    """Gate and shared expert read the unflattened input (Kolibri-1)."""
+
+    def forward(self, hidden_states):
+        _, weights, index = self.gate(hidden_states)
+        x = hidden_states.view(-1, hidden_states.shape[-1])
+        out = self.experts(x, index, weights).view(*hidden_states.shape)
+        return out + self.shared_expert(hidden_states)
+
+
+class MoEBlockSharedAveraged(MoEBlockSharedUnflattened):
+    """Averages the experts' and shared expert's outputs (Cohere2-MoE)."""
+
+    def forward(self, hidden_states):
+        return super().forward(hidden_states) / 2
+
+
+class MoEBlockScaled(MoEBlockSharedNoGate):
+    """Scales the experts' output by `routed_scaling_factor` (Laguna)."""
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        return self.experts(x, index, weights) * 2.5 + self.shared_expert(x)
+
+
+class MoEBlockScaledBiasedRouter(MoEBlockScaled):
+    """An unrecognized router with a scaled experts' output."""
+
+    def __init__(self):
+        super().__init__()
+        self.gate = BiasedRouter()
+
+
+class MoEBlockSharedBiasedRouter(MoEBlockSharedUnflattened):
+    """An unrecognized router with a shared expert."""
+
+    def __init__(self):
+        super().__init__()
+        self.gate = BiasedRouter()
+
+
+class MoEBlockGateExtraInput(MoEBlock):
+    """Passes the gate a block-owned bias the fused forward would drop (HYV3)."""
+
+    def __init__(self):
+        super().__init__(ExtraInputRouter)
+        self.register_buffer("e_score_correction_bias", torch.zeros(8))
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x, self.e_score_correction_bias)
+        return self.experts(x, index, weights)
+
+
+class MoEBlockCastCombine(MoEBlockSharedNoGate):
+    """Combines the outputs in fp32 (HYV3)."""
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        out = self.experts(x, index, weights).float() + self.shared_expert(x).float()
+        return out.to(x.dtype)
+
+
+class MoEBlockIdentityLatent(MoEBlockSharedNoGate):
+    """Wraps the experts in unused (identity) latent projections (NemotronH)."""
+
+    def __init__(self):
+        super().__init__()
+        self.fc1_latent_proj = nn.Identity()
+        self.fc2_latent_proj = nn.Identity()
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        out = self.fc2_latent_proj(
+            self.experts(self.fc1_latent_proj(x), index, weights)
+        )
+        return out + self.shared_expert(x)
+
+
 @pytest.mark.parametrize("sigmoid", [False, True])
 def test_moe_fuser_detects_router(sigmoid):
     with torch.device("meta"):
@@ -385,16 +506,6 @@ def test_moe_fuser_detects_shared_experts():
     assert fuser.shared_gate_name == "shared_expert_gate"
 
 
-def test_moe_fuser_skips_shared_detection_without_extra_children():
-    """With only experts and gate, shared-expert detection (and its block trace)
-    is skipped, so a gate-derived add is not misread as a shared expert."""
-    with torch.device("meta"):
-        block = MoEBlockNoShared()
-    fuser = MoEBlockFuser.match(block, "experts")
-    assert isinstance(fuser, MoEBlockFuser)
-    assert fuser.shared_name is None and fuser.shared_gate_name is None
-
-
 def test_moe_fuser_shared_without_gate():
     with torch.device("meta"):
         block = MoEBlockSharedNoGate()
@@ -422,6 +533,10 @@ def test_moe_fuser_detects_non_glu_shared_expert():
         MoEBlockTupleVar,  # tuple returned via a name binding, not a literal
         MoEBlockUnaccounted,  # weight-bearing child outside the fused dataflow
         MoEBlockUnaccountedBuffer,  # buffer-only child outside the fused dataflow
+        MoEBlockExpertsExtraArg,  # extra experts argument the fused forward drops
+        MoEBlockNoShared,  # gate-derived term added to the experts' output
+        MoEBlockSharedAveraged,  # experts' and shared expert's outputs averaged
+        MoEBlockGateExtraInput,  # gate input the fused forward drops
     ],
 )
 def test_moe_fuser_declines_unsupported(block_cls):
@@ -444,3 +559,53 @@ def test_moe_fuser_router_requires_connected_dataflow():
     with torch.device("meta"):
         block = MoEBlock(DisconnectedRouter)
     assert MoEBlockFuser.match(block, "experts") is None
+
+
+def test_moe_fuser_unrecognized_router_moves_shared_expert():
+    """Without vLLM routing, the shared expert still moves to vLLM."""
+    with torch.device("meta"):
+        block = MoEBlockSharedBiasedRouter()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert fuser.scoring_func is None
+    assert fuser.shared_name == "shared_expert"
+
+
+@pytest.mark.parametrize(
+    ("block_cls", "scoring_func"),
+    [
+        (MoEBlockSharedUnflattened, None),  # routing falls back to Transformers
+        (MoEBlock, "softmax"),  # no shared output to divide
+    ],
+)
+def test_moe_fuser_divided_shared_output(block_cls, scoring_func):
+    with torch.device("meta"):
+        fuser = MoEBlockFuser.match(block_cls(), "experts", divides_shared_output=True)
+    assert fuser.scoring_func == scoring_func
+
+
+def test_moe_fuser_scaled_output_needs_vllm_routing():
+    """Only vLLM routing applies the routed scale, so nothing else moves without it."""
+    with torch.device("meta"):
+        block = MoEBlockScaledBiasedRouter()
+    assert MoEBlockFuser.match(block, "experts") is None
+
+
+def test_moe_fuser_matches_experts_called_by_keyword():
+    with torch.device("meta"):
+        block = MoEBlockExpertsKeywords()
+    assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)
+
+
+def test_moe_fuser_scaled_output_with_vllm_routing():
+    """A scaled block fuses because vLLM routing applies `routed_scaling_factor`."""
+    with torch.device("meta"):
+        fuser = MoEBlockFuser.match(MoEBlockScaled(), "experts")
+    assert fuser.scoring_func == "softmax"
+    assert fuser.shared_name == "shared_expert"
+
+
+@pytest.mark.parametrize("block_cls", [MoEBlockCastCombine, MoEBlockIdentityLatent])
+def test_moe_fuser_looks_through_casts_and_identity(block_cls):
+    with torch.device("meta"):
+        fuser = MoEBlockFuser.match(block_cls(), "experts")
+    assert fuser.shared_name == "shared_expert"

vllm/model_executor/models/registry.py

@@ -717,6 +717,7 @@ _TRANSFORMERS_SUPPORTED_MODELS = {
     "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"),
     "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
     "Jais2ForCausalLM": ("transformers", "TransformersForCausalLM"),
+    "Kolibri1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
     "MellumForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
     "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"),
     "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"),

vllm/model_executor/models/transformers/fusers/moe.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing."""
+"""MoE fuser: run an HF MoE block's gate, experts and shared experts in `FusedMoE`."""
 
 import ast
+import functools
 import inspect
 import textwrap
 import types
@@ -11,20 +12,25 @@ from dataclasses import dataclass
 
 import torch
 from torch import fx, nn
+from typing_extensions import TypeIs
 
 from vllm.distributed import tensor_model_parallel_all_gather
+from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import GateLinear
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.models.transformers.fx_utils import (
     find_node,
     is_complete,
     is_op,
+    output_value,
     peel,
     trace,
 )
 from vllm.model_executor.models.transformers.utils import named_state
 from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk
 
+logger = init_logger(__name__)
+
 
 def _own_returns(node: ast.AST) -> Iterator[ast.Return]:
     """`return` statements in `node`'s own scope, not in nested functions."""
@@ -126,6 +132,35 @@ def _renormalizes(topk: fx.Node) -> bool:
     return False
 
 
+def _unwrap(node: object, moe_block: nn.Module) -> object:
+    """Strip reshapes, dtype casts and `nn.Identity` calls."""
+    while True:
+        node = peel(node)
+        if (
+            is_op(node, "reshape")
+            or is_op(node, "view")
+            or (
+                isinstance(node, fx.Node)
+                and node.op == "call_module"
+                and isinstance(moe_block.get_submodule(node.target), nn.Identity)
+            )
+        ):
+            node = node.args[0]
+        else:
+            return node
+
+
+def _reads_input(node: object, x: fx.Node, moe_block: nn.Module) -> TypeIs[fx.Node]:
+    """Is `node` a child module of `moe_block` called on just `x`?"""
+    return (
+        isinstance(node, fx.Node)
+        and node.op == "call_module"
+        and len(node.args) == 1
+        and not node.kwargs
+        and _unwrap(node.args[0], moe_block) is x
+    )
+
+
 class SharedExpertMLP(nn.Module):
     """Wraps an HF shared expert, applying the output gating it is paired with."""
 
@@ -162,7 +197,8 @@ class MoEBlockFuser:
     """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional)."""
 
     gate_name: str
-    scoring_func: str
+    scoring_func: str | None
+    """`None` if routing stays in the Transformers gate."""
     shared_name: str | None
     shared_gate_name: str | None
     router_dtype: torch.dtype | None = None
@@ -215,83 +251,142 @@ class MoEBlockFuser:
         scoring_func = "softmax" if is_softmax else "sigmoid"
         return scoring_func, _forced_dtype(logits_cone), renormalize
 
+    @staticmethod
+    def _match_experts(
+        graph: fx.Graph, experts_name: str, moe_block: nn.Module
+    ) -> tuple[fx.Node, fx.Node, fx.Node] | None:
+        """Matches `experts(x, ids, weights)` with `_, weights, ids = gate(x)`, `x`
+        being the block's input.
+
+        Returns the experts call, `x` and the gate call."""
+        experts = find_node(
+            graph, lambda n: n.op == "call_module" and n.target == experts_name
+        )
+        if experts is None:
+            return None
+        # Bind by signature so keyword calls match too. Fail closed on any other
+        # argument: the rewritten forward would drop it.
+        forward = inspect.signature(moe_block.get_submodule(experts_name).forward)
+        try:
+            bound = forward.bind(*experts.args, **experts.kwargs)
+        except TypeError:
+            return None
+        if len(bound.arguments) != 3:
+            return None
+        x, ids, weights = (_unwrap(a, moe_block) for a in bound.arguments.values())
+        if not (isinstance(x, fx.Node) and x.op == "placeholder"):
+            return None
+        if not (is_op(ids, "getitem") and is_op(weights, "getitem")):
+            return None
+        gate = ids.args[0]
+        if not (
+            _reads_input(gate, x, moe_block)
+            and ids.args == (gate, 2)
+            and weights.args == (gate, 1)
+        ):
+            return None
+        return experts, x, gate
+
     @staticmethod
     def _match_shared_experts(
-        graph: fx.Graph, experts: str
-    ) -> tuple[str | None, str | None]:
-        """Detects the shared expert and its optional gate by dataflow."""
-        experts_predicate = lambda n: n.op == "call_module" and n.target == experts
-        if (experts_node := find_node(graph, experts_predicate)) is None:
-            return None, None
-        from_experts = _reaches(experts_node, "users")
-        for add in graph.nodes:
-            if not is_op(add, "add"):
+        node: object, x: fx.Node, moe_block: nn.Module
+    ) -> tuple[str, str | None] | None:
+        """Matches `shared(x)` or `sigmoid(shared_gate(x)) * shared(x)`.
+
+        Returns the shared expert's name and its optional gate's name."""
+        if _reads_input(node, x, moe_block):
+            return node.target, None
+        if not (is_op(node, "mul") and len(node.args) == 2):
+            return None
+        for gated, shared in (node.args, node.args[::-1]):
+            gated, shared = peel(gated), peel(shared)
+            if not (is_op(gated, "sigmoid") and _reads_input(shared, x, moe_block)):
                 continue
-            operands = [a for a in add.args if isinstance(a, fx.Node)]
-            # Exactly one side is the experts' output; the other is the shared path.
-            sides = [a in from_experts for a in operands]
-            if len(operands) != 2 or sides.count(True) != 1:
-                continue
-            cone = _reaches(operands[sides.index(False)], "all_input_nodes")
-            modules = [n for n in cone if n.op == "call_module" and n.target != experts]
-            # A sigmoid wrapping one of those modules marks the shared-expert gate.
-            gate = next(
-                (
-                    src
-                    for n in cone
-                    if is_op(n, "sigmoid")
-                    and isinstance(src := peel(n.args[0]), fx.Node)
-                    and src in modules
-                ),
-                None,
-            )
-            shared = [n for n in modules if n is not gate]
-            if len(shared) != 1:
-                return None, None
-            return shared[0].target, (gate.target if gate is not None else None)
-        return None, None
+            gate = peel(gated.args[0])
+            if _reads_input(gate, x, moe_block) and _is_scalar_gate(
+                moe_block.get_submodule(gate.target)
+            ):
+                return shared.target, gate.target
+        return None
 
     @classmethod
-    def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None":
+    def match(
+        cls,
+        moe_block: nn.Module,
+        experts_name: str,
+        divides_shared_output: bool = False,
+    ) -> "MoEBlockFuser | None":
+        """Matches a block returning `experts(x, ids, weights) [* c] [+ shared(x)]`
+        with `_, weights, ids = gate(x)`, `x` being the block's input, up to
+        reshapes, casts and `nn.Identity`.
+
+        Args:
+            moe_block: The MoE block.
+            experts_name: The name of the experts in `moe_block`.
+            divides_shared_output: Whether vLLM routing would divide the shared
+                expert's output by `routed_scaling_factor` (fp16 overflow
+                protection), which Transformers decoder layers do not undo.
+
+        Returns:
+            A fuser if more than the experts can run in vLLM, else `None`.
+
+        """
         # Standard MoE block returns a single tensor.
         if _returns_tuple(type(moe_block)):
             return None
-        # Router: the child that scores + top-k selects.
-        gate_name = scoring_func = router_dtype = renormalize = None
-        for name, child in moe_block.named_children():
-            if (
-                name != experts_name
-                and (router := cls._match_router(child)) is not None
-            ):
-                gate_name = name
-                scoring_func, router_dtype, renormalize = router
-                break
-        if gate_name is None or scoring_func is None:
+        if (graph := trace(moe_block)) is None:
             return None
-        # Shared expert: a child the block adds to the experts' output.
-        shared_name = shared_gate_name = None
-        others = [
-            n
-            for n, _ in moe_block.named_children()
-            if n not in {experts_name, gate_name}
-        ]
-        if others:
-            graph = trace(moe_block)
-            if graph is None:
-                return None
-            shared_name, shared_gate_name = cls._match_shared_experts(
-                graph, experts_name
+        if (matched := cls._match_experts(graph, experts_name, moe_block)) is None:
+            return None
+        experts, x, gate = matched
+        unwrap = functools.partial(_unwrap, moe_block=moe_block)
+        gate_name = gate.target
+        router = cls._match_router(getattr(moe_block, gate_name))
+        scoring_func, router_dtype, renormalize = router or (None, None, None)
+
+        def is_routed(n: object) -> bool:
+            return n is experts or (
+                is_op(n, "mul")
+                and not isinstance(n.args[1], fx.Node)
+                and unwrap(n.args[0]) is experts
             )
-            if shared_gate_name is not None and not _is_scalar_gate(
-                getattr(moe_block, shared_gate_name)
-            ):
+
+        # Output: the experts' output, plus the shared expert's if any.
+        routed = out = unwrap(output_value(graph))
+        shared = None
+        if not is_routed(out):
+            if not (is_op(out, "add") and len(out.args) == 2):
                 return None
+            routed, shared = map(unwrap, out.args)
+            if is_routed(shared):
+                routed, shared = shared, routed
+            if not is_routed(routed):
+                return None
+        shared_name = shared_gate_name = None
+        if shared is not None:
+            if (match := cls._match_shared_experts(shared, x, moe_block)) is None:
+                return None
+            shared_name, shared_gate_name = match
         # Fail closed: `rewrite_forward` runs only the experts and the detected
         # shared expert, so any other stateful child would be dropped.
         accounted = {experts_name, gate_name, shared_name, shared_gate_name}
         for name, child in moe_block.named_children():
             if name not in accounted and next(named_state(child), None) is not None:
                 return None
+        if scoring_func is not None and shared_name and divides_shared_output:
+            logger.warning_once(
+                "%s could be fused but routing it in vLLM would apply "
+                "`routed_scaling_factor` by dividing the shared expert output, "
+                "which only fp16 overflow protection expects the decoder layer to "
+                "compensate for. Falling back to routing in Transformers; run in "
+                "bfloat16 to fuse it.",
+                type(moe_block).__name__,
+            )
+            scoring_func = router_dtype = renormalize = None
+        # Without vLLM routing, only a shared expert can move to vLLM, and only
+        # if the block does not scale the experts' output (vLLM routing does).
+        if scoring_func is None and (shared_name is None or routed is not experts):
+            return None
         return cls(
             gate_name,
             scoring_func,

vllm/model_executor/models/transformers/moe.py

@@ -127,6 +127,43 @@ direct_register_custom_op(
 )
 
 
+def _hf_topk_router_precomputed(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    moe_state: TransformersMoEState,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return `(topk_weights, topk_ids)` the HF block forward already computed:
+    weights arrive as `gating_output`, ids via `moe_state`."""
+    topk_weights = gating_output
+    topk_ids = moe_state.topk_ids
+    assert topk_ids is not None
+    # Handle all gather in expert parallel
+    if topk_ids.size(0) != hidden_states.size(0):
+        dp_metadata = get_forward_context().dp_metadata
+        assert dp_metadata is not None
+        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+        assert sizes is not None
+        is_sp = moe_state.is_sequence_parallel
+        group = get_ep_group() if is_sp else get_dp_group()
+        assert sizes[group.rank_in_group] == topk_ids.shape[0]
+        (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes)
+    return topk_weights, topk_ids
+
+
+def _hf_topk_router(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    gate: nn.Module,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Route with a Transformers gate returning `(_, topk_weights, topk_ids)`."""
+    _, topk_weights, topk_ids = gate(hidden_states)
+    return topk_weights, topk_ids
+
+
 class TransformersRoutedExperts(RoutedExperts):
     def get_expert_mapping(
         self,
@@ -310,26 +347,17 @@ class MoEMixin(MixtureOfExperts, Base):
                         is_fused_checkpoint_transposed=is_transposed,
                         routed_experts_cls=TransformersRoutedExperts,
                     )
-                    fuser = MoEBlockFuser.match(moe_block, experts_name)
                     # _maybe_apply_routed_scale_to_output edge case. Transformers
                     # decoder layers do not compensate for dividing by scaling factor.
-                    reaches_fp16_trick = (
+                    divides_shared_output = (
                         routed_scaling_factor != 1.0
                         and apply_routed_scale_to_output
                         and self.model_config.dtype == torch.float16
-                        and fuser is not None
-                        and fuser.shared_name is not None
                     )
-                    if reaches_fp16_trick:
-                        logger.warning_once(
-                            "%s could be fused but routing it in vLLM would apply "
-                            "`routed_scaling_factor` by dividing the shared expert "
-                            "output, which only fp16 overflow protection expects the "
-                            "decoder layer to compensate for. Falling back to routing "
-                            "in Transformers; run in bfloat16 to fuse it.",
-                            moe_block_cls,
-                        )
-                    if fuser is not None and not reaches_fp16_trick:
+                    fuser = MoEBlockFuser.match(
+                        moe_block, experts_name, divides_shared_output
+                    )
+                    if fuser is not None:
                         # MoE block forward is fully replaced.
                         # gate/router and shared expert (if any) runs in MoERunner.
                         shared_experts = fuser.shared_experts(moe_block, prefix)
@@ -341,39 +369,52 @@ class MoEMixin(MixtureOfExperts, Base):
                                 down_name = glu_fuser.down_name
                                 if down_name is not None:
                                     shared_down_projs.append((hf_shared, down_name))
-                        # Prefer config, otherwise read it from fuser.
-                        router_dtype = config_router_dtype or fuser.router_dtype
-                        gate = fuser.gate(moe_block, prefix, router_dtype)
+                        if scoring_func := fuser.scoring_func:
+                            # Internal routing kwargs.
+                            # Prefer config, otherwise read it from fuser.
+                            router_dtype = config_router_dtype or fuser.router_dtype
+                            gate = fuser.gate(moe_block, prefix, router_dtype)
+                            kwargs |= dict(
+                                scoring_func=scoring_func,
+                                renormalize=(
+                                    layer_renormalize
+                                    if fuser.renormalize is None
+                                    else fuser.renormalize
+                                ),
+                                gate=gate,
+                            )
+                            if router_dtype is not None:
+                                kwargs["router_logits_dtype"] = router_dtype
+                            if use_grouped_topk:
+                                kwargs |= grouped_topk_routing_kwargs
+                            if routed_scaling_factor != 1.0:
+                                kwargs |= routed_scaling_factor_kwargs
+                            bias = getattr(gate, "e_score_correction_bias", None)
+                            if bias is not None:
+                                kwargs["e_score_correction_bias"] = bias
+                        else:
+                            # External routing kwargs
+                            kwargs |= dict(
+                                num_expert_group=num_expert_group,
+                                topk_group=topk_group,
+                                custom_routing_function=partial(
+                                    _hf_topk_router,
+                                    gate=getattr(moe_block, fuser.gate_name),
+                                ),
+                            )
+
                         kwargs |= dict(
-                            scoring_func=fuser.scoring_func,
-                            renormalize=(
-                                layer_renormalize
-                                if fuser.renormalize is None
-                                else fuser.renormalize
-                            ),
                             is_sequence_parallel=(
                                 self.parallel_config.use_sequence_parallel_moe
                             ),
-                            gate=gate,
                             shared_experts=shared_experts,
                         )
-                        if router_dtype is not None:
-                            kwargs["router_logits_dtype"] = router_dtype
-                        if use_grouped_topk:
-                            kwargs |= grouped_topk_routing_kwargs
-                        if routed_scaling_factor != 1.0:
-                            kwargs |= routed_scaling_factor_kwargs
-                        bias = getattr(gate, "e_score_correction_bias", None)
-                        if bias is not None:
-                            kwargs["e_score_correction_bias"] = bias
                         fuser.rewrite_forward(moe_block)
-                        routed = "gate + experts"
-                        if fuser.shared_name:
-                            routed += " + shared experts"
                         logger.info_once(
-                            "Fused: %s (%s) -> MoERunner (internal routing)",
-                            routed,
+                            "Fused: gate + experts%s (%s) -> MoERunner (%s routing)",
+                            " + shared experts" if fuser.shared_name else "",
                             moe_block_cls,
+                            "external" if fuser.scoring_func is None else "internal",
                         )
                     else:
                         # MoE block forward is unmodified.
@@ -381,35 +422,11 @@ class MoEMixin(MixtureOfExperts, Base):
                         # We then smuggle the topk_ids in using a custom op.
                         moe_state = TransformersMoEState()
 
-                        def custom_routing_function(
-                            hidden_states: torch.Tensor,
-                            gating_output: torch.Tensor,
-                            topk: int,
-                            renormalize: bool,
-                            moe_state: TransformersMoEState,
-                        ):
-                            """Return `topk_weights` from `gating_output` and the
-                            `topk_ids` we stored in the layer earlier."""
-                            topk_weights = gating_output
-                            topk_ids = moe_state.topk_ids
-                            assert topk_ids is not None
-                            # Handle all gather in expert parallel
-                            if topk_ids.size(0) != hidden_states.size(0):
-                                dp_metadata = get_forward_context().dp_metadata
-                                assert dp_metadata is not None
-                                sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
-                                assert sizes is not None
-                                is_sp = moe_state.is_sequence_parallel
-                                group = get_ep_group() if is_sp else get_dp_group()
-                                assert sizes[group.rank_in_group] == topk_ids.shape[0]
-                                (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes)
-                            return topk_weights, topk_ids
-
                         kwargs |= dict(
                             num_expert_group=num_expert_group,
                             topk_group=topk_group,
                             custom_routing_function=partial(
-                                custom_routing_function, moe_state=moe_state
+                                _hf_topk_router_precomputed, moe_state=moe_state
                             ),
                             runner_cls=TransformersMoERunner,
                             runner_args={"moe_state": moe_state},