[None][feat] MoE scheduler: optional finalize-before-combine and FP4 scale-layout contract fix

tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md

@@ -118,7 +118,8 @@ comprehensive early-validation layer for every incompatible feature.
 ```text
 [EPLB start_wait_gpu] → routing → [EPLB done_wait_gpu + update_statistic + route]
   → [comm.prepare_dispatch (NVLink2-sided)] → quantize/dispatch (adaptive order)
-  → backend.run_moe → [EPLB start_set_cpu] → comm.combine → [EPLB done_set_cpu]
+  → backend.run_moe → [EPLB start_set_cpu] → [unfinalized_combine_fn]
+  → comm.combine → [EPLB done_set_cpu]
 
 Adaptive quantize/dispatch order (gated by comm.supports_post_quant_dispatch()):
   Post-quant flow: quantize_input() → comm.dispatch()   (send quantized data)
@@ -127,6 +128,8 @@ Adaptive quantize/dispatch order (gated by comm.supports_post_quant_dispatch()):
 
 EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`.
 
+`unfinalized_combine_fn` is a constructor argument (`create_moe(..., unfinalized_combine_fn=...)`) for a model that needs a per-`(token, expert)` transform applied to each expert output *before* the top-k sum. Once registered it runs on **every** `do_finalize=False` forward, with or without a comm strategy, so the forward's return type (dense tensor vs. unfinalized triple) depends only on static config. Under a comm strategy the scheduler runs it on the dispatched rows and the combine reduces already-finalized partials — equal to finalizing after the combine, because each term depends on one row alone and non-local slots contribute exact zeros. That reordering is only sound for a combine that is an unweighted row sum over a dense per-token tensor; `DeepEPLowLatency.combine` is not one (it consumes expert-major rows and applies the routing weights itself), which is what `Communication.supports_finalize_before_combine()` declares. `ConfigurableMoE.validate_backend` rejects at construction: a backend without `supports_unfinalized_output()`, a comm without `supports_finalize_before_combine()`, and `apply_router_weight_on_input`. With no hook registered, `do_finalize=False` plus a comm strategy stays a named runtime refusal instead of a crash inside the comm layer.
+
 ### Fused-comm execution flow (MegaMoE-style)
 
 `FusedCommMoEScheduler._forward_chunk` runs per chunk:

tensorrt_llm/_torch/moe/fused_moe/communication/base.py

@@ -123,6 +123,18 @@ def supports_post_quant_dispatch(self) -> bool:
         """
         return True
 
+    def supports_finalize_before_combine(self) -> bool:
+        """Whether ``combine`` is an unweighted row sum over a dense per-token tensor.
+
+        A model-owned finalize that runs before the combine
+        (``ConfigurableMoE.unfinalized_combine_fn``) is only correct when the
+        combine itself reduces the already-finalized partials without applying
+        the routing weights again. A strategy whose combine consumes a
+        different row layout or applies the weights itself answers ``False``,
+        and the incompatibility is rejected at construction.
+        """
+        return True
+
     def uses_internal_dispatch_quantization(self) -> bool:
         """Return whether dispatch itself quantizes BF16 activations.
 

tensorrt_llm/_torch/moe/fused_moe/communication/deep_ep_low_latency.py

