test-perf

stack-info: PR: #55, branch: drisspg/stack/14

Author: drisspg

transformer_nuggets/cute/fav4.py

@@ -0,0 +1,233 @@
+#!/usr/bin/env python3
+"""
+Benchmark comparing PyTorch native attention vs CUTE implementation.
+This benchmark has no dependencies on other flash-attn components.
+"""
+
+import math
+import time
+import torch
+import torch.nn.functional as F
+from typing import Tuple, Optional
+
+try:
+    from flash_attn.cute import flash_attn_func
+    CUTE_AVAILABLE = True
+except ImportError:
+    CUTE_AVAILABLE = False
+    print("Warning: CUTE implementation not available. Install with: pip install nvidia-cutlass-dsl>=4.1.0.dev0")
+
+
+def attention_pytorch(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    causal: bool = False,
+    softmax_scale: Optional[float] = None
+) -> torch.Tensor:
+    """
+    Standard PyTorch attention implementation.
+
+    Args:
+        q: Query tensor (batch, seqlen, nheads, headdim)
+        k: Key tensor (batch, seqlen, nheads_kv, headdim)
+        v: Value tensor (batch, seqlen, nheads_kv, headdim)
+        causal: Whether to apply causal masking
+        softmax_scale: Scale factor for attention scores
+
+    Returns:
+        output: (batch, seqlen, nheads, headdim)
+    """
+    batch_size, seqlen_q, nheads, headdim = q.shape
+    _, seqlen_k, nheads_kv, _ = k.shape
+
+    if softmax_scale is None:
+        softmax_scale = 1.0 / math.sqrt(headdim)
+
+    # Handle grouped query attention
+    if nheads != nheads_kv:
+        # Repeat k and v to match number of query heads
+        repeat_factor = nheads // nheads_kv
+        k = k.repeat_interleave(repeat_factor, dim=2)
+        v = v.repeat_interleave(repeat_factor, dim=2)
+
+    # Compute attention scores
+    scores = torch.einsum('bqhd,bkhd->bhqk', q, k) * softmax_scale
+
+    if causal:
+        causal_mask = torch.triu(
+            torch.full((seqlen_q, seqlen_k), float('-inf'), device=q.device, dtype=q.dtype),
+            diagonal=1
+        )
+        scores = scores + causal_mask
+
+    # Apply softmax
+    attn_weights = torch.softmax(scores, dim=-1)
+
+    # Apply attention to values
+    output = torch.einsum('bhqk,bkhd->bqhd', attn_weights, v)
+
+    return output
+
+
+def benchmark_function(func, *args, **kwargs):
+    """Benchmark a function with warmup and multiple runs."""
+    device = args[0].device
+
+    # Warmup
+    for _ in range(5):
+        with torch.no_grad():
+            _ = func(*args, **kwargs)
+        torch.cuda.synchronize(device)
+
+    # Benchmark
+    torch.cuda.synchronize(device)
+    start_time = time.time()
+
+    for _ in range(10):
+        with torch.no_grad():
+            result = func(*args, **kwargs)
+        torch.cuda.synchronize(device)
+
+    end_time = time.time()
+    avg_time = (end_time - start_time) / 10
+
+    return avg_time, result
+
+
+def compute_flops(batch_size: int, seqlen: int, nheads: int, headdim: int, causal: bool = False) -> int:
+    """Compute theoretical FLOPs for attention."""
+    # QK^T: batch_size * nheads * seqlen * seqlen * headdim
+    # Softmax: batch_size * nheads * seqlen * seqlen
+    # AV: batch_size * nheads * seqlen * seqlen * headdim
+    qk_flops = batch_size * nheads * seqlen * seqlen * headdim
+    softmax_flops = batch_size * nheads * seqlen * seqlen * 5  # approx for softmax
+    av_flops = batch_size * nheads * seqlen * seqlen * headdim
+
+    total_flops = qk_flops + softmax_flops + av_flops
+
+    # Causal attention reduces computation by ~half
+    if causal:
+        total_flops = total_flops // 2
+
+    return total_flops
+
+
+def run_benchmark():
+    """Run the benchmark comparing PyTorch vs CUTE."""
+
+    if not CUTE_AVAILABLE:
+        print("CUTE implementation not available. Exiting.")
+        return
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    if device == 'cpu':
+        print("CUDA not available. This benchmark requires GPU.")
+        return
+
+    dtype = torch.float16
+
+    # Test configurations
+    configs = [
+        # (batch_size, seqlen, nheads, headdim, causal)
+        (2, 512, 16, 64, False),
+        (2, 512, 16, 64, True),
+        (1, 1024, 16, 64, False),
+        (1, 1024, 16, 64, True),
+        (1, 2048, 16, 64, False),
+        (1, 2048, 16, 64, True),
+        (2, 512, 16, 128, False),
+        (2, 512, 16, 128, True),
+        (1, 1024, 16, 128, False),
+        (1, 1024, 16, 128, True),
+    ]
+
+    print("=" * 80)
+    print("PyTorch vs CUTE Flash Attention Benchmark")
+    print("=" * 80)
+    print(f"Device: {device}")
+    print(f"Dtype: {dtype}")
+    print()
+
+    results = []
+
+    for batch_size, seqlen, nheads, headdim, causal in configs:
+        print(f"Config: batch={batch_size}, seqlen={seqlen}, nheads={nheads}, headdim={headdim}, causal={causal}")
+
+        # Generate random inputs
+        q = torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype)
+        k = torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype)
+        v = torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype)
+
+        try:
+            # Benchmark PyTorch
+            pytorch_time, pytorch_output = benchmark_function(
+                attention_pytorch, q, k, v, causal=causal
+            )
+
+            # Benchmark CUTE
+            cute_time, (cute_output, _) = benchmark_function(
+                flash_attn_func, q, k, v, causal=causal
+            )
+
+            # Compute metrics
+            flops = compute_flops(batch_size, seqlen, nheads, headdim, causal)
+            pytorch_tflops = (flops / pytorch_time) / 1e12
+            cute_tflops = (flops / cute_time) / 1e12
+            speedup = pytorch_time / cute_time
+
+            # Check correctness (with some tolerance for numerical differences)
+            max_diff = torch.max(torch.abs(pytorch_output - cute_output)).item()
+            mean_diff = torch.mean(torch.abs(pytorch_output - cute_output)).item()
+
+            print(f"  PyTorch: {pytorch_time*1000:.2f}ms ({pytorch_tflops:.2f} TFLOPs/s)")
+            print(f"  CUTE:    {cute_time*1000:.2f}ms ({cute_tflops:.2f} TFLOPs/s)")
+            print(f"  Speedup: {speedup:.2f}x")
+            print(f"  Max diff: {max_diff:.2e}, Mean diff: {mean_diff:.2e}")
+
+            results.append({
+                'config': (batch_size, seqlen, nheads, headdim, causal),
+                'pytorch_time': pytorch_time,
+                'cute_time': cute_time,
+                'pytorch_tflops': pytorch_tflops,
+                'cute_tflops': cute_tflops,
+                'speedup': speedup,
+                'max_diff': max_diff,
+                'mean_diff': mean_diff
+            })
+
+        except Exception as e:
+            print(f"  Error: {e}")
+            results.append({
+                'config': (batch_size, seqlen, nheads, headdim, causal),
+                'error': str(e)
+            })
+
+        print()
+
+    # Summary
+    print("=" * 80)
+    print("SUMMARY")
+    print("=" * 80)
+
+    valid_results = [r for r in results if 'error' not in r]
+    if valid_results:
+        avg_speedup = sum(r['speedup'] for r in valid_results) / len(valid_results)
+        avg_cute_tflops = sum(r['cute_tflops'] for r in valid_results) / len(valid_results)
+        avg_pytorch_tflops = sum(r['pytorch_tflops'] for r in valid_results) / len(valid_results)
+
+        print(f"Average speedup: {avg_speedup:.2f}x")
+        print(f"Average CUTE performance: {avg_cute_tflops:.2f} TFLOPs/s")
+        print(f"Average PyTorch performance: {avg_pytorch_tflops:.2f} TFLOPs/s")
+
+        max_errors = [r['max_diff'] for r in valid_results]
+        print(f"Max numerical difference: {max(max_errors):.2e}")
+        print(f"Mean numerical difference: {sum(r['mean_diff'] for r in valid_results) / len(valid_results):.2e}")
+
+    error_count = len([r for r in results if 'error' in r])
+    if error_count > 0:
+        print(f"Failed configurations: {error_count}/{len(results)}")
+
+
+if __name__ == "__main__":
+    run_benchmark()