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

fix: stop folding global scale into float8 block scales

Browse Source 
The fold block_sf (float8) * global_sf (float32) -> float8 loses ~25% precision.
Product of ~56-448 block_sf * ~4.65e-05 global_sf lands in float8 low-precision
zone where step size is 25%. This makes model output garbage despite finite values.

Fix: keep block scales as original float8, return global scales separately as
float32 per-expert vectors. Apply global scale as per-expert GEMM alpha in
cutlass_grouped_nvfp4_gemm (already iterates per-expert). For L1 with separate
gate/up global scales, use gate_gs as alpha and apply up_correction ratio to
the up half post-GEMM.

weight_transform.py: no more _fold_global_scale, returns (w, sf, global_sf)
nvfp4_mega_moe.py: per-expert alpha = activation_gs * weight_gs
kernel.py: per_expert_alpha parameter in grouped GEMM
deepseek_v4.py: updated type hints and comments
This commit is contained in:
biondizzle
2026-05-15 12:42:53 +00:00
parent 56e62e916d
commit fd59222fc0
9 changed files with 955 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
vllm/patches/deepseek_v4.py
View File

@@ -346,8 +346,8 @@ class DeepseekV4MegaMoEExperts(nn.Module):
)
set_weight_attrs(self.w2_input_scale, weight_attrs)
self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
# Register in the static forward context so the custom-op wrapper
# can look up this module by name from within a torch.compile graph.
@@ -437,13 +437,15 @@ class DeepseekV4MegaMoEExperts(nn.Module):
from nvfp4_megamoe_kernel import transform_nvfp4_weights_for_mega_moe
# === Native NVFP4 path ===
# The DeepGEMM nvfp4 mega_moe kernel consumes NVFP4 directly:
# - E2M1 packed uint8 (same as checkpoint)
# The CUTLASS nvfp4 mega_moe kernel consumes NVFP4 directly:
# - E2M1 packed int8 (same as checkpoint)
# - UE4M3 block scales (float8_e4m3fn), group_size=16
# - float32 global scale folded into block scales
# No conversion to MXFP4. Experts stay NVFP4.
# - float32 global scales returned SEPARATELY (NOT folded into float8)
# Previous versions folded global scales into block scales via float8
# round-trip, which caused ~25% precision loss. Now, global scales
# are applied as per-expert GEMM alpha in float32 (exact).
# Fold global scales into block scales and transform for the kernel
# Transform weights — returns (w, sf, global_sf) tuples
self._transformed_l1_weights, self._transformed_l2_weights = (
transform_nvfp4_weights_for_mega_moe(
(self.w13_weight.data.contiguous(),