[TRTLLM-14955][refactor] Declare MoE backend behaviour instead of comparing classes

tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py

@@ -47,7 +47,6 @@
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
 from .communication import AllGatherReduceScatter, Communication, CommunicationFactory
-from .fused_moe_cute_dsl import CuteDslFusedMoE
 from .moe_scheduler import MoEScheduler, create_moe_scheduler
 
 # Attributes that ConfigurableMoE owns (computed in MoE.__init__ from real
@@ -410,8 +409,9 @@ def validate_config(self):
             )
 
     def _should_enable_dwdp(self) -> bool:
-        # DWDP is currently supported only for CuteDslFusedMoE with NVFP4 quantization.
-        if not isinstance(self.backend, CuteDslFusedMoE):
+        # DWDP is currently supported only by CuteDSL backends, and only with
+        # NVFP4 quantization.
+        if not self.backend.capabilities.supports_dwdp:
             return False
 
         quant_config = getattr(self.backend, "quant_config", None)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py

@@ -37,6 +37,7 @@
                       get_last_power_of_2_num_tokens_buckets,
                       last_positive_power_of_2)
 from .fused_moe_cutlass import CutlassFusedMoE
+from .impl_contract import MoERunContext, MoEStaticCapability, require_comm_plan
 from .quantization import MoEWeightLoadingMode, NVFP4CuteDslFusedMoEMethod
 from .routing import BaseMoeRoutingMethod
 
@@ -353,6 +354,14 @@ class CuteDslFusedMoE(CutlassFusedMoE):
         model_config (ModelConfig): Configuration object for the model.
     """
 
+    # ``supports_moe_lora`` is restated because CutlassFusedMoE declares True
+    # and the exact-class comparison it replaces answered False here.
+    # ``supports_dwdp`` is the capability this backend adds. CuteDslB12xFusedMoE
+    # derives from here and needs both, but must spell them out again: setting
+    # the attribute replaces the whole object rather than one field.
+    capabilities = MoEStaticCapability(supports_moe_lora=False,
+                                       supports_dwdp=True)
+
     @classmethod
     def can_implement(
         cls,
@@ -796,33 +805,28 @@ def run_moe_fp8_block_scales(
 
     def run_moe(
         self,
-        x: torch.Tensor,
-        token_selected_experts: torch.Tensor,
-        token_final_scales: Optional[torch.Tensor],
-        x_sf: Optional[torch.Tensor] = None,
-        moe_output: Optional[torch.Tensor] = None,
-        enable_alltoall: bool = False,
-        **kwargs,
+        ctx: MoERunContext,
+        *,
+        workspace: Optional[dict] = None,
     ) -> torch.Tensor:
         """
         Run MoE computation with CuteDSL backend.
 
         This method encapsulates the core MoE computation logic, handling different
         quantization schemes (fp8_block_scales and nvfp4).
 
-        Args:
-            # Standard MoE interface parameters:
-            x: Input hidden states (may be pre-quantized)
-            token_selected_experts: Expert IDs [num_tokens, top_k]. If EPLB is enabled,
-                                    this represents expert slots [num_tokens, top_k] instead.
-            token_final_scales: Final scaling factors for each token
-            x_sf: Input scale factors (optional, for certain quantization schemes)
-            moe_output: Pre-allocated MoE output buffer (optional, for NVLINK one-sided backend).
-            enable_alltoall: Whether alltoall communication is enabled.
-
         Returns:
             final_hidden_states tensor.
         """
+        del workspace  # CuteDSL kernels allocate their own intermediates.
+        plan = require_comm_plan(self, ctx)
+        x = ctx.x
+        token_selected_experts = ctx.token_selected_experts
+        token_final_scales = ctx.token_final_scales
+        x_sf = ctx.x_sf
+        moe_output = plan.moe_output
+        enable_alltoall = plan.enable_alltoall
+
         # Execute MoE computation
         if self.has_nvfp4:
             weight_view = self._build_local_weight_view()

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import replace
 from typing import Optional, Tuple, Union
 
 import torch
@@ -23,6 +24,7 @@
 from ...utils import ActivationType, Fp4QuantizedTensor
 from .fused_moe_cute_dsl import CuteDslFusedMoE
 from .fused_moe_cutlass import CutlassFusedMoE
+from .impl_contract import MoERunContext, MoEStaticCapability, require_comm_plan
 from .interface import _warn_and_return
 
 # Shared MoE output buffer pool, keyed by (max_num_tokens, hidden_size, dtype,
@@ -85,6 +87,11 @@ class on the MRO does not change which kernels execute, only where the
     ``create_moe.get_moe_cls``).
     """
 
