feat: full FP4 activations for mxf4nvf4 - E2M1 packed A side + UE4M3 scales

mxf4nvf4 requires BOTH A and B to be FP4 (E2M1 packed).
Changes:
- a_dtype_t: float_e4m3_t → float_e2m1_unpacksmem_t
- UMMA_K: 32 → 64 (FP4 MMA atom)
- L1 epilogue: FP8 quant → E2M1 FP4 quantization with nearest-neighbor
- L1 output SMEM: packed E2M1 (2 per byte), TMA store uint8
- TMA descriptors: adjusted for FP4 packing (K/2 bytes per row)
- SymmBuffer: uint8 activations, shape (M, K//2)
- Staging kernel: BF16 → E2M1 packed + UE4M3 block16 scales
This commit is contained in:
2026-05-11 20:29:08 +00:00
parent 2cd86ff5e7
commit b3d1aae038
4 changed files with 106 additions and 61 deletions

View File

@@ -312,7 +312,9 @@ def fp8_nvfp4_mega_moe(y: torch.Tensor,
"""NVFP4 mega MoE: uses kind::mxf4nvf4.block_scale.scale_vec::4X
with UE4M3 block scales (group_size=16).
Both activations AND weights are E2M1 packed (FP4×FP4).
Weight format: (uint8 E2M1 packed, int32 packed UTCCP UE4M3 scales)
Activation format: E2M1 packed uint8 + UE4M3 scales (computed by staging kernel)
Recipe: (1, 1, 16) — kGranK=16 for NVFP4 group_size=16.
"""
_C.fp8_nvfp4_mega_moe(