File size: 5,931 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 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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import os
import torch
import triton
from flash_attn import flash_attn_func
from torch.nn import functional as F
from fla.ops.comba import chunk_comba
from fla.ops.gated_delta_rule import chunk_gated_delta_rule
from fla.ops.generalized_delta_rule import chunk_dplr_delta_rule
from fla.ops.kda import chunk_kda
@triton.testing.perf_report(
triton.testing.Benchmark(
# argument names to use as an x-axis for the plot
x_names=['T'],
# different possible values for `x_name`
x_vals=[256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536],
# argument name whose value corresponds to a different line in the plot
line_arg='provider',
# possible values for `line_arg``
line_vals=['gdn', 'comba', 'kda', 'dplr', 'attn'],
# label name for the lines
line_names=['gdn', 'comba', 'kda', 'dplr', 'attn'],
# line styles
styles=[('blue', '-'), ('red', '-.'), ('green', '-'), ('orange', '-.'),
('purple', '-'), ('brown', '-.'), ('pink', '-'), ('gray', '-.')],
ylabel="Execution Time (ms)", # label name for the y-axis
# name for the plot. Used also as a file name for saving the plot.
plot_name="Performance",
args={},
),
)
def benchmark(T, provider):
from fla.utils import device
dtype = torch.bfloat16
B, H, D = 1, 16, 128
# Set TMA environment variable based on provider
original_tma_env = os.environ.get('FLA_USE_TMA', '0')
if provider.endswith('_no_tma'):
os.environ['FLA_USE_TMA'] = '0'
provider_base = provider.replace('_no_tma', '')
else:
os.environ['FLA_USE_TMA'] = '1'
provider_base = provider
quantiles = [0.5, 0.2, 0.8]
results = 0, 0, 0
do = torch.randn(B, T, H, D, dtype=dtype, device=device)
if provider_base == 'gdn':
q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
g = F.logsigmoid(torch.randn(B, T, H, dtype=dtype, device=device)).requires_grad_(True)
beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True)
results = triton.testing.do_bench(
lambda: chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
use_qk_l2norm_in_kernel=True,
)[0].backward(do),
quantiles=quantiles,
)
elif provider_base == 'attn':
q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
results = triton.testing.do_bench(
lambda: flash_attn_func(
q=q,
k=k,
v=v,
).backward(do),
quantiles=quantiles,
)
elif provider_base == 'comba':
q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
p = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
g = F.logsigmoid(torch.randn(B, T, H, dtype=torch.float, device=device)).requires_grad_(True)
beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True)
results = triton.testing.do_bench(
lambda: chunk_comba(
q=q,
k=k,
p=p,
v=v,
g=g,
beta=beta,
use_qk_l2norm_in_kernel=True,
)[0].backward(do),
quantiles=quantiles,
)
elif provider_base == 'kda':
q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
g = F.logsigmoid(torch.randn(B, T, H, D, dtype=dtype, device=device)).requires_grad_(True)
beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True)
results = triton.testing.do_bench(
lambda: chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
use_qk_l2norm_in_kernel=True,
)[0].backward(do),
quantiles=quantiles,
)
elif provider_base == 'dplr':
q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
a = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
b = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True)
g = F.logsigmoid(torch.randn(B, T, H, D, dtype=dtype, device=device)).requires_grad_(True)
beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True)
results = triton.testing.do_bench(
lambda: chunk_dplr_delta_rule(
q=q,
k=k,
v=v,
a=a,
b=b,
gk=g,
)[0].backward(do),
quantiles=quantiles,
)
# Restore original TMA environment variable
os.environ['FLA_USE_TMA'] = original_tma_env
return results
if __name__ == '__main__':
benchmark.run(print_data=True)
|