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

debug: fix gs scalar conversion + add traceback

Browse Source
This commit is contained in:
biondizzle
2026-05-15 18:27:44 +00:00
parent df916b87eb
commit 773967452f
1 changed files with 8 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
View File

@@ -403,9 +403,8 @@ def nvfp4_mega_moe_full(
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_global_sf[e0]: {l1_global_sf[e0].tolist()} shape={l1_global_sf[e0].shape}")
print(f"[GEMM-DEBUG] l1_alpha (igs): {l1_alpha:.6e}")
print(f"[GEMM-DEBUG] per_expert_alpha[{e0}]: {float(l1_alpha * l1_global_sf[e0]):.6e}")
# Dequantize activation
x_u8 = x_fp4[s0].view(torch.uint8)
@@ -433,15 +432,18 @@ def nvfp4_mega_moe_full(
ref_out = torch.nn.functional.linear(x_full.unsqueeze(0), w_full.T).squeeze(0) # (N,)
# Apply per-half global scale (gate_gs for first half, up_gs for second half)
gn = ref_out.shape[0] // 2
if gs.numel() == 1:
ref_out = ref_out * float(gs)
gs_vals = gs.detach().cpu().tolist()
if isinstance(gs_vals, float) or len(gs_vals) == 1:
ref_out = ref_out * (gs_vals if isinstance(gs_vals, float) else gs_vals[0])
else:
ref_out[:gn] = ref_out[:gn] * float(gs[0])
ref_out[gn:] = ref_out[gn:] * float(gs[1])
ref_out[:gn] = ref_out[:gn] * gs_vals[0]
ref_out[gn:] = ref_out[gn:] * gs_vals[1]
nvfp4_mega_moe_full._ref_l1 = (s0, e0, ref_out)
print(f"[BF16-REF-L1] expert={e0} amax={ref_out.abs().max():.4e} mean={ref_out.mean():.4e}")
except Exception as ex:
import traceback
print(f"[BF16-REF-L1] FAILED: {ex}")
traceback.print_exc()
# Step 2: L1 GEMM — slot-based, per-expert alpha
l1_slots, _ = nvfp4_mega_moe_l1(