nvfp4-megamoe-kernel

Files

biondizzle 794ebaf7e5 P4: Fused RMSNorm + NVFP4 quantize kernel (2 launches vs 6+)

- fused_rmsnorm_quantize.cu: two-kernel approach
  Kernel 1: rmsnorm_amax_gsa — compute RMS + amax of normalized output → gsa per row
  Kernel 2: rmsnorm_quantize_nvfp4 — normalize + quantize using GPU-computed gsa
- Python bridge: rmsnorm_quantize_nvfp4() in ops/quantize.py
- Python bridge: dequantize_nvfp4() in ops/quantize.py
- Unit test: test_fused_rmsnorm_quantize.py (production shapes: 7168 hidden)
- Eliminates ~488 kernel launches per token (122 sites × 4 launches saved)

2026-06-02 16:26:24 +00:00

attention

perf: skip MQA GQA expansion in FMHA (stride=0, no 128x K/V copy)

2026-06-02 03:54:03 +00:00

cache

fix: correct gather.py kernel_dir path

2026-05-30 21:12:09 +00:00

compressor

KV-1/KV-2/KV-3: NVFP4 compressed KV + FP8 indexer keys

2026-06-02 10:00:50 +00:00

cuda

P4: Fused RMSNorm + NVFP4 quantize kernel (2 launches vs 6+)