biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cleanup: remove dead _pack_ue4m3_to_uint32, fix data format docs

Browse Source 
weight_transform.py returns float8_e4m3fn scales, NOT packed uint32.
The _pack_ue4m3_to_uint32 function was never called. Removed it.
Updated README data formats to accurately reflect the pipeline:
- Weight scales: float8_e4m3fn (direct to CUTLASS, no unpack)
- Activation scales: uint32 packed (from staging kernel, unpacked to float8)
This commit is contained in:
biondizzle
2026-05-14 17:28:12 +00:00
parent 008f8cccbd
commit 1c15dadaa5
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
README.md
View File

@@ -115,7 +115,7 @@ nvfp4_megamoe_kernel/
| File | When it runs | What it does |
|------|-------------|--------------|
| `weight_transform.py` | Once at startup (weight loading) | Takes raw NVFP4 checkpoint weights, folds global scales into block scales, packs UE4M3 into uint32. Output: `((l1_w, l1_sf), (l2_w, l2_sf))` |
| `weight_transform.py` | Once at startup (weight loading) | Takes raw NVFP4 checkpoint weights, folds global scales into block scales. Returns scales as `float8_e4m3fn` (not packed uint32). Output: `((l1_w, l1_sf), (l2_w, l2_sf))` |
| `symm_buffer.py` | Once at startup (buffer alloc) | Pre-allocates GPU tensors for activations, scales, routing data, and all-reduce. These persist across forward passes. |
| `nvfp4_mega_moe.py` | Every forward pass | Orchestrates the MoE: reads from symm buffer → L1 GEMM → activation → re-quantize → L2 GEMM → output. Contains `stage_activation` (BF16→FP4 quantize for L1→L2) and `unpack_ue4m3_u32` (uint32 packed scales → float8). |
| `cutlass_nvfp4_gemm/kernel.py` | Every forward pass (called by nvfp4_mega_moe) | Per-expert loop: gather tokens for each expert, call CUTLASS GEMM, scatter results with routing weights. |
@@ -128,11 +128,11 @@ nvfp4_megamoe_kernel/
### Weights
- **Packed E2M1** (`int8`): 2 FP4 values per byte. Shape: `(E_per_rank, N, K//2)`, K-major layout.
- **UE4M3 block scales** (`float8_e4m3fn`): 1 scale per 16 FP4 values (group_size=16). Shape: `(E_per_rank, N, K//16)`.
- **UE4M3 block scales** (`float8_e4m3fn`): 1 scale per 16 FP4 values (group_size=16). Shape: `(E_per_rank, N, K//16)`. Returned as `float8_e4m3fn` from `weight_transform.py` — NOT packed uint32. The CUTLASS GEMM consumes float8 directly.
### Activations (after staging kernel)
- **Packed E2M1** (`int8`): Shape: `(num_tokens, K//2)`.
- **UE4M3 scales** (`uint32`): 4 UE4M3 values packed per uint32. Shape: `(num_tokens, K//64)`.
- **UE4M3 scales** (`uint32`): 4 UE4M3 values packed per uint32. Shape: `(num_tokens, K//64)`. Unpacked to `float8_e4m3fn` via `unpack_ue4m3_u32` before reaching the CUTLASS GEMM.
### GEMM dimensions (DeepSeek-V4-Pro)
- **L1 (gate_up_proj):** M×6144×7168 (per expert)

8
src/nvfp4_megamoe_kernel/weight_transform.py
View File

@@ -4,9 +4,13 @@ NVFP4 Weight Transformation for CUTLASS mega_moe kernel.
Converts raw NVFP4 checkpoint weights (uint8 E2M1 + float8_e4m3fn UE4M3 + float32 global scale)
into the format expected by the CUTLASS block-scaled GEMM kernel:
- Packed FP4 weights (int8, K-major)
- UE4M3 block scales (float8_e4m3fn, row-major — CUTLASS SF remap handles layout)
- UE4M3 block scales (float8_e4m3fn, row-major — CUTLASS SF remap handles interleaving)
This replaces deep_gemm.mega.transform_nvfp4_weights_for_mega_moe.
Weight scales are returned as float8_e4m3fn (NOT packed uint32). The CUTLASS GEMM
consumes float8 scales directly; only activation scales from the staging kernel come
as uint32 and need unpack_ue4m3_u32.
This replaces deep_gemma.mega.transform_nvfp4_weights_for_mega_moe.
Call signature matches the nightly vLLM deepseek_v4.py finalize_weights:
transform_nvfp4_weights_for_mega_moe(