From 9810de7109c58ed464ae4511b3840788ce272fb3 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 07:45:34 +0000 Subject: [PATCH] more debug --- src/nvfp4_megamoe_kernel/weight_transform.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py index 0e01033e..cadc0589 100644 --- a/src/nvfp4_megamoe_kernel/weight_transform.py +++ b/src/nvfp4_megamoe_kernel/weight_transform.py @@ -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.