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

61 lines
2.5 KiB
Markdown

# 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