[None][revert] - Revert "[TRTLLM-9108][feat] refactor MoE unit tests: add unified ConfigurableMoE test framework"

tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py

@@ -154,7 +154,7 @@ def create_strategy(
             logger.debug(f"NVLinkTwoSided not available: {e}")
 
         # Try DeepEP (if enabled and weight dtype is bfloat16)
-        if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "1") == "1" and act_dtype == torch.bfloat16:
+        if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "0") == "1" and act_dtype == torch.bfloat16:
             try:
                 strategy = DeepEP(
                     mapping,

tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py

@@ -81,6 +81,8 @@ def is_platform_supported() -> bool:
         """
         Check if DeepEP is supported on the current platform
         """
+        if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "0") != "1":
+            return False
         return deep_ep_installed
 
     @staticmethod
@@ -143,13 +145,6 @@ def dispatch(
         """
         all_rank_max_num_tokens = max(all_rank_num_tokens)
 
-        # DeepEP C++ kernel requires topk_weights (token_final_scales) to be float32,
-        # but downstream backends (e.g. TRTLLM) may require the original dtype.
-        # Convert to float32 for dispatch, then restore afterward.
-        original_scales_dtype = token_final_scales.dtype if token_final_scales is not None else None
-        if token_final_scales is not None and token_final_scales.dtype != torch.float32:
-            token_final_scales = token_final_scales.to(torch.float32)
-
         if not self.supports_post_quant_dispatch():
             # Pre-quant dispatch (unquantized data)
             (
@@ -220,14 +215,6 @@ def dispatch(
                 "padded": padded,
             }
 
-        # Restore token_final_scales to original dtype for downstream consumers
-        if (
-            token_final_scales is not None
-            and original_scales_dtype is not None
-            and token_final_scales.dtype != original_scales_dtype
-        ):
-            token_final_scales = token_final_scales.to(original_scales_dtype)
-
         return hidden_states, hidden_states_sf, token_selected_slots, token_final_scales
 
     def combine(

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

@@ -37,15 +37,6 @@ class DeepEPLowLatency(Communication):
     DeepEP Low Latency strategy supporting both pre-quant and post-quant
     """
 
-    SUPPORTED_HIDDEN_SIZES = {2048, 2560, 3584, 4096, 5120, 6144, 7168}
-    """set[int]: Hidden sizes supported by the low-latency DeepEP kernel (SWITCH_HIDDEN in launch.cuh)."""
-
-    SUPPORTED_HIDDEN_SIZES_EXTENSION = {4096, 6144, 7168}
-    """set[int]: Hidden sizes supported by extension kernels (nvfp4 post-quant/low-precision combine).
-
-    Sourced from SWITCH_HIDDEN_FOR_EXTENSION_KERNELS in extension_kernels.cu.
-    """
-
     def __init__(
         self,
         mapping: Mapping,
@@ -60,13 +51,6 @@ def __init__(
     ):
         super().__init__(mapping)
 
-        # Validate hidden_size against kernel constraints
-        if hidden_size not in self.SUPPORTED_HIDDEN_SIZES:
-            raise RuntimeError(
-                f"DeepEPLowLatency does not support hidden_size={hidden_size}. "
-                f"Supported hidden sizes: {sorted(self.SUPPORTED_HIDDEN_SIZES)}"
-            )
-
         # Store needed parameters
         self.num_slots = num_slots
         self.hidden_size = hidden_size
@@ -102,42 +86,24 @@ def is_platform_supported() -> bool:
         """
         Check if DeepEP Low Latency is supported on the current platform
         """
+        if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "0") != "1":
+            return False
         if not deep_ep_installed:
             return False
         return True
 
     def supports_post_quant_dispatch(self) -> bool:
         """
         DeepEP Low Latency supports post-quant for: fp8_qdq, nvfp4, w4afp8
-
-        Note: nvfp4 post-quant dispatch uses extension kernels which require
-        hidden_size in SUPPORTED_HIDDEN_SIZES_EXTENSION.
-
-        Note: fp8_qdq and w4afp8 post-quant dispatch views fp8 (1 byte) as
-        bf16 (2 bytes) via .view(torch.bfloat16), halving the hidden dimension.
-        The halved dimension must be in SUPPORTED_HIDDEN_SIZES for the dispatch
-        kernel (SWITCH_HIDDEN in internode_ll.cu) to work.
         """
         if not self.enable_postquant_alltoall:
             return False
-        if self._has_nvfp4():
-            # nvfp4 dispatch uses extension kernels with stricter hidden_size requirement
-            return self.hidden_size in self.SUPPORTED_HIDDEN_SIZES_EXTENSION
-        if self._has_fp8_qdq() or self._has_w4afp8():
-            # fp8/w4afp8 post-quant dispatch views fp8 (1 byte) as bf16 (2 bytes),
-            # halving the hidden dimension. The kernel must support the halved size.
-            return (self.hidden_size // 2) in self.SUPPORTED_HIDDEN_SIZES
-        return False
+        return self._has_nvfp4() or self._has_fp8_qdq() or self._has_w4afp8()
 
     def supports_low_precision_combine(self) -> bool:
         """
         DeepEP Low Latency supports low-precision combine for: fp8_qdq, nvfp4, w4afp8
-
-        Note: low-precision combine uses extension kernels which require
-        hidden_size in SUPPORTED_HIDDEN_SIZES_EXTENSION.
         """
-        if self.hidden_size not in self.SUPPORTED_HIDDEN_SIZES_EXTENSION:
-            return False
         return self._has_nvfp4() or self._has_fp8_qdq() or self._has_w4afp8()
 
     def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool:

tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py

@@ -100,25 +100,25 @@ def can_implement(
         cls,
         quant_algo,
         dtype_activation: torch.dtype = torch.bfloat16,
-        swiglu_gptoss_style: bool = False,
+        gptoss_style: bool = False,
     ):
         """
         ConfigurableMoE is a wrapper class that delegates to specific backends.
 
         To check capability, query the specific backend class directly:
-        - CutlassFusedMoE.can_implement(quant_algo, dtype_activation, swiglu_gptoss_style)
-        - TRTLLMGenFusedMoE.can_implement(quant_algo, dtype_activation, swiglu_gptoss_style)
+        - CutlassFusedMoE.can_implement(quant_algo, dtype_activation, gptoss_style)
+        - TRTLLMGenFusedMoE.can_implement(quant_algo, dtype_activation, gptoss_style)
         - etc.
 
         Args:
             quant_algo: The quantization algorithm to check (None for unquantized)
             dtype_activation: The activation data type
-            swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled
+            gptoss_style: Whether gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled
 
         Returns:
             Tuple[bool, Optional[str]]: Always returns (False, reason)
         """
-        del quant_algo, dtype_activation, swiglu_gptoss_style  # Unused - wrapper class
+        del quant_algo, dtype_activation, gptoss_style  # Unused - wrapper class
         return False, (
             "ConfigurableMoE is a wrapper class. "
             "Query the specific backend (CutlassFusedMoE, TRTLLMGenFusedMoE, etc.) directly."

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py

@@ -318,7 +318,7 @@ def can_implement(
         cls,
         quant_algo: Optional[QuantAlgo],
         dtype_activation: torch.dtype = torch.bfloat16,
-        swiglu_gptoss_style: bool = False,
+        gptoss_style: bool = False,
     ) -> Tuple[bool, Optional[str]]:
         """
         Check if CuteDslFusedMoE can implement the given quantization algorithm.
@@ -327,14 +327,14 @@ def can_implement(
         - NVFP4: SM in {100, 103}
 
         Does NOT support unquantized mode. Output dtype is hardcoded to bfloat16.
-        Does NOT support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit).
+        Does NOT support gptoss_style (bias/swiglu with custom alpha/beta/limit).
 
         Args:
             quant_algo: The quantization algorithm to check (None for unquantized)
             dtype_activation: The activation input data type. Only bfloat16 is supported
                 because output dtype is hardcoded to bfloat16 (input/output dtype must match).
-            swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
-                CuteDslFusedMoE does NOT support swiglu_gptoss_style.
+            gptoss_style: Whether gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
+                CuteDslFusedMoE does NOT support gptoss_style.
 
         Returns:
             Tuple[bool, Optional[str]]: (can_implement, skip_reason)
@@ -360,10 +360,10 @@ def can_implement(
             return _warn_and_return(
                 "CuteDslFusedMoE does not support unquantized mode")
 
-        # CuteDslFusedMoE does NOT support swiglu_gptoss_style
-        if swiglu_gptoss_style:
+        # CuteDslFusedMoE does NOT support gptoss_style
+        if gptoss_style:
             return _warn_and_return(
-                "CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)"
+                "CuteDslFusedMoE does not support gptoss_style (bias/swiglu with custom alpha/beta/limit)"
             )
 
         # NVFP4 - SM in {100, 103}

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

@@ -113,15 +113,15 @@ class CutlassFusedMoE(MoE):
         },
     }
 
+    # Quantization algorithms that support gptoss_style
     _GPTOSS_SUPPORTED_ALGOS = {QuantAlgo.W4A8_MXFP4_MXFP8}
-    """set[QuantAlgo]: Quantization algorithms that support swiglu_gptoss_style."""
 
     @classmethod
     def can_implement(
         cls,
         quant_algo: Optional[QuantAlgo],
         dtype_activation: torch.dtype = torch.bfloat16,
-        swiglu_gptoss_style: bool = False,
+        gptoss_style: bool = False,
     ) -> Tuple[bool, Optional[str]]:
         """
         Check if CutlassFusedMoE can implement the given quantization algorithm.
@@ -145,8 +145,8 @@ def can_implement(
                 - FP8/FP8_BLOCK_SCALES/W4A8_MXFP4_FP8: float16, bfloat16, float32
                 - NVFP4: float16, bfloat16, float8_e4m3fn
                 - W4A16_MXFP4/W4A8_AWQ/W8A16/W4A8_MXFP4_MXFP8: float16, bfloat16
-            swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
-                CutlassFusedMoE only supports swiglu_gptoss_style for W4A8_MXFP4_MXFP8 quantization.
+            gptoss_style: Whether gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
+                CutlassFusedMoE only supports gptoss_style for W4A8_MXFP4_MXFP8 quantization.
 
         Returns:
             Tuple[bool, Optional[str]]: (can_implement, skip_reason)
@@ -160,10 +160,10 @@ def can_implement(
             return _warn_and_return(
                 f"CutlassFusedMoE requires SM >= 80, got SM{sm_version}")
 
-        # Check swiglu_gptoss_style support
-        if swiglu_gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS:
+        # Check gptoss_style support
+        if gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS:
             return _warn_and_return(
-                f"CutlassFusedMoE swiglu_gptoss_style only supports W4A8_MXFP4_MXFP8 "
+                f"CutlassFusedMoE gptoss_style only supports W4A8_MXFP4_MXFP8 "
                 f"(got quant_algo={quant_algo})")
 
         # Check if quant_algo is supported

tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py

@@ -382,7 +382,7 @@ def can_implement(
         cls,
         quant_algo: Optional[QuantAlgo],
         dtype_activation: torch.dtype = torch.bfloat16,
-        swiglu_gptoss_style: bool = False,
+        gptoss_style: bool = False,
     ) -> Tuple[bool, Optional[str]]:
         """
         Check if DeepGemmFusedMoE can implement the given quantization algorithm.
@@ -391,15 +391,15 @@ def can_implement(
         - FP8_BLOCK_SCALES: SM in {100, 103}
 
         Does NOT support unquantized mode. Output dtype is hardcoded to bfloat16.
-        Does NOT support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit).
+        Does NOT support gptoss_style (bias/swiglu with custom alpha/beta/limit).
 
         Args:
             quant_algo: The quantization algorithm to check (None for unquantized)
             dtype_activation: The activation input data type. Supported types are
                 float32, bfloat16, and float16 (required by moe_permute_op kernel).
                 Note: Output dtype is always bfloat16 regardless of input dtype.
-            swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
-                DeepGemmFusedMoE does NOT support swiglu_gptoss_style.
+            gptoss_style: Whether gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
+                DeepGemmFusedMoE does NOT support gptoss_style.
 
         Returns:
             Tuple[bool, Optional[str]]: (can_implement, skip_reason)
@@ -425,10 +425,10 @@ def can_implement(
             return _warn_and_return(
                 "DeepGemmFusedMoE does not support unquantized mode")
 
-        # DeepGemmFusedMoE does NOT support swiglu_gptoss_style
-        if swiglu_gptoss_style:
+        # DeepGemmFusedMoE does NOT support gptoss_style
+        if gptoss_style:
             return _warn_and_return(
-                "DeepGemmFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)"
+                "DeepGemmFusedMoE does not support gptoss_style (bias/swiglu with custom alpha/beta/limit)"
             )
 
         # Only FP8_BLOCK_SCALES is supported

tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py

@@ -1283,12 +1283,12 @@ def can_implement(
         cls,
         quant_algo: Optional["QuantAlgo"],
         dtype_activation: torch.dtype = torch.bfloat16,
-        swiglu_gptoss_style: bool = False,
+        gptoss_style: bool = False,
     ) -> Tuple[bool, Optional[str]]:
         """
         Check if TritonFusedMoE can implement the given quantization algorithm.
 
-        TritonFusedMoE supports (SM90 only, swiglu_gptoss_style=True only):
+        TritonFusedMoE supports (SM90 only, gptoss_style=True only):
         - Unquantized (BF16 only)
         - FP8 per-tensor (QDQ)
         - W4A8_MXFP4_FP8
@@ -1298,8 +1298,8 @@ def can_implement(
             quant_algo: The quantization algorithm to check (None for unquantized)
             dtype_activation: The activation data type. In unquantized mode, activation,
                 weight, and output dtypes must all match (only bfloat16 supported).
-            swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
-                TritonFusedMoE ONLY supports swiglu_gptoss_style=True.
+            gptoss_style: Whether gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
+                TritonFusedMoE ONLY supports gptoss_style=True.
 
         Returns:
             Tuple[bool, Optional[str]]: (can_implement, skip_reason)
@@ -1316,10 +1316,10 @@ def can_implement(
             return _warn_and_return(
                 f"TritonFusedMoE only supports SM90, got SM{sm_version}")
 
-        # TritonFusedMoE ONLY supports swiglu_gptoss_style=True
-        if not swiglu_gptoss_style:
+        # TritonFusedMoE ONLY supports gptoss_style=True
+        if not gptoss_style:
             return _warn_and_return(
-                "TritonFusedMoE only supports swiglu_gptoss_style=True")
+                "TritonFusedMoE only supports gptoss_style=True")
 
         # Unquantized mode - only bfloat16 is supported
         if quant_algo is None:

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

@@ -87,7 +87,7 @@ class TRTLLMGenFusedMoE(MoE):
         QuantAlgo.W4A8_MXFP4_MXFP8,
     }
 
-    # Quantization algorithms that support swiglu_gptoss_style
+    # Quantization algorithms that support gptoss_style
     _GPTOSS_SUPPORTED_ALGOS = {
         QuantAlgo.NVFP4,
         QuantAlgo.W4A16_MXFP4,
@@ -100,7 +100,7 @@ def can_implement(
         cls,
         quant_algo: Optional[QuantAlgo],
         dtype_activation: torch.dtype = torch.bfloat16,
-        swiglu_gptoss_style: bool = False,
+        gptoss_style: bool = False,
     ) -> Tuple[bool, Optional[str]]:
         """
         Check if TRTLLMGenFusedMoE can implement the given quantization algorithm.
@@ -119,7 +119,7 @@ def can_implement(
             quant_algo: The quantization algorithm to check (None for unquantized)
             dtype_activation: The activation input data type. Only bfloat16 is supported.
                 See: forward_impl() assert x.dtype == torch.bfloat16 (line 722).
-            swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
+            gptoss_style: Whether gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
                 Only supported for nvfp4 and mxfp4 variants.
 
         Returns:
@@ -151,10 +151,10 @@ def can_implement(
             return _warn_and_return(
                 f"TRTLLMGenFusedMoE does not support quant_algo={quant_algo}")
 
-        # Check swiglu_gptoss_style support: only supported for nvfp4 and mxfp4 variants
-        if swiglu_gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS:
+        # Check gptoss_style support: only supported for nvfp4 and mxfp4 variants
+        if gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS:
             return _warn_and_return(
-                f"TRTLLMGenFusedMoE supports swiglu_gptoss_style (bias/swiglu) only for nvfp4 and mxfp4 variants, "
+                f"TRTLLMGenFusedMoE supports gptoss_style (bias/swiglu) only for nvfp4 and mxfp4 variants, "
                 f"got quant_algo={quant_algo}")
 
         return True, None

tensorrt_llm/_torch/modules/fused_moe/interface.py

@@ -1,18 +1,3 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import os
 import weakref
 from abc import abstractmethod
@@ -173,7 +158,7 @@ def can_implement(
         cls,
         quant_algo: Optional[QuantAlgo],
         dtype_activation: torch.dtype = torch.bfloat16,
-        swiglu_gptoss_style: bool = False,
+        gptoss_style: bool = False,
     ) -> Tuple[bool, Optional[str]]:
         """
         Check if this MoE backend can implement the given quantization algorithm.
@@ -191,7 +176,7 @@ def can_implement(
         Args:
             quant_algo: The quantization algorithm to check (None for unquantized)
             dtype_activation: The activation data type.
-            swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
+            gptoss_style: Whether gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled.
 
         Returns:
             Tuple[bool, Optional[str]]: (can_implement, skip_reason)