File size: 4,930 Bytes
fd6abd3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import torch
import time
import snapkitty_flash_attention
import json
import csv
from typing import List, Dict, Tuple
class FlashAttentionBenchmark:
def __init__(self, device='cuda', dtype=torch.half):
self.device = device
self.dtype = dtype
self.results = []
def benchmark_config(self, batch, heads, seq_len, dim, num_warmup=10, num_iters=100, causal=False):
Q = torch.randn(batch, heads, seq_len, dim, dtype=self.dtype, device=self.device)
K = torch.randn(batch, heads, seq_len, dim, dtype=self.dtype, device=self.device)
V = torch.randn(batch, heads, seq_len, dim, dtype=self.dtype, device=self.device)
for _ in range(num_warmup):
_ = snapkitty_flash_attention.flash_attention_fwd(Q, K, V)
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(num_iters):
O = snapkitty_flash_attention.flash_attention_fwd(Q, K, V)
torch.cuda.synchronize()
fwd_time = (time.perf_counter() - start) / num_iters * 1000
O = snapkitty_flash_attention.flash_attention_fwd(Q, K, V)
loss = O.sum()
for _ in range(num_warmup):
loss.backward()
Q.grad.zero_(); K.grad.zero_(); V.grad.zero_()
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(num_iters):
loss.backward()
Q.grad.zero_(); K.grad.zero_(); V.grad.zero_()
torch.cuda.synchronize()
bwd_time = (time.perf_counter() - start) / num_iters * 1000
torch.cuda.reset_peak_memory_stats()
_ = snapkitty_flash_attention.flash_attention_fwd(Q, K, V)
torch.cuda.synchronize()
peak_mem = torch.cuda.max_memory_allocated() / 1024**2
flops = 2 * batch * heads * seq_len * seq_len * dim + 2 * batch * heads * seq_len * dim
tflops = flops / (fwd_time / 1000) / 1e12
result = {
'batch': batch, 'heads': heads, 'seq_len': seq_len, 'dim': dim,
'fwd_ms': fwd_time, 'bwd_ms': bwd_time, 'total_ms': fwd_time + bwd_time,
'peak_mem_mb': peak_mem, 'tflops': tflops, 'causal': causal
}
self.results.append(result)
return result
def run_sweep(self, configs):
for config in configs:
print(f"Benchmarking: B={config[0]}, H={config[1]}, N={config[2]}, D={config[3]}")
try:
result = self.benchmark_config(*config)
print(f" FWD: {result['fwd_ms']:.2f} ms, BWD: {result['bwd_ms']:.2f} ms, TFLOPs: {result['tflops']:.2f}")
except Exception as e:
print(f" FAILED: {e}")
return self.results
def save_results(self, filepath):
with open(filepath + '.json', 'w') as f:
json.dump(self.results, f, indent=2)
if self.results and 'error' not in str(self.results[0]):
with open(filepath + '.csv', 'w', newline='') as f:
writer = csv.DictWriter(f, fieldnames=self.results[0].keys())
writer.writeheader()
writer.writerows(self.results)
def print_summary(self):
print("\n" + "="*100)
print(f"{'Config':<30} {'FWD (ms)':>10} {'BWD (ms)':>10} {'Total (ms)':>12} {'TFLOPs':>8} {'Mem (MB)':>10}")
print("-"*100)
for r in self.results:
cfg = f"B{r['batch']}H{r['heads']}N{r['seq_len']}D{r['dim']}"
print(f"{cfg:<30} {r['fwd_ms']:>10.2f} {r['bwd_ms']:>10.2f} {r['total_ms']:>12.2f} {r['tflops']:>8.2f} {r['peak_mem_mb']:>10.0f}")
print("="*100)
def compare_with_pytorch():
print("\n=== Correctness vs PyTorch SDPA ===")
torch.manual_seed(42)
configs = [(2,8,512,64), (2,8,1024,64), (2,8,2048,64), (4,16,1024,64)]
for batch, heads, seq_len, dim in configs:
Q = torch.randn(batch, heads, seq_len, dim, dtype=torch.half, device='cuda')
K = torch.randn(batch, heads, seq_len, dim, dtype=torch.half, device='cuda')
V = torch.randn(batch, heads, seq_len, dim, dtype=torch.half, device='cuda')
O_fa = snapkitty_flash_attention.flash_attention(Q, K, V)
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
O_sdpa = torch.nn.functional.scaled_dot_product_attention(Q, K, V)
max_diff = (O_fa - O_sdpa).abs().max().item()
print(f"B={batch} H={heads} N={seq_len} D={dim}: max_diff={max_diff:.6f}")
assert max_diff < 1e-3, f"Numerical mismatch: {max_diff}"
print("All correctness tests passed!")
if __name__ == '__main__':
bench = FlashAttentionBenchmark()
configs = [
(1,32,512,64), (1,32,1024,64), (1,32,2048,64), (1,32,4096,64),
(2,16,1024,64), (4,8,1024,64), (4,32,1024,64), (8,32,1024,64),
]
bench.run_sweep(configs)
bench.print_summary()
bench.save_results('benchmark_results')
compare_with_pytorch() |