diff --git a/tests/unit/test_fmha_mixed_fp8_debug.py b/tests/unit/test_fmha_mixed_fp8_debug.py new file mode 100644 index 00000000..5e200627 --- /dev/null +++ b/tests/unit/test_fmha_mixed_fp8_debug.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Minimal debug test for B1 mixed FP8 FMHA — compare per-step with BF16 reference. + +Tests a single head with small N to isolate the precision issue. +""" +import sys +import math +import torch +import torch.nn.functional as F + +def main(): + torch.manual_seed(42) + HD = 512; NOPE = 448; ROPE = 64 + H = 1; B = 1; T = 1 + N = 128 # small + scale = 1.0 / math.sqrt(HD) + + print(f"=== B1 Minimal Debug: N={N} H={H} HD={HD} ===") + + # Generate synthetic Q and KV + q_fp32 = torch.randn(B, H, T, HD, dtype=torch.float32) * 0.5 + k_fp32 = torch.randn(N, HD, dtype=torch.float32) * 0.5 + + q_bf16 = q_fp32.bfloat16().cuda() + + # Split KV + k_nope_fp32 = k_fp32[:, :NOPE].contiguous() + k_rope_fp32 = k_fp32[:, NOPE:].contiguous() + + # Quantize noPE to FP8 (same method as the production path) + amax = k_nope_fp32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + k_nope_scale = (amax / 448.0).squeeze(-1) # (N,) FP32 + k_nope_fp8 = (k_nope_fp32 / k_nope_scale.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn).view(torch.uint8) + + k_nope_fp8 = k_nope_fp8.cuda() + k_nope_scale = k_nope_scale.cuda() + k_rope_bf16 = k_rope_fp32.bfloat16().cuda() + + # Reference: BF16 SDPA + k_nope_dequant = k_nope_fp8.cpu().view(torch.float8_e4m3fn).bfloat16() * k_nope_scale.cpu().unsqueeze(-1).bfloat16() + k_full_bf16 = torch.cat([k_nope_dequant, k_rope_fp32.bfloat16()], dim=-1).cuda() + v_full_bf16 = k_full_bf16.clone() + + q_3d = q_bf16.squeeze(0) # (H, 1, HD) + k_3d = k_full_bf16.unsqueeze(0) # (1, N, HD) + v_3d = v_full_bf16.unsqueeze(0) # (1, N, HD) + + o_ref = F.scaled_dot_product_attention( + q_3d.float(), k_3d.unsqueeze(0).float(), v_3d.unsqueeze(0).float(), scale=scale + ).bfloat16() # (1, H, 1, HD) + o_ref = o_ref.squeeze(0) # (H, 1, HD) + + print(f"Reference: |o|={o_ref.abs().max().item():.6f} mean={o_ref.float().mean().item():.6f}") + print(f" o[0,0,:8]={o_ref[0,0,:8].float().tolist()}") + print(f" o[0,0,440:448]={o_ref[0,0,440:448].float().tolist()}") + + # Mixed FP8 kernel + from dsv4.kernels.attention.fmha_mixed_fp8_op import fmha_mixed_fp8_decode_raw + q_4d = q_bf16 # (B, H, T, HD) + o_mixed, lse = fmha_mixed_fp8_decode_raw( + q_4d, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE) + # o_mixed: (B, H, T, HD) + o_mixed_3d = o_mixed.squeeze(0) # (H, 1, HD) + + print(f"Mixed FP8: |o|={o_mixed.abs().max().item():.6f} mean={o_mixed.float().mean().item():.6f}") + print(f" o[0,0,:8]={o_mixed_3d[0,0,:8].float().tolist()}") + print(f" o[0,0,440:448]={o_mixed_3d[0,0,440:448].float().tolist()}") + + # Cosine + cos = F.cosine_similarity(o_ref.flatten().float(), o_mixed.flatten().float(), dim=0).item() + print(f"\nCosine: {cos:.6f}") + + # LSE comparison + # Reference LSE: log(sum(exp(scores))) + q_f = q_3d.float() # (H, 1, HD) + k_f = k_3d.unsqueeze(0).float() # (1, 1, N, HD) + scores = torch.matmul(q_f, k_f.transpose(-2, -1)) * scale # (H, 1, 1, N) + ref_lse = torch.logsumexp(scores, dim=-1) # (H, 1, 1) + print(f"Ref LSE: {ref_lse[0,0,0].item():.6f}") + print(f"Mixed LSE: {lse[0,0,0].item():.6f}") + + # Score distribution + print(f"\nScores: min={scores.min().item():.4f} max={scores.max().item():.4f} mean={scores.mean().item():.4f}") + + # Check if the noPE vs RoPE contributions are correct + q_nope_f = q_f[:, :, :NOPE] # (H, 1, NOPE) + q_rope_f = q_f[:, :, NOPE:] # (H, 1, ROPE) + k_nope_f = k_3d.unsqueeze(0).float()[:, :, :, :NOPE] # (1, 1, N, NOPE) + k_rope_f = k_3d.unsqueeze(0).float()[:, :, :, NOPE:] # (1, 1, N, ROPE) + + scores_nope = torch.matmul(q_nope_f, k_nope_f.transpose(-2, -1)) * scale + scores_rope = torch.matmul(q_rope_f, k_rope_f.transpose(-2, -1)) * scale + print(f"noPE scores: [{scores_nope.min().item():.4f}, {scores_nope.max().item():.4f}]") + print(f"RoPE scores: [{scores_rope.min().item():.4f}, {scores_rope.max().item():.4f}]") + + if cos < 0.999: + print(f"\n!!! COSINE TOO LOW ({cos:.6f}) — B1 KERNEL IS BROKEN !!!") + sys.exit(1) + else: + print(f"\nPASS: cosine {cos:.6f}") + sys.exit(0) + + +if __name__ == "__main__": + main()