Skip to content

Prepare impl. #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aoxy wants to merge 23 commits into
base: main
Choose a base branch
from
Open

Prepare impl. #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
2ff2a9d
Prepare impl.
Dec 20, 2025
f479d90
Right for some cases.
aoxy Jan 16, 2026
e9a60cb
First bwd.
aoxy Jan 19, 2026
1ecd4be
Useless optimization.
aoxy Jan 20, 2026
284d2d7
Another useless optimization.
aoxy Jan 20, 2026
4bbedfd
Clean code.
aoxy Jan 20, 2026
d1ea663
Update test.
aoxy Jan 20, 2026
8ddd284
Support bf16 sink.
aoxy Jan 20, 2026
42d1c15
Fix num_splits_dynamic.
aoxy Jan 21, 2026
bd815d6
Fix again.
aoxy Jan 21, 2026
14d3d19
Remove log.
aoxy Jan 21, 2026
63a1931
Update scripts.
aoxy Jan 22, 2026
4f31c27
Fix.
aoxy Jan 22, 2026
1c0dd9c
reduce chance of build oom (#2079)
Qubitium Jan 21, 2026
9ff82f4
Fix bf16.
aoxy Jan 28, 2026
0fc53d8
Fix bugs.
aoxy Jan 29, 2026
a17af5c
Fix bugs.
aoxy Jan 29, 2026
91624e5
Add use_fa4_sink option.
aoxy May 6, 2026
b6cb595
Add test.
aoxy May 7, 2026
05f9cfe
Update test.
aoxy May 8, 2026
0aa05b6
Use_fa4_sink to template.
aoxy May 22, 2026
58ecfef
fix v1.
aoxy May 23, 2026
37d3734
Add test.
aoxy May 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 72 additions & 0 deletions benchmarks/benchmark_flash_attention_install_res.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 259.99 TFLOPs/s, bwd: 197.78 TFLOPs/s, fwd + bwd: 212.29 TFLOPs/s
Pytorch fwd: 51.75 TFLOPs/s, bwd: 68.31 TFLOPs/s, fwd + bwd: 62.59 TFLOPs/s
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 296.25 TFLOPs/s, bwd: 243.13 TFLOPs/s, fwd + bwd: 256.26 TFLOPs/s
Pytorch fwd: 62.75 TFLOPs/s, bwd: 76.88 TFLOPs/s, fwd + bwd: 72.23 TFLOPs/s
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 313.91 TFLOPs/s, bwd: 270.73 TFLOPs/s, fwd + bwd: 281.81 TFLOPs/s
Pytorch fwd: 60.32 TFLOPs/s, bwd: 88.63 TFLOPs/s, fwd + bwd: 78.15 TFLOPs/s
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 316.63 TFLOPs/s, bwd: 284.33 TFLOPs/s, fwd + bwd: 292.86 TFLOPs/s
Pytorch fwd: 53.34 TFLOPs/s, bwd: 93.26 TFLOPs/s, fwd + bwd: 76.83 TFLOPs/s
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 319.05 TFLOPs/s, bwd: 294.69 TFLOPs/s, fwd + bwd: 301.26 TFLOPs/s
Pytorch fwd: 46.75 TFLOPs/s, bwd: 95.30 TFLOPs/s, fwd + bwd: 73.49 TFLOPs/s
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 314.26 TFLOPs/s, bwd: 299.77 TFLOPs/s, fwd + bwd: 303.77 TFLOPs/s
Pytorch fwd: 79.55 TFLOPs/s, bwd: 96.19 TFLOPs/s, fwd + bwd: 90.76 TFLOPs/s
### causal=False, headdim=128, batch_size=32, seqlen=512 ###
Flash2 fwd: 307.11 TFLOPs/s, bwd: 200.12 TFLOPs/s, fwd + bwd: 222.24 TFLOPs/s
Pytorch fwd: 74.11 TFLOPs/s, bwd: 105.25 TFLOPs/s, fwd + bwd: 93.97 TFLOPs/s
### causal=False, headdim=128, batch_size=16, seqlen=1024 ###
Flash2 fwd: 349.17 TFLOPs/s, bwd: 244.76 TFLOPs/s, fwd + bwd: 267.62 TFLOPs/s
Pytorch fwd: 100.38 TFLOPs/s, bwd: 130.96 TFLOPs/s, fwd + bwd: 120.47 TFLOPs/s
### causal=False, headdim=128, batch_size=8, seqlen=2048 ###
Flash2 fwd: 369.36 TFLOPs/s, bwd: 273.33 TFLOPs/s, fwd + bwd: 295.26 TFLOPs/s
Pytorch fwd: 107.31 TFLOPs/s, bwd: 159.98 TFLOPs/s, fwd + bwd: 140.31 TFLOPs/s
### causal=False, headdim=128, batch_size=4, seqlen=4096 ###
Flash2 fwd: 376.18 TFLOPs/s, bwd: 287.64 TFLOPs/s, fwd + bwd: 308.38 TFLOPs/s
Pytorch fwd: 100.58 TFLOPs/s, bwd: 175.80 TFLOPs/s, fwd + bwd: 144.85 TFLOPs/s
### causal=False, headdim=128, batch_size=2, seqlen=8192 ###
Flash2 fwd: 374.07 TFLOPs/s, bwd: 295.46 TFLOPs/s, fwd + bwd: 314.33 TFLOPs/s
Pytorch fwd: 90.88 TFLOPs/s, bwd: 184.54 TFLOPs/s, fwd + bwd: 142.56 TFLOPs/s
### causal=False, headdim=128, batch_size=1, seqlen=16384 ###
Flash2 fwd: 374.84 TFLOPs/s, bwd: 300.44 TFLOPs/s, fwd + bwd: 318.50 TFLOPs/s
Pytorch fwd: 155.80 TFLOPs/s, bwd: 188.79 TFLOPs/s, fwd + bwd: 178.02 TFLOPs/s
### causal=True, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 177.64 TFLOPs/s, bwd: 131.56 TFLOPs/s, fwd + bwd: 142.09 TFLOPs/s
Pytorch fwd: 16.06 TFLOPs/s, bwd: 34.13 TFLOPs/s, fwd + bwd: 25.83 TFLOPs/s
### causal=True, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 228.44 TFLOPs/s, bwd: 183.14 TFLOPs/s, fwd + bwd: 194.14 TFLOPs/s
Pytorch fwd: 18.16 TFLOPs/s, bwd: 38.52 TFLOPs/s, fwd + bwd: 29.18 TFLOPs/s
### causal=True, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 263.28 TFLOPs/s, bwd: 228.36 TFLOPs/s, fwd + bwd: 237.35 TFLOPs/s
Pytorch fwd: 17.65 TFLOPs/s, bwd: 44.29 TFLOPs/s, fwd + bwd: 30.94 TFLOPs/s
### causal=True, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 281.96 TFLOPs/s, bwd: 257.20 TFLOPs/s, fwd + bwd: 263.82 TFLOPs/s
Pytorch fwd: 15.48 TFLOPs/s, bwd: 46.60 TFLOPs/s, fwd + bwd: 29.60 TFLOPs/s
### causal=True, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 292.99 TFLOPs/s, bwd: 275.96 TFLOPs/s, fwd + bwd: 280.62 TFLOPs/s
Pytorch fwd: 14.19 TFLOPs/s, bwd: 47.69 TFLOPs/s, fwd + bwd: 28.48 TFLOPs/s
### causal=True, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 294.80 TFLOPs/s, bwd: 294.13 TFLOPs/s, fwd + bwd: 294.32 TFLOPs/s
Pytorch fwd: 18.97 TFLOPs/s, bwd: 48.12 TFLOPs/s, fwd + bwd: 33.44 TFLOPs/s
### causal=True, headdim=128, batch_size=32, seqlen=512 ###
Flash2 fwd: 203.80 TFLOPs/s, bwd: 145.04 TFLOPs/s, fwd + bwd: 158.06 TFLOPs/s
Pytorch fwd: 25.72 TFLOPs/s, bwd: 52.55 TFLOPs/s, fwd + bwd: 40.49 TFLOPs/s
### causal=True, headdim=128, batch_size=16, seqlen=1024 ###
Flash2 fwd: 262.21 TFLOPs/s, bwd: 199.92 TFLOPs/s, fwd + bwd: 214.48 TFLOPs/s
Pytorch fwd: 31.55 TFLOPs/s, bwd: 65.48 TFLOPs/s, fwd + bwd: 50.09 TFLOPs/s
### causal=True, headdim=128, batch_size=8, seqlen=2048 ###
Flash2 fwd: 302.12 TFLOPs/s, bwd: 245.33 TFLOPs/s, fwd + bwd: 259.25 TFLOPs/s
Pytorch fwd: 32.76 TFLOPs/s, bwd: 80.04 TFLOPs/s, fwd + bwd: 56.67 TFLOPs/s
### causal=True, headdim=128, batch_size=4, seqlen=4096 ###
Flash2 fwd: 328.82 TFLOPs/s, bwd: 277.86 TFLOPs/s, fwd + bwd: 290.73 TFLOPs/s
Pytorch fwd: 29.79 TFLOPs/s, bwd: 87.92 TFLOPs/s, fwd + bwd: 56.45 TFLOPs/s
### causal=True, headdim=128, batch_size=2, seqlen=8192 ###
Flash2 fwd: 339.26 TFLOPs/s, bwd: 293.81 TFLOPs/s, fwd + bwd: 305.50 TFLOPs/s
Pytorch fwd: 27.65 TFLOPs/s, bwd: 92.30 TFLOPs/s, fwd + bwd: 55.33 TFLOPs/s
### causal=True, headdim=128, batch_size=1, seqlen=16384 ###
Flash2 fwd: 346.64 TFLOPs/s, bwd: 323.29 TFLOPs/s, fwd + bwd: 329.63 TFLOPs/s
Pytorch fwd: 36.81 TFLOPs/s, bwd: 94.13 TFLOPs/s, fwd + bwd: 65.14 TFLOPs/s
72 changes: 72 additions & 0 deletions benchmarks/benchmark_flash_attention_res.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 260.69 TFLOPs/s, bwd: 197.79 TFLOPs/s, fwd + bwd: 212.44 TFLOPs/s
Pytorch fwd: 51.82 TFLOPs/s, bwd: 68.35 TFLOPs/s, fwd + bwd: 62.64 TFLOPs/s
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 296.89 TFLOPs/s, bwd: 243.43 TFLOPs/s, fwd + bwd: 256.63 TFLOPs/s
Pytorch fwd: 62.77 TFLOPs/s, bwd: 76.92 TFLOPs/s, fwd + bwd: 72.26 TFLOPs/s
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 314.14 TFLOPs/s, bwd: 271.00 TFLOPs/s, fwd + bwd: 282.07 TFLOPs/s
Pytorch fwd: 60.35 TFLOPs/s, bwd: 88.63 TFLOPs/s, fwd + bwd: 78.17 TFLOPs/s
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 319.20 TFLOPs/s, bwd: 285.11 TFLOPs/s, fwd + bwd: 294.08 TFLOPs/s
Pytorch fwd: 53.36 TFLOPs/s, bwd: 93.30 TFLOPs/s, fwd + bwd: 76.86 TFLOPs/s
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 315.76 TFLOPs/s, bwd: 294.74 TFLOPs/s, fwd + bwd: 300.46 TFLOPs/s
Pytorch fwd: 46.76 TFLOPs/s, bwd: 95.32 TFLOPs/s, fwd + bwd: 73.51 TFLOPs/s
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 316.32 TFLOPs/s, bwd: 299.80 TFLOPs/s, fwd + bwd: 304.35 TFLOPs/s
Pytorch fwd: 79.00 TFLOPs/s, bwd: 96.19 TFLOPs/s, fwd + bwd: 90.56 TFLOPs/s
### causal=False, headdim=128, batch_size=32, seqlen=512 ###
Flash2 fwd: 308.31 TFLOPs/s, bwd: 200.43 TFLOPs/s, fwd + bwd: 222.69 TFLOPs/s
Pytorch fwd: 74.27 TFLOPs/s, bwd: 105.22 TFLOPs/s, fwd + bwd: 94.02 TFLOPs/s
### causal=False, headdim=128, batch_size=16, seqlen=1024 ###
Flash2 fwd: 349.57 TFLOPs/s, bwd: 244.91 TFLOPs/s, fwd + bwd: 267.82 TFLOPs/s
Pytorch fwd: 100.55 TFLOPs/s, bwd: 131.12 TFLOPs/s, fwd + bwd: 120.64 TFLOPs/s
### causal=False, headdim=128, batch_size=8, seqlen=2048 ###
Flash2 fwd: 369.68 TFLOPs/s, bwd: 273.63 TFLOPs/s, fwd + bwd: 295.57 TFLOPs/s
Pytorch fwd: 107.39 TFLOPs/s, bwd: 160.01 TFLOPs/s, fwd + bwd: 140.36 TFLOPs/s
### causal=False, headdim=128, batch_size=4, seqlen=4096 ###
Flash2 fwd: 377.16 TFLOPs/s, bwd: 289.45 TFLOPs/s, fwd + bwd: 310.06 TFLOPs/s
Pytorch fwd: 100.63 TFLOPs/s, bwd: 175.76 TFLOPs/s, fwd + bwd: 144.86 TFLOPs/s
### causal=False, headdim=128, batch_size=2, seqlen=8192 ###
Flash2 fwd: 374.11 TFLOPs/s, bwd: 295.43 TFLOPs/s, fwd + bwd: 314.31 TFLOPs/s
Pytorch fwd: 90.94 TFLOPs/s, bwd: 184.66 TFLOPs/s, fwd + bwd: 142.65 TFLOPs/s
### causal=False, headdim=128, batch_size=1, seqlen=16384 ###
Flash2 fwd: 374.57 TFLOPs/s, bwd: 300.65 TFLOPs/s, fwd + bwd: 318.62 TFLOPs/s
Pytorch fwd: 156.04 TFLOPs/s, bwd: 188.74 TFLOPs/s, fwd + bwd: 178.08 TFLOPs/s
### causal=True, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 178.98 TFLOPs/s, bwd: 132.40 TFLOPs/s, fwd + bwd: 143.03 TFLOPs/s
Pytorch fwd: 16.10 TFLOPs/s, bwd: 34.19 TFLOPs/s, fwd + bwd: 25.88 TFLOPs/s
### causal=True, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 229.35 TFLOPs/s, bwd: 184.20 TFLOPs/s, fwd + bwd: 195.18 TFLOPs/s
Pytorch fwd: 18.17 TFLOPs/s, bwd: 38.52 TFLOPs/s, fwd + bwd: 29.18 TFLOPs/s
### causal=True, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 263.28 TFLOPs/s, bwd: 228.01 TFLOPs/s, fwd + bwd: 237.08 TFLOPs/s
Pytorch fwd: 17.65 TFLOPs/s, bwd: 44.31 TFLOPs/s, fwd + bwd: 30.95 TFLOPs/s
### causal=True, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 283.39 TFLOPs/s, bwd: 258.83 TFLOPs/s, fwd + bwd: 265.40 TFLOPs/s
Pytorch fwd: 15.50 TFLOPs/s, bwd: 46.61 TFLOPs/s, fwd + bwd: 29.62 TFLOPs/s
### causal=True, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 292.54 TFLOPs/s, bwd: 275.42 TFLOPs/s, fwd + bwd: 280.11 TFLOPs/s
Pytorch fwd: 14.19 TFLOPs/s, bwd: 47.70 TFLOPs/s, fwd + bwd: 28.48 TFLOPs/s
### causal=True, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 296.96 TFLOPs/s, bwd: 294.93 TFLOPs/s, fwd + bwd: 295.51 TFLOPs/s
Pytorch fwd: 18.97 TFLOPs/s, bwd: 48.08 TFLOPs/s, fwd + bwd: 33.43 TFLOPs/s
### causal=True, headdim=128, batch_size=32, seqlen=512 ###
Flash2 fwd: 203.08 TFLOPs/s, bwd: 144.93 TFLOPs/s, fwd + bwd: 157.85 TFLOPs/s
Pytorch fwd: 25.74 TFLOPs/s, bwd: 52.65 TFLOPs/s, fwd + bwd: 40.54 TFLOPs/s
### causal=True, headdim=128, batch_size=16, seqlen=1024 ###
Flash2 fwd: 260.47 TFLOPs/s, bwd: 200.44 TFLOPs/s, fwd + bwd: 214.57 TFLOPs/s
Pytorch fwd: 31.59 TFLOPs/s, bwd: 65.50 TFLOPs/s, fwd + bwd: 50.13 TFLOPs/s
### causal=True, headdim=128, batch_size=8, seqlen=2048 ###
Flash2 fwd: 304.26 TFLOPs/s, bwd: 245.34 TFLOPs/s, fwd + bwd: 259.71 TFLOPs/s
Pytorch fwd: 32.77 TFLOPs/s, bwd: 80.08 TFLOPs/s, fwd + bwd: 56.70 TFLOPs/s
### causal=True, headdim=128, batch_size=4, seqlen=4096 ###
Flash2 fwd: 327.53 TFLOPs/s, bwd: 268.25 TFLOPs/s, fwd + bwd: 282.88 TFLOPs/s
Pytorch fwd: 29.69 TFLOPs/s, bwd: 87.90 TFLOPs/s, fwd + bwd: 56.35 TFLOPs/s
### causal=True, headdim=128, batch_size=2, seqlen=8192 ###
Flash2 fwd: 338.84 TFLOPs/s, bwd: 298.21 TFLOPs/s, fwd + bwd: 308.79 TFLOPs/s
Pytorch fwd: 27.67 TFLOPs/s, bwd: 92.33 TFLOPs/s, fwd + bwd: 55.36 TFLOPs/s
### causal=True, headdim=128, batch_size=1, seqlen=16384 ###
Flash2 fwd: 343.30 TFLOPs/s, bwd: 324.38 TFLOPs/s, fwd + bwd: 329.57 TFLOPs/s
Pytorch fwd: 36.80 TFLOPs/s, bwd: 94.49 TFLOPs/s, fwd + bwd: 65.26 TFLOPs/s
27 changes: 19 additions & 8 deletions hopper/benchmark_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,10 +228,10 @@ def run(*args, **kwargs):
deterministic = False
batch_size = 2
# seqlen = 2048
seqlen = 8192
seqlen = 4096
# seqlen = 4096
# seqlen = 2047
dim = 2048
nheads = 24
# headdim = 128
# headdim = 64
headdim = 256
Expand All @@ -240,7 +240,7 @@ def run(*args, **kwargs):
# bs_seqlen_vals = [(16, 1024), (8, 2048), (4, 4096), (2, 8192), (1, 16384)]
# bs_seqlen_vals = [(32, 512), (16, 1024)]
# bs_seqlen_vals = [(2, 64 * 132)]
bs_seqlen_vals = [(2, 8192)]
bs_seqlen_vals = [(2, 4096)]
# bs_seqlen_vals = [(1, 16 * 1024)]
time_f = {}
time_b = {}
Expand All @@ -251,16 +251,16 @@ def run(*args, **kwargs):
# for headdim in [64, 96, 128]:
# for headdim in [64, 128, 256]:
# for headdim in [64, 96, 128, 192, 256]:
for headdim in [128]:
nheads = dim // headdim
for headdim in [128, 256]:
dim = nheads * headdim
# nheads = 128
# headdim = 64
# batch_size = 64
# seqlen = 512
# nheads = 8
# headdim = 128
nheads_kv = nheads
# nheads_kv = nheads // 4
# nheads_kv = nheads
nheads_kv = nheads // 12
# nheads_kv = 1
headdim_v = headdim
# headdim_v = 512
Expand All @@ -280,6 +280,7 @@ def run(*args, **kwargs):
k = torch.randn(batch_size, seqlen, nheads_kv, headdim, device=device, dtype=dtype_gen, requires_grad=True)
v = torch.randn(batch_size, seqlen, nheads_kv, headdim_v, device=device, dtype=dtype_gen, requires_grad=True)
q, k, v = [x.detach().to(dtype).requires_grad_() for x in [q, k, v]]
sink = torch.randn((nheads,), dtype=dtype_gen, device=device, requires_grad=True)
v_colmajor = v.detach().transpose(-1, -3).contiguous().transpose(-1, -3).requires_grad_()
v_fa3 = v if not V_colmajor else v_colmajor
qv = torch.randn(batch_size, seqlen_q, nheads, headdim_v, device=device, dtype=dtype_gen) if has_qv else None
Expand Down Expand Up @@ -309,7 +310,7 @@ def run(*args, **kwargs):

for causal in [False, True]:
# for causal in [True]:
print(f"\n### {headdim = }, {causal = }, {seqlen = } ###")
print(f"\n### {headdim = }, {causal = }, {seqlen = }, {nheads = }, {nheads_kv = } ###")
nFLOPS = flops(batch_size, nheads, seqlen_q, seqlen, headdim if not has_qv else headdim + headdim_v, headdim_v, causal=causal, window_size=window_size)
if cudnn is not None:
# if False:
Expand Down Expand Up @@ -360,20 +361,28 @@ def run(*args, **kwargs):
if not varlen:
# m1 = time_fwd(flash_attn_func_v3, q, k if page_size is None else k_paged, v_fa3 if page_size is None else v_paged, cache_leftpad = leftpad_k, page_table=page_table, causal=causal, window_size=window_size, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, repeats=repeats, verbose=verbose, desc='Fav3')
m1 = time_fwd(flash_attn_func_v3, q, k if page_size is None else k_paged, v_fa3 if page_size is None else v_paged, qv=qv, causal=causal, window_size=window_size, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, repeats=repeats, verbose=verbose, desc='Fav3')
m1_sink = time_fwd(flash_attn_func_v3, q, k if page_size is None else k_paged, v_fa3 if page_size is None else v_paged, qv=qv, causal=causal, window_size=window_size, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, learnable_sink=sink, repeats=repeats, verbose=verbose, desc='Fav3')
# pytorch_profiler(flash_attn_func_v3, q, k if page_size is None else k_paged, v_fa3 if page_size is None else v_paged, page_table=page_table, causal=causal, window_size=window_size, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa)
else:
m1 = time_fwd(flash_attn_varlen_func_v3, q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen, causal=causal, window_size=window_size, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, repeats=repeats, verbose=verbose, desc='Fav3')
m1_sink = time_fwd(flash_attn_varlen_func_v3, q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen, causal=causal, window_size=window_size, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, learnable_sink=sink, repeats=repeats, verbose=verbose, desc='Fav3')
# pytorch_profiler(flash_attn_varlen_func_v3, q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen, causal=causal, window_size=window_size, softcap=softcap, num_splits=num_splits)
time_f[(causal, headdim, batch_size, seqlen), "Flash3"] = m1.mean
time_f[(causal, headdim, batch_size, seqlen), "Flash3Sink"] = m1_sink.mean
if dtype != torch.float8_e4m3fn and headdim == headdim_v and not DISABLE_BACKWARD:
time.sleep(1)
if not varlen:
_, m1b = benchmark_backward(flash_attn_func_v3, q, k, v, causal=causal, window_size=window_size, softcap=softcap, deterministic=deterministic,
repeats=repeats, verbose=False, desc='Fav3')
_, m1b_sink = benchmark_backward(flash_attn_func_v3, q, k, v, causal=causal, window_size=window_size, softcap=softcap, deterministic=deterministic, learnable_sink=sink,
repeats=repeats, verbose=False, desc='Fav3Sink')
else:
_, m1b = benchmark_backward(flash_attn_varlen_func_v3, q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen, causal=causal, window_size=window_size, softcap=softcap, deterministic=deterministic,
repeats=repeats, verbose=False, desc='Fav3')
_, m1b_sink = benchmark_backward(flash_attn_varlen_func_v3, q_unpad, k_unpad, v_unpad, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen, causal=causal, window_size=window_size, softcap=softcap, deterministic=deterministic, learnable_sink=sink,
repeats=repeats, verbose=False, desc='Fav3Sink')
time_b[(causal, headdim, batch_size, seqlen), "Flash3"] = m1b.mean
time_b[(causal, headdim, batch_size, seqlen), "Flash3Sink"] = m1b_sink.mean
# time.sleep(1)
# if not varlen:
# pytorch_profiler(flash_attn_func_v3, q, k, v, causal=causal, deterministic=deterministic, backward=True)
Expand All @@ -394,8 +403,10 @@ def run(*args, **kwargs):
print(f'CuDNN fwd: {m2.mean * 1e3:.3f}ms, {(nFLOPS / m2.mean * 1e-12):.1f} TFLOPS')
print(f'CuDNN bwd: {m2b.mean * 1e3:.3f}ms, {(2.5 * nFLOPS / m2b.mean * 1e-12):.1f} TFLOPS')
print(f'Fav3 fwd: {m1.mean * 1e3:.3f}ms, {(nFLOPS / m1.mean * 1e-12):.1f} TFLOPS')
print(f'Fav3Sink fwd: {m1_sink.mean * 1e3:.3f}ms, {(nFLOPS / m1.mean * 1e-12):.1f} TFLOPS')
if dtype != torch.float8_e4m3fn and headdim == headdim_v and not DISABLE_BACKWARD:
print(f'Fav3 bwd: {m1b.mean * 1e3:.3f}ms, {(2.5 * nFLOPS / m1b.mean * 1e-12):.1f} TFLOPS')
print(f'Fav3Sink bwd: {m1b_sink.mean * 1e3:.3f}ms, {(2.5 * nFLOPS / m1b_sink.mean * 1e-12):.1f} TFLOPS')
# benchmark_forward(torch.square, k)
# print(f'cuBLAS: {m5.mean * 1e3:.3f}ms, {(nFLOPS_matmul / m5.mean * 1e-12):.1f} TFLOPS')
# print(time_f)
Expand Down
16 changes: 16 additions & 0 deletions hopper/benchmark_attn_install_res.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@

### headdim = 128, causal = False, seqlen = 8192 ###
Fav2 fwd: 2.796ms, 393.2 TFLOPS
Fav2 bwd: 9.343ms, 294.2 TFLOPS
CuDNN fwd: 1.692ms, 649.7 TFLOPS
CuDNN bwd: 4.990ms, 550.9 TFLOPS
Fav3 fwd: 1.569ms, 700.7 TFLOPS
Fav3 bwd: 4.328ms, 635.1 TFLOPS

### headdim = 128, causal = True, seqlen = 8192 ###
Fav2 fwd: 1.579ms, 348.2 TFLOPS
Fav2 bwd: 4.614ms, 297.9 TFLOPS
CuDNN fwd: 0.927ms, 593.3 TFLOPS
CuDNN bwd: 2.837ms, 484.5 TFLOPS
Fav3 fwd: 0.804ms, 684.2 TFLOPS
Fav3 bwd: 2.330ms, 589.8 TFLOPS
16 changes: 16 additions & 0 deletions hopper/benchmark_attn_res.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@

### headdim = 128, causal = False, seqlen = 8192 ###
Fav2 fwd: 2.798ms, 393.0 TFLOPS
Fav2 bwd: 9.343ms, 294.2 TFLOPS
CuDNN fwd: 1.692ms, 649.7 TFLOPS
CuDNN bwd: 4.985ms, 551.4 TFLOPS
Fav3 fwd: 1.570ms, 700.4 TFLOPS
Fav3 bwd: 4.320ms, 636.3 TFLOPS

### headdim = 128, causal = True, seqlen = 8192 ###
Fav2 fwd: 1.580ms, 348.0 TFLOPS
Fav2 bwd: 4.605ms, 298.4 TFLOPS
CuDNN fwd: 0.927ms, 593.0 TFLOPS
CuDNN bwd: 2.840ms, 483.9 TFLOPS
Fav3 fwd: 0.804ms, 684.0 TFLOPS
Fav3 bwd: 2.336ms, 588.4 TFLOPS
40 changes: 40 additions & 0 deletions hopper/benchmark_attn_sink.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@

### headdim = 128, causal = False, seqlen = 4096, nheads = 24, nheads_kv = 2 ###
Fav2 fwd: 1.072ms, 384.5 TFLOPS
Fav2 bwd: 3.610ms, 285.6 TFLOPS
CuDNN fwd: 0.638ms, 646.1 TFLOPS
CuDNN bwd: 2.000ms, 515.4 TFLOPS
Fav3 fwd: 0.970ms, 425.1 TFLOPS
Fav3Sink fwd: 0.971ms, 425.1 TFLOPS
Fav3 bwd: 4.526ms, 227.8 TFLOPS
Fav3Sink bwd: 5.617ms, 183.5 TFLOPS

### headdim = 128, causal = True, seqlen = 4096, nheads = 24, nheads_kv = 2 ###
Fav2 fwd: 0.633ms, 325.4 TFLOPS
Fav2 bwd: 1.963ms, 262.5 TFLOPS
CuDNN fwd: 0.370ms, 557.4 TFLOPS
CuDNN bwd: 1.259ms, 409.2 TFLOPS
Fav3 fwd: 0.503ms, 410.2 TFLOPS
Fav3Sink fwd: 0.503ms, 410.2 TFLOPS
Fav3 bwd: 2.432ms, 211.9 TFLOPS
Fav3Sink bwd: 2.922ms, 176.4 TFLOPS

### headdim = 256, causal = False, seqlen = 4096, nheads = 24, nheads_kv = 2 ###
Fav2 fwd: 2.451ms, 336.5 TFLOPS
Fav2 bwd: 8.963ms, 230.0 TFLOPS
CuDNN fwd: 1.191ms, 692.3 TFLOPS
CuDNN bwd: 6.106ms, 337.6 TFLOPS
Fav3 fwd: 1.082ms, 762.3 TFLOPS
Fav3Sink fwd: 1.083ms, 762.3 TFLOPS
Fav3 bwd: 4.590ms, 449.2 TFLOPS
Fav3Sink bwd: 5.750ms, 358.5 TFLOPS

### headdim = 256, causal = True, seqlen = 4096, nheads = 24, nheads_kv = 2 ###
Fav2 fwd: 1.341ms, 307.5 TFLOPS
Fav2 bwd: 4.384ms, 235.1 TFLOPS
CuDNN fwd: 0.679ms, 606.8 TFLOPS
CuDNN bwd: 3.185ms, 323.7 TFLOPS
Fav3 fwd: 0.571ms, 722.3 TFLOPS
Fav3Sink fwd: 0.571ms, 722.3 TFLOPS
Fav3 bwd: 2.521ms, 408.8 TFLOPS
Fav3Sink bwd: 2.977ms, 346.2 TFLOPS
Loading