From d0ed3d84a81cb79f53401ce5109889c9e296f51e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 13:21:25 +0000 Subject: [PATCH] debug: add L2, SiLU, and scatter pipeline prints --- src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 1daadfcf..206c777f 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -195,6 +195,7 @@ def nvfp4_mega_moe_l2( slot_expert_ids, per_expert_alpha=per_expert_alpha, ) + print(f"[L2-GEMM-OUT] slots={slot_out.shape[0]} N={N} amax={slot_out.abs().max().item():.4e} mean={slot_out.float().mean().item():.4e} nan={torch.isnan(slot_out).any().item()}") return slot_out # (num_slots, HIDDEN) bfloat16 @@ -385,14 +386,12 @@ def nvfp4_mega_moe_full( # Step 3: SiLU + Mul PER SLOT — nonlinearity before combining paths gate, up = l1_slots.chunk(2, dim=-1) + print(f"[L1-SPLIT] gate amax={gate.abs().max().item():.4e} mean={gate.float().mean().item():.4e} | up amax={up.abs().max().item():.4e} mean={up.float().mean().item():.4e}") activated = torch.nn.functional.silu(gate) * up + print(f"[SILU-ACT] amax={activated.abs().max().item():.4e} mean={activated.float().mean().item():.4e} nan={torch.isnan(activated).any().item()}") if activation_clamp is not None: activated = activated.clamp(max=activation_clamp) - if MEGA_MOE_DEBUG: - print(f"[silu] nan={torch.isnan(activated).any().item()} " - f"abs_max={activated.abs().max().item():.4e}") - # Step 4: Quantize activated slots → FP4 l1_fp4, l1_sf_out, l2_global_scale = stage_activation(activated) @@ -427,3 +426,4 @@ def nvfp4_mega_moe_full( slot_token, l2_slots * slot_weight.to(l2_slots.dtype).unsqueeze(1), ) + print(f"[SCATTER] y amax={y.abs().max().item():.4e} mean={y.float().mean().item():.4e} nan={torch.isnan(y).any().item()} slots={num_slots}")