| 26 | return attn @ value |
| 27 | |
| 28 | def test_standard_attention(seq_len, hidden_size, num_heads, batch_size=8, qkv=None): |
| 29 | # set random seed |
| 30 | q, k, v = qkv.clone().cuda().transpose(2,3).requires_grad_(True) |
| 31 | # clean memory record |
| 32 | torch.cuda.reset_peak_memory_stats() |
| 33 | # profile time and memory |
| 34 | with torch.autograd.profiler.profile(use_cuda=True, profile_memory=True) as prof: |
| 35 | out = attention(q, k, v) |
| 36 | out.sum().backward() |
| 37 | print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10)) |
| 38 | # peak memory usage |
| 39 | print(torch.cuda.max_memory_allocated() / 1024 ** 3) |
| 40 | |
| 41 | # clean cuda cache |
| 42 | del q, k, v |
| 43 | torch.cuda.empty_cache() |
| 44 | return out.transpose(1,2) |
| 45 | |
| 46 | def test_mixin(): |
| 47 | with torch.no_grad(): |