[None][feat] Add MegaMoECuteDsl NVFP4 MoE backend

Add the MegaMoECuteDsl backend (ported swap-AB MoE kernel) with NVFP4 quantization for the PyTorch fused-MoE path, wired through ConfigurableMoE / create_moe and the cute_dsl_megamoe_nvfp4_blackwell custom op.

Key points:
- Per-expert NVFP4 scales (fc1_alpha / fc2_alpha / fc1_norm_const) threaded through the v2 kernel ABI, derived from each expert raw input scales.
- FC1/FC2 weights are stored in the natural (slots, N, K_bytes) layout with K (hidden//2 / intermediate//2) innermost, and presented to the kernel as non-contiguous transpose(1, 2) views (slots, K_bytes, N) so the swap-AB kernel reads K-major with K stride-1.
  Materializing the transpose with .contiguous() would move K off the innermost axis and corrupt the GEMM (per-token cosine ~0).
- local_workspace is passed with a static cute layout (not mark_layout_dynamic), matching the kernel region-offset contract.
- AutoTuner input pre-hook only sanitizes the regenerated dynamic inputs and leaves the static per-expert weight scales / alphas untouched.
- deepgemm-graph NVFP4 reference (NVFP4RefMegaMoECuteDsl) for accuracy checks.

Verified on GB200: single-GPU test_moe_backend / test_moe_module megamoe_cutedsl
cases pass (per-token cosine ~1.0), and multi-GPU DEP (EP) cases pass.

Signed-off-by: xxi <xxi@nvidia.com>

Author: xxi-nv

.claude/skills/trtllm-moe-develop/SKILL.md

@@ -268,6 +268,26 @@ Checklist:
 - Existing legacy `forward` methods can be read for compatibility context, but
   they are not the default pattern for new backend work.
 
+### Imported Kernel ABI Checklist
+
+When importing or wrapping an upstream kernel, derive the TRT-LLM adapter
+contract from the lowest-level kernel consumer. Comments, docs, design notes,
+and parameter names are useful hints, but they are not proof of the runtime ABI.
+
+- Derive weight shape and layout from the kernel entrypoint, `make_layout`, TMA,
+  MMA/GEMM transforms, and stride usage. Record required tensor shape, stride,
+  physical storage layout, and boundary view layout.
+- Derive alpha and scale semantics from kernel consumption points. Trace where
+  alpha, norm constants, block scales, activation scales, and weight scales are
+  loaded and multiplied before deciding how upper layers compute or pack them.
+  Treat weight bytes, block scales/SF, and global alpha/norm constants as
+  separate contracts.
+- Design the upper-layer adapter from the kernel ABI upward. Map each kernel
+  input/output to an adapter responsibility: storage tensor, view/transposition,
+  dtype reinterpretation, padding, scale packing, workspace ownership,
+  synchronization, and output reduction. Validate parity with upstream
+  invocation dumps, not just final output.
+
 ### Quantization And Weights
 
 Role:

.pre-commit-config.yaml

@@ -302,6 +302,22 @@ common-files: &common_files |
         tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py |
         tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py |
         tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/config.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dispatch_kernel.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/iket_compat.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sf_swizzle.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py |
         tensorrt_llm/_torch/cute_dsl_utils.py |
         tensorrt_llm/_torch/debug/__init__.py |
         tensorrt_llm/_torch/debug/debug_hook.py |
@@ -1658,6 +1674,22 @@ legacy-files: &legacy_files |
         tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py |
         tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py |
         tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/config.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dispatch_kernel.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/iket_compat.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sf_swizzle.py |
+        tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py |
         tensorrt_llm/_torch/cute_dsl_utils.py |
         tensorrt_llm/_torch/debug/__init__.py |
         tensorrt_llm/_torch/debug/debug_hook.py |

legacy-files.txt

@@ -294,6 +294,22 @@ tensorrt_llm/_torch/cute_dsl_kernels/blackwell/__init__.py
 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py
 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py
 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/config.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dispatch_kernel.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/iket_compat.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sf_swizzle.py
+tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py
 tensorrt_llm/_torch/cute_dsl_utils.py
 tensorrt_llm/_torch/debug/__init__.py
 tensorrt_llm/_torch/debug/debug_hook.py

pyproject.toml

@@ -352,6 +352,22 @@ exclude = [
     "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py",
     "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py",
     "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/config.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dispatch_kernel.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/iket_compat.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sf_swizzle.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py",
     "tensorrt_llm/_torch/cute_dsl_utils.py",
     "tensorrt_llm/_torch/debug/__init__.py",
     "tensorrt_llm/_torch/debug/debug_hook.py",

ruff-legacy.toml

@@ -311,6 +311,22 @@ include = [
     "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py",
     "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py",
     "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/config.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dispatch_kernel.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/iket_compat.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sf_swizzle.py",
+    "tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py",
     "tensorrt_llm/_torch/cute_dsl_utils.py",
     "tensorrt_llm/_torch/debug/__init__.py",
     "tensorrt_llm/_torch/debug/debug_hook.py",

tensorrt_llm/_torch/autotuner.py

@@ -1468,8 +1468,7 @@ def _create_tensor_like(self, origin_tensor: torch.Tensor,
         if dtype == torch.float4_e2m1fn_x2:
             return (torch.rand(shapes, device=device) * 10 - 5).to(
                 torch.uint8).view(dtype)
-        else:
-            return (torch.rand(shapes, device=device) * 10 - 5).to(dtype)
+        return (torch.rand(shapes, device=device) * 10 - 5).to(dtype)
 
     def _prepare_input_tensors(
             self, profile: OptimizationProfile,

tensorrt_llm/_torch/custom_ops/__init__.py

@@ -58,6 +58,18 @@ def inplace_slice_copy(dest: torch.Tensor, src: torch.Tensor, dim1_start: int,
         'cute_dsl_nvfp4_dense_gemm_swiglu_fp4out_blackwell',
     ]
 
+    # MegaMoE NVFP4 op probes a strict superset of IS_CUTLASS_DSL_AVAILABLE
+    # (cutlass.torch + cutlass._mlir + cute_nvgpu MMA atoms + the ported
+    # CuteDSL kernel package). The cute_dsl_megamoe_custom_op module
+    # sets ``IS_MEGAMOE_OP_AVAILABLE`` based on its own try/except probe;
+    # importing the module is safe regardless of the result -- it just
+    # logs and leaves ``IS_MEGAMOE_OP_AVAILABLE = False`` on partial
+    # cutlass-dsl installs so callers can fall back via the factory.
+    from .cute_dsl_megamoe_custom_op import IS_MEGAMOE_OP_AVAILABLE
+    if IS_MEGAMOE_OP_AVAILABLE:
+        from .cute_dsl_megamoe_custom_op import cute_dsl_megamoe_nvfp4_blackwell
+        __all__ += ['cute_dsl_megamoe_nvfp4_blackwell']
+
 if IS_CUDA_TILE_AVAILABLE:
     from .cuda_tile_custom_ops import (cuda_tile_rms_norm,
                                        cuda_tile_rms_norm_fuse_residual_)