# 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 |