@@ -161,6 +161,16 @@ def supports_post_quant_dispatch(self) -> bool:
             return (self.hidden_size // 2) in self.SUPPORTED_HIDDEN_SIZES
         return False
 
+    def supports_finalize_before_combine(self) -> bool:
+        """``combine`` here is not an unweighted row sum over a per-token tensor.
+
+        It consumes expert-major rows (``[local_experts, ep_size * max_tokens,
+        hidden]``) and applies the routing weights itself, and the post-dispatch
+        ``token_final_scales`` are unit placeholders. A finalize that already
+        applied the real weights per token cannot feed it.
+        """
+        return False
+
     def supports_low_precision_combine(self) -> bool:
         """
         DeepEP Low Latency supports low-precision combine for: fp8_qdq, nvfp4, w4afp8

tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py

@@ -29,7 +29,7 @@
 
 import copy
 from contextlib import contextmanager
-from typing import Dict, List, Optional, Type, Union
+from typing import Callable, Dict, List, Optional, Type, Union
 
 import torch
 
@@ -159,6 +159,7 @@ def __init__(
         override_quant_config: Optional["QuantConfig"] = None,
         moe_cls: Optional[Type] = None,
         communication_method: Optional[str] = None,
+        unfinalized_combine_fn: Optional[Callable[..., torch.Tensor]] = None,
         **kwargs,
     ):
         super().__init__(
@@ -207,6 +208,26 @@ def __init__(
         # ========== Create Communication Strategy ==========
         self.comm = self._create_comm_strategy_auto()
 
+        # ========== Model-owned finalize for do_finalize=False ==========
+        # A model that needs a per-(token, expert) transform applied to each
+        # expert output before the top-k sum registers its finalize here, and
+        # ``ExternalCommMoEScheduler`` runs it on every ``do_finalize=False``
+        # forward -- on the dispatched rows, before any communication combine
+        # -- so the forward's return type depends only on this static config,
+        # never on which comm strategy was auto-selected. Signature:
+        #
+        #     fn(*, gemm2_output, expanded_idx_to_permuted_idx,
+        #        token_selected_slots, routing_weights, num_tokens)
+        #       -> [num_tokens, hidden_size]
+        #
+        # It must be computed from one ``(token, expert)`` row at a time, so
+        # that combining the per-rank finalized partials equals finalizing
+        # after the combine. ``validate_backend`` rejects every combination
+        # this contract cannot hold under (backend without the unfinalized
+        # triple, a combine that is not an unweighted row sum, router weights
+        # folded into the activations).
+        self.unfinalized_combine_fn = unfinalized_combine_fn
+
         # ========== Chunking Configuration ==========
         # moe_max_num_tokens is set in ModelConfig.__post_init__ if not specified
         # The default value is max_num_tokens * dp_size
@@ -690,6 +711,10 @@ def validate_backend(self, backend: MoE | MoEImplBase | None) -> None:
           1. ``backend`` is not None.
           2. If EPLB is enabled, the backend must support routing
              separation (``backend._supports_load_balancer()``).
+          3. If ``unfinalized_combine_fn`` is registered, the combination
+             must be able to honor it: the backend returns the unfinalized
+             triple, the combine (if any) is an unweighted row sum, and the
+             routing weights are not folded into the activations.
 
         Backend-specific checks are delegated to
         ``backend.validate_configurable_moe(self)``; backends with extra
@@ -714,6 +739,36 @@ def validate_backend(self, backend: MoE | MoEImplBase | None) -> None:
                 f"Either disable EPLB or use a backend that supports load balancer."
             )
 
+        if self.unfinalized_combine_fn is not None:
+            # Runs after ``self.comm`` is assigned, so the three refusals
+            # below cover the exact configuration the forwards will run with.
+            if not callable(self.unfinalized_combine_fn):
+                raise ValueError(
+                    f"unfinalized_combine_fn must be callable, got "
+                    f"{type(self.unfinalized_combine_fn).__name__}."
+                )
+            if not backend.supports_unfinalized_output():
+                raise ValueError(
+                    f"unfinalized_combine_fn is registered but backend "
+                    f"{backend.__class__.__name__} cannot return the unfinalized "
+                    "(gemm2_output, expert_weights, expanded_idx_to_permuted_idx) "
+                    "triple for do_finalize=False; there is nothing for the "
+                    "finalize to run on."
+                )
+            if self.comm is not None and not self.comm.supports_finalize_before_combine():
+                raise ValueError(
+                    f"unfinalized_combine_fn is registered but "
+                    f"{type(self.comm).__name__}.combine is not an unweighted row "
+                    "sum over a dense per-token tensor, so partials finalized "
+                    "before the combine would not reduce to the finalized result."
+                )
+            if self.apply_router_weight_on_input:
+                raise ValueError(
+                    "unfinalized_combine_fn needs the routing weights, but "
+                    "apply_router_weight_on_input folds them into the activations "
+                    "and leaves only placeholders. The two are not combinable."
+                )
+
         backend.validate_configurable_moe(self)
 
     def create_weights(self):

tensorrt_llm/_torch/moe/fused_moe/create_moe.py

@@ -1,6 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional
 
 import torch
 
@@ -281,6 +281,7 @@ def create_moe(
     layer_idx: Optional[int] = None,
     activation: MoEActivation = DEFAULT_MOE_ACTIVATION,
     communication_method: Optional[str] = None,
+    unfinalized_combine_fn: Optional[Callable[..., torch.Tensor]] = None,
     allow_backend_degradation: bool = True,
 ) -> MoE | VanillaMoE:
     """
@@ -304,6 +305,11 @@ def create_moe(
             ``SwigluActivation(clamp=7.0)`` or
             ``SiTuActivation(gate_softcap=..., linear_softcap=...)``
         communication_method: Optional ConfigurableMoE communication method
+        unfinalized_combine_fn: Optional model-owned finalize that
+            ``ConfigurableMoE`` runs on every ``do_finalize=False`` forward,
+            before any communication combine. See the signature and contract on
+            ``ConfigurableMoE.__init__``; validated against the resolved
+            backend and comm strategy in ``validate_backend``.
         allow_backend_degradation: When False, a requested backend that cannot
             serve this layer raises with the rejection trail instead of falling
             back. For callers that must know they got the backend they asked
@@ -377,12 +383,15 @@ def create_moe(
             bias=bias,
             activation=activation,
             communication_method=communication_method,
+            unfinalized_combine_fn=unfinalized_combine_fn,
         )
 
     # TritonFusedMoE and VanillaMoE are not wrapped by ConfigurableMoE
     # and own their communication and forward paths.
     if communication_method is not None:
         raise ValueError("communication_method requires ConfigurableMoE.")
+    if unfinalized_combine_fn is not None:
+        raise ValueError("unfinalized_combine_fn requires ConfigurableMoE.")
     return create_moe_backend(
         moe_cls=moe_cls,
         routing_method=routing_method,

tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py

@@ -72,6 +72,18 @@ def supports_moe_output_in_alltoall_workspace(self) -> bool:
         """
         return False
 
+    def supports_unfinalized_output(self) -> bool:
+        """Whether ``run_moe(do_finalize=False)`` returns the unfinalized triple.
+
+        ``True`` means the kernel honors ``do_finalize=False`` and hands back
+        ``(gemm2_output, expert_weights, expanded_idx_to_permuted_idx)`` for
+        the caller to finalize. The default is ``False``: most kernels either
+        assert ``do_finalize`` or finalize regardless, and a registered
+        ``ConfigurableMoE.unfinalized_combine_fn`` must be rejected at
+        construction rather than fed a dense tensor at runtime.
+        """
+        return False
+
     @property
     def uses_locality_domain(self) -> bool:
         """True when this impl runs partitioned across locality domains."""

tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py

@@ -337,6 +337,12 @@ def run_fp4_block_scale_moe(
         tune_max_num_tokens=8192,
         use_dp=False,
     ):
+        # The native runners take the block scales as one flat linear-layout
+        # tensor (``fp4BlockScaleMoe.cpp`` checks ``dim() == 1``); the
+        # framework hands them over token-major 2-D as ``quantize_input``
+        # produced them. Flattening is this provider's ABI detail.
+        if hidden_states_scale is not None:
+            hidden_states_scale = hidden_states_scale.flatten()
         hidden_size = gemm1_weights.shape[-1] * 2
         if gated_act_type == int(ActType_TrtllmGen.SiTu):
             # Fused SiTu FC1 cubins exist for two input formats: NVFP4
@@ -626,14 +632,17 @@ def fp4_quantize(
         is_sf_8x4_layout: bool = False,
         enable_pdl: Optional[bool] = None,
     ):
+        # Keyword binding: the provider inserts new options (e.g.
+        # ``is_global_scale_inversed``) ahead of ``enable_pdl`` in the
+        # positional signature, so positional calls misbind across releases.
         return self._fp4_quantize(
             input,
-            global_scale,
-            sf_vec_size,
-            sf_use_ue8m0,
-            is_sf_swizzled_layout,
-            is_sf_8x4_layout,
-            enable_pdl,
+            global_scale=global_scale,
+            sf_vec_size=sf_vec_size,
+            sf_use_ue8m0=sf_use_ue8m0,
+            is_sf_swizzled_layout=is_sf_swizzled_layout,
+            is_sf_8x4_layout=is_sf_8x4_layout,
+            enable_pdl=enable_pdl,
         )
 
     def mxfp8_quantize(
@@ -777,6 +786,9 @@ def run_fp4_block_scale_moe(
         tune_max_num_tokens=8192,
         use_dp=False,
     ):
+        # The framework hands the scales over token-major 2-D as
+        # ``quantize_input`` produced them, which is exactly what the
+        # FlashInfer API asserts (``shape[0] == num_tokens``). No reshape.
         if router_logits is not None:
             outputs = self._fused_moe.trtllm_fp4_block_scale_moe(
                 router_logits,