+    # Restated rather than inherited: the LoRA gate this replaces compared the
+    # exact class and answered False here, while the DWDP gate used isinstance
+    # and answered True through CuteDslFusedMoE.
+    capabilities = MoEStaticCapability(supports_moe_lora=False, supports_dwdp=True)
+
     # SM versions on which the FlashInfer b12x NVFP4 MoE kernel is available.
     # SM120 = desktop Blackwell (RTX 5090 / GB202); SM121 = GB10 / DGX Spark.
     _SUPPORTED_SM_VERSIONS = frozenset({120, 121})
@@ -216,29 +223,21 @@ def quantize_input(
     @nvtx_range("[b12x] run_moe")
     def run_moe(
         self,
-        x: torch.Tensor,
-        token_selected_experts: torch.Tensor,
-        token_final_scales: torch.Tensor,
-        x_sf: Optional[torch.Tensor] = None,
-        is_sf_swizzled: bool = True,
-        output_dtype: Optional[torch.dtype] = None,
-        tuner_num_tokens: Optional[int] = None,
-        tuner_top_k: Optional[int] = None,
-        moe_output: Optional[torch.Tensor] = None,
-        enable_alltoall: Optional[bool] = None,
+        ctx: MoERunContext,
+        *,
+        workspace: Optional[dict] = None,
     ) -> torch.Tensor:
+        plan = require_comm_plan(self, ctx)
+        x = ctx.x
         if self._route_to_cutlass(x):
             # ``CutlassFusedMoE.run_moe`` forwards ``output_dtype`` straight
             # into the C++ ``trtllm::fused_moe`` op, which requires a concrete
             # high-precision ``ScalarType`` (uint8 / FP4-packed activations are
             # rejected at the kernel epilogue with "Invalid output type Byte").
-            # Schedulers that drive ``run_moe`` directly (the KV-cache capacity
-            # probe, for one) leave ``output_dtype`` unset, so fall back to
-            # ``x.dtype`` if it is a real compute dtype, else bf16. Mirrors the
-            # ``forward_chunk`` convention while staying safe for the FP4
-            # quant-input path (``x`` is uint8 after ``quantize_input``).
+            # ``ConfigurableMoE.forward`` always fills ``output_dtype``, so this
+            # only narrows the type for anything driving ``run_moe`` without it.
             _HIGH_PRECISION = {torch.float16, torch.bfloat16, torch.float32}
-            cutlass_output_dtype = output_dtype
+            cutlass_output_dtype = ctx.output_dtype
             if cutlass_output_dtype is None:
                 cutlass_output_dtype = (
                     x.dtype
@@ -247,17 +246,13 @@ def run_moe(
                 )
             return CutlassFusedMoE.run_moe(
                 self,
-                x,
-                token_selected_experts=token_selected_experts,
-                token_final_scales=token_final_scales,
-                x_sf=x_sf,
-                is_sf_swizzled=is_sf_swizzled,
-                output_dtype=cutlass_output_dtype,
-                tuner_num_tokens=tuner_num_tokens,
-                tuner_top_k=tuner_top_k,
-                moe_output=moe_output,
-                enable_alltoall=enable_alltoall,
+                replace(ctx, output_dtype=cutlass_output_dtype),
+                workspace=workspace,
             )
+        token_selected_experts = ctx.token_selected_experts
+        token_final_scales = ctx.token_final_scales
+        x_sf = ctx.x_sf
+        moe_output = plan.moe_output
         if self.b12x_wrapper is None or self._b12x_weights is None:
             raise RuntimeError(
                 "CuteDslB12xFusedMoE.run_moe called before process_weights_after_loading completed."

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

@@ -28,6 +28,8 @@
 from ...peft.lora.validation import has_moe_lora_targets
 from ...utils import (ActivationType, AuxStreamType, EventType,
                       Fp4QuantizedTensor)
+from .impl_contract import (MoEInputRequirement, MoERunContext,
+                            MoEStaticCapability, require_comm_plan)
 from .interface import MoE
 from .quantization import UnquantizedFusedMoEMethod
 
@@ -85,6 +87,13 @@ class CutlassFusedMoE(MoE):
             equals to: dynamic quant + routing(topK, etc.) [+ fp4_allgather] + scatter + gemm1 + swiglu + gemm2 + finalizeMoeRoute [no allreduce] + reducescatter
     """
 
+    # Routed-expert MoE LoRA is fused into this backend's op only; the
+    # subclasses below each restate ``supports_moe_lora=False``.
+    capabilities = MoEStaticCapability(supports_moe_lora=True)
+
+    # Inherited by every subclass, matching the isinstance check this replaces.
+    input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32)
+
     # Quantization algorithm support table for can_implement()
     # Format: quant_algo -> {sm_constraint, dtypes}
     # sm_constraint types:
@@ -850,42 +859,53 @@ def create_weights(self):
     def supports_moe_output_in_alltoall_workspace(self):
         return True
 
+    def _tuner_shapes(
+        self,
+        ctx: MoERunContext,
+        enable_alltoall: Optional[bool],
+    ) -> Tuple[Optional[int], Optional[int]]:
+        """Token/top-k shapes the profiling tuner should key on.
+
+        Only meaningful under alltoall: the tuner must see pre-alltoall token
+        counts so tactics cached during the no-alltoall warmup still apply at
+        runtime. Without alltoall the kernel derives both from ``x`` itself.
+        """
+        if not enable_alltoall:
+            return None, None
+        if ctx.all_rank_num_tokens is not None:
+            tuner_num_tokens = sum(ctx.all_rank_num_tokens)
+        else:
+            tuner_num_tokens = ctx.x.shape[0] * self.mapping.tp_size
+        return tuner_num_tokens, self.routing_method.top_k
+
     def run_moe(
         self,
-        x: torch.Tensor,
-        token_selected_experts: torch.Tensor,
-        token_final_scales: torch.Tensor,
-        x_sf: Optional[torch.Tensor] = None,
-        is_sf_swizzled: bool = True,
-        output_dtype: Optional[torch.dtype] = None,
-        tuner_num_tokens: Optional[int] = None,
-        tuner_top_k: Optional[int] = None,
-        moe_output: Optional[torch.Tensor] = None,
-        enable_alltoall: Optional[bool] = None,
-        lora_params: Optional[Dict] = None,
+        ctx: MoERunContext,
+        *,
+        workspace: Optional[dict] = None,
     ) -> torch.Tensor:
         """
         Run MoE computation with Cutlass backend.
 
         This method encapsulates the core MoE computation logic, handling different
         quantization schemes.
 
-        Args:
-            x: Input hidden states (may be pre-quantized)
-            token_selected_experts: Expert IDs or expert slots [num_tokens, top_k]
-                                   If EPLB is enabled, represents expert slots; otherwise expert IDs
-            token_final_scales: Final scaling factors for each token
-            x_sf: Input scale factors (optional, for certain quantization schemes)
-            is_sf_swizzled: Whether scaling factors are swizzled
-            output_dtype: Output data type (optional)
-            tuner_num_tokens: Number of tokens for profiling tuner (optional)
-            tuner_top_k: Top-k value for profiling tuner (optional)
-            moe_output: Pre-allocated output buffer (optional)
-            enable_alltoall: Whether alltoall communication is enabled (optional). If None, defaults to self.enable_alltoall.
-
         Returns:
             final_hidden_states: Output tensor from MoE computation
         """
+        del workspace  # Cutlass allocates its own intermediates.
+        plan = require_comm_plan(self, ctx)
+        x = ctx.x
+        token_selected_experts = ctx.token_selected_experts
+        token_final_scales = ctx.token_final_scales
+        x_sf = ctx.x_sf
+        output_dtype = ctx.output_dtype
+        lora_params = ctx.lora_params
+        is_sf_swizzled = plan.input_sf_swizzled
+        moe_output = plan.moe_output
+        enable_alltoall = plan.enable_alltoall
+        tuner_num_tokens, tuner_top_k = self._tuner_shapes(ctx, enable_alltoall)
+
         # W4A16 NVFP4 fallback (SM<100).
         if isinstance(self.quant_method, W4A16NVFP4CutlassFusedMoEMethod):
             return self._run_moe_w4a16_nvfp4(
@@ -904,8 +924,7 @@ def run_moe(
         if self.has_deepseek_fp8_block_scales and get_sm_version() == 120:
             from .fused_moe_triton_fp8_block_scale import \
                 run_triton_fp8_block_scale_moe
-            _use_alltoall = (enable_alltoall if enable_alltoall is not None else
-                             self.enable_alltoall)
+
             # forward_chunk sets token_final_scales=None when
             # apply_router_weight_on_input=True (weights already folded into x);
             # substitute ones so the Triton kernel's per-token scaling is a no-op.
@@ -918,7 +937,7 @@ def run_moe(
             # (0 .. expert_size_per_partition-1), so remap and zero-scale any
             # non-local token-expert pairs to suppress their contribution.
             local_n = self.expert_size_per_partition
-            if _use_alltoall:
+            if enable_alltoall:
                 # After alltoall dispatch, IDs are already local; padding = local_n
                 local_ids = token_selected_experts.clamp(0, local_n - 1)
                 is_local = token_selected_experts < local_n
@@ -959,9 +978,6 @@ def run_moe(
             elif self.has_w4a16_mxfp4:
                 weight_dtype = torch.uint8
 
-        if enable_alltoall is None:
-            enable_alltoall = self.enable_alltoall
-
         use_dynamic_fc2_scale = (self.has_nvfp4 and getattr(
             self, 'force_dynamic_quantization', False)
                                  and hasattr(self, 'fc2_weight_scale_2'))
@@ -1034,16 +1050,20 @@ def _run_moe_w4a16_nvfp4(
         tuner_num_tokens: Optional[int] = None,
         tuner_top_k: Optional[int] = None,
         moe_output: Optional[torch.Tensor] = None,
-        enable_alltoall: Optional[bool] = None,
+        *,
+        enable_alltoall: bool,
     ) -> torch.Tensor:
         """W4A16 fallback for NVFP4 MoE on SM<100. Active-mask dequant into
         a static [E_total, N, K] bf16 workspace, then bf16 fused_moe with the
         original (global) token_selected_experts. CUDA-graph capturable.
+
+        ``enable_alltoall`` has no default because it picks the expert-id remap
+        below, and either default is silently wrong for half the callers: the
+        ids are local after an alltoall dispatch and global otherwise, so a
+        wrong guess shifts every id by ``slot_start`` without failing.
         """
         assert isinstance(self.quant_method, W4A16NVFP4CutlassFusedMoEMethod)
 
-        if enable_alltoall is None:
-            enable_alltoall = self.enable_alltoall
         if output_dtype is None:
             output_dtype = x.dtype