nvfp4-megamoe-kernel/tests/unit/test_b1_mixed_fp8_prefill.py

#!/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()