test: add CuTeDSL NVFP4 GEMM test using reference ScaledGroupedGemmKernel

Tests the NVIDIA reference kernel with our quantization pipeline:
1. Quantize BF16 → NVFP4 (our stage_activation logic)
2. Pad and swizzle scale factors (to_blocked)
3. Run ScaledGroupedGemmKernel (2Dx3D scenario)
4. Compare against BF16 matmul reference

Also adds cutedsl/moe.py module for the future pipeline integration.

src/nvfp4_megamoe_kernel/cutedsl/__init__.py

@@ -0,0 +1,14 @@
+"""
+NVFP4 MoE kernel using NVIDIA's CuTeDSL ScaledGroupedGemmKernel.
+
+This replaces the broken C++ CUTLASS kernel. The CuTeDSL kernel handles:
+- NVFP4 (Float4E2M1FN + Float8E4M3FN, sf_vec_size=16) natively
+- Block-scaled SF layouts (no manual remap needed)
+- Full Blackwell pipeline (TMA → MMA → Epilogue overlap)
+- Per-expert global scales for NVFP4
+
+We just need to:
+1. Quantize activations to FP4 (stage_activation)
+2. Call the kernel with the right tensor layout
+3. Apply MoE routing (gate/up fusion, SiLU, scatter)
+"""

src/nvfp4_megamoe_kernel/cutedsl/moe.py

