biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
102 Commits 1 Branch 17 Tags
c7f6a1dc4d4e5f9deadfc4e017cdb6dcf64cd439
Commit Graph

7 Commits

Author SHA1 Message Date
biondizzle
c7f6a1dc4d fix: transpose B and SFB on the Python side at weight-load time, and adjust the SFB remap kernel to read from column-major source layout 2026-05-15 04:35:45 +00:00
biondizzle
ff6bb32684 Plumb global scale as GEMM alpha instead of folding into UE4M3 
stage_activation now returns (x_fp4, x_sf, input_global_scale).
The global scale is applied as the CUTLASS GEMM alpha parameter
in the epilogue: D = alpha * A @ B, avoiding the fp32→UE4M3
round-trip that folding would introduce.

Changes:
- stage_activation: returns global scale as 3rd value
- cutlass_nvfp4_gemm C++ binding: alpha param (was hardcoded 1.0)
- cutlass_grouped_nvfp4_gemm: passes alpha to per-expert GEMM
- nvfp4_mega_moe_l1/l2: accept alpha, pass to grouped GEMM
- nvfp4_moe_full: reads symm_buffer.input_global_scale for L1,
  uses stage_activation's returned global scale for L2
- SymmBuffer: added input_global_scale field
- vllm patch: stores global scale from stage_activation
 2026-05-15 03:32:19 +00:00
biondizzle
ce4c4b6fcb debug empty output 2026-05-14 22:13:32 +00:00
biondizzle
c016e66e23 Add CUDA sync + NaN/Inf check after each expert GEMM in grouped kernel 2026-05-14 11:27:58 +00:00
biondizzle
869151d211 Fix kernel.py: remove broken expand on scale factors (was expanding sf to weight size) 2026-05-14 10:36:16 +00:00
biondizzle
25cbc85afe Replace kernel.py with thin wrapper around pre-compiled _C extension 2026-05-14 09:25:56 +00:00
biondizzle
f375c80bfe feat: CUTLASS NVFP4 block-scaled GEMM kernel (native SM100 Blackwell) 
- Native NVFP4 block-scaled MMA using CUTLASS MainloopSm100TmaUmmaWarpSpecializedBlockScaled
- Invokes mxf8f6f4.block_scale tensor core instructions (tcgen05.mma)
- E2M1 (packed int8) + UE4M3 (float8_e4m3fn) block-16 scales → BF16 output
- No dequantization: hardware block-scaled MMA avoids costly dequantize+BF16 path
- PyTorch CUDA extension with CollectiveBuilder auto-deduction
- Grouped expert GEMM for MoE dispatch (32 experts/rank, top-6 routing)
- Integrated into nvfp4_mega_moe.py as primary path with TileLang fallback
- Standalone C API (cutlass_nvfp4_gemm.cu) for direct B200 compilation
- Build script, setup.py, and test script for B200 deployment

Files:
  cutlass_nvfp4_gemm/ — Kernel source, PyTorch binding, build/test scripts
  nvfp4_mega_moe.py — Updated to use CUTLASS kernel when available
 2026-05-13 23:11:15 +00:00