File size: 1,576 Bytes
eafbe80 | 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 |
import torch
import triton
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=['B', 'T', 'H', 'chunk_size'],
x_vals=[
(b, t, h, c)
for b in [8]
for t in [2048, 4096, 8192]
for h in [16, 64]
for c in [16, 32, 64]
],
line_arg='provider',
line_vals=[
'solve_tril_tma',
],
line_names=[
'solve_tril_tma',
],
styles=[('green', '-'), ('green', '--')],
ylabel="Time (ms)",
plot_name="solve_tril_performance",
args={},
),
)
def benchmark(B, T, H, chunk_size, provider):
from fla.ops.utils.solve_tril import solve_tril
from fla.utils import device
requires_grad = True
dtype = torch.float32
k = torch.randn((B, H, T, 64), dtype=dtype, device=device, requires_grad=requires_grad)
k = torch.nn.functional.normalize(k, dim=-1)
padding_size = (chunk_size - T % chunk_size) % chunk_size
T_padded = T + padding_size
k_padded = torch.nn.functional.pad(k, (0, 0, 0, padding_size, 0, 0, 0, 0))
k_padded = k_padded.reshape(B, H, T_padded // chunk_size, chunk_size, 64)
A = (k_padded @ k_padded.transpose(-1, -2)).tril(-1)
A = A.permute(0, 2, 1, 3, 4).contiguous()
A = A.view(B, T_padded, H, chunk_size)
A = A[:, :T, :, :]
results = triton.testing.do_bench(
lambda: solve_tril(A),
quantiles=[0.5, 0.2, 0.8],
)
return results
if __name__ == '__main__':
benchmark.run(print_data=True, save_path=".")
|