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

fix the god damn projections

Browse Source
This commit is contained in:
biondizzle
2026-05-15 08:02:02 +00:00
parent 9810de7109
commit d493193d25
2 changed files with 30 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
src/nvfp4_megamoe_kernel/weight_transform.py
View File

@@ -26,17 +26,28 @@ import torch
def _fold_global_scale(
weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn
weight_scale_2: torch.Tensor, # (E, 1) or (E,) or scalar float32
weight_scale_2: torch.Tensor, # (E,) or (E, 2) or scalar float32
) -> torch.Tensor:
"""Fold global scale into block scales: UE4M3 * FP32 → float32.
Each expert has its own global scale. Broadcasts (E,1,1) → (E, N, K//16).
For fused projections (w13 = gate+up), weight_scale_2 is (E, 2):
scale_2[e, 0] applies to gate_proj rows, scale_2[e, 1] applies to up_proj rows.
N is split: gate = weight_scale[:, :N//2, :], up = weight_scale[:, N//2:, :]
For single projections (w2), weight_scale_2 is (E,) or scalar.
"""
sf_f32 = weight_scale.to(torch.float32)
gs = weight_scale_2.to(torch.float32)
if gs.numel() == 1:
sf_f32 = sf_f32 * gs
elif gs.dim() == 2 and gs.shape[1] == 2:
# Fused projection: (E, 2) — gate and up have separate global scales
# weight_scale is (E, N, K//16), N = gate_N + up_N
gate_N = sf_f32.shape[1] // 2
gs_gate = gs[:, 0].unsqueeze(-1) # (E, 1)
gs_up = gs[:, 1].unsqueeze(-1) # (E, 1)
sf_f32[:, :gate_N, :] = sf_f32[:, :gate_N, :] * gs_gate.unsqueeze(-1)
sf_f32[:, gate_N:, :] = sf_f32[:, gate_N:, :] * gs_up.unsqueeze(-1)
else:
# Per-expert global scale — broadcast multiply
while gs.dim() < sf_f32.dim():
@@ -87,13 +98,12 @@ def transform_nvfp4_weights_for_mega_moe(
print(f"[SF-DEBUG] raw l1_sf dtype={l1_weight_scale.dtype} range=[{l1_sf_f32_raw.min().item():.4e}, {l1_sf_f32_raw.max().item():.4e}] "
f"unique_raw={torch.unique(l1_weight_scale.view(torch.uint8)).numel()}")
if l1_gs_raw is not None:
print(f"[SF-DEBUG] l1_gs dtype={l1_weight_scale_2.dtype} shape={l1_weight_scale_2.shape} range=[{l1_gs_raw.min().item():.4e}, {l1_gs_raw.max().item():.4e}] "
print(f"[SF-DEBUG] l1_gs dtype={l1_weight_scale_2.dtype} shape={tuple(l1_weight_scale_2.shape)} "
f"range=[{l1_gs_raw.min().item():.4e}, {l1_gs_raw.max().item():.4e}] "
f"unique_gs={torch.unique(l1_gs_raw).numel()}")
# Show what happens after fold
folded = l1_sf_f32_raw[0] * l1_gs_raw[0] if l1_gs_raw.numel() > 1 else l1_sf_f32_raw[0] * l1_gs_raw
folded_u8 = folded.clamp(0, 448).to(torch.float8_e4m3fn)
print(f"[SF-DEBUG] after fold e=0: range=[{folded.min().item():.4e}, {folded.max().item():.4e}] "
f"unique_folded_u8={torch.unique(folded_u8.view(torch.uint8)).numel()}")
if l1_gs_raw.dim() == 2 and l1_gs_raw.shape[1] == 2:
print(f"[SF-DEBUG] gate gs unique={torch.unique(l1_gs_raw[:, 0]).numel()} "
f"up gs unique={torch.unique(l1_gs_raw[:, 1]).numel()}")
# Fold global scales into block scales
# The logical_widths branch was wrong: it treated gs as per-projection

16
vllm/patches/deepseek_v4.py
View File

@@ -295,9 +295,10 @@ class DeepseekV4MegaMoEExperts(nn.Module):
set_weight_attrs(self.w13_weight_scale, weight_attrs)
self.w13_weight_scale.quant_method = "block"
# NVFP4 global scales: float32, per-expert
# NVFP4 global scales: float32, per-expert, per-projection (gate, up)
# shape (num_local_experts, 2) — one scale for gate_proj, one for up_proj
self.w13_weight_scale_2 = nn.Parameter(
torch.zeros(num_local_experts, dtype=torch.float32),
torch.zeros(num_local_experts, 2, dtype=torch.float32),
requires_grad=False,
)
set_weight_attrs(self.w13_weight_scale_2, weight_attrs)
@@ -379,9 +380,16 @@ class DeepseekV4MegaMoEExperts(nn.Module):
f"param_shape={tuple(param.data[local_expert_id].shape)} "
f"loaded_absmax={loaded_weight.view(torch.int8).abs().max().item()}")
# Scalar params (weight_scale_2, input_scale): 1D per-expert
# Scalar params (weight_scale_2, input_scale): per-expert
if "weight_scale_2" in weight_name or "input_scale" in weight_name:
param.data[local_expert_id].copy_(loaded_weight)
if "w13_" in weight_name and "weight_scale_2" in weight_name:
# w13 is fused gate+up — store gate and up scales separately
# shard_id tells us which projection: w1=gate, w3=up
proj_idx = 0 if shard_id == "w1" else 1
param.data[local_expert_id, proj_idx].copy_(loaded_weight)
else:
# w2 or input_scale — single scalar per expert
param.data[local_expert_id].copy_(loaded_weight)
return True
expert_data = param.data[local_expert_id]