#!/usr/bin/env python3 """B1 mixed FP8/BF16 prefill FMHA — unit test. Tests the T>1 prefill kernel at production values: HD=512, NOPE=448, ROPE=64, H=128, T=1..64, N=128..2048. 1. T=1 prefill vs decode kernel (should be identical) 2. T>1 prefill vs PyTorch SDPA reference 3. T>1 with attention sinks 4. Large N (production context lengths) 5. Multi-batch No model weights needed — uses synthetic random data. """ import sys, math import torch import torch.nn.functional as F def quantize_fp8_e4m3(x_fp32): amax = x_fp32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) scale = amax / 448.0 fp8 = (x_fp32 / scale).clamp(-448, 448).to(torch.float8_e4m3fn) return fp8.view(torch.uint8), scale.squeeze(-1) def cosine(a, b): return F.cosine_similarity(a.flatten().float(), b.flatten().float(), dim=0).item() def main(): HD = 512; NOPE = 448; ROPE = 64; H = 128 scale = 1.0 / math.sqrt(HD) print("=" * 70) print("B1 MIXED FP8 PREFILL FMHA — UNIT TEST") print(f"Production values: HD={HD}, NOPE={NOPE}, ROPE={ROPE}, H={H}") print("=" * 70) results = {} # ---- Test 1: T=1 prefill vs decode kernel ---- print("\n" + "=" * 70) print("TEST 1: T=1 prefill vs T=1 decode (should be identical)") print("=" * 70) try: from dsv4.kernels.attention.fmha_mixed_fp8_op import fmha_mixed_fp8_decode_raw from dsv4.kernels.attention.fmha_mixed_fp8_prefill_op import fmha_mixed_fp8_prefill_raw torch.manual_seed(42) B = 1; T = 1; N = 256 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() k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_fp32[:, :NOPE]) k_rope_bf16 = k_fp32[:, NOPE:].bfloat16() k_nope_fp8 = k_nope_fp8.cuda() k_nope_scale = k_nope_scale.cuda() k_rope_bf16 = k_rope_bf16.cuda() o_decode, _ = fmha_mixed_fp8_decode_raw( q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE) o_prefill, _ = fmha_mixed_fp8_prefill_raw( q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE) cos_val = cosine(o_decode, o_prefill) print(f" T=1 decode vs prefill: cos={cos_val:.8f}") assert cos_val >= 0.999, f"T=1 decode vs prefill cos={cos_val:.6f} < 0.999" results["1_t1_vs_decode"] = True print(" PASS") except Exception as e: print(f" FAIL: {e}") results["1_t1_vs_decode"] = False # ---- Test 2: T>1 prefill vs SDPA reference ---- print("\n" + "=" * 70) print("TEST 2: T>1 prefill vs PyTorch SDPA") print("=" * 70) all_pass = True for T in [1, 2, 4, 8, 16, 32]: for N in [128, 512]: print(f"\n T={T} N={N}") try: from dsv4.kernels.attention.fmha_mixed_fp8_prefill_op import fmha_mixed_fp8_prefill_raw torch.manual_seed(42) q_fp32 = torch.randn(1, 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() k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_fp32[:, :NOPE]) k_rope_bf16 = k_fp32[:, NOPE:].bfloat16() k_nope_fp8 = k_nope_fp8.cuda() k_nope_scale = k_nope_scale.cuda() k_rope_bf16 = k_rope_bf16.cuda() o_prefill, lse = fmha_mixed_fp8_prefill_raw( q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE) # Reference: dequantize, run SDPA per query position nope_dequant = k_nope_fp8.view(torch.float8_e4m3fn).cpu().float() * k_nope_scale.cpu().unsqueeze(-1).float() k_full = torch.cat([nope_dequant, k_fp32[:, NOPE:]], dim=-1).bfloat16().cuda() k_4d = k_full.unsqueeze(0).unsqueeze(0).expand(1, 1, -1, -1) v_4d = k_4d.clone() o_ref = F.scaled_dot_product_attention(q_bf16, k_4d, v_4d, scale=scale) cos_val = cosine(o_prefill, o_ref) print(f" cos={cos_val:.6f} |prod|={o_prefill.float().abs().max().item():.4f} " f"|ref|={o_ref.float().abs().max().item():.4f}") if cos_val < 0.999: all_pass = False print(f" FAIL") else: print(f" PASS") except Exception as e: print(f" ERROR: {e}") all_pass = False results["2_t1_vs_sdpa"] = all_pass # ---- Test 3: T>1 with attention sinks ---- print("\n" + "=" * 70) print("TEST 3: T>1 with attention sinks") print("=" * 70) try: from dsv4.kernels.attention.fmha_mixed_fp8_prefill_op import fmha_mixed_fp8_prefill_raw T = 4; N = 256 torch.manual_seed(42) q_fp32 = torch.randn(1, 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() k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_fp32[:, :NOPE]) k_rope_bf16 = k_fp32[:, NOPE:].bfloat16() k_nope_fp8 = k_nope_fp8.cuda(); k_nope_scale = k_nope_scale.cuda(); k_rope_bf16 = k_rope_bf16.cuda() sink_bias = torch.randn(H, dtype=torch.float32) * 2.0 o_with, _ = fmha_mixed_fp8_prefill_raw( q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, attn_sink=sink_bias, rope_dim=ROPE) o_no, _ = fmha_mixed_fp8_prefill_raw( q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE) diff = (o_with - o_no).float().abs().max().item() print(f" Max diff with/without sink: {diff:.6f}") assert diff > 1e-4, "Sink bias has no effect" results["3_sinks"] = True print(" PASS") except Exception as e: print(f" FAIL: {e}") results["3_sinks"] = False # ---- Test 4: Large N ---- print("\n" + "=" * 70) print("TEST 4: Large N (production context)") print("=" * 70) all_pass = True for N in [1024, 2048, 4096]: for T in [4, 16]: print(f"\n T={T} N={N}") try: from dsv4.kernels.attention.fmha_mixed_fp8_prefill_op import fmha_mixed_fp8_prefill_raw torch.manual_seed(42) q_fp32 = torch.randn(1, 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() k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_fp32[:, :NOPE]) k_rope_bf16 = k_fp32[:, NOPE:].bfloat16() k_nope_fp8 = k_nope_fp8.cuda(); k_nope_scale = k_nope_scale.cuda(); k_rope_bf16 = k_rope_bf16.cuda() o_prefill, lse = fmha_mixed_fp8_prefill_raw( q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE) nope_dequant = k_nope_fp8.view(torch.float8_e4m3fn).cpu().float() * k_nope_scale.cpu().unsqueeze(-1).float() k_full = torch.cat([nope_dequant, k_fp32[:, NOPE:]], dim=-1).bfloat16().cuda() k_4d = k_full.unsqueeze(0).unsqueeze(0).expand(1, 1, -1, -1) v_4d = k_4d.clone() o_ref = F.scaled_dot_product_attention(q_bf16, k_4d, v_4d, scale=scale) cos_val = cosine(o_prefill, o_ref) print(f" cos={cos_val:.6f}") if cos_val < 0.999: all_pass = False print(f" FAIL") else: print(f" PASS") except Exception as e: print(f" ERROR: {e}") all_pass = False results["4_large_n"] = all_pass # ---- Summary ---- print("\n" + "=" * 70) print("SUMMARY") print("=" * 70) all_ok = True for name, passed in results.items(): status = "PASS" if passed else "FAIL" if not passed: all_ok = False print(f" {name}: {status}") print() sys.exit(0 if all_ok else 1) if __name__ == "__main__": main()