feat(moe): add unified CUTLASS BF16 and W4A16 runners

flashinfer/fused_moe/__init__.py

@@ -22,6 +22,8 @@
     BackendOptions,
     CuteDslConfig,
     CutlassConfig,
+    CutlassBf16Config,
+    CutlassW4A16Config,
     ExecutionConfig,
     ExpertConfig,
     MoEActivationPack,
@@ -40,6 +42,8 @@
 from .runners import (  # noqa: F401
     B12xNvfp4Runner,
     B12xW4A16Runner,
+    CutlassBf16Runner,
+    CutlassW4A16Runner,
     CuteDslNvfp4Runner,
     TrtllmFp4RoutedRunner,
     TrtllmFp8BlockRunner,
@@ -129,11 +133,15 @@
     "ActivationConfig",
     "B12xNvfp4Config",
     "B12xNvfp4Runner",
+    "CutlassBf16Runner",
+    "CutlassW4A16Runner",
     "B12xW4A16Config",
     "B12xW4A16Runner",
     "BackendOptions",
     "CuteDslConfig",
     "CutlassConfig",
+    "CutlassBf16Config",
+    "CutlassW4A16Config",
     "ExecutionConfig",
     "ExpertConfig",
     "CuteDslNvfp4Runner",

flashinfer/fused_moe/api.py

@@ -32,6 +32,7 @@
 
 import torch
 from torch import Tensor
+from typing_extensions import deprecated
 
 from ..tllm_enums import ActivationType, RoutingInputMode, RoutingMethodType
 
@@ -234,6 +235,13 @@ def __repr__(self) -> str:
 # compiles for major 12 as well, but those cubins fail at runtime on SM120/121.
 _TRTLLM_ROUTED_FP8_ARCHS = (100, 103)
 
+# Dense BF16 follows the architecture dispatch already exposed by the flat
+# CUTLASS API.
+_CUTLASS_BF16_ARCHS = (89, 90, 100, 103, 107, 110, 120, 121)
+
+# W4A16 uses Hopper-specific mixed-input weight and scale layouts.
+_CUTLASS_W4A16_ARCHS = (90,)
+
 
 @dataclass(frozen=True)
 class TrtllmFp4Config:
@@ -478,18 +486,120 @@ def __repr__(self) -> str:
         return "TrtllmMxInt4Config()"
 
 
+@deprecated(
+    "CutlassConfig is deprecated and non-runnable; use CutlassBf16Config or "
+    "CutlassW4A16Config instead."
+)
 @dataclass(frozen=True)
 class CutlassConfig:
-    """CUTLASS backend — broadest architecture support."""
+    """Legacy quantization-neutral CUTLASS configuration placeholder.
+
+    .. deprecated::
+        Use :class:`CutlassBf16Config` or :class:`CutlassW4A16Config` instead.
+
+    This type is preserved for source compatibility, but it is intentionally
+    not registered with :class:`MoELayer` and therefore is not runnable. Select
+    a concrete tensor contract such as :class:`CutlassBf16Config` or
+    :class:`CutlassW4A16Config` instead.
+    """
 
     @classmethod
     def supported(cls, arch: int) -> bool:
-        return True  # universal fallback
+        # Compatibility-only placeholder: it has no registered runner and must
+        # never be surfaced as a dispatch candidate by BackendOptions.valid_for().
+        return False
 
     def __repr__(self) -> str:
         return "CutlassConfig()"
 
 
+@dataclass(frozen=True)
+class CutlassBf16Config:
+    """CUTLASS BF16 backend for the unified MoE API.
+
+    Architecture coverage follows the dense-BF16 legacy flat API. The unified
+    GPU tests currently exercise SM90.
+
+    This backend supports packed precomputed routing with SwiGLU and requires
+    ``do_finalize=True``. Expert parallelism and shared experts are not
+    supported.
+    """
+
+    @classmethod
+    def supported(cls, arch: int) -> bool:
+        return arch in _CUTLASS_BF16_ARCHS
+
+    @staticmethod
+    def prepare_weights(
+        w1_bf16,
+        w2_bf16,
+        *,
+        num_local_experts: int,
+        hidden_size: int,
+        intermediate_size: int,
+        device=None,
+    ):
+        """Build the ``cutlass_bf16`` canonical BF16 weight view.
+
+        GEMM1 uses the public ``[up, gate]`` row convention.  Unlike TRTLLM's
+        BlockMajorK path, CUTLASS BF16 kernels consume these weights directly
+        and need no physical reordering.
+        """
+        from .prepare import prepare_cutlass_bf16_weights
+
+        return prepare_cutlass_bf16_weights(
+            w1_bf16,
+            w2_bf16,
+            num_local_experts=num_local_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            device=device,
+        )
+
+    def __repr__(self) -> str:
+        return "CutlassBf16Config()"
+
+
+@dataclass(frozen=True)
+class CutlassW4A16Config:
+    """CUTLASS MXFP4-weight x BF16-activation backend for SM90.
+
+    This backend supports packed precomputed routing with SwiGLU and requires
+    ``do_finalize=True``. Expert parallelism and shared experts are not
+    supported. Both ``hidden_size`` and ``intermediate_size`` must be divisible
+    by 128.
+    """
+
+    @classmethod
+    def supported(cls, arch: int) -> bool:
+        return arch in _CUTLASS_W4A16_ARCHS
+
+    @staticmethod
+    def prepare_weights(
+        w1_bf16,
+        w2_bf16,
+        *,
+        num_local_experts: int,
+        hidden_size: int,
+        intermediate_size: int,
+        device=None,
+    ):
+        """Quantize and interleave canonical BF16 weights for SM90 W4A16."""
+        from .prepare import prepare_cutlass_w4a16_weights
+
+        return prepare_cutlass_w4a16_weights(
+            w1_bf16,
+            w2_bf16,
+            num_local_experts=num_local_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            device=device,
+        )
+
+    def __repr__(self) -> str:
+        return "CutlassW4A16Config()"
+
+
 @dataclass(frozen=True)
 class CuteDslConfig:
     """CuteDSL NVFP4 backend — SM100 family only (Blackwell SM100, SM103).
@@ -624,6 +734,8 @@ def __repr__(self) -> str:
     TrtllmBf16Config,
     TrtllmMxInt4Config,
     CutlassConfig,
+    CutlassBf16Config,
+    CutlassW4A16Config,
     CuteDslConfig,
     B12xNvfp4Config,
     B12xW4A16Config,
@@ -636,6 +748,8 @@ def __repr__(self) -> str:
     TrtllmBf16Config,
     TrtllmMxInt4Config,
     CutlassConfig,
+    CutlassBf16Config,
+    CutlassW4A16Config,
     CuteDslConfig,
     B12xNvfp4Config,
     B12xW4A16Config,
@@ -690,7 +804,8 @@ def __iter__(self):
         TrtllmFp8PerTensorConfig(),
         TrtllmBf16Config(),
         TrtllmMxInt4Config(),
-        CutlassConfig(),
+        CutlassBf16Config(),
+        CutlassW4A16Config(),
         CuteDslConfig(),
     )
 )

flashinfer/fused_moe/core.py

@@ -434,6 +434,37 @@ def __init__(
 
             self.fused_moe_runner = MoERunner.runner_dict[instance_key]
 
+        def get_cache_key_extras(self, _inputs: List[torch.Tensor]) -> tuple:
+            # Stage profiling passes only activation and weight tensors, so the
+            # profile key captures their shapes but not constructor-fixed options
+            # such as top-k, parallel ranks, quantization mode, or activation.
+            # The in-memory runner hash distinguishes instances, but it is
+            # intentionally excluded from persisted file keys. Include those
+            # options here to prevent runners with identical tensor profiles from
+            # reusing incompatible saved tactics.
+            return (
+                self.x_dtype,
+                self.weight_dtype,
+                self.output_dtype,
+                self.top_k,
+                self.tp_size,
+                self.tp_rank,
+                self.ep_size,
+                self.ep_rank,
+                self.cluster_size,
+                self.cluster_rank,
+                self.enable_alltoall,
+                self.use_deepseek_fp8_block_scale,
+                self.use_w4_group_scaling,
+                self.use_mxfp8_act_scaling,
+                self.use_wfp4afp8_humming,
+                self.min_latency_mode,
+                self.enable_pdl,
+                int(self.activation_type),
+                self.use_packed_weights,
+                self.use_fused_finalize,
+            )
+
         def get_valid_tactics(
             self,
             inputs: List[torch.Tensor],
@@ -878,6 +909,7 @@ def _cutlass_fused_moe_workspace_size(
 
     # Register the module
     return SimpleNamespace(
+        MoERunner=MoERunner,
         cutlass_fused_moe=cutlass_fused_moe,
         cutlass_fused_moe_workspace_size=_cutlass_fused_moe_workspace_size,
         interleave_moe_weights_for_sm90_mixed_gemm=(

flashinfer/fused_moe/layer.py

@@ -30,6 +30,8 @@
 from .api import (
     B12xNvfp4Config,
     B12xW4A16Config,
+    CutlassBf16Config,
+    CutlassW4A16Config,
     CuteDslConfig,
     MoEActivationPack,
     MoEConfig,
@@ -43,6 +45,8 @@
 from .runners import (
     B12xNvfp4Runner,
     B12xW4A16Runner,
+    CutlassBf16Runner,
+    CutlassW4A16Runner,
     CuteDslNvfp4Runner,
     TrtllmBf16RoutedRunner,
     TrtllmFp4RoutedRunner,
@@ -57,6 +61,8 @@
 # backend_key / tuning_config / pack_inputs as attributes or class members;
 # typing the list with this Union gives mypy the visibility it needs.
 _RunnerT = Union[
+    CutlassBf16Runner,
+    CutlassW4A16Runner,
     CuteDslNvfp4Runner,
     TrtllmFp4RoutedRunner,
     TrtllmBf16RoutedRunner,
@@ -69,6 +75,8 @@
 
 # Map backend-config class -> runner class
 _BACKEND_RUNNERS: Dict[type, Type[_RunnerT]] = {
+    CutlassBf16Config: CutlassBf16Runner,
+    CutlassW4A16Config: CutlassW4A16Runner,
     CuteDslConfig: CuteDslNvfp4Runner,
     TrtllmFp4Config: TrtllmFp4RoutedRunner,
     TrtllmBf16Config: TrtllmBf16RoutedRunner,
@@ -117,6 +125,7 @@ def __init__(self, config: MoEConfig, device: Optional[torch.device] = None):
                 runner.check_support()
             except (NotImplementedError, ValueError, RuntimeError):
                 continue
+            runner.build()
             self.runners.append(runner)
 
         if not self.runners: