feat(moe): add CuTe DSL NVFP4 W4A16 MegaMoE

docs/design_docs/moe_ep_architecture.md

@@ -29,12 +29,16 @@ owns dispatch, expert compute, and combine; output is always BF16
 
 | Backend (alias) | Activation | Weight | Output | Arch | Tuning |
 |---|---|---|---|---|---|
+| `sm100_bf16_nvfp4_bf16_cutedsl` | BF16 | NVFP4 (block-16 E4M3, optional per-expert FP32 global scales) | BF16 | SM100/SM103 | same `knobs` surface as the NVFP4 backend |
 | `sm100_nvfp4_nvfp4_bf16_cutedsl` (`nvfp4_cutedsl`) | NVFP4 (block-16) | NVFP4 (block-16) | BF16 | SM100 family | `knobs=None` → token-count heuristic; `knobs=dict` → pinned; `knobs="auto"` → collective compile+time sweep at first forward (never in serving); winners cacheable via `FLASHINFER_MOE_EP_KNOB_CACHE` |
 | `sm100_mxfp8_mxfp8_bf16_cutedsl` (`mxfp8_cutedsl`) | MXFP8 (block-32 UE8M0) | MXFP8 (block-32 UE8M0) | BF16 | SM100 family | same `knobs` surface as the NVFP4 backend |
 | `sm100_fp8_fp4_bf16_deepgemm` (`deep_gemm_mega`) | FP8 (E4M3, block-32 UE8M0) | FP4 (int8-packed, block-32) | BF16 | SM100 family | — (DeepGEMM selects its own JIT configs internally) |
 | `sm90_fp8_fp8_bf16_pull_cutedsl` (`sm90_pull_fp8`) | FP8 (E4M3/E5M2; per-tensor or DeepGEMM-style blockwise scales) | FP8 (same `fp8_scale_mode`) | BF16 | SM90 exactly | explicit geometry knobs on the config (`swap_ab`, `mma_tiler_mnk`); no tuner/knob-cache yet |
 | `sm90_fp8_fp8_bf16_push_cuda` (`sm90_push_fp8`) | FP8 (E4M3) | FP8 (E4M3) | BF16 | SM90 | — (static dimensions/protocol choices only) |
 
+W4A16 returns BF16 expert terms and combines them in a separate kernel,
+applying routing scores and accumulating in FP32 in fixed top-k order.
+
 The SM90 pull-style CuTeDSL tree is process-exclusive with the SM100 CuTeDSL
 tree (module names collide). Weight inputs are canonical BF16 `MoEWeightPack`
 by default (the backend quantizes at `preprocess_weights`); kernel-ready
@@ -317,12 +321,21 @@ classDiagram
 | Split kernel | `identity` | `IdentityConfig` — comm-only; `dummy_moe_weights` OK |
 | Split kernel | `fused_moe` | `FusedMoeKernelConfig(moe_config=...)` — bridges to `flashinfer.fused_moe`; BF16 + W4A4/W4A8/W4A16; LL EXPERT_MAJOR / RANK_MAJOR / HT FLAT |
 | Mega kernel | `sm100_fp8_fp4_bf16_deepgemm` | `Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig` — FP8/FP4, sm_100+ |
+| Mega kernel | `sm100_bf16_nvfp4_bf16_cutedsl` | `Sm100_Bf16_Nvfp4_Bf16_Cutedsl_MegaMoeConfig` — BF16/NVFP4, SM100/SM103 |
 | Mega kernel | `sm100_nvfp4_nvfp4_bf16_cutedsl` | `Sm100_Nvfp4_Nvfp4_Bf16_Cutedsl_MegaMoeConfig` — NVFP4, sm_100+ |
 | Mega kernel | `sm100_mxfp8_mxfp8_bf16_cutedsl` | `Sm100_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig` — MXFP8 (`kind` e4m3/e5m2), sm_100+ |
 
 **Mega weights:** with `preprocess_weights=True` (default), canonical bf16 or pre-quantized `MoEWeightPack` is transformed at init. With `preprocess_weights=False`, supply `MegaConfig.transformed_weights` (from `preprocess_*_mega_weights`).
 
