diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 6f3bb07d..9e053269 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -397,6 +397,16 @@ def nvfp4_mega_moe_full( try: s0 = slot_token[0].item() e0 = slot_expert_local[0].item() + # Dump raw GEMM inputs for expert e0 + print(f"[GEMM-DEBUG] expert={e0} s0={s0}") + print(f"[GEMM-DEBUG] x_fp4[s0] first 8 bytes: {x_fp4[s0].view(torch.uint8)[:8].tolist()}") + print(f"[GEMM-DEBUG] x_sf[s0] first 8: {x_sf[s0].to(torch.float32)[:8].tolist()}") + print(f"[GEMM-DEBUG] l1_w[e0] first 8 bytes: {l1_w[e0].view(torch.uint8).flatten()[:8].tolist()}") + print(f"[GEMM-DEBUG] l1_sf[e0] first 8: {l1_sf[e0].to(torch.float32).flatten()[:8].tolist()}") + print(f"[GEMM-DEBUG] l1_global_sf[e0]: {l1_global_sf[e0].tolist()}") + print(f"[GEMM-DEBUG] l1_alpha (igs): {l1_alpha:.6e}") + print(f"[GEMM-DEBUG] per_expert_alpha[{e0}]: {float(per_expert_alpha[e0]):.6e}") + # Dequantize activation x_u8 = x_fp4[s0].view(torch.uint8) lo = (x_u8 & 0x0F).long() @@ -457,6 +467,9 @@ def nvfp4_mega_moe_full( cos = torch.nn.functional.cosine_similarity(nvfp4_out.unsqueeze(0), ref_f.unsqueeze(0)).item() mse = (nvfp4_out - ref_f).pow(2).mean().item() print(f"[COSINE-L1] expert={e0} cosine={cos:.6f} mse={mse:.4e} nvfp4_amax={nvfp4_out.abs().max():.4e} ref_amax={ref_f.abs().max():.4e}") + # Dump first 8 output values from each + print(f"[NVFP4-OUT-8] {nvfp4_out[:8].tolist()}") + print(f"[REF-OUT-8] {ref_f[:8].tolist()}") except Exception as ex: print(f"[COSINE-L1] FAILED: {ex}")