diff --git a/CORRRECTNESS_BACKLOG.md b/CORRRECTNESS_BACKLOG.md index 14d00b94..7200eeb3 100644 --- a/CORRRECTNESS_BACKLOG.md +++ b/CORRRECTNESS_BACKLOG.md @@ -98,6 +98,3 @@ Let me check what seq_len the FMHA is seeing. At L1 during prefill of the first ``` SO SINCE WE HAD TO TOUCH FMHA ANYWAY IN PART B. WE DID THAT FIRST AND TRIED TO GET THAT CORRECT BEFORE WE REVISTED THIS ISSUE!!! - -## Suggested sequence (we shouldve already tried all of these) -A1 (stop set) → A2 (penalty test) → if still broken: A3 (visible-range parity vs reference) → A4 (inverse-RoPE check). Then \ No newline at end of file diff --git a/tests/unit/test_fmha_mixed_fp8_cosine.py b/tests/unit/test_fmha_mixed_fp8_cosine.py new file mode 100644 index 00000000..5ab87e67 --- /dev/null +++ b/tests/unit/test_fmha_mixed_fp8_cosine.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +"""Cosine verification test for B1 mixed FP8/BF16 decode FMHA. + +Generates synthetic Q and KV in DSV4 storage format (FP8 noPE + BF16 RoPE), +runs the mixed FP8 decode kernel and the BF16 reference, and compares +per-head cosine similarity. + +Production sizes: HD=512, NOPE=448, ROPE=64, N=128..2048, H=128. +""" +import sys +import math +import torch +import torch.nn.functional as F + +def quantize_fp8_e4m3(x_fp32): + """Quantize FP32 tensor to FP8_E4M3 with per-row scale.""" + # x_fp32: (rows, cols) + amax = x_fp32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + scale = amax / 448.0 # E4M3 max representable + scaled = x_fp32 / scale + fp8 = scaled.to(torch.float8_e4m3fn) + return fp8.view(torch.uint8), scale.squeeze(-1) + + +def run_mixed_fp8_decode(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=64): + """Run the B1 mixed FP8 decode FMHA kernel.""" + from dsv4.kernels.attention.fmha_mixed_fp8_op import fmha_mixed_fp8_decode_raw + + B, H, T, HD = q_bf16.shape + q4 = q_bf16.permute(0, 2, 1, 3).contiguous() # (B, T, H, HD) -> need (B, H, T, HD) + q4 = q_bf16 # already (B, H, T, HD) + + o, lse = fmha_mixed_fp8_decode_raw( + q4, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=rope_dim) + return o # (B, H, T, HD) BF16 + + +def run_bf16_reference(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=64): + """Run BF16 reference FMHA by dequantizing FP8 noPE to BF16.""" + B, H, T, HD = q_bf16.shape + NOPE = HD - rope_dim + N = k_nope_fp8.shape[0] + + # Dequantize FP8 noPE → BF16 + k_nope_flat = k_nope_fp8.view(torch.float8_e4m3fn) + k_nope_bf16 = k_nope_flat.bfloat16() # (N, NOPE) + # Apply per-row scale + k_nope_bf16 = k_nope_bf16 * k_nope_scale.unsqueeze(-1).bfloat16() + + # Concat noPE + RoPE into full KV + k_full = torch.cat([k_nope_bf16, k_rope_bf16], dim=-1) # (N, HD) + + # V = K for MQA (self-attention decode) + v_full = k_full.clone() + + # Run BF16 FMHA + from dsv4.kernels.attention.production import dsv4_attention + q_3d = q_bf16.squeeze(2) # (B, H, HD) -> need (H, T, HD) per batch + results = [] + for b in range(B): + q_b = q_3d[b].transpose(0, 1) # (H, 1, HD) -> (H, T=1, HD) + # dsv4_attention expects (n_q, T, hd) or (batch, n_q, T, hd) + o_b = dsv4_attention(q_b.unsqueeze(0), k_full.unsqueeze(0).unsqueeze(0), + v_full.unsqueeze(0).unsqueeze(0), scale) + results.append(o_b) + o = torch.cat(results, dim=0) # (B, H, T, HD) + return o + + +def test_cosine(N_values, H=128, HD=512, rope_dim=64, B=1, seed=42): + """Test cosine similarity between mixed FP8 and BF16 reference FMHA.""" + torch.manual_seed(seed) + NOPE = HD - rope_dim + scale = 1.0 / math.sqrt(HD) + + all_pass = True + for N in N_values: + print(f"\n--- N={N} H={H} HD={HD} ---") + + # Generate synthetic Q (BF16) + q_fp32 = torch.randn(B, H, 1, HD, dtype=torch.float32) * 0.5 + q_bf16 = q_fp32.bfloat16().cuda() + + # Generate synthetic KV — split into noPE (FP8) + RoPE (BF16) + k_fp32 = torch.randn(N, HD, dtype=torch.float32) * 0.5 + k_nope_fp32 = k_fp32[:, :NOPE].contiguous() + k_rope_fp32 = k_fp32[:, NOPE:].contiguous() + + # Quantize noPE to FP8 + k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_nope_fp32) + k_nope_fp8 = k_nope_fp8.cuda() + k_nope_scale = k_nope_scale.cuda() + + # RoPE stays BF16 + k_rope_bf16 = k_rope_fp32.bfloat16().cuda() + + # Run mixed FP8 decode + try: + o_mixed = run_mixed_fp8_decode(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim) + except Exception as e: + print(f" MIXED FP8 FAILED: {e}") + all_pass = False + continue + + # Run BF16 reference + try: + o_ref = run_bf16_reference(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim) + except Exception as e: + print(f" BF16 REF FAILED: {e}") + all_pass = False + continue + + # Compare + o_mixed_f = o_mixed.float() + o_ref_f = o_ref.float() + + # Global cosine + cos_global = F.cosine_similarity(o_mixed_f.flatten(), o_ref_f.flatten(), dim=0).item() + + # Per-head cosine (averaged) + # o shape: (B, H, 1, HD) -> per-head: (B, H, HD) + o_mixed_h = o_mixed_f.squeeze(2) # (B, H, HD) + o_ref_h = o_ref_f.squeeze(2) + per_head_cos = F.cosine_similarity(o_mixed_h, o_ref_h, dim=-1) # (B, H) + min_cos = per_head_cos.min().item() + mean_cos = per_head_cos.mean().item() + + # Magnitude check + mixed_max = o_mixed_f.abs().max().item() + ref_max = o_ref_f.abs().max().item() + + pass_threshold = 0.999 + passed = cos_global >= pass_threshold + status = "PASS" if passed else "FAIL" + + print(f" {status}: cos_global={cos_global:.6f} min_head_cos={min_cos:.6f} " + f"mean_head_cos={mean_cos:.6f}") + print(f" |mixed|={mixed_max:.4f} |ref|={ref_max:.4f} " + f"ratio={mixed_max/ref_max:.4f}" if ref_max > 0 else " |ref|=0") + + if not passed: + all_pass = False + # Print worst heads + worst = per_head_cos[0].argsort()[:5] + print(f" Worst heads: {worst.tolist()} cos={per_head_cos[0][worst].tolist()}") + + return all_pass + + +if __name__ == "__main__": + # Production-scale test: N from 128 to 2048 + N_values = [128, 256, 512, 1024, 2048] + if len(sys.argv) > 1: + N_values = [int(x) for x in sys.argv[1].split(',')] + + print("=" * 70) + print("B1 Mixed FP8/BF16 FMHA Cosine Test") + print("Production sizes: HD=512, NOPE=448, ROPE=64, H=128") + print("=" * 70) + + all_pass = test_cosine(N_values) + + print("\n" + "=" * 70) + if all_pass: + print("ALL TESTS PASSED") + else: + print("SOME TESTS FAILED") + sys.exit(1)