project_6/verify_ixformer_attn.py

#!/usr/bin/env python3
"""Test ixformer native attention with head_dim=256 on BI-V100.

Qwen3.5 uses head_dim=256 for full attention layers.
We bypassed ixformer because of "head_dim > 128 limit".
Test if that's actually true on this hardware.

Tests:
1. ixformer.scaled_dot_product_attention with head_dim=256
2. ixformer.flash_attn_func with head_dim=256
3. ixformer.vllm_single_query_cached_kv_attention with head_dim=256
4. Compare outputs vs PyTorch reference
"""
import sys
import torch
import torch.nn.functional as F


def pytorch_sdpa_ref(q, k, v, is_causal=True):
    """Reference: standard scaled dot-product attention."""
    scale = q.shape[-1] ** -0.5
    attn = torch.matmul(q * scale, k.transpose(-2, -1))
    if is_causal:
        L = q.shape[-2]
        S = k.shape[-2]
        mask = torch.triu(torch.ones(L, S, device=q.device, dtype=torch.bool), diagonal=S-L+1)
        attn = attn.masked_fill(mask, float('-inf'))
    attn = torch.softmax(attn, dim=-1)
    return torch.matmul(attn, v)


def main():
    print("=" * 60)
    print("BI-V100 ixformer attention head_dim=256 test")
    print("=" * 60)

    if not torch.cuda.is_available():
        print("FATAL: No CUDA")
        return 1

    try:
        import ixformer
        print(f"ixformer available: True")
    except ImportError:
        print("ixformer not available")
        return 1

    # Check what's available
    has_sdpa = hasattr(ixformer, 'scaled_dot_product_attention')
    has_flash = hasattr(ixformer, 'flash_attn_func')
    has_varlen = hasattr(ixformer, 'flash_attn_varlen_func')
    has_paged = hasattr(ixformer, 'vllm_single_query_cached_kv_attention')
    print(f"scaled_dot_product_attention: {has_sdpa}")
    print(f"flash_attn_func: {has_flash}")
    print(f"flash_attn_varlen_func: {has_varlen}")
    print(f"vllm_single_query_cached_kv_attention: {has_paged}")

    # Qwen3.5 full attention dims (TP=4):
    # num_heads=4, num_kv_heads=1, head_dim=256
    batch = 1
    num_heads = 4
    num_kv_heads = 1
    head_dim = 256
    torch.manual_seed(42)

    # --- Test 1: scaled_dot_product_attention ---
    if has_sdpa:
        for seq_len in [1, 4, 16, 64, 128]:
            q = torch.randn(batch, num_heads, seq_len, head_dim,
                            device="cuda", dtype=torch.float16)
            # GQA: kv has fewer heads
            k = torch.randn(batch, num_kv_heads, seq_len, head_dim,
                            device="cuda", dtype=torch.float16)
            v = torch.randn(batch, num_kv_heads, seq_len, head_dim,
                            device="cuda", dtype=torch.float16)
            # Expand kv to match q heads for reference
            k_exp = k.expand(-1, num_heads, -1, -1)
            v_exp = v.expand(-1, num_heads, -1, -1)
            ref = pytorch_sdpa_ref(q, k_exp, v_exp, is_causal=(seq_len > 1))

            try:
                # Try without causal first
                out = ixformer.scaled_dot_product_attention(
                    q, k_exp, v_exp, is_causal=(seq_len > 1))
                diff = (out.float() - ref.float()).abs().max().item()
                has_nan = out.isnan().any().item()
                print(f"\n  SDPA seq={seq_len}: diff={diff:.6f} nan={has_nan} "
                      f"{'PASS' if diff < 0.01 and not has_nan else 'FAIL'}")
            except Exception as e:
                print(f"\n  SDPA seq={seq_len}: EXCEPTION: {e}")

    # --- Test 2: flash_attn_func ---
    if has_flash:
        for seq_len in [1, 4, 16, 64]:
            # flash_attn expects (batch, seqlen, nheads, headdim)
            q = torch.randn(batch, seq_len, num_heads, head_dim,
                            device="cuda", dtype=torch.float16)
            k = torch.randn(batch, seq_len, num_kv_heads, head_dim,
                            device="cuda", dtype=torch.float16)
            v = torch.randn(batch, seq_len, num_kv_heads, head_dim,
                            device="cuda", dtype=torch.float16)
            try:
                out = ixformer.flash_attn_func(q, k, v, causal=True)
                has_nan = out.isnan().any().item()
                print(f"  flash_attn seq={seq_len}: shape={out.shape} nan={has_nan}")
            except Exception as e:
                print(f"  flash_attn seq={seq_len}: EXCEPTION: {e}")

    # --- Test 3: head_dim=128 (known to work) vs head_dim=256 ---
    if has_sdpa:
        print("\n--- Comparison: head_dim=128 vs head_dim=256 ---")
        for hd in [128, 256]:
            q = torch.randn(1, 4, 16, hd, device="cuda", dtype=torch.float16)
            k = torch.randn(1, 4, 16, hd, device="cuda", dtype=torch.float16)
            v = torch.randn(1, 4, 16, hd, device="cuda", dtype=torch.float16)
            try:
                out = ixformer.scaled_dot_product_attention(q, k, v, is_causal=True)
                print(f"  head_dim={hd}: OK shape={out.shape} nan={out.isnan().any().item()}")
            except Exception as e:
                print(f"  head_dim={hd}: EXCEPTION: {e}")

    # --- Test 4: Performance if it works ---
    if has_sdpa:
        print("\n--- Performance: SDPA head_dim=256 seq=64 ---")
        q = torch.randn(1, 4, 64, 256, device="cuda", dtype=torch.float16)
        k = torch.randn(1, 4, 64, 256, device="cuda", dtype=torch.float16)
        v = torch.randn(1, 4, 64, 256, device="cuda", dtype=torch.float16)

        import time
        # Warmup
        try:
            for _ in range(5):
                ixformer.scaled_dot_product_attention(q, k, v, is_causal=True)
            torch.cuda.synchronize()

            N = 50
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            for _ in range(N):
                ixformer.scaled_dot_product_attention(q, k, v, is_causal=True)
            torch.cuda.synchronize()
            ix_ms = (time.perf_counter() - t0) / N * 1000

            torch.cuda.synchronize()
            t0 = time.perf_counter()
            for _ in range(N):
                pytorch_sdpa_ref(q, k, v, is_causal=True)
            torch.cuda.synchronize()
            pt_ms = (time.perf_counter() - t0) / N * 1000

            print(f"  ixformer: {ix_ms:.3f} ms")
            print(f"  PyTorch:  {pt_ms:.3f} ms")
            print(f"  Speedup:  {pt_ms/ix_ms:.2f}x")
        except Exception as e:
            print(f"  Performance test failed: {e}")

    print("\n" + "=" * 60)
    return 0


if __name__ == "__main__":
    sys.exit(main())