nvfp4-megamoe-kernel/NVFP4-1.1_APPROACH.md

NVFP4-1.1 Approach Update (2026-05-28)

CuTeDSL Float-to-Int Limitation

CuTeDSL CANNOT convert Float32 to Int32. Both cutlass.Int32(float_val) and float_val.to(cutlass.Int32) fail with "LLVM ERROR: unsupported operation" during PTX lowering. The MLIR arith.FloatToSIOp is generated but the LLVM backend cannot lower it.

This blocks in-kernel FP4 pack, which requires:

1. FP8 E4M3 bit pattern computation (exponent + mantissa as integers)
2. E2M1 nibble index computation (half_step → index as integer)
3. Nibble packing into bytes (bit shifts and OR on integers)

What works in CuTeDSL

• cute.arch.fmax, cute.arch.fmin — float min/max ✅
• cute.floor — floor function ✅
• cute.absf — float abs ✅
• cute.arch.load / cute.arch.store — scalar GMEM I/O ✅
• cute.arch.cvt_i8_bf16 — int8 → BF16 (one-way) ✅
• cute.arch.cvt_f4e2m1_f16 — FP4 → BF16 (one-way) ✅
• Float arithmetic (+, -, *, /) ✅
• Type casts: Float32 ↔ BFloat16 ✅
• Int32 arithmetic (shifts, OR, AND) ✅ (but can't create Int32 from Float32!)

What does NOT work

• Float32 → Int32 conversion (any method)
• Inline PTX / inline assembly
• BF16 → Int8 (reverse of cvt_i8_bf16)
• BF16 → FP4 (reverse of cvt_f4e2m1_f16)

Revised approach: Compact SwiGLU output

Since in-kernel FP4 pack is blocked, the best optimization within CuTeDSL's capabilities is:

Modify the SwiGLU epilogue to skip writing gate subtiles to GMEM.

Current flow:

L1 GEMM → BF16 interleaved [gate*8, swiglu*8, ...] → 2*intermediate BF16
  → deinterleave_quantize_nvfp4_cuda → FP4 + SF

Target flow:

L1 GEMM → BF16 compact [swiglu, swiglu, ...] → intermediate BF16  (HALF the write!)
  → quantize_nvfp4_gpu → FP4 + SF  (simpler kernel, no deinterleave)

Wins:

• 50% less BF16 GMEM written (skip gate columns)
• Simpler quantization kernel (no deinterleave needed)
• quantize_nvfp4_gpu is already tested and proven

The full FP4 fusion can be revisited when CuTeDSL adds float-to-int support or when the attention final-epilogue is rewritten in CUTLASS C++ (ROADMAP Priority 2).

Implementation plan

1. Add compact_output tensor parameter to kernel (shape: (tokens, intermediate) BF16)
2. In epilogue: gate subtiles → skip SMEM write + TMA store
3. Up subtiles → write to compact_output via TMA store (not the interleaved C tensor)
4. This requires a new TMA atom and descriptor for the compact output
5. Update runner and MoE layer to use quantize_nvfp4_gpu instead of deinterleave_quantize_nvfp4_cuda