DeepGEMM/README_NVFP4.md

DeepGEMM NVFP4 Mega MoE Kernel

Overview

This branch adds native NVFP4 support to DeepGEMM's mega MoE kernel. NVFP4 uses E2M1 packed weights with UE4M3 block scales (group_size=16), as opposed to MXFP4 which uses UE8M0 scales (group_size=32).

HARD RULE: MoE experts stay in NVFP4. Never convert to MXFP4.

What Changed

CUDA Kernel (sm100_fp8_nvfp4_mega_moe.cuh)

Parameter	MXFP4 (original)	NVFP4 (new)
kGranK	32	16
Scale factor type	float_ue8m0_t	float_ue4m3_t
PTX instruction	kind::mxf8f6f4.block_scale	kind::mxf4nvf4.block_scale.scale_vec::4X
SF vector size	block32 / 2X	block16 / 4X
TMEM SF layout	2 sub-columns per UMMA	4 sub-columns per UMMA
UTCCP col stride	i * 4	i * 8
kNumSFATmemCols	SF_BLOCK_M / 32	SF_BLOCK_M / 32 * 4
kNumSFBTmemCols	SF_BLOCK_N / 32	SF_BLOCK_N / 32 * 4
Activation SF format	UE8M0 (>> 23)	UE4M3 (float→e4m3 cast)
kNumSFUint32	kHidden / 128	kHidden / 64
recipe	(1, 1, 32)	(1, 1, 16)

PTX Wrappers (tcgen05.cuh)

Added:

• SM100_MMA_MXF4NVF4_2x1SM_SS — 2-CTA NVFP4 block-scaled MMA
• SM100_MMA_MXF4NVF4_SS — 1-CTA NVFP4 block-scaled MMA

Python API (mega/__init__.py)

• fp8_nvfp4_mega_moe() — main entry point
• transform_nvfp4_weights_for_mega_moe() — converts checkpoint weights to kernel format
• _pack_nvfp4_sf_for_utccp() — packs UE4M3 scales into int32 UTCCP layout

C++ Bindings

• csrc/apis/mega_nvfp4.hpp — SymmBuffer with NVFP4 SF strides (K/16)
• csrc/jit_kernels/impls/sm100_fp8_nvfp4_mega_moe.hpp — JIT kernel

Architecture

NVFP4 Checkpoint                    Kernel Format
┌─────────────────────┐            ┌────────────────────────┐
│ weight: uint8       │───────────>│ uint8 (same, E2M1)    │
│ (E2M1, 2 per byte)  │            │ packed, interleaved    │
├─────────────────────┤            ├────────────────────────┤
│ weight_scale:       │            │ int32 (UTCCP layout)   │
│ float8_e4m3fn       │──pack────>│ 4 UE4M3 bytes → 1 i32 │
│ (UE4M3, group=16)   │  +UTCCP   │ then 4x32 transpose    │
├─────────────────────┤            └────────────────────────┘
│ weight_scale_2:     │
│ float32 (global)    │──────> Applied at weight load time
│                     │        (multiply into block scales)
└─────────────────────┘

Scale Factor Flow

Checkpoint: weight_scale (UE4M3) × weight_scale_2 (FP32) = dequantized BF16

Kernel path:
  1. pre-scale: weight_scale_2 * weight_scale → float32 block scales
  2. pack to UE4M3: float32 → cutlass::float_e4m3_t → uint8 → pack 4→int32
  3. UTCCP transpose for TMA consumption
  4. Tensor Core reads UE4M3 scales via mxf4nvf4 instruction

Important Notes

scale_format_ constraint

The CUTLASS instruction descriptor has a single scale_format_ bit (0=E4M3, 1=E8M0) that applies to BOTH A and B scale factors. This means both activation (SFA) and weight (SFB) scales must use the same format. The L1 epilogue outputs UE4M3 activation scales to match the NVFP4 weight scales.

Weight scale_2 handling

The NVFP4 checkpoint has a dual-level scaling scheme:

• weight_scale: per-block UE4M3 (group_size=16)
• weight_scale_2: per-tensor float32 global scale

The weight_scale_2 must be multiplied into the block scales before packing for the kernel. This is done in transform_nvfp4_weights_for_mega_moe().

Remaining Work

• Test compilation on B200 (SM100)
• Verify UTCCP 4X column stride (i*8) — may need adjustment based on TMEM layout diagrams
• Verify SF packing: UE4M3 bytes → int32 layout matches what the UTCCP instruction expects
• Verify the L1 epilogue UE4M3 conversion (float → e4m3 cast + sign bit clear)
• Validate scale_format_ bit value: currently set by make_instr_desc_block_scaled<float_ue4m3_t> which sets scale_format_=0 (E4M3)
• Verify kNumSFATmemCols and kNumSFBTmemCols calculations for 4X layout
• Integration with vLLM DeepseekV4MegaMoEExperts class
• Weight loading: map NVFP4 checkpoint params to DeepseekV4MegaMoEExperts
• End-to-end quality test: compare NVFP4 mega_moe output vs FlashInfer FP4 MoE