@@ -0,0 +1,171 @@
+"""
+NVFP4 MoE pipeline using CuTeDSL ScaledGroupedGemmKernel.
+
+Replaces the broken C++ CUTLASS path. Uses NVIDIA's official MoE scaled
+grouped GEMM kernel from the CUTLASS CuTeDSL examples.
+
+Usage:
+    from nvfp4_megamoe_kernel.cutedsl.moe import nvfp4_mega_moe_full
+"""
+
+import sys
+import os
+import torch
+import cutlass
+import cutlass.cute as cute
+import cutlass.torch as cutlass_torch
+import cutlass.utils as utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+
+# Add the CuTeDSL examples to the path so we can import the kernel
+_CUTLASS_ROOT = os.environ.get("CUTLASS_ROOT", "/root/cutlass")
+_CUTEDSL_EXAMPLES = os.path.join(_CUTLASS_ROOT, "examples/python/CuTeDSL")
+if _CUTEDSL_EXAMPLES not in sys.path:
+    sys.path.insert(0, _CUTEDSL_EXAMPLES)
+
+from cute.blackwell.kernel.moe.torch_scaled_grouped_mm import ScaledGroupedGemmKernel
+
+from nvfp4_megamoe_kernel.nvfp4_mega_moe import (
+    stage_activation,
+    _quantize_to_e2m1,
+)
+
+# ── Module-level compiled kernel cache ──
+_compiled_l1_kernel = None
+_compiled_l2_kernel = None
+_l1_kernel_config = None
+_l2_kernel_config = None
+
+
+def _get_torch_dtype(cutlass_dtype):
+    """Convert CUTLASS dtype to PyTorch dtype."""
+    mapping = {
+        cutlass.Float4E2M1FN: torch.float4_e2m1fn_x2,
+        cutlass.Float8E4M3FN: torch.float8_e4m3fn,
+        cutlass.Float8E8M0FNU: torch.float8_e8m0fnu,
+        cutlass.BFloat16: torch.bfloat16,
+        cutlass.Float16: torch.float16,
+        cutlass.Float32: torch.float32,
+    }
+    return mapping.get(cutlass_dtype)
+
+
+def _torch_tensor_to_cute(torch_tensor: torch.Tensor) -> cute.Tensor:
+    """Convert a PyTorch GPU tensor to a CuTe tensor with dynamic layout."""
+    cute_tensor = cutlass_torch.from_dlpack(torch_tensor)
+    leading_dim = cutlass_torch.get_leading_dim(torch_tensor)
+    cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim)
+    return cute_tensor
+
+
+def _compile_kernel_once(kernel, sample_tensors, global_scale_a=None, global_scale_b=None):
+    """Compile the CuTeDSL kernel on first call, cache the result."""
+    import cuda.bindings.driver as cuda
+
+    a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute = sample_tensors
+
+    gsa_cute = _torch_tensor_to_cute(global_scale_a) if global_scale_a is not None else None
+    gsb_cute = _torch_tensor_to_cute(global_scale_b) if global_scale_b is not None else None
+
+    cluster_size = kernel.cluster_shape_mn[0] * kernel.cluster_shape_mn[1]
+    max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+    compiled = cute.compile(
+        kernel,
+        a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
+        max_active_clusters, stream,
+        global_scale_a=gsa_cute,
+        global_scale_b=gsb_cute,
+    )
+    return compiled
+
+
+def run_scaled_grouped_gemm(
+    mat_a: torch.Tensor,        # (tokens_sum, K_packed) float4_e2m1fn_x2 — row-major (K-major for CuTe)
+    mat_b: torch.Tensor,        # (experts, K_packed, N) float4_e2m1fn_x2 — K-major
+    scale_a: torch.Tensor,      # (tokens_sum, K_sf) float8_e4m3fn — row-major
+    scale_b: torch.Tensor,      # (experts, K_sf, N) float8_e4m3fn — K-major after transpose
+    expert_offsets: torch.Tensor,  # (experts,) int32 — cumulative token offsets
+    global_scale_a: torch.Tensor = None,  # (experts,) float32 — NVFP4 per-expert activation scale
+    global_scale_b: torch.Tensor = None,  # (experts,) float32 — NVFP4 per-expert weight scale
+    mma_tiler_mn: tuple = (128, 128),
+    cluster_shape_mn: tuple = (1, 1),
+) -> torch.Tensor:
+    """Run the CuTeDSL NVFP4 scaled grouped GEMM.
+    
+    2Dx3D scenario: A(tokens, K) x B(experts, K, N) -> C(tokens, N)
+    
+    Args:
+        mat_a: Activation tensor (tokens_sum, K_packed) in FP4
+        mat_b: Weight tensor (experts, K_packed, N) in FP4
+        scale_a: Activation block scales (tokens_sum, K_sf) in E4M3
+        scale_b: Weight block scales (experts, K_sf, N) in E4M3
+        expert_offsets: Cumulative token end offsets per expert
+        global_scale_a: Per-expert float32 activation global scale (NVFP4)
+        global_scale_b: Per-expert float32 weight global scale (NVFP4)
+    
+    Returns:
+        Output tensor (tokens_sum, N) in BF16
+    """
+    global _compiled_l1_kernel, _l1_kernel_config
+
+    tokens_sum = mat_a.shape[0]
+    k_packed = mat_a.shape[1]
+    num_experts = mat_b.shape[0]
+    n_dim = mat_b.shape[2]
+    k_dim = k_packed * 2  # 2 FP4 values per byte
+
+    # Output tensor
+    out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
+
+    # Create kernel config
+    kernel = ScaledGroupedGemmKernel(
+        scenario="2Dx3D",
+        sf_vec_size=16,
+        accumulate_on_output=False,
+        separate_tensormap_init=True,
+        consistent_token_padding=False,
+        mma_tiler_mnk=(*mma_tiler_mn, 256),
+        cluster_shape_mnk=(*cluster_shape_mn, 1),
+    )
+
+    # Convert to CuTe tensors
+    a_cute = _torch_tensor_to_cute(mat_a)
+    b_cute = _torch_tensor_to_cute(mat_b)
+    sfa_cute = _torch_tensor_to_cute(scale_a)
+    sfb_cute = _torch_tensor_to_cute(scale_b)
+    c_cute = _torch_tensor_to_cute(out)
+    offs_cute = _torch_tensor_to_cute(expert_offsets)
+
+    # Workspace
+    workspace_size = kernel.get_workspace_size(num_experts)
+    workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=mat_a.device)
+    ws_cute = _torch_tensor_to_cute(workspace)
+
+    gsa_cute = _torch_tensor_to_cute(global_scale_a) if global_scale_a is not None else None
+    gsb_cute = _torch_tensor_to_cute(global_scale_b) if global_scale_b is not None else None
+
+    import cuda.bindings.driver as cuda
+    cluster_size = kernel.cluster_shape_mn[0] * kernel.cluster_shape_mn[1]
+    max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+    # Compile and run
+    compiled = cute.compile(
+        kernel,
+        a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
+        max_active_clusters, stream,
+        global_scale_a=gsa_cute,
+        global_scale_b=gsb_cute,
+    )
+
+    compiled(
+        a_cute, b_cute, sfa_cute, sfb_cute, c_cute, offs_cute, ws_cute,
+        stream,
+        global_scale_a=gsa_cute,
+        global_scale_b=gsb_cute,
+    )
+    torch.cuda.synchronize()
+
+    return out