From 879adc324d99263d84f0296b8ee394826eabd162 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 14 May 2026 14:17:44 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20=5Ffold=5Fglobal=5Fscale=20=E2=80=94=20r?= =?UTF-8?q?emove=20broken=20logical=5Fwidths=20branch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The logical_widths branch took expert 0 and 1's global scales and applied them to ALL experts. For L1 with logical_widths=[3072,3072], every expert got expert-0's scale on its gate half and expert-1's scale on its up half. All other experts' global scales were discarded. The else branch correctly broadcasts each expert's own (E,1) global scale across (E, N, K//16). Removed the dead logical_widths code. --- src/nvfp4_megamoe_kernel/weight_transform.py | 25 +++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py index ed4fd6a5..6d4e74ad 100644 --- a/src/nvfp4_megamoe_kernel/weight_transform.py +++ b/src/nvfp4_megamoe_kernel/weight_transform.py @@ -22,25 +22,19 @@ import torch def _fold_global_scale( weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn - weight_scale_2: torch.Tensor, # (E, num_logical) or (E, 1) or scalar float32 - logical_widths: list[int] = None, + weight_scale_2: torch.Tensor, # (E, 1) or (E,) or scalar float32 ) -> torch.Tensor: - """Fold global scale into block scales: UE4M3 * FP32 → float32.""" + """Fold global scale into block scales: UE4M3 * FP32 → float32. + + Each expert has its own global scale. Broadcasts (E,1,1) → (E, N, K//16). + """ 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.numel() > 1 and logical_widths is not None: - expanded = [] - for i, w in enumerate(logical_widths): - if i < len(gs.flatten()): - expanded.append(gs.flatten()[i].expand(w)) - if expanded: - global_scale = torch.cat(expanded).unsqueeze(1) - sf_f32 = sf_f32 * global_scale else: - # Per-expert or per-row global scale — broadcast multiply + # Per-expert global scale — broadcast multiply while gs.dim() < sf_f32.dim(): gs = gs.unsqueeze(-1) sf_f32 = sf_f32 * gs.expand_as(sf_f32) @@ -82,9 +76,12 @@ def transform_nvfp4_weights_for_mega_moe( l2_weight, l2_weight_scale = l2_tuple # Fold global scales into block scales + # Both L1 and L2 use per-expert global scales (shape (E,1) or (E,)). + # The logical_widths branch was wrong: it treated gs as per-projection + # scalars and only used experts 0 and 1's scales for ALL experts. + # The else branch correctly broadcasts each expert's own global scale. if l1_weight_scale_2 is not None: - l1_sf_folded = _fold_global_scale(l1_weight_scale, l1_weight_scale_2, - logical_widths=[3072, 3072]) + l1_sf_folded = _fold_global_scale(l1_weight_scale, l1_weight_scale_2) else: l1_sf_folded = l1_weight_scale.to(torch.float32)