From 1c15dadaa59f826b4699d60f679787c7005afbe0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 14 May 2026 17:28:12 +0000 Subject: [PATCH] cleanup: remove dead _pack_ue4m3_to_uint32, fix data format docs 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) --- README.md | 6 +++--- src/nvfp4_megamoe_kernel/weight_transform.py | 8 ++++++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index c8080d0b..4a4e930e 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py index 893ea0ee..48c0c1f9 100644 --- a/src/nvfp4_megamoe_kernel/weight_transform.py +++ b/src/nvfp4_megamoe_kernel/weight_transform.py @@ -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(