From d493193d251bf7f7f4ccac9c90e071c90fe9fa8e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 08:02:02 +0000 Subject: [PATCH] fix the god damn projections --- src/nvfp4_megamoe_kernel/weight_transform.py | 26 ++++++++++++++------ vllm/patches/deepseek_v4.py | 16 +++++++++--- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py index cadc0589..6794a39d 100644 --- a/src/nvfp4_megamoe_kernel/weight_transform.py +++ b/src/nvfp4_megamoe_kernel/weight_transform.py @@ -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 diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index ad906734..93692955 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -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]