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

debug: dump raw GEMM inputs + first 8 output values

Browse Source
This commit is contained in:
biondizzle
2026-05-15 17:02:40 +00:00
parent 9159cb6bb3
commit de8acc7965
1 changed files with 13 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
View File

@@ -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}")