diff --git a/verify_ixformer_attn.py b/verify_ixformer_attn.py new file mode 100644 index 00000000..71129bd3 --- /dev/null +++ b/verify_ixformer_attn.py @@ -0,0 +1,162 @@ +#!/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())