B1: add mixed FP8 FMHA cosine verification test (HD=512, N=128-2048)

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

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)