nvfp4-megamoe-kernel/REWRITE_PLAN.md

NVFP4 MegaMoE Kernel Rewrite — CuTeDSL

Decision

The C++ CUTLASS kernel is fundamentally broken. The GEMM produces cosine 0.05 with real data (despite SF remap = 0 errors and uniform tests passing). The root cause is likely in how CUTLASS's C++ API handles FP4 packing/tiling internally — something we can't easily debug or fix.

We're replacing it with NVIDIA's CuTeDSL approach (Python-based CUTLASS kernels compiled via MLIR → PTX). This is what the NVIDIA CUTLASS team recommends for Blackwell, and they have a working reference:

cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/moe/torch_scaled_grouped_mm.py

This is a 3900-line MoE scaled grouped GEMM that supports NVFP4 out of the box.

Reference Files (copied to reference/)

• reference/moe_torch_scaled_grouped_mm.py — the kernel (3900 lines)
• reference/moe_moe_utils.py — tensormap constructor
• reference/moe_moe_persistent_scheduler.py — MoE tile scheduler
• reference/moe_moe_sched_extension.py — scheduler extension for scaled GEMM
• reference/grouped_blockscaled_gemm.py — simpler grouped blockscaled reference
• reference/dense_blockscaled_gemm_persistent.py — dense (non-grouped) reference
• reference/blockscaled_layout.py — SF layout utilities

Architecture

The Reference Kernel: ScaledGroupedGemmKernel

7-warp persistent kernel:

• Warps 0-3: Epilogue (TMEM → RMEM → SMEM → GMEM)
• Warp 4: MMA (tcgen05.mma.block_scale with SFA/SFB in TMEM)
• Warp 5: TMA load (A, B, SFA, SFB from GMEM → SMEM)
• Warp 6: Scheduler (MoE persistent tile scheduler)

Supports:

• NVFP4 (Float4E2M1FN + Float8E4M3FN + sf_vec_size=16)
• 2Dx3D scenario (our case): A(tokens, K) x B(experts, K, N) → C(tokens, N)
• PyTorch interface via torch.nn.functional.scaled_grouped_mm

Our Adaptation

We need to:

1. Use ScaledGroupedGemmKernel directly (or with minimal adaptation)
2. Replace our Python pipeline (nvfp4_mega_moe.py) to call the CuTeDSL kernel
3. Handle the MoE-specific stuff (routing, gate/up fusion, SiLU, scatter)

The CuTeDSL kernel handles A/B/SFA/SFB tiling and MMA internally. We NO LONGER need:

• Our custom SF remap kernel (CuTeDSL handles SF layouts natively)
• Our C++ CUTLASS extension (cutlass_nvfp4_gemm.cu)
• transform_nvfp4_weights_for_mega_moe (CuTeDSL uses the natural layout)

We STILL need:

• stage_activation (BF16 → FP4 quantization)
• The MoE routing logic (top-k selection, slot mapping)
• Gate/up fusion with up_correction
• SiLU activation
• Scatter with routing weights

Plan

Phase 1: Get the reference kernel running standalone

• Set up CuTeDSL build environment on B200
• Run the reference example with NVFP4 config
• Verify it produces correct output

Phase 2: Integrate into our pipeline

• Create a Python module that wraps ScaledGroupedGemmKernel
• Replace cutlass_nvfp4_gemm calls with CuTeDSL kernel calls
• Handle weight format (the CuTeDSL kernel expects raw FP4 + SF, no transpose)

Phase 3: Full MoE pipeline

• Wire up stage_activation → L1 GEMM → SiLU → stage_activation → L2 GEMM → scatter
• Test with layertest.py (should get cosine ~0.995)

Phase 4: vLLM integration

• Update the vLLM patch to use the CuTeDSL kernel
• Remove the C++ extension build from the Dockerfile
• Test full inference

Key Differences from Current Approach

Current (C++ CUTLASS)	New (CuTeDSL)
C++ .cu file	Python .py file
Manual SF remap kernel	CuTe handles SF natively
Manual weight transpose	CuTe handles layout
pip install C++ extension	cute.compile() JIT
Broken FP4 handling	Reference-verified
No Blackwell pipeline	Full TMA/MMA/Epilogue overlap