+W4A16 additionally accepts optional FP32 `[local_experts]`
+`w13_global_scale` / `w2_global_scale` (omitted means one), applied after FP32
+GEMM accumulation. Its prepared `(weight, scale, alpha)` triples share the
+W4A4 weight/scale layout. Layer construction rejects these fields for other
+backends.
+
 **Mega activations:** with `quantize_input=True` (default), bf16 `[T, hidden]` is quantized into symm workspace at forward. Non-bf16 with `quantize_input=True` raises `MoEEpConfigError`; use `quantize_input=False` and pre-quantized activations plus `MoEEpTensors.scales`.
+The BF16-activation CuTeDSL backends, including W4A16, instead copy BF16
+inputs with `quantize_input=True` and reject pre-quantized activation inputs.
 
 ## Runtime
 

flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py

@@ -438,6 +438,21 @@ def _setup_transform_partitions(
             transform_tiler,
         )
 
+    @cute.jit
+    def _finish_transform_stage(
+        self,
+        a_load2trans_pipeline: pipeline.PipelineTmaAsync,
+        trans2mma_pipeline: pipeline.PipelineAsyncUmma,
+        cur_a_load2trans_consumer_state: pipeline.PipelineState,
+        trans2mma_producer_state: pipeline.PipelineState,
+    ) -> None:
+        if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM):
+            cute.arch.fence_view_async_tmem_store()
+        else:
+            cute.arch.fence_proxy("async.shared", space="cta")
+        a_load2trans_pipeline.consumer_release(cur_a_load2trans_consumer_state)
+        trans2mma_pipeline.producer_commit(trans2mma_producer_state)
+
     @cute.jit
     def _transform_tile(
         self,
@@ -537,14 +552,12 @@ def _transform_tile(
                 tAsA_transform[a_transform_stage_coord],
                 dst_copy_a,
             )
-            if cutlass.const_expr(
-                self.transform_a_source == tcgen05.OperandSource.TMEM
-            ):
-                cute.arch.fence_view_async_tmem_store()
-            else:
-                cute.arch.fence_proxy("async.shared", space="cta")
-            a_load2trans_pipeline.consumer_release(cur_a_load2trans_consumer_state)
-            trans2mma_pipeline.producer_commit(trans2mma_producer_state)
+            self._finish_transform_stage(
+                a_load2trans_pipeline,
+                trans2mma_pipeline,
+                cur_a_load2trans_consumer_state,
+                trans2mma_producer_state,
+            )
             trans2mma_producer_state.advance()
             if trans2mma_producer_state.count < k_tile_cnt:
                 peek_trans2mma_empty_status = trans2mma_pipeline.producer_try_acquire(

flashinfer/moe_ep/__init__.py

@@ -11,16 +11,19 @@
         mega/
           kernel/           fused comm + local MoE kernels
       modes/                split and mega orchestration layers
+      cute_dsl/             FlashInfer-maintained CuTe DSL implementations
       kernel_src/           vendored kernel drops (verbatim src/ + shim/)
 
 Import layering (strict, one direction)::
 
-    layer / modes / core  -->  backends  -->  kernel_src.<drop> shim  -->  src/
+    layer / modes / core  -->  backends  -->  cute_dsl or kernel_src.<drop>
+    cute_dsl  -->  kernel_src.<drop> public helpers  -->  shim/  -->  src/
 
 - Only a drop's ``shim/`` may import that drop's vendored ``src/`` tree;
   nothing else imports ``src/``, ever.
-- Only ``backends/`` may import a drop's shim, and only through the drop's
-  package ``__init__`` (``kernel_src.<drop>``), never shim submodules.
+- Only ``backends/`` and ``cute_dsl/`` may import a drop's shim, and only
+  through the drop's package ``__init__`` (``kernel_src.<drop>``), never shim
+  submodules.
 - The layer, ``modes/``, ``core/``, and everything above use backend APIs
   only (config classes + the ``core.kernel.registry``) — no ``kernel_src``,
   no shim.
@@ -69,6 +72,10 @@
     Sm100_Bf16_Bf16_Bf16_RankMajorCuda_MegaMoeConfig,
     preprocess_mega_weights as preprocess_bf16_rank_major_cuda_mega_weights,
 )
+from .backends.mega.kernel.sm100.bf16_nvfp4_bf16_cutedsl import (
+    Sm100_Bf16_Nvfp4_Bf16_Cutedsl_MegaMoeConfig,
+    preprocess_mega_weights as preprocess_w4a16_cutedsl_mega_weights,
+)
 from .backends.mega.kernel.sm100.mxfp8_mxfp8_bf16_cutedsl import (
     Sm100_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig,
     preprocess_mega_weights as preprocess_mxfp8_cutedsl_mega_weights,
@@ -179,6 +186,7 @@
     "Bf16CutedslMegaMoeConfig",
     "Sm100_Bf16_Bf16_Bf16_Cutedsl_MegaMoeConfig",
     "Sm100_Bf16_Bf16_Bf16_RankMajorCuda_MegaMoeConfig",
+    "Sm100_Bf16_Nvfp4_Bf16_Cutedsl_MegaMoeConfig",
     "CombineInputParams",
     "CombineOutput",
     "Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig",
@@ -250,6 +258,7 @@
     "preprocess_mega_weights",
     "preprocess_bf16_cutedsl_mega_weights",
     "preprocess_bf16_rank_major_cuda_mega_weights",
+    "preprocess_w4a16_cutedsl_mega_weights",
     "preprocess_mxfp8_cutedsl_mega_weights",
     "preprocess_nvfp4_cutedsl_mega_weights",
     "preprocess_sm120_mxfp8_cutedsl_mega_weights",

flashinfer/moe_ep/backends/mega/kernel/sm100/bf16_nvfp4_bf16_cutedsl/__init__.py

@@ -0,0 +1,11 @@
+"""W4A16 CuTe DSL MegaMoE backend."""
+
+from .backend import W4A16CutedslMegaKernelBackend
+from .config import Sm100_Bf16_Nvfp4_Bf16_Cutedsl_MegaMoeConfig
+from .weights import preprocess_mega_weights
+
+__all__ = [
+    "W4A16CutedslMegaKernelBackend",
+    "Sm100_Bf16_Nvfp4_Bf16_Cutedsl_MegaMoeConfig",
+    "preprocess_mega_weights",
+]

flashinfer/moe_ep/backends/mega/kernel/sm100/bf16_nvfp4_bf16_cutedsl/backend.py

@@ -0,0 +1,229 @@
+"""Fused EP with packed NVFP4 weights and BF16 token transport."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from ......config import BootstrapConfig, FleetParams
+from ......core.kernel.base import MegaKernelBackend
+from ......core.kernel.registry import register_mega_kernel
+from ......core.kernel.workspace_pool import knobs_pool_key
+from ......core.runtime import bf16_cutedsl_runtime_requirements
+from ......core.validation.common import (
+    MoEEpArchError,
+    MoEEpConfigError,
+    validate_mega_fleet_params,
+    validate_mega_forward_inputs,
+)
+from ......weights import MoEWeightPack
+from .staging import stage_mega_moe_inputs
+from .config import Sm100_Bf16_Nvfp4_Bf16_Cutedsl_MegaMoeConfig
+from .weights import (
+    TransformedMegaWeights,
+    preprocess_mega_weights,
+    validate_transformed_mega_weights,
+)
+
+if TYPE_CHECKING:
+    from ......tensors import MoEEpTensors
+
+
+@register_mega_kernel("sm100_bf16_nvfp4_bf16_cutedsl")
+class W4A16CutedslMegaKernelBackend(MegaKernelBackend):
+    supports_global_weight_scales = True
+
+    @classmethod
+    def kernel_name(cls) -> str:
+        return "sm100_bf16_nvfp4_bf16_cutedsl"
+
+    def __init__(self, config: Sm100_Bf16_Nvfp4_Bf16_Cutedsl_MegaMoeConfig) -> None:
+        super().__init__(config)
+        self._kernel_config = config
+        self._autotune_pending = config.knobs == "auto"
+        self._autotune_winner: dict | None = None
+
+    def runtime_requirements(self, bootstrap: BootstrapConfig) -> frozenset[str]:
+        return bf16_cutedsl_runtime_requirements(bootstrap)
+
+    def validate_init(
+        self, bootstrap: BootstrapConfig, fleet_params: FleetParams
+    ) -> None:
+        if torch.cuda.is_available():
+            cc = torch.cuda.get_device_capability()
+            if cc not in ((10, 0), (10, 3)):
+                raise MoEEpArchError("W4A16 MegaMoE requires SM100 or SM103")
+        config = self._kernel_config
+        validate_mega_fleet_params(
+            fleet_params,
+            bootstrap.world_size,
+            intermediate_size=config.intermediate_size,
+            top_k=config.top_k,
+            alignment=32,
+        )
+        if config.intermediate_size % 64:
+            raise MoEEpConfigError(
+                "W4A16 MegaMoE requires intermediate size divisible by 64"
+            )
+        if config.top_k > min(32, fleet_params.num_experts):
+            raise MoEEpConfigError(
+                "W4A16 MegaMoE top_k must not exceed 32 or num_experts"
+            )
+
+    def preprocess_weights(
+        self, weights: MoEWeightPack, fleet_params: FleetParams
+    ) -> TransformedMegaWeights:
+        return preprocess_mega_weights(
+            weights,
+            intermediate_size=self._kernel_config.intermediate_size,
+            hidden_size=fleet_params.token_hidden_size,
+        )
+
+    def validate_transformed_weights(
+        self,
+        transformed_weights: TransformedMegaWeights,
+        bootstrap: BootstrapConfig,
+        fleet_params: FleetParams,
+    ) -> None:
+        validate_transformed_mega_weights(
+            transformed_weights,
+            intermediate_size=self._kernel_config.intermediate_size,
+            hidden_size=fleet_params.token_hidden_size,
+            world_size=bootstrap.world_size,
+            num_experts=fleet_params.num_experts,
+        )
+
+    def _allocate_workspace(self, fleet_params: FleetParams) -> Any:
+        from ......cute_dsl.megamoe.nvfp4_w4a16 import (
+            get_symm_buffer_for_w4a16_mega_moe,
+        )
+
+        config = self._kernel_config
+        return get_symm_buffer_for_w4a16_mega_moe(
+            fleet_params.num_experts,
+            fleet_params.max_tokens_per_rank,
+            config.top_k,
+            fleet_params.token_hidden_size,
+            config.intermediate_size,
+            self.ep_rank,
+            self.ep_world_size,
+            gate_up_clamp=config.gate_up_clamp,
+            knobs=config.knobs if isinstance(config.knobs, dict) else None,
+        )
+
+    def _workspace_pool_key(self, fleet_params: FleetParams) -> Any:
+        config = self._kernel_config
+        if config.knobs == "auto":
+            # Tuning mutates this session's frontend; never share it with a
+            # separately tuned layer (same rule as the NVFP4 Mega backend).
+            return None
+        return (
+            self.kernel_name(),
+            torch.cuda.current_device(),
+            self.ep_rank,
+            self.ep_world_size,
+            id(self.ep_comm_group),
+            fleet_params.num_experts,
+            fleet_params.max_tokens_per_rank,
+            fleet_params.token_hidden_size,
+            config.intermediate_size,
+            config.top_k,
+            config.gate_up_clamp,
+            knobs_pool_key(config.knobs),
+        )
+
+    def validate_forward(
+        self, t: MoEEpTensors, fleet_params: FleetParams, *, quantize_input: bool
+    ) -> None:
+        if not quantize_input:
+            raise MoEEpConfigError(
+                "W4A16 MegaMoE accepts BF16 activations; keep MegaConfig.quantize_input=True"
+            )
+        if any(
+            v is not None
+            for v in (t.scales, t.fc1_alpha, t.fc2_alpha, t.fc1_norm_const)
+        ):
+            raise MoEEpConfigError(
+                "W4A16 MegaMoE does not accept activation quantization fields; "
+                "pass weight global scales in PrequantizedMoEWeights"
+            )
+        if t.hidden_states.ndim != 2 or t.topk_ids.ndim != 2:
+            raise MoEEpConfigError("W4A16 activations and routing must be 2D")
+        if t.hidden_states.dtype != torch.bfloat16:
+            raise MoEEpConfigError("W4A16 MegaMoE hidden_states must be BF16")
+        if t.topk_ids.dtype not in (torch.int32, torch.int64):
+            raise MoEEpConfigError("W4A16 MegaMoE topk_ids must be int32 or int64")
+        if t.topk_weights.dtype != torch.float32:
+            raise MoEEpConfigError("W4A16 MegaMoE topk_weights must be FP32")
+        if (
+            not t.hidden_states.is_cuda
+            or t.topk_ids.device != t.hidden_states.device
+            or t.topk_weights.device != t.hidden_states.device
+        ):
+            raise MoEEpConfigError(
+                "W4A16 activations and routing must share a CUDA device"
+            )
+        validate_mega_forward_inputs(
+            t.hidden_states,
+            t.topk_ids,
+            t.topk_weights,
+            fleet_params,
+            top_k=self._kernel_config.top_k,
+            quantize_input=True,
+        )
+
+    def validate_capture_ready(
+        self, workspace: Any, transformed_weights: TransformedMegaWeights
+    ) -> None:
+        mega = workspace._frontend._mega
+        if mega is None or mega.compiled is None:
+            raise RuntimeError(
+                "MegaMoE workspace is not warmed for CUDA graph capture; "
+                "call layer.warmup(..., workspace=workspace) first"
+            )
+
+    def stage_inputs(
+        self, t: MoEEpTensors, workspace: Any, *, quantize_input: bool
+    ) -> None:
+        stage_mega_moe_inputs(
+            t.hidden_states,
+            t.topk_weights,
+            t.topk_ids,
+            workspace.x,
+            workspace.topk_idx,
+            workspace.topk_weights,
+        )
+
+    def compute(
+        self,
+        workspace: Any,
+        transformed_weights: TransformedMegaWeights,
+        *,
+        output: torch.Tensor,
+    ) -> torch.Tensor:
+        from ......cute_dsl.megamoe.nvfp4_w4a16 import w4a16_mega_moe
+
+        if self._autotune_pending:
+            from ......cute_dsl.megamoe.nvfp4_w4a16 import autotune_w4a16_mega_moe
+
+            self._autotune_winner = dict(
+                autotune_w4a16_mega_moe(
+                    output,
+                    transformed_weights[0],
+                    transformed_weights[1],
+                    workspace,
+                    num_tokens=output.shape[0],
+                    gate_up_clamp=self._kernel_config.gate_up_clamp,
+                )
+            )
+            self._autotune_pending = False
+        w4a16_mega_moe(
+            output,
+            transformed_weights[0],
+            transformed_weights[1],
+            workspace,
+            num_tokens=output.shape[0],
+            gate_up_clamp=self._kernel_config.gate_up_clamp,
+        )
+        return output