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

fix: MoE and shared expert global scale — gsb=ws2, gsa=input_scale (same bug as Nvfp4Linear)

Browse Source
This commit is contained in:
biondizzle
2026-06-01 02:31:12 +00:00
parent f4b444b456
commit 66b98e5794
1 changed files with 17 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
single_shot_inference.py
View File

@@ -454,7 +454,8 @@ def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg):
l1_fp4_list.append(torch.cat([gw, uw], dim=0).to(dev))
if gws is not None and uws is not None: l1_sf_list.append(torch.cat([gws, uws], dim=0).to(dev))
gs = gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)
l1_gs_list.append(gs)
l1_gs_list.append(1.0) # gsb base — ws2 will be folded in by _ensure_stacked
l1_gsa_list.append(gs) # gsa = input_scale
# weight_scale_2: scalar, folded into global_scale_b
l1_ws2_list.append(gws2.to(dev) if gws2 is not None else None)
dw, dws, dws2, disc = get_nvfp4_weight(all_w, ep, 'down_proj')
@@ -462,7 +463,8 @@ def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg):
l2_fp4_list.append(dw.to(dev))
if dws is not None: l2_sf_list.append(dws.to(dev))
gs2 = disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)
l2_gs_list.append(gs2)
l2_gs_list.append(1.0) # gsb base
l2_gsa_list.append(gs2) # gsa = input_scale
l2_ws2_list.append(dws2.to(dev) if dws2 is not None else None)
if not l1_fp4_list: log.warning(f"L{li}: No expert weights found"); return
l1_stacked = torch.stack(l1_fp4_list).to(dev)
@@ -471,6 +473,11 @@ def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg):
l2_sf_stacked = torch.stack(l2_sf_list).to(dev) if l2_sf_list else None
del l1_fp4_list, l1_sf_list, l2_fp4_list, l2_sf_list
moe.prepare_weights_from_stacked(l1_stacked, l1_sf_stacked, l1_gs_list, l2_stacked, l2_sf_stacked, l2_gs_list)
# Set activation global scales from input_scale (gsa = input_scale, gsb = weight_scale_2)
if l1_gsa_list:
moe._l1_activation_global_scale = l1_gsa_list[0] # Use first expert's input_scale
if l2_gsa_list:
moe._l2_activation_global_scale = l2_gsa_list[0] # Use first expert's input_scale
moe.l1_ws2 = l1_ws2_list
moe.l2_ws2 = l2_ws2_list
@@ -481,13 +488,17 @@ def _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg):
if gw is not None and uw is not None:
se.l1_fp4 = [torch.cat([gw, uw], dim=0).to(dev)]
se.l1_sf = [torch.cat([gws, uws], dim=0).to(dev)] if gws is not None and uws is not None else [torch.zeros(1, device=dev, dtype=torch.float8_e4m3fn)]
se.l1_gs = [gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)]
se.l1_ws2 = [gws2.to(dev) if gws2 is not None else None] # weight_scale_2
l1_isc = gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)
se.l1_gs = [1.0] # gsb base — ws2 will be folded in by finalize_weights
se.l1_ws2 = [gws2.to(dev) if gws2 is not None else None]
se._l1_activation_global_scale = l1_isc # gsa = input_scale
if dw is not None:
se.l2_fp4 = [dw.to(dev)]
se.l2_sf = [dws.to(dev)] if dws is not None else [torch.zeros(1, device=dev, dtype=torch.float8_e4m3fn)]
se.l2_gs = [disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)]
se.l2_ws2 = [dws2.to(dev) if dws2 is not None else None] # weight_scale_2
l2_isc = disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)
se.l2_gs = [1.0] # gsb base
se.l2_ws2 = [dws2.to(dev) if dws2 is not None else None]
se._l2_activation_global_scale = l2_isc # gsa = input_scale
def _cache_layer_weights_no_experts(all_w, n_layers, devices):
cached = {}