|
4 | 4 | import torch |
5 | 5 | from torch.nn import functional as F |
6 | 6 | from torch.utils.cpp_extension import load |
| 7 | +from functools import partial |
| 8 | +from typing import Optional |
7 | 9 |
|
8 | 10 | torch.set_grad_enabled(False) |
9 | 11 | # Load the CUDA kernel as a python module |
10 | | -custom_flash_attn = load(name='custom_flash_attn', |
11 | | - sources=[ |
12 | | - 'flash_attn.cc', |
13 | | - 'flash_attn_1_fwd_f32.cu', |
14 | | - 'flash_attn_2_fwd_f32.cu' |
15 | | - ], |
16 | | - extra_cuda_cflags=['-O2']) |
| 12 | +lib = load(name='flash_attn_lib', |
| 13 | + sources=['flash_attn_1_fwd_f32.cu'], |
| 14 | + extra_cuda_cflags=[ |
| 15 | + "-O3", |
| 16 | + "-U__CUDA_NO_HALF_OPERATORS__", |
| 17 | + "-U__CUDA_NO_HALF_CONVERSIONS__", |
| 18 | + "-U__CUDA_NO_HALF2_OPERATORS__", |
| 19 | + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", |
| 20 | + "--expt-relaxed-constexpr", |
| 21 | + "--expt-extended-lambda", |
| 22 | + "--use_fast_math" |
| 23 | + ], |
| 24 | + extra_cflags=['-std=c++17']) |
17 | 25 |
|
18 | 26 | # Use small model params, otherwise slower than manual attention. See caveats in README. |
19 | | -batch_size = 16 |
20 | | -n_head = 12 |
21 | | -seq_len = 64 |
22 | | -head_embd = 64 |
23 | 27 |
|
24 | | -q = torch.randn(batch_size, n_head, seq_len, head_embd).float().cuda() |
25 | | -k = torch.randn(batch_size, n_head, seq_len, head_embd).float().cuda() |
26 | | -v = torch.randn(batch_size, n_head, seq_len, head_embd).float().cuda() |
27 | | -q.requires_grad = False |
28 | | -k.requires_grad = False |
29 | | -v.requires_grad = False |
30 | | -print('=== profiling manual attention ===') |
31 | | - |
32 | | -def manual_attn(q, k, v): |
| 28 | +# un-fused naive attn |
| 29 | +def manual_attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): |
33 | 30 | att = (q @ k.transpose(-2, -1) * (1.0 / math.sqrt(k.size(-1)))) |
34 | 31 | att = F.softmax(att, dim=-1) |
35 | 32 | y = att @ v |
36 | 33 | return y |
37 | 34 |
|
38 | | -for _ in range(2): |
39 | | - manual_result = manual_attn(q, k, v) # warmup |
40 | | - |
41 | | -torch.cuda.synchronize() |
42 | | -with torch.autograd.profiler.profile(use_cuda=True, with_flops=True) as prof: |
43 | | - with torch.autograd.profiler.record_function("manual_attn"): |
44 | | - manual_result = manual_attn(q, k, v) |
45 | | -print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=10)) |
46 | 35 |
|
47 | | -for _ in range(2): |
48 | | - custom_result = custom_flash_attn.flash_attn_1_fwd_f32(q, k, v) # warmup |
49 | | -print('=== profiling flash_attn_1_fwd_f32 attention === ') |
50 | | -with torch.autograd.profiler.profile(use_cuda=True, with_flops=True) as prof: |
51 | | - with torch.autograd.profiler.record_function("flash_attn_1_fwd_f32"): |
52 | | - custom_result = custom_flash_attn.flash_attn_1_fwd_f32(q, k, v) |
53 | | -print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=10)) |
54 | | -print('attn values sanity check:', torch.allclose(custom_result, manual_result, rtol=0, atol=1e-02)) |
55 | | - |
56 | | -# Why custom flash attn is slow than naive attn in for loop test ? |
57 | | -REPEAT = 10 |
58 | | -manual_result = manual_attn(q, k, v) # warmup |
59 | | -st = time.time() |
60 | | -for _ in range(REPEAT): |
61 | | - manual_result = manual_attn(q, k, v) |
| 36 | +def run_benchmark(perf_func: callable, |
| 37 | + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, |
| 38 | + tag: str, out: Optional[torch.Tensor] = None, |
| 39 | + warmup: int = 10, iters: int = 200, |
| 40 | + show_all: bool = False): |
| 41 | + if out is not None: |
| 42 | + out.fill_(0) |
| 43 | + if out is not None: |
| 44 | + for i in range(warmup): |
| 45 | + perf_func(q, k, v, out) |
| 46 | + else: |
| 47 | + for i in range(warmup): |
| 48 | + _ = perf_func(q, k, v) |
| 49 | + |
62 | 50 | torch.cuda.synchronize() |
63 | | -print(f"manual attention mean time(ms): {((time.time() - st) * 1000) / REPEAT}") |
64 | | -custom_result = custom_flash_attn.flash_attn_1_fwd_f32(q, k, v) # warmup |
65 | | -st = time.time() |
66 | | -for _ in range(REPEAT): |
67 | | - custom_result = custom_flash_attn.flash_attn_1_fwd_f32(q, k, v) |
| 51 | + start = time.time() |
| 52 | + # iters |
| 53 | + if out is not None: |
| 54 | + for i in range(iters): |
| 55 | + perf_func(q, k, v, out) |
| 56 | + else: |
| 57 | + for i in range(iters): |
| 58 | + out = perf_func(q, k, v) |
68 | 59 | torch.cuda.synchronize() |
69 | | -print(f"flash_attn_1_fwd_f32 mean time(ms): {((time.time() - st) * 1000) / REPEAT}") |
| 60 | + end = time.time() |
| 61 | + total_time = (end - start) * 1000 # ms |
| 62 | + mean_time = total_time / iters |
| 63 | + out_info = f"out_{tag}" |
| 64 | + out_val = out.flatten().detach().cpu().numpy().tolist()[:3] |
| 65 | + out_val = [round(v, 8) for v in out_val] |
| 66 | + print(f"{out_info:>17}: {out_val}, time:{mean_time:.8f}ms") |
| 67 | + if show_all: print(out[0, 0, 0, :]) |
| 68 | + return out.clone(), mean_time |
| 69 | + |
70 | 70 |
|
| 71 | +print("-" * 80) |
| 72 | +# batch_size, n_head, seq_len, head_dim (B,nh,N,d) |
| 73 | +B, nh, N, d = 16, 12, 64, 64 |
| 74 | +q = torch.randn(B, nh, N, d).float().cuda().contiguous() |
| 75 | +k = torch.randn(B, nh, N, d).float().cuda().contiguous() |
| 76 | +v = torch.randn(B, nh, N, d).float().cuda().contiguous() |
| 77 | +o = torch.randn(B, nh, N, d).float().cuda().contiguous() |
| 78 | +q.requires_grad = False |
| 79 | +k.requires_grad = False |
| 80 | +v.requires_grad = False |
| 81 | +o.requires_grad = False |
| 82 | +run_benchmark(lib.flash_attn_1_fwd_f32, q, k, v, "fa1fwdf32") |
| 83 | +run_benchmark(lib.flash_attn_1_fwd_f32_v2, q, k, v, "fa1fwdf32(v2)", o) |
| 84 | +run_benchmark(manual_attn, q, k, v, "attnf32_th") |
| 85 | +print("-" * 80) |
0 commit comments