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

more debug

Browse Source
This commit is contained in:
biondizzle
2026-05-15 07:45:34 +00:00
parent 1a37b66922
commit 9810de7109
1 changed files with 16 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/nvfp4_megamoe_kernel/weight_transform.py
View File

@@ -79,6 +79,22 @@ def transform_nvfp4_weights_for_mega_moe(
l1_weight, l1_weight_scale = l1_tuple
l2_weight, l2_weight_scale = l2_tuple
# DEBUG: check raw scales before folding
l1_sf_f32_raw = l1_weight_scale.to(torch.float32)
l1_gs_raw = l1_weight_scale_2.to(torch.float32) if l1_weight_scale_2 is not None else None
if not getattr(transform_nvfp4_weights_for_mega_moe, '_sf_debug', False):
transform_nvfp4_weights_for_mega_moe._sf_debug = True
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}] "
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()}")
# Fold global scales into block scales
# 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.