Unify MoE CuTe DSL dispatch to be dtype agnostic

benchmarks/bench_cute_dsl_moe_distributed.py

@@ -881,10 +881,11 @@ def _create_distributed_moe_layer(
     )
     weight_pack = MoEWeightPack()
     weight_pack.prepare_for(
-        "cute_dsl_nvfp4",
+        "cute_dsl",
         CuteDslConfig.prepare_weights(
             w13,
             w2,
+            variant=moe_config.quant.variant,
             num_local_experts=num_local_experts,
             hidden_size=CFG.hidden_size,
             intermediate_size=intermediate_size,

benchmarks/bench_moe_deepseek.py

@@ -339,7 +339,7 @@ def bench_cute_dsl(
 
     Args:
         use_wrapper: If True, use CuteDslMoEWrapper API (recommended for CUDA graph).
-                    If False, use cute_dsl_fused_moe_nvfp4 functional API.
+                    If False, use cute_dsl_fused_moe functional API.
         do_autotune: If True, run the pre-warm pass under autotune(True) so the
                     autotuner profiles all buckets and populates its cache. The
                     measurement loop runs OUTSIDE the autotune context so that
@@ -499,7 +499,7 @@ def run(x, x_sf, router_logits, routing_bias, topk_values, topk_indices):
             )
     else:
         # Use functional API
-        from flashinfer import cute_dsl_fused_moe_nvfp4
+        from flashinfer import cute_dsl_fused_moe
 
         def run(x, x_sf, router_logits, routing_bias, topk_values, topk_indices):
             x, x_sf, per_token_scale = prepare_activation(x, x_sf)
@@ -513,7 +513,7 @@ def run(x, x_sf, router_logits, routing_bias, topk_values, topk_indices):
                 topk_values=topk_values,
                 topk_indices=topk_indices,
             )
-            return cute_dsl_fused_moe_nvfp4(
+            return cute_dsl_fused_moe(
                 x=x,
                 x_sf=x_sf,
                 token_selected_experts=topk_indices,

benchmarks/routines/moe.py

@@ -249,7 +249,7 @@ def parse_moe_args(line, parser):
         default=False,
         help=(
             "Use the functional MoE API instead of the wrapper class: "
-            "cute_dsl_fused_moe_nvfp4 vs CuteDslMoEWrapper for "
+            "cute_dsl_fused_moe vs CuteDslMoEWrapper for "
             "cute_dsl_fp4_block_scale_moe, and b12x_fused_moe vs B12xMoEWrapper for "
             "b12x_fused_moe. Useful for verifying that the wrapper's workspace cache "
             "eliminates per-call allocation overhead."
@@ -1380,7 +1380,7 @@ def testCuteDslFp4BlockScaleMoe(args):
 
     This test:
     1. Creates NVFP4-quantized weights and fp4-quantized inputs for CuTe DSL kernels
-    2. Runs MoE via CuteDslMoEWrapper (or cute_dsl_fused_moe_nvfp4 when
+    2. Runs MoE via CuteDslMoEWrapper (or cute_dsl_fused_moe when
        ``--use_functional_api`` is set). SwiGLU only.
     3. Measures performance metrics (TFLOPS, TB/sec)
 
@@ -1460,17 +1460,17 @@ def testCuteDslFp4BlockScaleMoe(args):
 
     if use_functional:
         from functools import partial
-        from flashinfer import cute_dsl_fused_moe_nvfp4
+        from flashinfer import cute_dsl_fused_moe
 
         # Pre-allocate output buffer to avoid per-call allocation
         moe_output = torch.empty(
             num_tokens, hidden_size, dtype=torch.bfloat16, device=device
         )
 
         if args.verbose >= 1:
-            print("[INFO] Using CuTe DSL functional API (cute_dsl_fused_moe_nvfp4)")
+            print("[INFO] Using CuTe DSL functional API (cute_dsl_fused_moe)")
         runner = partial(
-            cute_dsl_fused_moe_nvfp4,
+            cute_dsl_fused_moe,
             num_experts=num_experts,
             top_k=top_k,
             num_local_experts=local_num_experts,
@@ -2591,7 +2591,7 @@ def testUnifiedNvfp4Moe(args):
     num_active_experts = int(local_topk_ids.unique().numel())
 
     weight_pack = MoEWeightPack()
-    weight_pack.prepare_for("cute_dsl_nvfp4", cute_dsl_view)
+    weight_pack.prepare_for("cute_dsl", cute_dsl_view)
     weight_pack.prepare_for("trtllm_fp4_routed", trtllm_view)
 
     # ---- MoELayer config --------------------------------------------------

docs/api/fused_moe.rst

@@ -108,8 +108,7 @@ The CuteDSL backends are conditionally available when the
 .. autosummary::
     :toctree: ../generated
 
-    cute_dsl_fused_moe_nvfp4
-    cute_dsl_fused_moe_mxfp8_mxfp4
+    cute_dsl_fused_moe
     b12x_fused_moe
 
 .. autoclass:: CuteDslMoEWrapper
@@ -119,12 +118,10 @@ The CuteDSL backends are conditionally available when the
 
     .. automethod:: __init__
 
-.. autoclass:: CuteDslMxfp8Mxfp4MoEWrapper
-    :members:
-    :inherited-members:
-    :show-inheritance:
-
-    .. automethod:: __init__
+The dtype-specific ``cute_dsl_fused_moe_nvfp4``,
+``cute_dsl_fused_moe_mxfp8_mxfp4``, and ``CuteDslMxfp8Mxfp4MoEWrapper``
+names are deprecated compatibility aliases. Use ``cute_dsl_fused_moe`` or
+``CuteDslMoEWrapper`` with the corresponding ``quant_mode`` instead.
 
 .. autoclass:: B12xMoEWrapper
     :members:

docs/design_docs/moe_ep_architecture.md

@@ -60,8 +60,7 @@ backend packs them); the kernel backend computes on this rank's expert shard.
 | `identity` (`IdentityConfig`) | passthrough | none | dispatch tensor unchanged | any | — |
 | `fused_moe` (`FusedMoeKernelConfig`) with `TrtllmBf16Config` | BF16 | BF16 | BF16 | SM100 family | inner `MoELayer` AutoTuner: per-runner tactic search + cross-backend winner per token bucket, up to `ExecutionConfig.tune_max_num_tokens` |
 | `fused_moe` with `TrtllmFp4Config` | NVFP4 (block-16, quantized post-dispatch in the bridge) | NVFP4 (block-16) | BF16 | SM100/SM103/SM107 | same `MoELayer` AutoTuner |
-| `fused_moe` with `CuteDslConfig` | NVFP4 (block-16) or BF16 (W4A16) | NVFP4 or 4-bit (W4A16) | BF16 | SM100/SM103 | same `MoELayer` AutoTuner |
-| `sm100_mxfp8_mxfp4_bf16_cutedsl` (`Sm100_Mxfp8_Mxfp4_Bf16_Cutedsl_SplitConfig`) | MXFP8 (block-32 UE8M0; quantized post-dispatch, or pre-dispatch packed payload with `mxfp8_dispatch=True`, nccl_ep only, bit-identical) | MXFP4 (block-32 UE8M0) | BF16 | SM100/SM103 | `tactic=` pins a kernel tactic; unpinned runs the AutoTuner default tactic (tactic sweep only under an autotune context — not tuned in the shipped benchmarks) |
+| `fused_moe` with `CuteDslConfig` | NVFP4 (W4A4), MXFP8 (W4A8), or BF16 (W4A16); W4A8 may use pre-dispatch packed payloads with `mxfp8_dispatch=True` | NVFP4/MXFP4 | BF16 | SM100/SM103 (W4A4/W4A16 also SM107) | same `MoELayer` AutoTuner |
 
 `fused_moe` accepts exactly one backend candidate per `MoEConfig` (weight
 views are prepared for the first match only). The W4A8 split kernel requires
@@ -186,7 +185,7 @@ moe_ep/
   config.py, tensors.py, weights.py, layer.py, algo_knobs.py, errors.py
   core/comm, core/kernel, core/runtime, core/validation, core/bootstrap_utils.py
   backends/split/comm/{nccl_ep,nixl_ep}
-  backends/split/kernel/{identity,fused_moe,sm100/mxfp8_mxfp4_bf16_cutedsl}
+  backends/split/kernel/{identity,fused_moe}
   backends/mega/kernel/sm100/{bf16_bf16_bf16_cutedsl,nvfp4_nvfp4_bf16_cutedsl,mxfp8_mxfp8_bf16_cutedsl,fp8_fp4_bf16_deepgemm}
   backends/mega/kernel/sm90/{fp8_fp8_bf16_pull_cutedsl,fp8_fp8_bf16_push_cuda}
   kernel_src/cutedsl_megamoe/  ← Blackwell CuTeDSL kernel src (kernel team) + FI shim
@@ -239,7 +238,7 @@ Kernels register via `@register_split_kernel` / `@register_mega_kernel` when `ba
 - LL **EXPERT_MAJOR** — `[num_local_experts, cap, hidden]` (`cap = max_tokens_per_rank * world`), each row pre-assigned to one expert; the bridge synthesizes `top_k=1` / `final_scales=1` and **combine owns the real top-k reweight**.
 - LL **RANK_MAJOR** / **HT FLAT** — `[world, max_tokens_per_rank, hidden]` carrying received `topk_idx` / `topk_weights`; the runner uses the real `top_k` with non-local picks masked to weight 0, and combine just sums across ranks.
 
-Both bf16 and NVFP4 are supported in the compute path (`MoEConfig.quant.variant`); NVFP4 activations are quantized in the bridge (linear SF layout). W4A8 (MXFP8 activations × MXFP4 weights) is a separate dedicated split kernel backend, `sm100_mxfp8_mxfp4_bf16_cutedsl`, with its own routing synthesis mirroring this bridge and an optional MXFP8-packed dispatch payload (see **Available backends**).
+BF16, W4A4, W4A8, and W4A16 are supported through the unified compute path (`MoEConfig.quant.variant`); quantized activations are prepared in the bridge, and W4A8 optionally packs its MXFP8 payload before dispatch (see **Available backends**).
 
 **Mega:** pass `MegaConfig(megakernel=...)`. Weights required as the layer's `weights` argument. Workspace allocated on first forward. Output is bf16 `[num_tokens, token_hidden_size]` where `num_tokens = MoEEpTensors.num_tokens` (may be `< max_tokens_per_rank`). `fleet_knobs` are ignored. NIXL-EP split layers require `BootstrapConfig.tcp_store` at init.
 

flashinfer/__init__.py

@@ -1,4 +1,4 @@
 """
 Copyright (c) 2023 by FlashInfer team.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -108,14 +108,26 @@
 # CuteDSL high-level APIs (conditionally if cute_dsl available)
 with contextlib.suppress(ImportError):
     from .fused_moe import (
-        cute_dsl_fused_moe_nvfp4 as cute_dsl_fused_moe_nvfp4,
+        cute_dsl_fused_moe as cute_dsl_fused_moe,
         CuteDslMoEWrapper as CuteDslMoEWrapper,
-        cute_dsl_fused_moe_mxfp8_mxfp4 as cute_dsl_fused_moe_mxfp8_mxfp4,
-        CuteDslMxfp8Mxfp4MoEWrapper as CuteDslMxfp8Mxfp4MoEWrapper,
         b12x_fused_moe as b12x_fused_moe,
         B12xMoEWrapper as B12xMoEWrapper,
     )
     from .gdn_prefill import chunk_gated_delta_rule as chunk_gated_delta_rule
+
+
+def __getattr__(name: str):
+    if name in {
+        "cute_dsl_fused_moe_nvfp4",
+        "cute_dsl_fused_moe_mxfp8_mxfp4",
+        "CuteDslMxfp8Mxfp4MoEWrapper",
+    }:
+        from . import fused_moe as _fused_moe
+
+        return getattr(_fused_moe, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 # The fused GDN decode step is surfaced here like the other GDN APIs; the
 # code lives under flashinfer/gdn_kernels/experimental/ (see its README),
 # but "experimental" describes the file location, not the import path.

flashinfer/fused_moe/__init__.py

@@ -14,6 +14,8 @@
 limitations under the License.
 """
 
+import warnings
+
 # Unified MoE API
 from .api import (  # noqa: F401
     # Typed activation values
@@ -75,7 +77,7 @@
     CutlassNvfp4Runner,
     CutlassW4A16Runner,
     CutlassW4A8Runner,
-    CuteDslNvfp4Runner,
+    CuteDslRunner,
     TrtllmBf16RoutedRunner,
     TrtllmFp4RoutedRunner,
     TrtllmFp8BlockRunner,
@@ -159,10 +161,8 @@
 # CuteDSL MoE APIs (conditionally imported if cute_dsl available)
 try:
     from .cute_dsl import (
-        cute_dsl_fused_moe_nvfp4,
+        cute_dsl_fused_moe,
         CuteDslMoEWrapper,
-        cute_dsl_fused_moe_mxfp8_mxfp4,
-        CuteDslMxfp8Mxfp4MoEWrapper,
         b12x_fused_moe,
         B12xMoEWrapper,
     )
@@ -171,6 +171,26 @@
 except ImportError:
     _cute_dsl_available = False
 
+
+def __getattr__(name: str):
+    if name == "CuteDslNvfp4Runner":
+        warnings.warn(
+            "CuteDslNvfp4Runner is deprecated; use CuteDslRunner instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return CuteDslRunner
+    if _cute_dsl_available and name in {
+        "cute_dsl_fused_moe_nvfp4",
+        "cute_dsl_fused_moe_mxfp8_mxfp4",
+        "CuteDslMxfp8Mxfp4MoEWrapper",
+    }:
+        from .cute_dsl import fused_moe as _cute_dsl_fused_moe
+
+        return getattr(_cute_dsl_fused_moe, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     # Typed activation values
     "ActivationConfig",
@@ -211,6 +231,7 @@
     "CutlassW4A8Runner",
     "ExecutionConfig",
     "ExpertConfig",
+    "CuteDslRunner",
     "CuteDslNvfp4Runner",
     "MoEActivationPack",
     "RoutingInputMode",
@@ -286,9 +307,10 @@
 # Add CuteDSL exports if available
 if _cute_dsl_available:
     __all__ += [
+        "cute_dsl_fused_moe",
         "cute_dsl_fused_moe_nvfp4",
-        "CuteDslMoEWrapper",
         "cute_dsl_fused_moe_mxfp8_mxfp4",
+        "CuteDslMoEWrapper",
         "CuteDslMxfp8Mxfp4MoEWrapper",
         "b12x_fused_moe",
         "B12xMoEWrapper",

flashinfer/fused_moe/api.py

@@ -1113,7 +1113,7 @@ def __repr__(self) -> str:
 
 @dataclass(frozen=True)
 class CuteDslConfig:
-    """CuteDSL NVFP4 backend — SM100 family only (Blackwell SM100, SM103).
+    """CuteDSL FP4 backend for W4A4, W4A8, and W4A16.
 
     The underlying CuteDSL kernel throws at launch on SM120/SM121/SM130.
     """
@@ -1128,22 +1128,23 @@ def prepare_weights(
         w1_bf16,
         w2_bf16,
         *,
+        variant: QuantVariant = QuantVariant.NVFP4,
         num_local_experts: int,
         hidden_size: int,
         intermediate_size: int,
         activation: Optional[ActivationConfig] = None,
         device=None,
     ):
-        """Build the ``cute_dsl_nvfp4`` weight view from canonical bf16 weights.
+        """Build the ``cute_dsl`` weight view from canonical BF16 weights.
 
-        Register the result with ``MoEWeightPack.prepare_for("cute_dsl_nvfp4", ...)``.
-        See :func:`flashinfer.fused_moe.prepare.prepare_cute_dsl_nvfp4_weights`.
+        Register the result with ``MoEWeightPack.prepare_for("cute_dsl", ...)``.
         """
-        from .prepare import prepare_cute_dsl_nvfp4_weights
+        from .prepare import prepare_cute_dsl_weights
 
-        return prepare_cute_dsl_nvfp4_weights(
+        return prepare_cute_dsl_weights(
             w1_bf16,
             w2_bf16,
+            variant=variant,
             num_local_experts=num_local_experts,
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,

flashinfer/fused_moe/cute_dsl/__init__.py

@@ -11,37 +11,44 @@
 # 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.
-"""
-CuteDSL-based Fused MoE Kernels for NVFP4 and MXFP8 x MXFP4.
-"""
+"""CuteDSL-based Fused MoE Kernels for block-scaled FP4 compute."""
 
 from ...cute_dsl.utils import is_cute_dsl_available
 
 # Conditionally import CuTe-DSL kernels
 if is_cute_dsl_available():
+    from . import fused_moe as _fused_moe
     from .fused_moe import (
-        cute_dsl_fused_moe_nvfp4,
+        cute_dsl_fused_moe,
         CuteDslMoEWrapper,
     )
-    from .fused_moe_mxfp8_mxfp4 import (
-        cute_dsl_fused_moe_mxfp8_mxfp4,
-        CuteDslMxfp8Mxfp4MoEWrapper,
-    )
     from .b12x_moe import (
         b12x_fused_moe,
         B12xMoEWrapper,
     )
 
+_DEPRECATED_APIS = (
+    "cute_dsl_fused_moe_nvfp4",
+    "cute_dsl_fused_moe_mxfp8_mxfp4",
+    "CuteDslMxfp8Mxfp4MoEWrapper",
+)
+
+
+def __getattr__(name: str):
+    if is_cute_dsl_available() and name in _DEPRECATED_APIS:
+        return getattr(_fused_moe, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "is_cute_dsl_available",
 ]
 
 if is_cute_dsl_available():
     __all__ += [
-        "cute_dsl_fused_moe_nvfp4",
+        "cute_dsl_fused_moe",
         "CuteDslMoEWrapper",
-        "cute_dsl_fused_moe_mxfp8_mxfp4",
-        "CuteDslMxfp8Mxfp4MoEWrapper",
+        *_DEPRECATED_APIS,
         "b12x_fused_moe",
         "B12xMoEWrapper",
     ]

flashinfer/fused_moe/cute_dsl/_inputs_helper.py

@@ -35,7 +35,7 @@
 Reference: ports trt-llm's GroupedGemmInputsHelper.generate_num_tokens_per_expert
 and generate_token_selected_experts from
 tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (lines 31-167) and the
-``inputs_pre_hook`` mechanism from CuteDslFusedMoENvfp4InputsHelper in
+``inputs_pre_hook`` mechanism from the CuteDSL fused-MoE runner in
 tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py.
 """