biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
2,281 Commits 1 Branch 17 Tags
04cf8ca848d796aebc19f2688cdac068bb2de9ec
Commit Graph

4 Commits

Author SHA1 Message Date
biondizzle
7ef6402936 KV-1/KV-2/KV-3: NVFP4 compressed KV + FP8 indexer keys 
Architecture:
- Compressed KV: stored as NVFP4 (E2M1 + E4M3 + FP32 gsa)
  - Write path: compress→FP32 → FP32 RoPE → quantize FP32→NVFP4
  - Read path: dequant_nvfp4/dequant_nvfp4_selective → BF16 for FMHA
  - No BF16 intermediate in the write path
- Indexer keys: stored as FP8_E4M3 (1 byte + per-row scale)
  - Write path: compress→FP32 → quantize FP32→FP8_E4M3
  - Read path: dequant_fp8_e4m3 → BF16 for scoring
- SWA: remains BF16 (8MB total, fits in L2)

New kernels in kv_quantize.cu:
- compute_amax_gsa_fp32: per-row gsa from FP32 input
- quantize_nvfp4_from_fp32: FP32→NVFP4 with GPU gsa buffer
- quantize_fp8_e4m3_from_fp32: FP32→FP8_E4M3 for indexer keys
- dequant_fp8_e4m3 / dequant_fp8_e4m3_selective: FP8→BF16
- rope_fp32: FP32 GPT-J interleaved RoPE (no BF16)

Proven two-kernel pattern (same as quantize_nvfp4_gpu_fused):
  Kernel 1: amax_gsa (GPU-only)
  Kernel 2: quantize from buffer (GPU gsa)
No shared memory bugs. No cross-CTA race conditions.

KVCache updated:
- comp_kv_fp4/sf/gsa: NVFP4 storage (3.5× smaller than BF16)
- comp_idx_fp8/scale: FP8_E4M3 storage (1.9× smaller than BF16)
- comp_kv property: dequant NVFP4→BF16 on demand
- comp_kv_selective: dequant only top-k entries (bandwidth savings)
- comp_idx_kv property: dequant FP8→BF16 on demand

Removed: compressor_reduce_quant.cu (buggy single-kernel approach)
 2026-06-02 10:00:50 +00:00
biondizzle
f23320b5b2 KV-1/KV-2: Fused compress+NVFP4 quantize kernels + dequant 
- compressor_reduce_quant.cu: Single-kernel CSA/HCA compress + RMSNorm + NVFP4 quantize.
  No intermediate BF16. FP32 → E2M1 + E4M3 + FP32 gsa in one kernel.
  Shared memory: ~2.5KB per CTA (FP32 staging + nibble buffer).

- dequant_nvfp4.cu: NVFP4 → BF16 dequantization kernels.
  Full dequant (HCA dense gather) and selective dequant (CSA top-k gather).
  Single kernel launch per gather operation.

- production_compress.py: Added csa_compress_production_nvfp4() and
  hca_compress_production_nvfp4() — production path for KV-1/KV-2.

- loader.py: Preload dequant_nvfp4 and compressor_reduce_quant modules.

- test_kv_compress_quant.py: Unit tests verifying cos >= 0.999
  between BF16 reference and NVFP4 round-trip path.
 2026-06-02 09:37:53 +00:00
biondizzle
62041b78bf fix: import torch.utils.cpp_extension explicitly in production_compress 2026-06-01 05:20:44 +00:00
biondizzle
b380028c49 feat: production compressor/indexer — NVFP4 GEMM + CUDA softmax/reduce kernel 
- New compressor_reduce.cu: CSA/HCA token-level softmax + weighted sum + kv_norm
  One block per compressed entry, 128 threads, FP32 accumulation
  CSA: overlapping Ca/Cb streams (2m tokens per block)
  HCA: single stream (m tokens per block)
  Includes apply_kv_norm kernel (unweighted RMSNorm + weight)

- New production_compress.py: Python wrapper for CUDA kernels

- single_shot_inference.py: Compressor/Indexer now use production Nvfp4Linear
  for kv_proj, gate_proj, q_b_proj, weights_proj projections
  Then CUDA reduce kernel for softmax + weighted sum
  No more PyTorch reference nvfp4_linear_ref in compressor/indexer path
 2026-06-01 05:18:59 +00:00