debug: BF16 reference GEMM + cosine comparison for L1

src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py

@@ -391,6 +391,48 @@ def nvfp4_mega_moe_full(
     assert slot_token.numel() == num_slots
     assert slot_weight.numel() == num_slots
 
+    # BF16 reference: dequantize and run BF16 GEMM for the first slot to compare
+    if not getattr(nvfp4_mega_moe_full, '_ref_diag', False):
+        nvfp4_mega_moe_full._ref_diag = True
+        try:
+            s0 = slot_token[0].item()
+            e0 = slot_expert_local[0].item()
+            # Dequantize activation
+            x_u8 = x_fp4[s0].view(torch.uint8)
+            lo = (x_u8 & 0x0F).long()
+            hi = ((x_u8 >> 4) & 0x0F).long()
+            x_nib = torch.stack([lo, hi], dim=-1).reshape(-1)
+            x_signs = (x_nib >> 3).float() * -2 + 1
+            x_mags = _E2M1_MAGNITUDES.to(device=x_u8.device)[(x_nib & 0x07)]
+            x_deq = x_signs * x_mags
+            sf_exp = x_sf[s0].to(torch.float32).repeat_interleave(16, dim=-1)
+            igs = float(l1_global_scale) if not isinstance(l1_global_scale, float) else l1_global_scale
+            x_bf16 = (x_deq * sf_exp * igs).to(torch.bfloat16)
+            # Dequantize L1 weight for expert e0
+            w_u8 = l1_w[e0].view(torch.uint8)
+            wlo = (w_u8 & 0x0F).long()
+            whi = ((w_u8 >> 4) & 0x0F).long()
+            w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0], -1)
+            w_signs = (w_nib >> 3).float() * -2 + 1
+            w_mags = _E2M1_MAGNITUDES.to(device=w_u8.device)[(w_nib & 0x07)]
+            w_deq = w_signs * w_mags
+            w_sf_exp = l1_sf[e0].to(torch.float32).repeat_interleave(16, dim=0)
+            gs = l1_global_sf[e0]
+            if gs.dim() == 0:
+                w_bf16 = (w_deq * w_sf_exp * gs.item()).to(torch.bfloat16)
+            else:
+                w_bf16 = (w_deq * w_sf_exp).to(torch.bfloat16)
+            ref_out = torch.nn.functional.linear(x_bf16.unsqueeze(0), w_bf16.T).squeeze(0)
+            if gs.dim() > 0:
+                gn = ref_out.shape[0] // 2
+                ref_out[:gn] = ref_out[:gn] * gs[0].item()
+                ref_out[gn:] = ref_out[gn:] * gs[1].item()
+            ref_out = ref_out * igs
+            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:
+            print(f"[BF16-REF-L1] FAILED: {ex}")
+
     # Step 2: L1 GEMM — slot-based, per-expert alpha
     l1_slots, _ = nvfp4_mega_moe_l1(
         x_fp4, x_sf, l1_w, l1_sf,
@@ -399,6 +441,19 @@ def nvfp4_mega_moe_full(
         alpha=l1_alpha,
     )  # (num_slots, 2*INTER) bfloat16
 
+    # Compare L1 NVFP4 output to BF16 reference
+    if hasattr(nvfp4_mega_moe_full, '_ref_l1') and not getattr(nvfp4_mega_moe_full, '_ref_comp', False):
+        nvfp4_mega_moe_full._ref_comp = True
+        try:
+            s0, e0, ref = nvfp4_mega_moe_full._ref_l1
+            nvfp4_out = l1_slots[0].float()
+            ref_f = ref.float()
+            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}")
+        except Exception as ex:
+            print(f"[COSINE-L1] FAILED: {ex}")
+
     # Post-L1 shape asserts
     assert l1_slots.shape[0] == num_slots