From 2ff2a9d49f48239dbfb2f90f5668698d7d019c44 Mon Sep 17 00:00:00 2001 From: jerryao Date: Sat, 20 Dec 2025 18:10:25 +0800 Subject: [PATCH 01/23] Prepare impl. --- .../benchmark_flash_attention_install_res.txt | 72 +++++++++++++++++++ benchmarks/benchmark_flash_attention_res.txt | 72 +++++++++++++++++++ hopper/benchmark_attn_install_res.txt | 16 +++++ hopper/benchmark_attn_res.txt | 16 +++++ reinstall.sh | 14 ++++ 5 files changed, 190 insertions(+) create mode 100644 benchmarks/benchmark_flash_attention_install_res.txt create mode 100644 benchmarks/benchmark_flash_attention_res.txt create mode 100644 hopper/benchmark_attn_install_res.txt create mode 100644 hopper/benchmark_attn_res.txt create mode 100755 reinstall.sh diff --git a/benchmarks/benchmark_flash_attention_install_res.txt b/benchmarks/benchmark_flash_attention_install_res.txt new file mode 100644 index 00000000000..07262abe2e5 --- /dev/null +++ b/benchmarks/benchmark_flash_attention_install_res.txt @@ -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 diff --git a/benchmarks/benchmark_flash_attention_res.txt b/benchmarks/benchmark_flash_attention_res.txt new file mode 100644 index 00000000000..07ff23e0da1 --- /dev/null +++ b/benchmarks/benchmark_flash_attention_res.txt @@ -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 diff --git a/hopper/benchmark_attn_install_res.txt b/hopper/benchmark_attn_install_res.txt new file mode 100644 index 00000000000..08579e82290 --- /dev/null +++ b/hopper/benchmark_attn_install_res.txt @@ -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 diff --git a/hopper/benchmark_attn_res.txt b/hopper/benchmark_attn_res.txt new file mode 100644 index 00000000000..152b73f5e57 --- /dev/null +++ b/hopper/benchmark_attn_res.txt @@ -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 diff --git a/reinstall.sh b/reinstall.sh new file mode 100755 index 00000000000..d5b0bd78cf6 --- /dev/null +++ b/reinstall.sh @@ -0,0 +1,14 @@ +python setup.py bdist_wheel + +cd hopper +python setup.py bdist_wheel + +cd ../csrc/fused_dense_lib +python setup.py bdist_wheel + +cd ../.. +cp dist/*.whl ./whls/ +cp hopper/dist/*.whl ./whls/ +cp csrc/fused_dense_lib/dist/*.whl ./whls/ + +pip install ./whls/*.whl From f479d9022e36b5bb243a76e97f1ecaffeac4f018 Mon Sep 17 00:00:00 2001 From: aoxy Date: Fri, 16 Jan 2026 14:12:34 +0800 Subject: [PATCH 02/23] Right for some cases. --- hopper/epilogue_fwd.hpp | 7 +-- hopper/flash.h | 2 + hopper/flash_api.cpp | 46 +++++++++++++++++-- hopper/flash_attn_interface.py | 52 +++++++++++++++++----- hopper/flash_fwd_combine_kernel.h | 21 ++++++++- hopper/flash_fwd_combine_launch_template.h | 19 ++++---- hopper/flash_fwd_kernel_sm90.h | 12 +++-- hopper/flash_fwd_launch_template.h | 19 ++++---- hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp | 9 +++- hopper/softmax.h | 25 +++++++++-- hopper/static_switch.h | 10 +++++ hopper/test_util.py | 10 ++++- reinstall.sh | 29 +++++++----- warp_reinstall.sh | 1 + 14 files changed, 204 insertions(+), 58 deletions(-) create mode 100755 warp_reinstall.sh diff --git a/hopper/epilogue_fwd.hpp b/hopper/epilogue_fwd.hpp index 69102e8c4e6..7012425b623 100644 --- a/hopper/epilogue_fwd.hpp +++ b/hopper/epilogue_fwd.hpp @@ -410,7 +410,8 @@ struct CollectiveEpilogueFwd { store_zero( Params const& params, int thread_idx, - cute::tuple const& block_coord + cute::tuple const& block_coord, + float lse_val=-INFINITY ) { static constexpr int kBlockM = get<0>(TileShape_MNK_PV{}); auto [m_block, bidh, bidb, split_idx] = block_coord; @@ -437,13 +438,13 @@ struct CollectiveEpilogueFwd { if (thread_idx < kBlockM) { const int row = m_block * kBlockM + thread_idx; if constexpr (!PackGQA) { - if (row < seqlen_o) { mLSE(row) = -INFINITY; } + if (row < seqlen_o) { mLSE(row) = lse_val; } } else { if (row < seqlen_o * qhead_per_khead) { int m_idx, h_idx; m_idx = params.qhead_per_khead_divmod.divmod(h_idx, row); // mLSE has shape ((qhead_per_khead, seqlen_q)) and it's unhappy with just 1 "make_coord" - mLSE(make_coord(make_coord(h_idx, m_idx))) = -INFINITY; + mLSE(make_coord(make_coord(h_idx, m_idx))) = lse_val; } } } diff --git a/hopper/flash.h b/hopper/flash.h index bee89e5f054..30c1b9db311 100644 --- a/hopper/flash.h +++ b/hopper/flash.h @@ -159,6 +159,7 @@ struct Flash_fwd_params : public Qkv_params { int arch; int num_sm; + void *__restrict__ learnable_sink_ptr; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -205,6 +206,7 @@ struct Flash_bwd_params : public Flash_fwd_params { bool deterministic; index_t dq_accum_split_stride; + void *__restrict__ dsink_ptr; }; //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp index 58188137777..e76451972d4 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -680,7 +680,8 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq std::optional &scheduler_metadata_, // (b + 1) int num_splits, std::optional pack_gqa_, - int const sm_margin + int const sm_margin, + std::optional learnable_sink_ ) { auto dprops = at::cuda::getCurrentDeviceProperties(); @@ -1045,6 +1046,17 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq params.kv_batch_idx = reinterpret_cast(kv_batch_idx.data_ptr()); } + if (learnable_sink_.has_value()) { + at::Tensor learnable_sink = learnable_sink_.value(); + TORCH_CHECK(learnable_sink.dtype() == torch::kFloat32, "Learnable sink must have dtype fp32"); + CHECK_DEVICE(learnable_sink); CHECK_CONTIGUOUS(learnable_sink); + TORCH_CHECK(learnable_sink.stride(-1) == 1, "Learnable sink tensor must have contiguous last dimension"); + CHECK_SHAPE(learnable_sink, num_heads); + params.learnable_sink_ptr = learnable_sink.data_ptr(); + } else { + params.learnable_sink_ptr = nullptr; + } + at::Tensor out_accum, softmax_lse_accum; auto outaccum_type = at::ScalarType::Float; if (params.num_splits > 1) { @@ -1233,7 +1245,9 @@ std::vector mha_bwd( int window_size_right, float const softcap, bool const deterministic, - int const sm_margin) { + int const sm_margin, + std::optional learnable_sink_, + std::optional dsink_) { #ifdef FLASHATTENTION_DISABLE_BACKWARD TORCH_CHECK(false, "This flash attention build does not support backward."); @@ -1414,7 +1428,7 @@ std::vector mha_bwd( auto opts = q.options(); // Need softmax_d to have total_q_padded_rounded since we want its address to be aligned by 16/8 bytes for TMA / LDG.64 - at::Tensor softmax_d, softmax_lse_log2; + at::Tensor softmax_d, dsink, softmax_lse_log2; if (!is_varlen) { // Need softmax_d to have seqlen_q_rounded since we want its address to be aligned by 16/8 bytes for TMA / LDG.64 softmax_d = torch::empty({batch_size, num_heads, seqlen_q_rounded}, opts.dtype(at::kFloat)); @@ -1471,6 +1485,30 @@ std::vector mha_bwd( params.dv = head_size_v; params.dv_rounded = head_size_v_rounded; + + if (learnable_sink_.has_value()) { + at::Tensor learnable_sink = learnable_sink_.value(); + TORCH_CHECK(learnable_sink.dtype() == torch::kFloat32, "Learnable sink must have dtype fp32"); + CHECK_DEVICE(learnable_sink); CHECK_CONTIGUOUS(learnable_sink); + TORCH_CHECK(learnable_sink.stride(-1) == 1, "Learnable sink tensor must have contiguous last dimension"); + CHECK_SHAPE(learnable_sink, num_heads); + if (dsink_.has_value()) { + dsink = dsink_.value(); + TORCH_CHECK(dsink.dtype() == torch::kFloat32, "dsink must have dtype fp32"); + CHECK_DEVICE(dsink); CHECK_CONTIGUOUS(dsink); + TORCH_CHECK(dsink.stride(-1) == 1, "dsink tensor must have contiguous last dimension"); + CHECK_SHAPE(dsink, num_heads); + } else { + dsink = torch::empty_like(learnable_sink); + } + dsink.zero_(); + params.learnable_sink_ptr = learnable_sink.data_ptr(); + params.dsink_ptr = dsink.data_ptr(); + } else { + params.learnable_sink_ptr = nullptr; + params.dsink_ptr = nullptr; + } + // auto tile_count_semaphore = (params.is_causal || params.is_local) ? torch::zeros({1}, opts.dtype(torch::kInt32)) : torch::empty({1}, opts.dtype(torch::kInt32)); // params.tile_count_semaphore = tile_count_semaphore.data_ptr(); // Will be zero'ed out in the backward preprocess kernel @@ -1504,7 +1542,7 @@ std::vector mha_bwd( softmax_d.zero_(); } - return { dq, dk, dv, softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum }; + return { dq, dk, dv, softmax_d, dsink, softmax_lse_log2, dq_accum, dk_accum, dv_accum }; } std::vector diff --git a/hopper/flash_attn_interface.py b/hopper/flash_attn_interface.py index 06782fa409b..e41e20eb1e1 100644 --- a/hopper/flash_attn_interface.py +++ b/hopper/flash_attn_interface.py @@ -49,7 +49,8 @@ def _flash_attn_forward( scheduler_metadata=None, num_splits=1, pack_gqa=None, - sm_margin=0): + sm_margin=0, + learnable_sink=None): q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ @@ -96,6 +97,7 @@ def _flash_attn_forward( num_splits, pack_gqa, sm_margin, + learnable_sink, ) return out, softmax_lse, *rest @@ -122,10 +124,12 @@ def _flash_attn_backward( softcap=0.0, deterministic=False, sm_margin=0, + learnable_sink=None, + dsink=None, ): # dq, dk, dv are allocated by us so they should already be contiguous dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - dq, dk, dv, softmax_d, *rest = flash_attn_3_cuda.bwd( + dq, dk, dv, softmax_d, dsink, *rest = flash_attn_3_cuda.bwd( dout, q, k, @@ -148,8 +152,10 @@ def _flash_attn_backward( softcap, deterministic, sm_margin, + learnable_sink, + dsink, ) - return dq, dk, dv, softmax_d + return dq, dk, dv, softmax_d, dsink class FlashAttnQKVPackedFunc(torch.autograd.Function): @@ -166,6 +172,7 @@ def forward( deterministic=False, num_heads_q=None, sm_margin=0, + learnable_sink=None, ): if softmax_scale is None: softmax_scale = qkv.shape[-1] ** (-0.5) @@ -199,7 +206,7 @@ def forward( sm_margin=sm_margin, ) # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) - ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse, learnable_sink) ctx.softmax_scale = softmax_scale ctx.causal = causal ctx.window_size = window_size @@ -213,7 +220,7 @@ def forward( @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse = ctx.saved_tensors + q, k, v, out, softmax_lse, learnable_sink = ctx.saved_tensors assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" if ctx.ndim == 5: qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) @@ -225,6 +232,7 @@ def backward(ctx, dout, *args): qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:]) dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2) + dsink = None if learnable_sink is None else torch.empty_like(learnable_sink) _flash_attn_backward( dout, q, @@ -244,9 +252,11 @@ def backward(ctx, dout, *args): ctx.softcap, ctx.deterministic, ctx.sm_margin, + learnable_sink, + dsink, ) dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None, None, None + return dqkv, dsink, None, None, None, None, None, None, None, None, None, None, None class FlashAttnFunc(torch.autograd.Function): @@ -268,6 +278,7 @@ def forward( pack_gqa=None, deterministic=False, sm_margin=0, + learnable_sink=None, ): if softmax_scale is None: softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) @@ -293,9 +304,10 @@ def forward( num_splits=num_splits, pack_gqa=pack_gqa, sm_margin=sm_margin, + learnable_sink=learnable_sink, ) # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) - ctx.save_for_backward(q, k, v, out, softmax_lse) + ctx.save_for_backward(q, k, v, out, softmax_lse, learnable_sink) ctx.softmax_scale = softmax_scale ctx.causal = causal ctx.window_size = window_size @@ -307,9 +319,10 @@ def forward( @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse = ctx.saved_tensors + q, k, v, out, softmax_lse, learnable_sink = ctx.saved_tensors assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + dsink = None if learnable_sink is None else torch.empty_like(learnable_sink) _flash_attn_backward( dout, q, @@ -329,11 +342,13 @@ def backward(ctx, dout, *args): ctx.softcap, ctx.deterministic, ctx.sm_margin, + learnable_sink, + dsink, ) dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + return dq, dk, dv, dsink, None, None, None, None, None, None, None, None, None, None, None, None, None, None class FlashAttnVarlenFunc(torch.autograd.Function): @@ -361,6 +376,7 @@ def forward( pack_gqa=None, deterministic=False, sm_margin=0, + learnable_sink=None, ): if softmax_scale is None: softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) @@ -390,9 +406,10 @@ def forward( num_splits=num_splits, pack_gqa=pack_gqa, sm_margin=sm_margin, + learnable_sink=learnable_sink, ) # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) - ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink) ctx.max_seqlen_q = max_seqlen_q ctx.max_seqlen_k = max_seqlen_k ctx.softmax_scale = softmax_scale @@ -406,9 +423,10 @@ def forward( @staticmethod def backward(ctx, dout, *args): - q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink = ctx.saved_tensors assert ctx.attention_chunk == 0, "FA3 backward does not support attention_chunk" dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + dsink = None if learnable_sink is None else torch.empty_like(learnable_sink) _flash_attn_backward( dout, q, @@ -431,11 +449,13 @@ def backward(ctx, dout, *args): ctx.softcap, ctx.deterministic, ctx.sm_margin, + learnable_sink, + dsink, ) dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + return dq, dk, dv, dsink, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None def flash_attn_qkvpacked_func( @@ -449,6 +469,7 @@ def flash_attn_qkvpacked_func( deterministic=False, num_heads_q=None, sm_margin=0, + learnable_sink=None, ): """dropout_p should be set to 0.0 during evaluation If Q, K, V are already stacked into 1 tensor, this function will be faster than @@ -495,6 +516,7 @@ def flash_attn_qkvpacked_func( deterministic, num_heads_q, sm_margin, + learnable_sink, ) @@ -513,6 +535,7 @@ def flash_attn_func( pack_gqa=None, deterministic=False, sm_margin=0, + learnable_sink=None, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads @@ -574,6 +597,7 @@ def flash_attn_func( pack_gqa, deterministic, sm_margin, + learnable_sink, ) @@ -598,6 +622,7 @@ def flash_attn_varlen_func( pack_gqa=None, deterministic=False, sm_margin=0, + learnable_sink=None, ): return FlashAttnVarlenFunc.apply( q, @@ -620,6 +645,7 @@ def flash_attn_varlen_func( pack_gqa, deterministic, sm_margin, + learnable_sink, ) @@ -658,6 +684,7 @@ def flash_attn_with_kvcache( pack_gqa=None, # Can be tuned for speed sm_margin=0, # Can be tuned if some SMs are used for communication return_softmax_lse=False, + learnable_sink=None, ): """ If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from @@ -785,6 +812,7 @@ def flash_attn_with_kvcache( num_splits=num_splits, pack_gqa=pack_gqa, sm_margin=sm_margin, + learnable_sink=learnable_sink, ) # return (out, softmax_lse) if return_softmax_lse else out return (out, softmax_lse, *rest) if return_softmax_lse else out diff --git a/hopper/flash_fwd_combine_kernel.h b/hopper/flash_fwd_combine_kernel.h index a22e05969d9..e9cda0b869c 100644 --- a/hopper/flash_fwd_combine_kernel.h +++ b/hopper/flash_fwd_combine_kernel.h @@ -22,7 +22,7 @@ namespace flash { using namespace cute; template + bool Is_even_K, bool Varlen, class Element, class ElementPartial, class ArchTag_, bool Has_sink> class FlashAttnFwdCombine { public: @@ -146,6 +146,7 @@ class FlashAttnFwdCombine { int const* const seqused = nullptr; int const* const num_splits_dynamic_ptr = nullptr; int* const semaphore_to_reset = nullptr; + float const* const ptr_Sink = nullptr; }; // Kernel entry point API @@ -165,6 +166,7 @@ class FlashAttnFwdCombine { int const* const seqused = nullptr; int const* const num_splits_dynamic_ptr = nullptr; int* const semaphore_to_reset = nullptr; + float const* const ptr_Sink = nullptr; }; // Convert to underlying arguments. In this case, a simple copy for the aliased type. @@ -187,7 +189,8 @@ class FlashAttnFwdCombine { args.cu_seqlens, args.seqused, args.num_splits_dynamic_ptr, - args.semaphore_to_reset + args.semaphore_to_reset, + args.ptr_Sink }; } @@ -369,6 +372,20 @@ class FlashAttnFwdCombine { } SumOp sum_op; lse_sum_cur = Allreduce::run(lse_sum_cur, sum_op); + if constexpr (Has_sink) { + int mi = int(get<1>(ts2rcLSE(_0{}, _0{}, m))); + int idx = m_block * kBlockM + mi; + if (idx < max_idx) { + int m_idx, bidh; + if constexpr (!Varlen) { + bidh = params.seqlen_divmod.divmod(m_idx, idx); + } else { + bidh = seqlen_divmod_dynamic.divmod(m_idx, idx); + } + const float sink_val_exp = expf((params.ptr_Sink)[bidh] - lse_max); + lse_sum_cur += sink_val_exp; + } + } lse_sum(m) = logf(lse_sum_cur) + lse_max; float inv_sum = (lse_sum_cur == 0.f || lse_sum_cur != lse_sum_cur) ? 0.f : 1.f / lse_sum_cur; #pragma unroll diff --git a/hopper/flash_fwd_combine_launch_template.h b/hopper/flash_fwd_combine_launch_template.h index 11d422924b4..090dc500850 100644 --- a/hopper/flash_fwd_combine_launch_template.h +++ b/hopper/flash_fwd_combine_launch_template.h @@ -17,12 +17,12 @@ using namespace cute; -template +template void run_flash_fwd_combine(Flash_fwd_params ¶ms, cudaStream_t stream, bool enable_pdl) { using ArchTag = std::conditional_t= 90, cutlass::arch::Sm90, cutlass::arch::Sm80>; using TileShape_MK = cute::Shape, Int>; using CombineKernel = flash::FlashAttnFwdCombine; + IsEvenK, Varlen, Element, ElementPartial, ArchTag, Has_sink>; typename CombineKernel::Arguments args { static_cast(params.oaccum_ptr), @@ -35,7 +35,8 @@ void run_flash_fwd_combine(Flash_fwd_params ¶ms, cudaStream_t stream, bool e {params.o_row_stride, _1{}, params.o_head_stride, !Varlen ? params.o_batch_stride : 0}, // stride_O static_cast(params.softmax_lse_ptr), {_1{}, !Varlen ? params.seqlen_q : params.total_q, !Varlen ? params.h * params.seqlen_q : 0}, // stride_LSE - params.cu_seqlens_q, params.seqused_q, params.num_splits_dynamic_ptr, params.tile_count_semaphore + params.cu_seqlens_q, params.seqused_q, params.num_splits_dynamic_ptr, params.tile_count_semaphore, + reinterpret_cast(params.learnable_sink_ptr) }; typename CombineKernel::Params kernel_params = CombineKernel::to_underlying_arguments(args); @@ -60,21 +61,23 @@ void run_mha_fwd_combine_(Flash_fwd_params ¶ms, cudaStream_t stream, bool en static constexpr int kBlockM = kBlockK % 128 == 0 ? 8 : (kBlockK % 64 == 0 ? 16 : 32); ARCH_SWITCH(params.arch, Arch, [&] { BOOL_SWITCH(params.cu_seqlens_q || params.seqused_q, Varlen, [&] { + SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { if constexpr (kBlockM >= 16) { // If kBlockM == 8 then the minimum number of splits is 32. if (params.num_splits <= 16) { - run_flash_fwd_combine(params, stream, enable_pdl); + run_flash_fwd_combine(params, stream, enable_pdl); return; } } if (params.num_splits <= 32) { - run_flash_fwd_combine(params, stream, enable_pdl); + run_flash_fwd_combine(params, stream, enable_pdl); } else if (params.num_splits <= 64) { - run_flash_fwd_combine(params, stream, enable_pdl); + run_flash_fwd_combine(params, stream, enable_pdl); } else if (params.num_splits <= 128) { - run_flash_fwd_combine(params, stream, enable_pdl); + run_flash_fwd_combine(params, stream, enable_pdl); } else { - run_flash_fwd_combine(params, stream, enable_pdl); + run_flash_fwd_combine(params, stream, enable_pdl); } + }); }); }); } diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index 47b3817cd28..61cc4b05d24 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -49,6 +49,7 @@ class FlashAttnFwdSm90 { static constexpr int NumProducerThreads = CollectiveMainloop::NumProducerThreads; static constexpr bool SameHeadDim = CollectiveMainloop::SameHeadDim; static constexpr bool LargeHeadDimV = CollectiveMainloop::LargeHeadDimV; + static constexpr bool Has_sink = CollectiveMainloop::Has_sink; static_assert(CollectiveMainloop::LargeHeadDimV == CollectiveEpilogue::LargeHeadDimV); using SeqlenInfo_t = typename CollectiveMainloop::SeqlenInfo_t; @@ -113,7 +114,7 @@ class FlashAttnFwdSm90 { alignas(16) typename CollectiveMainloop::MainloopPipelineKVNew::SharedStorage pipeline_v_new; alignas(16) typename TileScheduler::SharedStorage smem_scheduler; } pipelines; - + // float sink_val; }; static constexpr int SharedStorageSize = sizeof(SharedStorage); @@ -406,14 +407,16 @@ class FlashAttnFwdSm90 { } // If there's tanh softcap, the scaling will be done before tanh. float softmax_scale_log2 = params.mainloop.softmax_scale_log2; + int const bidh = get<1>(block_coord); + float sink_val = !Has_sink ? -INFINITY : params.mainloop.ptr_Sink[bidh]; if constexpr (Is_FP8 && !Has_softcap) { - int const bidh = get<1>(block_coord); int const bidh_kv = !PackGQA ? params.mainloop.qhead_per_khead_divmod.divide(bidh) : bidh; float const q_descale = params.mainloop.ptr_q_descale == nullptr ? 1.0f : params.mainloop.ptr_q_descale[bidb * get<0>(params.mainloop.stride_q_descale) + bidh_kv * get<1>(params.mainloop.stride_q_descale)]; float const k_descale = params.mainloop.ptr_k_descale == nullptr ? 1.0f : params.mainloop.ptr_k_descale[bidb * get<0>(params.mainloop.stride_k_descale) + bidh_kv * get<1>(params.mainloop.stride_k_descale)]; softmax_scale_log2 *= q_descale * k_descale; } - flash::Softmax softmax(softmax_scale_log2); + if (threadIdx.x == 128) { printf("=====LargeHeadDimV=%d=get<1>(block_coord)=%d, softmax_scale_log2=%f, shared_storage.sink_val=%f, sink_val=%f\n", int(LargeHeadDimV), get<1>(block_coord), softmax_scale_log2, sink_val, sink_val); } + flash::Softmax softmax(softmax_scale_log2, sink_val); // Attention output (GEMM-II) accumulator. Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_MNK_PV{})); bool tile_valid; @@ -445,7 +448,8 @@ class FlashAttnFwdSm90 { threadIdx.x - MmaThreadOffset, block_coord); } else { // Write 0 to gO and -inf to gLSE. - epilogue.store_zero(params.epilogue, threadIdx.x - MmaThreadOffset, block_coord); + // const float lse_val = Has_sink ? sink_val : -INFINITY; + epilogue.store_zero(params.epilogue, threadIdx.x - MmaThreadOffset, block_coord, sink_val); } } epilogue.store_tail(); diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index b8af2977f11..adb290dee08 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -26,7 +26,7 @@ using namespace cute; template + bool PackGQA, bool Split, bool V_colmajor, bool Has_sink> void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { static_assert(!(Is_causal && Is_local), "Causal and Local cannot be enabled at the same time"); static_assert(!(AppendKV && V_colmajor), "AppendKV and V_colmajor cannot be enabled at the same time"); @@ -51,7 +51,7 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { using ClusterShape = cute::Shape, _1, _1>; using CollectiveMainloop = std::conditional_t< Arch >= 90, - flash::CollectiveMainloopFwdSm90, + flash::CollectiveMainloopFwdSm90, flash::CollectiveMainloopFwdSm80 >; using CollectiveEpilogue = flash::CollectiveEpilogueFwd; @@ -126,7 +126,8 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { params.kv_batch_idx, params.cu_seqlens_q, params.cu_seqlens_k, params.cu_seqlens_knew, params.seqused_q, params.seqused_k, - params.leftpad_k, params.seqlens_rotary + params.leftpad_k, params.seqlens_rotary, + reinterpret_cast(params.learnable_sink_ptr) }; typename CollectiveEpilogue::Arguments epilogue_args { static_cast(params.o_ptr), @@ -169,10 +170,10 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { dim3 grid_dims = AttnKernel::get_grid_shape(kernel_params); dim3 block_dims = AttnKernel::get_block_shape(); int smem_size = AttnKernel::SharedStorageSize; - // int smem_size_q = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_q)); - // int smem_size_k = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k)); - // int smem_size_v = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v)); - // printf("smem_size = %d, q = %d, k = %d, v = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v); + int smem_size_q = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_q)); + int smem_size_k = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k)); + int smem_size_v = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v)); + printf("========== smem_size = %d, q = %d, k = %d, v = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v); // Get the ptr to kernel function. if constexpr (size(ClusterShape{}) > 1) { void const* kernel = (void const*) cutlass::device_kernel; @@ -210,10 +211,12 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { BOOL_SWITCH(params.qv_ptr, HasQV_, [&] { static constexpr bool HasQv = HasQV_ && Arch == 90 && !Is_FP8 && kHeadDim == 64 && kHeadDimV >= 256; APPENDKV_SWITCH(params.knew_ptr, AppendKV, [&] { + SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { // Only use Cluster if number of tiles along seqlen_q is even and not varlen CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; - run_flash_fwd(params, stream); + run_flash_fwd(params, stream); + }); }); }); }); diff --git a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp index ba699f17105..4129eb34d1e 100644 --- a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp @@ -30,7 +30,7 @@ using namespace cute; template + bool MmaPV_is_RS, bool IntraWGOverlap, bool PackGQA_, bool Split_, bool V_colmajor_, bool Has_sink_> struct CollectiveMainloopFwdSm90 { static constexpr int kStages = Stages; @@ -52,6 +52,7 @@ struct CollectiveMainloopFwdSm90 { static constexpr bool PackGQA = PackGQA_; static constexpr bool Split = Split_; static constexpr bool V_colmajor = V_colmajor_; + static constexpr bool Has_sink = Has_sink_; static constexpr bool Transpose_V = Is_FP8 && !V_colmajor; static constexpr bool Use_TMA_Q = !PackGQA; static constexpr bool Use_TMA_KV = !PagedKVNonTMA; @@ -396,6 +397,7 @@ struct CollectiveMainloopFwdSm90 { int const* const seqused_k = nullptr; int const* const leftpad_k = nullptr; int const* const seqlens_rotary = nullptr; + float const* const ptr_Sink = nullptr; }; // Device side kernel params @@ -453,6 +455,7 @@ struct CollectiveMainloopFwdSm90 { int const* const seqused_k = nullptr; int const* const leftpad_k = nullptr; int const *const seqlens_rotary = nullptr; + float const* const ptr_Sink = nullptr; }; static Params @@ -564,7 +567,7 @@ struct CollectiveMainloopFwdSm90 { !Split ? 1 : args.num_splits, args.kv_batch_idx, args.cu_seqlens_q, args.cu_seqlens_k, args.cu_seqlens_k_new, - args.seqused_q, args.seqused_k, args.leftpad_k, args.seqlens_rotary}; + args.seqused_q, args.seqused_k, args.leftpad_k, args.seqlens_rotary, args.ptr_Sink}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -610,6 +613,8 @@ struct CollectiveMainloopFwdSm90 { params.window_size_left, params.window_size_right, params.attention_chunk_divmod, params.qhead_per_khead_divmod); // It's possible to have n_block_max <= n_block_min. Loading K can cause illegal memory access. + // shared_storage.sink_val = !Has_sink ? -INFINITY : params.ptr_Sink[bidh]; + if constexpr (Is_causal || Is_local || Varlen || Split) { if (n_block_max <= n_block_min) { scheduler_prefetch(); diff --git a/hopper/softmax.h b/hopper/softmax.h index 8fcdb6bd070..7c004801b03 100644 --- a/hopper/softmax.h +++ b/hopper/softmax.h @@ -89,14 +89,17 @@ __forceinline__ __device__ void scale_apply_exp2(Tensor &tenso //////////////////////////////////////////////////////////////////////////////////////////////////// -template +template struct Softmax { using TensorT = decltype(make_tensor(Shape>{})); TensorT row_max, row_sum; float const softmax_scale_log2; + float const softmax_scale; + float const sink_val; - CUTLASS_DEVICE Softmax(float const softmax_scale_log2_) : softmax_scale_log2(softmax_scale_log2_) {}; + CUTLASS_DEVICE Softmax(float const softmax_scale_log2_, float const sink_val_ = -INFINITY) + : softmax_scale_log2(softmax_scale_log2_), softmax_scale(softmax_scale_log2_ * float(M_LN2)), sink_val(sink_val_) {}; template __forceinline__ __device__ TensorT max_get_scale(Tensor0 &acc_s) { @@ -105,7 +108,14 @@ struct Softmax { static_assert(CUTE_STATIC_V(size<0>(scores)) == kNRows); TensorT scores_scale; if constexpr (Is_first) { - flash::template reduce_max(scores, row_max); + if constexpr (Has_sink) { + const float sink_scaled = sink_val / softmax_scale; + #pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { row_max(mi) = sink_scaled; } + flash::template reduce_max(scores, row_max); + } else { + flash::template reduce_max(scores, row_max); + } cute::fill(scores_scale, 1.f); } else { Tensor scores_max_prev = make_fragment_like(row_max); @@ -141,6 +151,13 @@ struct Softmax { #pragma unroll for (int mi = 0; mi < size(row_sum); ++mi) { float sum = row_sum(mi); + if constexpr (Has_sink) { + static constexpr float max_offset = float(Max_offset); + // const float max_scaled = row_max(mi) == -INFINITY ? 0.f : row_max(mi) * softmax_scale; + // sum += exp2f(sink_val * float(M_LOG2E) - max_scaled); + const float max_scaled = max(mi) == -INFINITY ? 0.f : (row_max(mi) * softmax_scale_log2) - max_offset; + sum += exp2f(sink_val * float(M_LOG2E) - max_scaled); + } float inv_sum = (sum == 0.f || sum != sum) ? 0.f : 1.f / sum; scores_scale(mi) = inv_sum * final_scale; // For FP8, we might have scaled the output of exp by 2**8 so we need to divide sum by that amount. @@ -148,7 +165,7 @@ struct Softmax { static constexpr float sum_scale = 1.f / float(1 << Max_offset); sum *= sum_scale; } - row_sum(mi) = (sum == 0.f || sum != sum) ? -INFINITY : row_max(mi) * (softmax_scale_log2 * float(M_LN2)) + __logf(sum); + row_sum(mi) = (sum == 0.f || sum != sum) ? -INFINITY : row_max(mi) * softmax_scale + __logf(sum); } return scores_scale; }; diff --git a/hopper/static_switch.h b/hopper/static_switch.h index 5e13b5f93a8..3fb11ea223b 100644 --- a/hopper/static_switch.h +++ b/hopper/static_switch.h @@ -179,3 +179,13 @@ return __VA_ARGS__(); \ } \ }() + +#ifdef FLASHATTENTION_DISABLE_SINK + #define SINK_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + constexpr static bool CONST_NAME = false; \ + return __VA_ARGS__(); \ + }() +#else + #define SINK_SWITCH BOOL_SWITCH +#endif diff --git a/hopper/test_util.py b/hopper/test_util.py index 7331ea62ca1..21eb0373b78 100644 --- a/hopper/test_util.py +++ b/hopper/test_util.py @@ -243,6 +243,7 @@ def attention_ref( upcast=True, reorder_ops=False, intermediate_dtype=None, + learnable_sink=None, ): """ Arguments: @@ -322,7 +323,14 @@ def attention_ref( scores.masked_fill_(local_mask, float("-inf")) if attn_bias is not None: scores = scores + attn_bias - attention = torch.softmax(scores, dim=-1).to(v.dtype) + if learnable_sink is None: + attention = torch.softmax(scores, dim=-1).to(v.dtype) + else: + learnable_sink = learnable_sink.to(q.dtype) + sinks = repeat(learnable_sink, 'h -> b h n 1', b=scores.shape[0], n=scores.shape[2]) + scores = torch.cat([scores, sinks], dim=-1) + attention = torch.softmax(scores, dim=-1).to(v.dtype) + attention = attention[..., :-1] # We want to mask here so that the attention matrix doesn't have any NaNs # Otherwise we'll get NaN in dV if query_padding_mask is not None: diff --git a/reinstall.sh b/reinstall.sh index d5b0bd78cf6..56ac3d1a510 100755 --- a/reinstall.sh +++ b/reinstall.sh @@ -1,14 +1,23 @@ -python setup.py bdist_wheel +#!/bin/bash +set -e -cd hopper -python setup.py bdist_wheel +export FLASH_ATTENTION_DISABLE_SM80="TRUE" +export FLASH_ATTENTION_DISABLE_FP8="TRUE" +export FLASH_ATTENTION_DISABLE_SOFTCAP="TRUE" +# export FLASH_ATTENTION_DISABLE_LOCAL="TRUE" +# export FLASH_ATTENTION_DISABLE_BACKWARD="TRUE" +export FLASH_ATTENTION_DISABLE_APPENDKV="TRUE" -cd ../csrc/fused_dense_lib -python setup.py bdist_wheel +export FLASH_ATTENTION_FORCE_BUILD="TRUE" -cd ../.. -cp dist/*.whl ./whls/ -cp hopper/dist/*.whl ./whls/ -cp csrc/fused_dense_lib/dist/*.whl ./whls/ +# export FLASH_ATTENTION_DISABLE_HDIM64="TRUE" +# export FLASH_ATTENTION_DISABLE_HDIM96="TRUE" +# export FLASH_ATTENTION_DISABLE_HDIM128="TRUE" +# export FLASH_ATTENTION_DISABLE_HDIM192="TRUE" -pip install ./whls/*.whl +export MAX_JOBS=32 + +cd hopper +# rm -rf __pycache__/ flash_attn_3.egg-info/ build/ dist/ +python setup.py bdist_wheel +pip install --no-deps --force-reinstall dist/*.whl diff --git a/warp_reinstall.sh b/warp_reinstall.sh new file mode 100755 index 00000000000..3e74bd28792 --- /dev/null +++ b/warp_reinstall.sh @@ -0,0 +1 @@ +time ./reinstall.sh > my_install.log 2>&1 From e9a60cb699034bbc9f4500d43e028a0066ff481f Mon Sep 17 00:00:00 2001 From: aoxy Date: Mon, 19 Jan 2026 21:08:30 +0800 Subject: [PATCH 03/23] First bwd. --- hopper/epilogue_bwd.hpp | 27 ++++++++++++++++++------ hopper/flash_api.cpp | 1 + hopper/flash_attn_interface.py | 6 +++--- hopper/flash_bwd_kernel_sm90.h | 5 +++-- hopper/flash_bwd_launch_template.h | 16 ++++++++------ hopper/flash_fwd_launch_template.h | 7 +++--- hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp | 20 +++++++++++++++--- warp_reinstall.sh | 2 ++ 8 files changed, 61 insertions(+), 23 deletions(-) diff --git a/hopper/epilogue_bwd.hpp b/hopper/epilogue_bwd.hpp index 6d9b5f4f596..b64ae854878 100644 --- a/hopper/epilogue_bwd.hpp +++ b/hopper/epilogue_bwd.hpp @@ -19,7 +19,7 @@ namespace flash { using namespace cute; template + int NumEpilogueThreads_, bool Varlen_, bool dKV_swapAB_, int AtomLayoutKdKV=1, bool Has_sink=false> struct CollectiveEpilogueBwd { using TileShape_MNK = TileShape_MNK_; @@ -114,6 +114,7 @@ struct CollectiveEpilogueBwd { int* dv_semaphore; int const* cu_seqlens; int const* seqused; + float* dsink_ptr; }; // Device side kernel params @@ -127,6 +128,7 @@ struct CollectiveEpilogueBwd { TMA_dKV tma_store_dK, tma_store_dV; int const* cu_seqlens = nullptr; int const* seqused = nullptr; + float* dsink_ptr; }; static Params @@ -148,7 +150,7 @@ struct CollectiveEpilogueBwd { } }(); return {args.ptr_dK, args.shape_dK, args.stride_dK, args.ptr_dV, args.shape_dV, args.stride_dV, - tma_store_dK, tma_store_dV, args.cu_seqlens, args.seqused}; + tma_store_dK, tma_store_dV, args.cu_seqlens, args.seqused, args.dsink_ptr}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -168,7 +170,8 @@ struct CollectiveEpilogueBwd { SharedStorage& shared_storage, TiledMma tiled_mma, int thread_idx, - cute::tuple const& block_coord + cute::tuple const& block_coord, + float dsink_val=0.0f ) { auto [n_block, bidh, bidb] = block_coord; @@ -267,6 +270,11 @@ struct CollectiveEpilogueBwd { gmem_tiled_copy_dKV, tdKVrdK, tdKVgdK, tdKVcdKV, tdKVpdK, std::min(seqlen_info.seqlen - n_block * kBlockN, kBlockN) ); } + + if constexpr (Has_sink) { + float* dsink_ptr = reinterpret_cast(params.dsink_ptr); + atomicAdd(dsink_ptr + bidh, -dsink_val); + } } CUTLASS_DEVICE void @@ -318,7 +326,7 @@ struct CollectiveEpilogueBwd { }; template + int NumEpilogueThreads_, bool Varlen_, bool Deterministic, bool Has_sink=false> struct CollectiveEpilogueBwdGQA { using TileShape_MNK = TileShape_MNK_; @@ -374,6 +382,7 @@ struct CollectiveEpilogueBwdGQA { int* dv_semaphore; int const* cu_seqlens; int const* seqused; + float* dsink_ptr; }; // Device side kernel params @@ -389,6 +398,7 @@ struct CollectiveEpilogueBwdGQA { int* dv_semaphore; int const* cu_seqlens = nullptr; int const* seqused = nullptr; + float* dsink_ptr; }; static Params @@ -400,7 +410,7 @@ struct CollectiveEpilogueBwdGQA { return {args.ptr_dKaccum, args.shape_dKaccum, args.stride_dKaccum, args.ptr_dVaccum, args.shape_dVaccum, args.stride_dVaccum, cutlass::FastDivmod(cute::ceil_div(args.num_heads_q, get<1>(args.shape_dKaccum))), args.dk_semaphore, args.dv_semaphore, - args.cu_seqlens, args.seqused}; + args.cu_seqlens, args.seqused, args.dsink_ptr}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -416,7 +426,8 @@ struct CollectiveEpilogueBwdGQA { SharedStorage& shared_storage, TiledMma tiled_mma, int thread_idx, - cute::tuple const& block_coord + cute::tuple const& block_coord, + float dsink_val=0.0f ) { auto [n_block, bidh, bidb] = block_coord; @@ -507,6 +518,10 @@ struct CollectiveEpilogueBwdGQA { #pragma unroll for (int i = 0; i < size(tdKrdK_atomic); ++i) { atomicAdd(&tdKgdK_atomic(i), tdKrdK_atomic(i)); } } + if constexpr (Has_sink) { + float* dsink_ptr = reinterpret_cast(params.dsink_ptr); + atomicAdd(dsink_ptr + bidh, -dsink_val); + } if constexpr (Deterministic) { Barrier::arrive_inc(lock_ptr, thread_idx, n_block * num_batch * num_head_kv); } diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp index e76451972d4..4def6028d28 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -1053,6 +1053,7 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq TORCH_CHECK(learnable_sink.stride(-1) == 1, "Learnable sink tensor must have contiguous last dimension"); CHECK_SHAPE(learnable_sink, num_heads); params.learnable_sink_ptr = learnable_sink.data_ptr(); + params.pack_gqa = false; // Disable pack_gqa if learnable sink is provided } else { params.learnable_sink_ptr = nullptr; } diff --git a/hopper/flash_attn_interface.py b/hopper/flash_attn_interface.py index e41e20eb1e1..f1e76e8e4cb 100644 --- a/hopper/flash_attn_interface.py +++ b/hopper/flash_attn_interface.py @@ -256,7 +256,7 @@ def backward(ctx, dout, *args): dsink, ) dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, dsink, None, None, None, None, None, None, None, None, None, None, None + return dqkv, None, None, None, None, None, None, None, None, None, None, dsink, None class FlashAttnFunc(torch.autograd.Function): @@ -348,7 +348,7 @@ def backward(ctx, dout, *args): dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, dsink, None, None, None, None, None, None, None, None, None, None, None, None, None, None + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink, None class FlashAttnVarlenFunc(torch.autograd.Function): @@ -455,7 +455,7 @@ def backward(ctx, dout, *args): dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, dsink, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink, None def flash_attn_qkvpacked_func( diff --git a/hopper/flash_bwd_kernel_sm90.h b/hopper/flash_bwd_kernel_sm90.h index b93a0219161..aed5e9fcf89 100644 --- a/hopper/flash_bwd_kernel_sm90.h +++ b/hopper/flash_bwd_kernel_sm90.h @@ -261,12 +261,13 @@ class FlashAttnBwdSm90 { // dK and dV output accumulator. Tensor tdKrdK = partition_fragment_C(tiled_mma_dKV, select(TileShape_MNK{})); Tensor tdVrdV = partition_fragment_C(tiled_mma_dKV, select(TileShape_MNK{})); + float dsink_val = 0.0f; bool tile_valid = mainloop.mma( params.mainloop, pipeline_q, pipeline_do, smem_pipe_read, smem_pipe_read_do, - tdKrdK, tdVrdV, threadIdx.x - NumCopyThreads, work_idx, block_coord, shared_storage); + tdKrdK, tdVrdV, threadIdx.x - NumCopyThreads, work_idx, block_coord, shared_storage, dsink_val); if (tile_valid) { epilogue.store(params.epilogue, tdKrdK, tdVrdV, shared_storage, tiled_mma_dKV, - threadIdx.x - NumCopyThreads, block_coord); + threadIdx.x - NumCopyThreads, block_coord, dsink_val); } else { epilogue.store_zero(params.epilogue, threadIdx.x - NumCopyThreads, block_coord); } diff --git a/hopper/flash_bwd_launch_template.h b/hopper/flash_bwd_launch_template.h index b6e8810b25f..34495cbc55b 100644 --- a/hopper/flash_bwd_launch_template.h +++ b/hopper/flash_bwd_launch_template.h @@ -28,7 +28,7 @@ template + bool V_in_regs=false, bool Has_sink=false> void run_flash_bwd(Flash_bwd_params ¶ms, cudaStream_t stream) { static_assert(!(Is_causal && Is_local), "Is_causal and Is_local cannot be true at the same time."); using ElementAccum = float; @@ -83,15 +83,15 @@ void run_flash_bwd(Flash_bwd_params ¶ms, cudaStream_t stream) { Arch >= 90, flash::CollectiveMainloopBwdSm90, + SdP_swapAB, dKV_swapAB, dQ_swapAB, NumMmaWarpGroups, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ, V_in_regs, Has_sink>, flash::CollectiveMainloopBwdSm80 >; using CollectiveEpilogue = std::conditional_t< !GQA, - flash::CollectiveEpilogueBwd= 90 ? 1 : cutlass::NumWarpsPerWarpGroup) / AtomLayoutNdKV>, - flash::CollectiveEpilogueBwdGQA + flash::CollectiveEpilogueBwd= 90 ? 1 : cutlass::NumWarpsPerWarpGroup) / AtomLayoutNdKV, Has_sink>, + flash::CollectiveEpilogueBwdGQA >; using Scheduler = std::conditional_t< Is_causal && !Varlen, @@ -131,7 +131,8 @@ void run_flash_bwd(Flash_bwd_params ¶ms, cudaStream_t stream) { params.b, params.dq_semaphore, params.cu_seqlens_q, params.cu_seqlens_k, - params.seqused_q, params.seqused_k + params.seqused_q, params.seqused_k, + reinterpret_cast(params.learnable_sink_ptr) }; // The case work with GQA is ugly but idk how to fix it. typename CollectiveEpilogue::Arguments epilogue_args { @@ -170,6 +171,7 @@ void run_flash_bwd(Flash_bwd_params ¶ms, cudaStream_t stream) { params.dv_semaphore, params.cu_seqlens_k, params.seqused_k, + reinterpret_cast(params.dsink_ptr), }; int num_blocks_n = cutlass::ceil_div(params.seqlen_k, get<1>(TileShape_MNK{})); @@ -301,11 +303,13 @@ template(params, stream); - run_flash_bwd(params, stream); + run_flash_bwd(params, stream); // }); }); + }); }); } diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index adb290dee08..16eb7c8a48b 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -173,7 +173,7 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { int smem_size_q = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_q)); int smem_size_k = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k)); int smem_size_v = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v)); - printf("========== smem_size = %d, q = %d, k = %d, v = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v); + // printf("========== smem_size = %d, q = %d, k = %d, v = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v); // Get the ptr to kernel function. if constexpr (size(ClusterShape{}) > 1) { void const* kernel = (void const*) cutlass::device_kernel; @@ -212,10 +212,11 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { static constexpr bool HasQv = HasQV_ && Arch == 90 && !Is_FP8 && kHeadDim == 64 && kHeadDimV >= 256; APPENDKV_SWITCH(params.knew_ptr, AppendKV, [&] { SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { + static constexpr bool PackGQA_Sink = PackGQA && !Has_sink; // Only use Cluster if number of tiles along seqlen_q is even and not varlen - CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { + CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; - run_flash_fwd(params, stream); + run_flash_fwd(params, stream); }); }); }); diff --git a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp index ec34e20eca1..b0a1f229d90 100644 --- a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp @@ -31,7 +31,7 @@ template + bool Mma_dP_is_RS=false, bool Has_sink=false> struct CollectiveMainloopBwdSm90 { static constexpr int kStages = Stages; @@ -320,6 +320,7 @@ struct CollectiveMainloopBwdSm90 { int const* const cu_seqlens_k = nullptr; int const* const seqused_q = nullptr; int const* const seqused_k = nullptr; + float const* const ptr_Sink = nullptr; }; // Device side kernel params @@ -350,6 +351,7 @@ struct CollectiveMainloopBwdSm90 { int const* const cu_seqlens_k = nullptr; int const* const seqused_q = nullptr; int const* const seqused_k = nullptr; + float const* const ptr_Sink = nullptr; }; static Params @@ -406,7 +408,7 @@ struct CollectiveMainloopBwdSm90 { args.window_size_left, args.window_size_right, attention_chunk_divmod, !Has_softcap ? 0.f : args.softmax_scale / args.softcap_val, args.num_batch, args.dq_semaphore, - args.cu_seqlens_q, args.cu_seqlens_k, args.seqused_q, args.seqused_k}; + args.cu_seqlens_q, args.cu_seqlens_k, args.seqused_q, args.seqused_k, args.ptr_Sink}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -683,7 +685,8 @@ struct CollectiveMainloopBwdSm90 { int thread_idx, int &work_idx, cute::tuple block_coord, - SharedStorage& shared_storage + SharedStorage& shared_storage, + float &dsink_val=0.0f ) { static_assert(is_rmem::value, "dK and dV tensor must be rmem resident."); @@ -790,6 +793,8 @@ struct CollectiveMainloopBwdSm90 { int const seqlen_q = seqlen_info.seqlen_q; int const seqlen_k = seqlen_info.seqlen_k; + const float sink_val = !Has_sink ? -INFINITY : reinterpret_cast(params.ptr_Sink)[bidh]; + // For the case where we do atomicAdd directly to gdQaccum instead of using TMA bool const is_varlen = Varlen && params.cu_seqlens_q; Tensor mdQaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.ptr_dQaccum)), @@ -876,15 +881,24 @@ struct CollectiveMainloopBwdSm90 { Tensor dS = make_tensor(tdPrdP.data(), scores.layout()); #pragma unroll for (int mi = 0; mi < size<0>(dS); ++mi) { + float dsink_val_cols = 0.0f; float const dP_sum_cur = [&] { if constexpr (!ShuffledPsum) return tLSErdPsum(mi); else return __shfl_sync(0xffffffff, tLSErdPsum(mi / 8), (mi % 8) * 4 + (thread_idx % 4)); }(); #pragma unroll for (int ni = 0; ni < size<1>(dS); ++ni) { + if constexpr (Has_sink) { dsink_val_cols += scores(mi, ni) * dS(mi, ni); } dS(mi, ni) = scores(mi, ni) * (dS(mi, ni) - dP_sum_cur); if constexpr (Has_softcap) { dS(mi, ni) *= dtanh(mi, ni); } } + if constexpr (Has_sink) { + float const lse_scaled = [&] { + if constexpr (!ShuffleLSE) return tLSErLSE(mi); + else return __shfl_sync(0xffffffff, tLSErLSE(mi / 8), (mi % 8) * 4 + (thread_idx % 4)); + }(); + dsink_val += dsink_val_cols * exp2f(sink_val * float(M_LOG2E) - lse_scaled); + } } // Convert scores from fp32 to fp16/bf16 diff --git a/warp_reinstall.sh b/warp_reinstall.sh index 3e74bd28792..2c1a511309c 100755 --- a/warp_reinstall.sh +++ b/warp_reinstall.sh @@ -1 +1,3 @@ time ./reinstall.sh > my_install.log 2>&1 +grep -i "error" my_install.log +tail -n 10 my_install.log \ No newline at end of file From 1ecd4be3f77528d4ed28dbc95a257d68bef32ba0 Mon Sep 17 00:00:00 2001 From: aoxy Date: Tue, 20 Jan 2026 16:31:05 +0800 Subject: [PATCH 04/23] Useless optimization. --- hopper/benchmark_attn.py | 25 ++++++--- hopper/epilogue_bwd.hpp | 68 ++++++++++++++++++++---- hopper/flash_bwd_kernel_sm90.h | 15 +++++- hopper/flash_bwd_launch_template.h | 1 + hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp | 3 +- 5 files changed, 94 insertions(+), 18 deletions(-) diff --git a/hopper/benchmark_attn.py b/hopper/benchmark_attn.py index 33e5d282716..3e42c73511b 100644 --- a/hopper/benchmark_attn.py +++ b/hopper/benchmark_attn.py @@ -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 @@ -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 @@ -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=torch.float32, 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 @@ -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: @@ -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) @@ -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) diff --git a/hopper/epilogue_bwd.hpp b/hopper/epilogue_bwd.hpp index b64ae854878..2bd361f9af6 100644 --- a/hopper/epilogue_bwd.hpp +++ b/hopper/epilogue_bwd.hpp @@ -29,6 +29,7 @@ struct CollectiveEpilogueBwd { static constexpr bool Varlen = Varlen_; static constexpr bool dKV_swapAB = dKV_swapAB_; static constexpr bool Use_TMA = !Varlen && ArchTag::kMinComputeCapability >= 90; + static constexpr uint32_t SinkHeads = 24; static_assert(ArchTag::kMinComputeCapability >= 80); @@ -115,6 +116,7 @@ struct CollectiveEpilogueBwd { int const* cu_seqlens; int const* seqused; float* dsink_ptr; + int const batch_size; }; // Device side kernel params @@ -129,6 +131,8 @@ struct CollectiveEpilogueBwd { int const* cu_seqlens = nullptr; int const* seqused = nullptr; float* dsink_ptr; + int const num_heads_q; + int const batch_size; }; static Params @@ -150,7 +154,7 @@ struct CollectiveEpilogueBwd { } }(); return {args.ptr_dK, args.shape_dK, args.stride_dK, args.ptr_dV, args.shape_dV, args.stride_dV, - tma_store_dK, tma_store_dV, args.cu_seqlens, args.seqused, args.dsink_ptr}; + tma_store_dK, tma_store_dV, args.cu_seqlens, args.seqused, args.dsink_ptr, args.num_heads_q, args.batch_size}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -272,14 +276,34 @@ struct CollectiveEpilogueBwd { } if constexpr (Has_sink) { - float* dsink_ptr = reinterpret_cast(params.dsink_ptr); - atomicAdd(dsink_ptr + bidh, -dsink_val); + SumOp sum_op; + dsink_val = Allreduce<32>::run(dsink_val, sum_op); + if (thread_idx % 32 == 0) { + if (dsink_val != 0.0f) { + Layout smem_layout = make_layout(make_shape(Int{})); + Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); + atomicAdd(&dsink_h_sum(bidh), -dsink_val); + } + } } } + template CUTLASS_DEVICE void - store_tail() { - // if constexpr (Use_TMA) { tma_store_wait<0>(); } + store_tail(Params const& params, + SharedStorage& shared_storage, + int thread_idx + ) { + if constexpr (Has_sink) { + if (thread_idx < SinkHeads) { + Layout smem_layout = make_layout(make_shape(Int{})); + Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); + float local_sum = dsink_h_sum(thread_idx); + if (local_sum != 0.0f) { + atomicAdd(params.dsink_ptr + thread_idx, local_sum); + } + } + } } // Write 0 to dK and dV @@ -335,6 +359,7 @@ struct CollectiveEpilogueBwdGQA { static constexpr int NumEpilogueThreads = NumEpilogueThreads_; static constexpr bool Varlen = Varlen_; static constexpr bool Use_TMA = ArchTag::kMinComputeCapability >= 90; + static constexpr uint32_t SinkHeads = 24; static_assert(ArchTag::kMinComputeCapability >= 80); @@ -383,6 +408,7 @@ struct CollectiveEpilogueBwdGQA { int const* cu_seqlens; int const* seqused; float* dsink_ptr; + int const batch_size; }; // Device side kernel params @@ -399,6 +425,8 @@ struct CollectiveEpilogueBwdGQA { int const* cu_seqlens = nullptr; int const* seqused = nullptr; float* dsink_ptr; + int const num_heads_q; + int const batch_size; }; static Params @@ -410,7 +438,7 @@ struct CollectiveEpilogueBwdGQA { return {args.ptr_dKaccum, args.shape_dKaccum, args.stride_dKaccum, args.ptr_dVaccum, args.shape_dVaccum, args.stride_dVaccum, cutlass::FastDivmod(cute::ceil_div(args.num_heads_q, get<1>(args.shape_dKaccum))), args.dk_semaphore, args.dv_semaphore, - args.cu_seqlens, args.seqused, args.dsink_ptr}; + args.cu_seqlens, args.seqused, args.dsink_ptr, args.num_heads_q, args.batch_size}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -519,8 +547,16 @@ struct CollectiveEpilogueBwdGQA { for (int i = 0; i < size(tdKrdK_atomic); ++i) { atomicAdd(&tdKgdK_atomic(i), tdKrdK_atomic(i)); } } if constexpr (Has_sink) { - float* dsink_ptr = reinterpret_cast(params.dsink_ptr); - atomicAdd(dsink_ptr + bidh, -dsink_val); + SumOp sum_op; + dsink_val = Allreduce<32>::run(dsink_val, sum_op); + if (thread_idx % 32 == 0) { + if (dsink_val != 0.0f) { + Layout smem_layout = make_layout(make_shape(Int{})); + Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); + atomicAdd(&dsink_h_sum(bidh), -dsink_val); + } + } + } if constexpr (Deterministic) { Barrier::arrive_inc(lock_ptr, thread_idx, n_block * num_batch * num_head_kv); @@ -529,8 +565,22 @@ struct CollectiveEpilogueBwdGQA { // flash::named_barrier_arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp, static_cast(BwdNamedBarriers::KVEmpty) /*id*/); } + template CUTLASS_DEVICE void - store_tail() { + store_tail(Params const& params, + SharedStorage& shared_storage, + int thread_idx + ) { + if constexpr (Has_sink) { + if (thread_idx < SinkHeads) { + Layout smem_layout = make_layout(make_shape(Int{})); + Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); + float local_sum = dsink_h_sum(thread_idx); + if (local_sum != 0.0f) { + atomicAdd(params.dsink_ptr + thread_idx, local_sum); + } + } + } } // Write 0 to dK and dV diff --git a/hopper/flash_bwd_kernel_sm90.h b/hopper/flash_bwd_kernel_sm90.h index aed5e9fcf89..87c2a7f52ec 100644 --- a/hopper/flash_bwd_kernel_sm90.h +++ b/hopper/flash_bwd_kernel_sm90.h @@ -42,6 +42,7 @@ class FlashAttnBwdSm90 { using MainloopArguments = typename CollectiveMainloop::Arguments; using MainloopParams = typename CollectiveMainloop::Params; static constexpr bool dKV_swapAB = CollectiveMainloop::dKV_swapAB; + static constexpr bool Has_sink = CollectiveMainloop::Has_sink; // Epilogue derived types using CollectiveEpilogue = CollectiveEpilogue_; @@ -68,6 +69,9 @@ class FlashAttnBwdSm90 { // static constexpr uint32_t LoadRegisterRequirement = 40; // static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 232 : 152; + static constexpr uint32_t SinkHeads = 24; // 24 is the maximum number of heads that can be processed in a single block? + using TensorDSink = std::conditional_t, cute::array>; + // Kernel level shared memory storage struct SharedStorage { struct TensorStorage : cute::aligned_struct<128> { @@ -84,6 +88,7 @@ class FlashAttnBwdSm90 { alignas(16) typename TileScheduler::SharedStorage smem_scheduler; } pipelines; + TensorDSink dsink_h_sum; }; static constexpr int SharedStorageSize = sizeof(SharedStorage); @@ -198,6 +203,14 @@ class FlashAttnBwdSm90 { CollectiveMainloop mainloop; CollectiveEpilogue epilogue; + if constexpr (Has_sink) { + Layout smem_layout = make_layout(make_shape(Int{})); + Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); + if (threadIdx.x < SinkHeads) { + dsink_h_sum(threadIdx.x) = 0.f; + } + } + // We need this to guarantee that the Pipeline init is visible to all producers and consumer blocks in the Cluster if constexpr (size(ClusterShape{}) > 1) { cute::cluster_arrive_relaxed(); @@ -273,7 +286,7 @@ class FlashAttnBwdSm90 { } } - epilogue.store_tail(); + epilogue.store_tail(params.epilogue, shared_storage, threadIdx.x - NumCopyThreads); } } diff --git a/hopper/flash_bwd_launch_template.h b/hopper/flash_bwd_launch_template.h index 34495cbc55b..c9d34b6911d 100644 --- a/hopper/flash_bwd_launch_template.h +++ b/hopper/flash_bwd_launch_template.h @@ -172,6 +172,7 @@ void run_flash_bwd(Flash_bwd_params ¶ms, cudaStream_t stream) { params.cu_seqlens_k, params.seqused_k, reinterpret_cast(params.dsink_ptr), + params.b, }; int num_blocks_n = cutlass::ceil_div(params.seqlen_k, get<1>(TileShape_MNK{})); diff --git a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp index b0a1f229d90..1001aea2a99 100644 --- a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp @@ -31,7 +31,7 @@ template + bool Mma_dP_is_RS=false, bool Has_sink_=false> struct CollectiveMainloopBwdSm90 { static constexpr int kStages = Stages; @@ -53,6 +53,7 @@ struct CollectiveMainloopBwdSm90 { static constexpr bool SdP_swapAB = SdP_swapAB_; static constexpr bool dKV_swapAB = dKV_swapAB_; static constexpr bool dQ_swapAB = dQ_swapAB_; + static constexpr bool Has_sink = Has_sink_; static constexpr bool Q_dO_same_stages = kStages == kStages_dO; From 284d2d79162fecbf16c5f9bbf5d1e2bc5ac7e564 Mon Sep 17 00:00:00 2001 From: aoxy Date: Tue, 20 Jan 2026 19:49:59 +0800 Subject: [PATCH 05/23] Another useless optimization. --- hopper/epilogue_bwd.hpp | 96 ++++++++++++++++++++-------------- hopper/flash_bwd_kernel_sm90.h | 31 ++++++----- 2 files changed, 74 insertions(+), 53 deletions(-) diff --git a/hopper/epilogue_bwd.hpp b/hopper/epilogue_bwd.hpp index 2bd361f9af6..36fbf9c6523 100644 --- a/hopper/epilogue_bwd.hpp +++ b/hopper/epilogue_bwd.hpp @@ -29,7 +29,7 @@ struct CollectiveEpilogueBwd { static constexpr bool Varlen = Varlen_; static constexpr bool dKV_swapAB = dKV_swapAB_; static constexpr bool Use_TMA = !Varlen && ArchTag::kMinComputeCapability >= 90; - static constexpr uint32_t SinkHeads = 24; + static constexpr uint32_t SinkHeads = 25; static_assert(ArchTag::kMinComputeCapability >= 80); @@ -275,32 +275,40 @@ struct CollectiveEpilogueBwd { ); } - if constexpr (Has_sink) { - SumOp sum_op; - dsink_val = Allreduce<32>::run(dsink_val, sum_op); - if (thread_idx % 32 == 0) { - if (dsink_val != 0.0f) { - Layout smem_layout = make_layout(make_shape(Int{})); - Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); - atomicAdd(&dsink_h_sum(bidh), -dsink_val); - } - } - } + // if constexpr (Has_sink) { + // SumOp sum_op; + // dsink_val = Allreduce<32>::run(dsink_val, sum_op); + // if (thread_idx % 32 == 0) { + // if (dsink_val != 0.0f) { + // Layout smem_layout = make_layout(make_shape(Int{})); + // Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); + // atomicAdd(&dsink_h_sum(bidh), -dsink_val); + // } + // } + // } } - template + template CUTLASS_DEVICE void store_tail(Params const& params, SharedStorage& shared_storage, int thread_idx ) { - if constexpr (Has_sink) { - if (thread_idx < SinkHeads) { - Layout smem_layout = make_layout(make_shape(Int{})); - Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); - float local_sum = dsink_h_sum(thread_idx); - if (local_sum != 0.0f) { - atomicAdd(params.dsink_ptr + thread_idx, local_sum); + if constexpr (Has_sink) { + if (thread_idx == 0) { + float final_sum[SinkHeads] = {0}; + #pragma unroll + for (int w = 0; w < NumMmaWarp; ++w) { + #pragma unroll + for (int h = 0; h < SinkHeads; ++h) { + final_sum[h] += shared_storage.dsink_warp_acc[h][w]; + } + } + #pragma unroll + for (int h = 0; h < params.num_heads_q; ++h) { + if (final_sum[h] != 0.0f) { + atomicAdd(params.dsink_ptr + h, -final_sum[h]); + } } } } @@ -359,7 +367,7 @@ struct CollectiveEpilogueBwdGQA { static constexpr int NumEpilogueThreads = NumEpilogueThreads_; static constexpr bool Varlen = Varlen_; static constexpr bool Use_TMA = ArchTag::kMinComputeCapability >= 90; - static constexpr uint32_t SinkHeads = 24; + static constexpr uint32_t SinkHeads = 25; static_assert(ArchTag::kMinComputeCapability >= 80); @@ -546,18 +554,18 @@ struct CollectiveEpilogueBwdGQA { #pragma unroll for (int i = 0; i < size(tdKrdK_atomic); ++i) { atomicAdd(&tdKgdK_atomic(i), tdKrdK_atomic(i)); } } - if constexpr (Has_sink) { - SumOp sum_op; - dsink_val = Allreduce<32>::run(dsink_val, sum_op); - if (thread_idx % 32 == 0) { - if (dsink_val != 0.0f) { - Layout smem_layout = make_layout(make_shape(Int{})); - Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); - atomicAdd(&dsink_h_sum(bidh), -dsink_val); - } - } + // if constexpr (Has_sink) { + // SumOp sum_op; + // dsink_val = Allreduce<32>::run(dsink_val, sum_op); + // if (thread_idx % 32 == 0) { + // if (dsink_val != 0.0f) { + // Layout smem_layout = make_layout(make_shape(Int{})); + // Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); + // atomicAdd(&dsink_h_sum(bidh), -dsink_val); + // } + // } - } + // } if constexpr (Deterministic) { Barrier::arrive_inc(lock_ptr, thread_idx, n_block * num_batch * num_head_kv); } @@ -565,19 +573,27 @@ struct CollectiveEpilogueBwdGQA { // flash::named_barrier_arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp, static_cast(BwdNamedBarriers::KVEmpty) /*id*/); } - template + template CUTLASS_DEVICE void store_tail(Params const& params, SharedStorage& shared_storage, int thread_idx ) { - if constexpr (Has_sink) { - if (thread_idx < SinkHeads) { - Layout smem_layout = make_layout(make_shape(Int{})); - Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); - float local_sum = dsink_h_sum(thread_idx); - if (local_sum != 0.0f) { - atomicAdd(params.dsink_ptr + thread_idx, local_sum); + if constexpr (Has_sink) { + if (thread_idx == 0) { + float final_sum[SinkHeads] = {0}; + #pragma unroll + for (int w = 0; w < NumMmaWarp; ++w) { + #pragma unroll + for (int h = 0; h < SinkHeads; ++h) { + final_sum[h] += shared_storage.dsink_warp_acc[h][w]; + } + } + #pragma unroll + for (int h = 0; h < params.num_heads_q; ++h) { + if (final_sum[h] != 0.0f) { + atomicAdd(params.dsink_ptr + h, -final_sum[h]); + } } } } diff --git a/hopper/flash_bwd_kernel_sm90.h b/hopper/flash_bwd_kernel_sm90.h index 87c2a7f52ec..76959d9b061 100644 --- a/hopper/flash_bwd_kernel_sm90.h +++ b/hopper/flash_bwd_kernel_sm90.h @@ -69,7 +69,7 @@ class FlashAttnBwdSm90 { // static constexpr uint32_t LoadRegisterRequirement = 40; // static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 232 : 152; - static constexpr uint32_t SinkHeads = 24; // 24 is the maximum number of heads that can be processed in a single block? + static constexpr uint32_t SinkHeads = 32; // Padded to avoid shared memory bank conflicts. Original max heads: 24. using TensorDSink = std::conditional_t, cute::array>; // Kernel level shared memory storage @@ -88,7 +88,8 @@ class FlashAttnBwdSm90 { alignas(16) typename TileScheduler::SharedStorage smem_scheduler; } pipelines; - TensorDSink dsink_h_sum; + // TensorDSink dsink_h_sum; + float dsink_warp_acc[SinkHeads][NumMmaWarpGroups * cutlass::NumWarpsPerWarpGroup]; }; static constexpr int SharedStorageSize = sizeof(SharedStorage); @@ -203,14 +204,13 @@ class FlashAttnBwdSm90 { CollectiveMainloop mainloop; CollectiveEpilogue epilogue; - if constexpr (Has_sink) { - Layout smem_layout = make_layout(make_shape(Int{})); - Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); - if (threadIdx.x < SinkHeads) { - dsink_h_sum(threadIdx.x) = 0.f; - } + int const mma_warp_idx = warp_idx - NumLoadWarpGroups * cutlass::NumWarpsPerWarpGroup; + int const mma_lane = (threadIdx.x - NumCopyThreads) % cutlass::NumThreadsPerWarp; + for (int i = mma_lane; i < SinkHeads; i+= cutlass::NumThreadsPerWarp) { + shared_storage.dsink_warp_acc[i][mma_warp_idx] = 0.0f; } + // We need this to guarantee that the Pipeline init is visible to all producers and consumer blocks in the Cluster if constexpr (size(ClusterShape{}) > 1) { cute::cluster_arrive_relaxed(); @@ -262,6 +262,7 @@ class FlashAttnBwdSm90 { mainloop.mma_init(); scheduler.init_consumer(); + int thread_idx = threadIdx.x - NumCopyThreads; int work_idx = 0; CUTLASS_PRAGMA_NO_UNROLL for (auto work_tile_info = scheduler.template get_initial_work(params.scheduler); @@ -277,16 +278,20 @@ class FlashAttnBwdSm90 { float dsink_val = 0.0f; bool tile_valid = mainloop.mma( params.mainloop, pipeline_q, pipeline_do, smem_pipe_read, smem_pipe_read_do, - tdKrdK, tdVrdV, threadIdx.x - NumCopyThreads, work_idx, block_coord, shared_storage, dsink_val); + tdKrdK, tdVrdV, thread_idx, work_idx, block_coord, shared_storage, dsink_val); + SumOp sum_op; + dsink_val = Allreduce::run(dsink_val, sum_op); + if (lane_predicate) { + shared_storage.dsink_warp_acc[bidh][mma_warp_idx] += dsink_val; + } if (tile_valid) { epilogue.store(params.epilogue, tdKrdK, tdVrdV, shared_storage, tiled_mma_dKV, - threadIdx.x - NumCopyThreads, block_coord, dsink_val); + thread_idx, block_coord, dsink_val); } else { - epilogue.store_zero(params.epilogue, threadIdx.x - NumCopyThreads, block_coord); + epilogue.store_zero(params.epilogue, thread_idx, block_coord); } - } - epilogue.store_tail(params.epilogue, shared_storage, threadIdx.x - NumCopyThreads); + epilogue.store_tail(params.epilogue, shared_storage, thread_idx); } } From 4bbedfd213d5a4a69cefebc2f504db6cd28f8bfa Mon Sep 17 00:00:00 2001 From: aoxy Date: Tue, 20 Jan 2026 21:37:50 +0800 Subject: [PATCH 06/23] Clean code. --- hopper/benchmark_attn.py | 2 +- hopper/benchmark_attn_sink.txt | 40 ++++++++++ hopper/epilogue_bwd.hpp | 98 +++++------------------- hopper/flash_bwd_kernel_sm90.h | 24 +----- hopper/flash_bwd_launch_template.h | 13 ++-- hopper/flash_fwd_kernel_sm90.h | 3 - hopper/flash_fwd_launch_template.h | 20 ++--- hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp | 2 - hopper/softmax.h | 2 - 9 files changed, 80 insertions(+), 124 deletions(-) create mode 100644 hopper/benchmark_attn_sink.txt diff --git a/hopper/benchmark_attn.py b/hopper/benchmark_attn.py index 3e42c73511b..80234b58bde 100644 --- a/hopper/benchmark_attn.py +++ b/hopper/benchmark_attn.py @@ -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 = {} diff --git a/hopper/benchmark_attn_sink.txt b/hopper/benchmark_attn_sink.txt new file mode 100644 index 00000000000..a669e83c9c9 --- /dev/null +++ b/hopper/benchmark_attn_sink.txt @@ -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 diff --git a/hopper/epilogue_bwd.hpp b/hopper/epilogue_bwd.hpp index 36fbf9c6523..1bce5e07ff4 100644 --- a/hopper/epilogue_bwd.hpp +++ b/hopper/epilogue_bwd.hpp @@ -29,7 +29,6 @@ struct CollectiveEpilogueBwd { static constexpr bool Varlen = Varlen_; static constexpr bool dKV_swapAB = dKV_swapAB_; static constexpr bool Use_TMA = !Varlen && ArchTag::kMinComputeCapability >= 90; - static constexpr uint32_t SinkHeads = 25; static_assert(ArchTag::kMinComputeCapability >= 80); @@ -116,7 +115,6 @@ struct CollectiveEpilogueBwd { int const* cu_seqlens; int const* seqused; float* dsink_ptr; - int const batch_size; }; // Device side kernel params @@ -131,8 +129,6 @@ struct CollectiveEpilogueBwd { int const* cu_seqlens = nullptr; int const* seqused = nullptr; float* dsink_ptr; - int const num_heads_q; - int const batch_size; }; static Params @@ -154,7 +150,7 @@ struct CollectiveEpilogueBwd { } }(); return {args.ptr_dK, args.shape_dK, args.stride_dK, args.ptr_dV, args.shape_dV, args.stride_dV, - tma_store_dK, tma_store_dV, args.cu_seqlens, args.seqused, args.dsink_ptr, args.num_heads_q, args.batch_size}; + tma_store_dK, tma_store_dV, args.cu_seqlens, args.seqused, args.dsink_ptr}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -275,43 +271,18 @@ struct CollectiveEpilogueBwd { ); } - // if constexpr (Has_sink) { - // SumOp sum_op; - // dsink_val = Allreduce<32>::run(dsink_val, sum_op); - // if (thread_idx % 32 == 0) { - // if (dsink_val != 0.0f) { - // Layout smem_layout = make_layout(make_shape(Int{})); - // Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); - // atomicAdd(&dsink_h_sum(bidh), -dsink_val); - // } - // } - // } + if constexpr (Has_sink) { + SumOp sum_op; + dsink_val = Allreduce::run(dsink_val, sum_op); + if (thread_idx % cutlass::NumThreadsPerWarp == 0 && dsink_val != 0.0f) { + atomicAdd(params.dsink_ptr + bidh, -dsink_val); + } + } } - template CUTLASS_DEVICE void - store_tail(Params const& params, - SharedStorage& shared_storage, - int thread_idx - ) { - if constexpr (Has_sink) { - if (thread_idx == 0) { - float final_sum[SinkHeads] = {0}; - #pragma unroll - for (int w = 0; w < NumMmaWarp; ++w) { - #pragma unroll - for (int h = 0; h < SinkHeads; ++h) { - final_sum[h] += shared_storage.dsink_warp_acc[h][w]; - } - } - #pragma unroll - for (int h = 0; h < params.num_heads_q; ++h) { - if (final_sum[h] != 0.0f) { - atomicAdd(params.dsink_ptr + h, -final_sum[h]); - } - } - } - } + store_tail() { + // if constexpr (Use_TMA) { tma_store_wait<0>(); } } // Write 0 to dK and dV @@ -367,7 +338,6 @@ struct CollectiveEpilogueBwdGQA { static constexpr int NumEpilogueThreads = NumEpilogueThreads_; static constexpr bool Varlen = Varlen_; static constexpr bool Use_TMA = ArchTag::kMinComputeCapability >= 90; - static constexpr uint32_t SinkHeads = 25; static_assert(ArchTag::kMinComputeCapability >= 80); @@ -416,7 +386,6 @@ struct CollectiveEpilogueBwdGQA { int const* cu_seqlens; int const* seqused; float* dsink_ptr; - int const batch_size; }; // Device side kernel params @@ -433,8 +402,6 @@ struct CollectiveEpilogueBwdGQA { int const* cu_seqlens = nullptr; int const* seqused = nullptr; float* dsink_ptr; - int const num_heads_q; - int const batch_size; }; static Params @@ -446,7 +413,7 @@ struct CollectiveEpilogueBwdGQA { return {args.ptr_dKaccum, args.shape_dKaccum, args.stride_dKaccum, args.ptr_dVaccum, args.shape_dVaccum, args.stride_dVaccum, cutlass::FastDivmod(cute::ceil_div(args.num_heads_q, get<1>(args.shape_dKaccum))), args.dk_semaphore, args.dv_semaphore, - args.cu_seqlens, args.seqused, args.dsink_ptr, args.num_heads_q, args.batch_size}; + args.cu_seqlens, args.seqused, args.dsink_ptr}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -554,18 +521,13 @@ struct CollectiveEpilogueBwdGQA { #pragma unroll for (int i = 0; i < size(tdKrdK_atomic); ++i) { atomicAdd(&tdKgdK_atomic(i), tdKrdK_atomic(i)); } } - // if constexpr (Has_sink) { - // SumOp sum_op; - // dsink_val = Allreduce<32>::run(dsink_val, sum_op); - // if (thread_idx % 32 == 0) { - // if (dsink_val != 0.0f) { - // Layout smem_layout = make_layout(make_shape(Int{})); - // Tensor dsink_h_sum = make_tensor(make_smem_ptr(shared_storage.dsink_h_sum.data()), smem_layout); - // atomicAdd(&dsink_h_sum(bidh), -dsink_val); - // } - // } - - // } + if constexpr (Has_sink) { + SumOp sum_op; + dsink_val = Allreduce::run(dsink_val, sum_op); + if (thread_idx % cutlass::NumThreadsPerWarp == 0 && dsink_val != 0.0f) { + atomicAdd(params.dsink_ptr + bidh, -dsink_val); + } + } if constexpr (Deterministic) { Barrier::arrive_inc(lock_ptr, thread_idx, n_block * num_batch * num_head_kv); } @@ -573,30 +535,8 @@ struct CollectiveEpilogueBwdGQA { // flash::named_barrier_arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp, static_cast(BwdNamedBarriers::KVEmpty) /*id*/); } - template CUTLASS_DEVICE void - store_tail(Params const& params, - SharedStorage& shared_storage, - int thread_idx - ) { - if constexpr (Has_sink) { - if (thread_idx == 0) { - float final_sum[SinkHeads] = {0}; - #pragma unroll - for (int w = 0; w < NumMmaWarp; ++w) { - #pragma unroll - for (int h = 0; h < SinkHeads; ++h) { - final_sum[h] += shared_storage.dsink_warp_acc[h][w]; - } - } - #pragma unroll - for (int h = 0; h < params.num_heads_q; ++h) { - if (final_sum[h] != 0.0f) { - atomicAdd(params.dsink_ptr + h, -final_sum[h]); - } - } - } - } + store_tail() { } // Write 0 to dK and dV diff --git a/hopper/flash_bwd_kernel_sm90.h b/hopper/flash_bwd_kernel_sm90.h index 76959d9b061..613715dad2a 100644 --- a/hopper/flash_bwd_kernel_sm90.h +++ b/hopper/flash_bwd_kernel_sm90.h @@ -69,8 +69,6 @@ class FlashAttnBwdSm90 { // static constexpr uint32_t LoadRegisterRequirement = 40; // static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 232 : 152; - static constexpr uint32_t SinkHeads = 32; // Padded to avoid shared memory bank conflicts. Original max heads: 24. - using TensorDSink = std::conditional_t, cute::array>; // Kernel level shared memory storage struct SharedStorage { @@ -88,8 +86,6 @@ class FlashAttnBwdSm90 { alignas(16) typename TileScheduler::SharedStorage smem_scheduler; } pipelines; - // TensorDSink dsink_h_sum; - float dsink_warp_acc[SinkHeads][NumMmaWarpGroups * cutlass::NumWarpsPerWarpGroup]; }; static constexpr int SharedStorageSize = sizeof(SharedStorage); @@ -204,12 +200,6 @@ class FlashAttnBwdSm90 { CollectiveMainloop mainloop; CollectiveEpilogue epilogue; - int const mma_warp_idx = warp_idx - NumLoadWarpGroups * cutlass::NumWarpsPerWarpGroup; - int const mma_lane = (threadIdx.x - NumCopyThreads) % cutlass::NumThreadsPerWarp; - for (int i = mma_lane; i < SinkHeads; i+= cutlass::NumThreadsPerWarp) { - shared_storage.dsink_warp_acc[i][mma_warp_idx] = 0.0f; - } - // We need this to guarantee that the Pipeline init is visible to all producers and consumer blocks in the Cluster if constexpr (size(ClusterShape{}) > 1) { @@ -262,7 +252,6 @@ class FlashAttnBwdSm90 { mainloop.mma_init(); scheduler.init_consumer(); - int thread_idx = threadIdx.x - NumCopyThreads; int work_idx = 0; CUTLASS_PRAGMA_NO_UNROLL for (auto work_tile_info = scheduler.template get_initial_work(params.scheduler); @@ -278,20 +267,15 @@ class FlashAttnBwdSm90 { float dsink_val = 0.0f; bool tile_valid = mainloop.mma( params.mainloop, pipeline_q, pipeline_do, smem_pipe_read, smem_pipe_read_do, - tdKrdK, tdVrdV, thread_idx, work_idx, block_coord, shared_storage, dsink_val); - SumOp sum_op; - dsink_val = Allreduce::run(dsink_val, sum_op); - if (lane_predicate) { - shared_storage.dsink_warp_acc[bidh][mma_warp_idx] += dsink_val; - } + tdKrdK, tdVrdV, threadIdx.x - NumCopyThreads, work_idx, block_coord, shared_storage, dsink_val); if (tile_valid) { epilogue.store(params.epilogue, tdKrdK, tdVrdV, shared_storage, tiled_mma_dKV, - thread_idx, block_coord, dsink_val); + threadIdx.x - NumCopyThreads, block_coord, dsink_val); } else { - epilogue.store_zero(params.epilogue, thread_idx, block_coord); + epilogue.store_zero(params.epilogue, threadIdx.x - NumCopyThreads, block_coord); } } - epilogue.store_tail(params.epilogue, shared_storage, thread_idx); + epilogue.store_tail(); } } diff --git a/hopper/flash_bwd_launch_template.h b/hopper/flash_bwd_launch_template.h index c9d34b6911d..ec0198b0468 100644 --- a/hopper/flash_bwd_launch_template.h +++ b/hopper/flash_bwd_launch_template.h @@ -172,7 +172,6 @@ void run_flash_bwd(Flash_bwd_params ¶ms, cudaStream_t stream) { params.cu_seqlens_k, params.seqused_k, reinterpret_cast(params.dsink_ptr), - params.b, }; int num_blocks_n = cutlass::ceil_div(params.seqlen_k, get<1>(TileShape_MNK{})); @@ -304,12 +303,12 @@ template(params, stream); - run_flash_bwd(params, stream); -// }); - }); + SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { +// BOOL_SWITCH(params.deterministic, Deterministic, [&] { + // run_flash_bwd(params, stream); + run_flash_bwd(params, stream); +// }); + }); }); }); } diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index 61cc4b05d24..184cf62e8b5 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -114,7 +114,6 @@ class FlashAttnFwdSm90 { alignas(16) typename CollectiveMainloop::MainloopPipelineKVNew::SharedStorage pipeline_v_new; alignas(16) typename TileScheduler::SharedStorage smem_scheduler; } pipelines; - // float sink_val; }; static constexpr int SharedStorageSize = sizeof(SharedStorage); @@ -415,7 +414,6 @@ class FlashAttnFwdSm90 { float const k_descale = params.mainloop.ptr_k_descale == nullptr ? 1.0f : params.mainloop.ptr_k_descale[bidb * get<0>(params.mainloop.stride_k_descale) + bidh_kv * get<1>(params.mainloop.stride_k_descale)]; softmax_scale_log2 *= q_descale * k_descale; } - if (threadIdx.x == 128) { printf("=====LargeHeadDimV=%d=get<1>(block_coord)=%d, softmax_scale_log2=%f, shared_storage.sink_val=%f, sink_val=%f\n", int(LargeHeadDimV), get<1>(block_coord), softmax_scale_log2, sink_val, sink_val); } flash::Softmax softmax(softmax_scale_log2, sink_val); // Attention output (GEMM-II) accumulator. Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_MNK_PV{})); @@ -448,7 +446,6 @@ class FlashAttnFwdSm90 { threadIdx.x - MmaThreadOffset, block_coord); } else { // Write 0 to gO and -inf to gLSE. - // const float lse_val = Has_sink ? sink_val : -INFINITY; epilogue.store_zero(params.epilogue, threadIdx.x - MmaThreadOffset, block_coord, sink_val); } } diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index 16eb7c8a48b..ce82062e581 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -170,10 +170,10 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { dim3 grid_dims = AttnKernel::get_grid_shape(kernel_params); dim3 block_dims = AttnKernel::get_block_shape(); int smem_size = AttnKernel::SharedStorageSize; - int smem_size_q = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_q)); - int smem_size_k = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k)); - int smem_size_v = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v)); - // printf("========== smem_size = %d, q = %d, k = %d, v = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v); + // int smem_size_q = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_q)); + // int smem_size_k = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k)); + // int smem_size_v = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v)); + // printf("smem_size = %d, q = %d, k = %d, v = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v); // Get the ptr to kernel function. if constexpr (size(ClusterShape{}) > 1) { void const* kernel = (void const*) cutlass::device_kernel; @@ -212,12 +212,12 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { static constexpr bool HasQv = HasQV_ && Arch == 90 && !Is_FP8 && kHeadDim == 64 && kHeadDimV >= 256; APPENDKV_SWITCH(params.knew_ptr, AppendKV, [&] { SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { - static constexpr bool PackGQA_Sink = PackGQA && !Has_sink; - // Only use Cluster if number of tiles along seqlen_q is even and not varlen - CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { - static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; - run_flash_fwd(params, stream); - }); + static constexpr bool PackGQA_Sink = PackGQA && !Has_sink; + // Only use Cluster if number of tiles along seqlen_q is even and not varlen + CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { + static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; + run_flash_fwd(params, stream); + }); }); }); }); diff --git a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp index 4129eb34d1e..984bab16242 100644 --- a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp @@ -613,8 +613,6 @@ struct CollectiveMainloopFwdSm90 { params.window_size_left, params.window_size_right, params.attention_chunk_divmod, params.qhead_per_khead_divmod); // It's possible to have n_block_max <= n_block_min. Loading K can cause illegal memory access. - // shared_storage.sink_val = !Has_sink ? -INFINITY : params.ptr_Sink[bidh]; - if constexpr (Is_causal || Is_local || Varlen || Split) { if (n_block_max <= n_block_min) { scheduler_prefetch(); diff --git a/hopper/softmax.h b/hopper/softmax.h index 7c004801b03..6ee01362c62 100644 --- a/hopper/softmax.h +++ b/hopper/softmax.h @@ -153,8 +153,6 @@ struct Softmax { float sum = row_sum(mi); if constexpr (Has_sink) { static constexpr float max_offset = float(Max_offset); - // const float max_scaled = row_max(mi) == -INFINITY ? 0.f : row_max(mi) * softmax_scale; - // sum += exp2f(sink_val * float(M_LOG2E) - max_scaled); const float max_scaled = max(mi) == -INFINITY ? 0.f : (row_max(mi) * softmax_scale_log2) - max_offset; sum += exp2f(sink_val * float(M_LOG2E) - max_scaled); } From d1ea6630ac66f9cc3c548a10459411c040385dae Mon Sep 17 00:00:00 2001 From: aoxy Date: Tue, 20 Jan 2026 22:31:29 +0800 Subject: [PATCH 07/23] Update test. --- hopper/setup.py | 2 + hopper/test_flash_attn.py | 77 ++++++++++++++++++++++++++++++++------- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/hopper/setup.py b/hopper/setup.py index 7ed8abce15f..845f6b6504e 100644 --- a/hopper/setup.py +++ b/hopper/setup.py @@ -61,6 +61,7 @@ DISABLE_HDIM192 = os.getenv("FLASH_ATTENTION_DISABLE_HDIM192", "FALSE") == "TRUE" DISABLE_HDIM256 = os.getenv("FLASH_ATTENTION_DISABLE_HDIM256", "FALSE") == "TRUE" DISABLE_SM8x = os.getenv("FLASH_ATTENTION_DISABLE_SM80", "FALSE") == "TRUE" +DISABLE_SINK = os.getenv("FLASH_ATTENTION_DISABLE_SINK", "FALSE") == "TRUE" ENABLE_VCOLMAJOR = os.getenv("FLASH_ATTENTION_ENABLE_VCOLMAJOR", "FALSE") == "TRUE" @@ -468,6 +469,7 @@ def nvcc_threads_args(): + (["-DFLASHATTENTION_DISABLE_HDIM256"] if DISABLE_HDIM256 else []) + (["-DFLASHATTENTION_DISABLE_SM8x"] if DISABLE_SM8x else []) + (["-DFLASHATTENTION_ENABLE_VCOLMAJOR"] if ENABLE_VCOLMAJOR else []) + + (["-DFLASHATTENTION_DISABLE_SINK"] if DISABLE_SINK else []) ) DTYPE_FWD_SM80 = ["bf16"] + (["fp16"] if not DISABLE_FP16 else []) diff --git a/hopper/test_flash_attn.py b/hopper/test_flash_attn.py index 80d4dc0c15c..f9167d46138 100644 --- a/hopper/test_flash_attn.py +++ b/hopper/test_flash_attn.py @@ -37,6 +37,7 @@ DISABLE_HDIM128 = os.getenv("FLASH_ATTENTION_DISABLE_HDIM128", "FALSE") == "TRUE" DISABLE_HDIM192 = os.getenv("FLASH_ATTENTION_DISABLE_HDIM192", "FALSE") == "TRUE" DISABLE_HDIM256 = os.getenv("FLASH_ATTENTION_DISABLE_HDIM256", "FALSE") == "TRUE" +DISABLE_SINK = os.getenv("FLASH_ATTENTION_DISABLE_SINK", "FALSE") == "TRUE" COMPILED_HDIMS = ( [] @@ -101,8 +102,9 @@ ], ) # @pytest.mark.parametrize('seqlen_q,seqlen_k', [(128, 128)]) +@pytest.mark.parametrize("has_learnable_sink", [False] + ([True] if not DISABLE_SINK else [])) def test_flash_attn_output( - seqlen_q, seqlen_k, d, causal, local, softcap, V_colmajor, deterministic, has_qv, mha_type, dtype + seqlen_q, seqlen_k, d, causal, local, softcap, V_colmajor, deterministic, has_qv, mha_type, dtype, has_learnable_sink ): if V_colmajor and (seqlen_k % 16 != 0 or dtype != torch.float8_e4m3fn): pytest.skip("V_colmajor requires seqlen_k to be a multiple of 16 and dtype to be float8_e4m3fn") @@ -133,6 +135,11 @@ def test_flash_attn_output( qv_ref = torch.randn(batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref).to(dtype).to(dtype_ref) else: qv_ref = None + if has_learnable_sink: + learnable_sink_ref = torch.randn(nheads, device=device, dtype=torch.float32, requires_grad=True) + learnable_sink = learnable_sink_ref.detach().requires_grad_() + else: + learnable_sink_ref, learnable_sink = None, None # Put window_size after QKV randn so that window_size changes from test to test window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) # window_size = (-1, -1) if not local else (16, 0) @@ -155,7 +162,8 @@ def test_flash_attn_output( q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, window_size=window_size, attention_chunk=attention_chunk, - softcap=softcap + softcap=softcap, + learnable_sink=learnable_sink_ref, ) out_pt, attn_pt = attention_ref( q_ref, @@ -172,6 +180,7 @@ def test_flash_attn_output( upcast=False, reorder_ops=True, intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, + learnable_sink=learnable_sink_ref, ) # qk = torch.einsum('bshd,bthd->bhst', q_ref, k_ref).float() @@ -203,7 +212,8 @@ def test_flash_attn_output( attention_chunk=attention_chunk, softcap=softcap, pack_gqa=pack_gqa, - num_splits=num_splits + num_splits=num_splits, + learnable_sink=learnable_sink, ) print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") @@ -243,7 +253,7 @@ def test_flash_attn_output( # deterministic, # 0, # sm_margin # ) - dq, dk, dv = torch.autograd.grad(out, (q, k, v), g) + dq, dk, dv = torch.autograd.grad(out, (q, k, v), g, retain_graph=has_learnable_sink) # print(f"dO_O max diff: {(softmax_d - do_o).abs().max().item()}") # assert (softmax_d - do_o).abs().max().item() <= 1e-5 # assert dq_accum.abs().max().item() == 0.0 @@ -256,8 +266,8 @@ def test_flash_attn_output( # dK = torch.einsum('bhts,bthd->bshd', dP, q.float()) # dq, dk, dv = torch.autograd.grad(out, (q, k, v), g) - dq_ref, dk_ref, dv_ref = torch.autograd.grad(out_ref, (q_ref, k_ref, v_ref), g) - dq_pt, dk_pt, dv_pt = torch.autograd.grad(out_pt, (q_ref, k_ref, v_ref), g) + dq_ref, dk_ref, dv_ref = torch.autograd.grad(out_ref, (q_ref, k_ref, v_ref), g, retain_graph=has_learnable_sink) + dq_pt, dk_pt, dv_pt = torch.autograd.grad(out_pt, (q_ref, k_ref, v_ref), g, retain_graph=has_learnable_sink) print(f"dQ max diff: {(dq - dq_ref).abs().max().item()}") print(f"dK max diff: {(dk - dk_ref).abs().max().item()}") print(f"dV max diff: {(dv - dv_ref).abs().max().item()}") @@ -277,6 +287,18 @@ def test_flash_attn_output( assert (dk - dk_ref).abs().max().item() <= rtol * (dk_pt - dk_ref).abs().max().item() + dk_atol dv_atol = 2 * (dv_ref + 0.3 - 0.3 - dv_ref).abs().max().item() + (0 if softcap == 0 else 3e-4) assert (dv - dv_ref).abs().max().item() <= rtol * (dv_pt - dv_ref).abs().max().item() + dv_atol + if has_learnable_sink: + (dsink,) = torch.autograd.grad(out, learnable_sink, g) + (dsink_ref,) = torch.autograd.grad(out_ref, learnable_sink_ref, g) + (dsink_pt,) = torch.autograd.grad(out_pt, learnable_sink_ref, g) + print(f"dSink max diff: {(dsink - dsink_ref).abs().max().item()}") + print(f"dSink mean diff: {(dsink - dsink_ref).abs().mean().item()}") + print(f"dSink Pytorch max diff: {(dsink_pt - dsink_ref).abs().max().item()}") + print(f"dSink Pytorch mean diff: {(dsink_pt - dsink_ref).abs().mean().item()}") + dsink_atol = 2 * (dsink_ref + 0.3 - 0.3 - dsink_ref).abs().max().item() + (0 if softcap == 0 else 3e-4) + assert (dsink - dsink_ref).abs().max().item() <= rtol * (dsink_pt - dsink_ref).abs().max().item() + dsink_atol + + # @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) @@ -330,8 +352,9 @@ def test_flash_attn_output( (2048, 2048), ], ) +@pytest.mark.parametrize("has_learnable_sink", [False] + ([True] if not DISABLE_SINK else [])) def test_flash_attn_varlen_output( - seqlen_q, seqlen_k, d, add_unused_qkv, causal, local, softcap, deterministic, has_qv, mha_type, dtype + seqlen_q, seqlen_k, d, add_unused_qkv, causal, local, softcap, deterministic, has_qv, mha_type, dtype, has_learnable_sink ): device = "cuda" # set seed @@ -360,6 +383,11 @@ def test_flash_attn_varlen_output( qv_ref = torch.randn(batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref).to(dtype).to(dtype_ref) else: qv_ref = None + if has_learnable_sink: + learnable_sink_ref = torch.randn(nheads, device=device, dtype=torch.float32, requires_grad=True) + learnable_sink = learnable_sink_ref.detach().requires_grad_() + else: + learnable_sink_ref, learnable_sink = None, None # Put window_size after QKV randn so that window_size changes from test to test window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) if dtype == torch.float8_e4m3fn: @@ -426,7 +454,8 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, window_size=window_size, attention_chunk=attention_chunk, - softcap=softcap + softcap=softcap, + learnable_sink=learnable_sink_ref, ) out_pt, attn_pt = attention_ref( q_ref, @@ -443,6 +472,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): upcast=False, reorder_ops=True, intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, + learnable_sink=learnable_sink_ref, ) @@ -476,6 +506,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): window_size=window_size, attention_chunk=attention_chunk, softcap=softcap, + learnable_sink=learnable_sink, ) out = output_pad_fn(out_unpad) if query_unused_mask is not None: @@ -523,7 +554,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): # deterministic, # 0, # sm_margin # ) - dq_unpad, dk_unpad, dv_unpad = torch.autograd.grad(out_unpad, (q_unpad, k_unpad, v_unpad), g_unpad) + dq_unpad, dk_unpad, dv_unpad = torch.autograd.grad(out_unpad, (q_unpad, k_unpad, v_unpad), g_unpad, retain_graph=has_learnable_sink) dq = dq_pad_fn(dq_unpad) dk = dk_pad_fn(dk_unpad) dv = dk_pad_fn(dv_unpad) @@ -549,8 +580,8 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): # dq, dk, dv = torch.autograd.grad(out, (q, k, v), g) - dq_ref, dk_ref, dv_ref = torch.autograd.grad(out_ref, (q_ref, k_ref, v_ref), g) - dq_pt, dk_pt, dv_pt = torch.autograd.grad(out_pt, (q_ref, k_ref, v_ref), g) + dq_ref, dk_ref, dv_ref = torch.autograd.grad(out_ref, (q_ref, k_ref, v_ref), g, retain_graph=has_learnable_sink) + dq_pt, dk_pt, dv_pt = torch.autograd.grad(out_pt, (q_ref, k_ref, v_ref), g, retain_graph=has_learnable_sink) print(f"dQ max diff: {(dq - dq_ref).abs().max().item()}") print(f"dK max diff: {(dk - dk_ref).abs().max().item()}") print(f"dV max diff: {(dv - dv_ref).abs().max().item()}") @@ -571,6 +602,16 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): dv_atol = 2 * (dv_ref + 0.3 - 0.3 - dv_ref).abs().max().item() + (0 if softcap == 0 else 3e-4) assert (dv - dv_ref).abs().max().item() <= rtol * (dv_pt - dv_ref).abs().max().item() + dv_atol + if has_learnable_sink: + (dsink,) = torch.autograd.grad(out_unpad, learnable_sink, g_unpad) + (dsink_ref,) = torch.autograd.grad(out_ref, learnable_sink_ref, g) + (dsink_pt,) = torch.autograd.grad(out_pt, learnable_sink_ref, g) + print(f"dSink max diff: {(dsink - dsink_ref).abs().max().item()}") + print(f"dSink mean diff: {(dsink - dsink_ref).abs().mean().item()}") + print(f"dSink Pytorch max diff: {(dsink_pt - dsink_ref).abs().max().item()}") + print(f"dSink Pytorch mean diff: {(dsink_pt - dsink_ref).abs().mean().item()}") + dsink_atol = 2 * (dsink_ref + 0.3 - 0.3 - dsink_ref).abs().max().item() + (0 if softcap == 0 else 3e-4) + assert (dsink - dsink_ref).abs().max().item() <= rtol * (dsink_pt - dsink_ref).abs().max().item() + dsink_atol # @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) @pytest.mark.parametrize("dtype", [torch.bfloat16] + ([torch.float8_e4m3fn] if not DISABLE_FP8 else [])) @@ -624,6 +665,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): ], ) # @pytest.mark.parametrize('seqlen_q,seqlen_k', [(256, 128)]) +@pytest.mark.parametrize("has_learnable_sink", [False] + ([True] if not DISABLE_SINK else [])) def test_flash_attn_kvcache( seqlen_q, seqlen_k, @@ -641,6 +683,7 @@ def test_flash_attn_kvcache( new_kv, mha_type, dtype, + has_learnable_sink, ): if page_size is not None and seqlen_k % page_size != 0: pytest.skip() @@ -674,6 +717,11 @@ def test_flash_attn_kvcache( qv = torch.randn(batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref).to(dtype).to(dtype_ref) else: qv = None + if has_learnable_sink: + learnable_sink_ref = torch.randn(nheads, device=device, dtype=torch.float32, requires_grad=True) + learnable_sink = learnable_sink_ref.detach().requires_grad_() + else: + learnable_sink_ref, learnable_sink = None, None if varlen_q: query_padding_mask = generate_random_padding_mask(seqlen_q, batch_size, device, mode="random") q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, *rest = unpad_input(q, query_padding_mask) @@ -817,6 +865,7 @@ def test_flash_attn_kvcache( window_size=window_size, attention_chunk=attention_chunk, key_leftpad=cache_leftpad, + learnable_sink=learnable_sink_ref, ) out_pt, _ = attention_ref( q_ro, @@ -831,7 +880,8 @@ def test_flash_attn_kvcache( upcast=False, reorder_ops=True, key_leftpad=cache_leftpad, - intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None + intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, + learnable_sink=learnable_sink_ref, ) q = q.to(dtype) q_unpad = q_unpad.to(dtype) if varlen_q else None @@ -894,7 +944,8 @@ def test_flash_attn_kvcache( rotary_interleaved=rotary_interleaved, scheduler_metadata=scheduler_metadata, num_splits=num_splits, - return_softmax_lse=True + return_softmax_lse=True, + learnable_sink=learnable_sink, ) if varlen_q: out = output_pad_fn(out) From 8ddd284882145d405e55e9efab80a13c3f039281 Mon Sep 17 00:00:00 2001 From: aoxy Date: Tue, 20 Jan 2026 22:43:13 +0800 Subject: [PATCH 08/23] Support bf16 sink. --- hopper/benchmark_attn.py | 2 +- hopper/flash_attn_interface.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/hopper/benchmark_attn.py b/hopper/benchmark_attn.py index 80234b58bde..9b96e071b3b 100644 --- a/hopper/benchmark_attn.py +++ b/hopper/benchmark_attn.py @@ -280,7 +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=torch.float32, device=device, requires_grad=True) + 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 diff --git a/hopper/flash_attn_interface.py b/hopper/flash_attn_interface.py index f1e76e8e4cb..678bcd85625 100644 --- a/hopper/flash_attn_interface.py +++ b/hopper/flash_attn_interface.py @@ -62,6 +62,7 @@ def _flash_attn_forward( ] rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] seqlens_rotary = maybe_contiguous(seqlens_rotary) + learnable_sink = None if learnable_sink is None else learnable_sink.float() out, softmax_lse, *rest = flash_attn_3_cuda.fwd( q, k, @@ -129,6 +130,9 @@ def _flash_attn_backward( ): # dq, dk, dv are allocated by us so they should already be contiguous dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + dsink_dtype = None if dsink is None else dsink.dtype + learnable_sink = None if learnable_sink is None else learnable_sink.float() + dsink = None if dsink is None else dsink.float() dq, dk, dv, softmax_d, dsink, *rest = flash_attn_3_cuda.bwd( dout, q, @@ -155,7 +159,7 @@ def _flash_attn_backward( learnable_sink, dsink, ) - return dq, dk, dv, softmax_d, dsink + return dq, dk, dv, softmax_d, (None if dsink is None else dsink.to(dsink_dtype)) class FlashAttnQKVPackedFunc(torch.autograd.Function): From 42d1c15905a7772c198e41895a13a26193c8def2 Mon Sep 17 00:00:00 2001 From: aoxy Date: Wed, 21 Jan 2026 15:01:42 +0800 Subject: [PATCH 09/23] Fix num_splits_dynamic. --- hopper/flash_fwd_combine_kernel.h | 2 +- hopper/flash_fwd_combine_launch_template.h | 28 +++++++++++----------- hopper/flash_fwd_kernel_sm90.h | 13 +++++++++- hopper/flash_fwd_launch_template.h | 2 +- hopper/test_flash_attn.py | 2 ++ 5 files changed, 30 insertions(+), 17 deletions(-) diff --git a/hopper/flash_fwd_combine_kernel.h b/hopper/flash_fwd_combine_kernel.h index e9cda0b869c..c6acfa3d6fa 100644 --- a/hopper/flash_fwd_combine_kernel.h +++ b/hopper/flash_fwd_combine_kernel.h @@ -382,7 +382,7 @@ class FlashAttnFwdCombine { } else { bidh = seqlen_divmod_dynamic.divmod(m_idx, idx); } - const float sink_val_exp = expf((params.ptr_Sink)[bidh] - lse_max); + const float sink_val_exp = expf((params.ptr_Sink)[bidh] - lse_max_cur); lse_sum_cur += sink_val_exp; } } diff --git a/hopper/flash_fwd_combine_launch_template.h b/hopper/flash_fwd_combine_launch_template.h index 090dc500850..728d92002a3 100644 --- a/hopper/flash_fwd_combine_launch_template.h +++ b/hopper/flash_fwd_combine_launch_template.h @@ -62,21 +62,21 @@ void run_mha_fwd_combine_(Flash_fwd_params ¶ms, cudaStream_t stream, bool en ARCH_SWITCH(params.arch, Arch, [&] { BOOL_SWITCH(params.cu_seqlens_q || params.seqused_q, Varlen, [&] { SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { - if constexpr (kBlockM >= 16) { // If kBlockM == 8 then the minimum number of splits is 32. - if (params.num_splits <= 16) { - run_flash_fwd_combine(params, stream, enable_pdl); - return; + if constexpr (kBlockM >= 16) { // If kBlockM == 8 then the minimum number of splits is 32. + if (params.num_splits <= 16) { + run_flash_fwd_combine(params, stream, enable_pdl); + return; + } + } + if (params.num_splits <= 32) { + run_flash_fwd_combine(params, stream, enable_pdl); + } else if (params.num_splits <= 64) { + run_flash_fwd_combine(params, stream, enable_pdl); + } else if (params.num_splits <= 128) { + run_flash_fwd_combine(params, stream, enable_pdl); + } else { + run_flash_fwd_combine(params, stream, enable_pdl); } - } - if (params.num_splits <= 32) { - run_flash_fwd_combine(params, stream, enable_pdl); - } else if (params.num_splits <= 64) { - run_flash_fwd_combine(params, stream, enable_pdl); - } else if (params.num_splits <= 128) { - run_flash_fwd_combine(params, stream, enable_pdl); - } else { - run_flash_fwd_combine(params, stream, enable_pdl); - } }); }); }); diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index 184cf62e8b5..4959a99b6a2 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -407,7 +407,18 @@ class FlashAttnFwdSm90 { // If there's tanh softcap, the scaling will be done before tanh. float softmax_scale_log2 = params.mainloop.softmax_scale_log2; int const bidh = get<1>(block_coord); - float sink_val = !Has_sink ? -INFINITY : params.mainloop.ptr_Sink[bidh]; + float sink_val = -INFINITY; + if constexpr (Has_sink) { + if constexpr (Split && Varlen) { + uint32_t num_splits_dynamic_u = reinterpret_cast(get<3>(block_coord)) >> 16; // first 16 bits are for num_splits + int num_splits_dynamic = reinterpret_cast(num_splits_dynamic_u); + if (num_splits_dynamic <= 1) { + sink_val = params.mainloop.ptr_Sink[bidh]; + } + } else { + sink_val = params.mainloop.ptr_Sink[bidh]; + } + } if constexpr (Is_FP8 && !Has_softcap) { int const bidh_kv = !PackGQA ? params.mainloop.qhead_per_khead_divmod.divide(bidh) : bidh; float const q_descale = params.mainloop.ptr_q_descale == nullptr ? 1.0f : params.mainloop.ptr_q_descale[bidb * get<0>(params.mainloop.stride_q_descale) + bidh_kv * get<1>(params.mainloop.stride_q_descale)]; diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index ce82062e581..43b91192126 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -216,7 +216,7 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { // Only use Cluster if number of tiles along seqlen_q is even and not varlen CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; - run_flash_fwd(params, stream); + run_flash_fwd(params, stream); }); }); }); diff --git a/hopper/test_flash_attn.py b/hopper/test_flash_attn.py index f9167d46138..ceeec938339 100644 --- a/hopper/test_flash_attn.py +++ b/hopper/test_flash_attn.py @@ -506,6 +506,8 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): window_size=window_size, attention_chunk=attention_chunk, softcap=softcap, + pack_gqa=pack_gqa, + num_splits=num_splits, learnable_sink=learnable_sink, ) out = output_pad_fn(out_unpad) From bd815d67cb63967ebe263a66db567eff590a4946 Mon Sep 17 00:00:00 2001 From: aoxy Date: Wed, 21 Jan 2026 19:53:11 +0800 Subject: [PATCH 10/23] Fix again. --- hopper/flash_fwd_kernel_sm90.h | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index 4959a99b6a2..fd2e2e923ba 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -409,11 +409,14 @@ class FlashAttnFwdSm90 { int const bidh = get<1>(block_coord); float sink_val = -INFINITY; if constexpr (Has_sink) { - if constexpr (Split && Varlen) { - uint32_t num_splits_dynamic_u = reinterpret_cast(get<3>(block_coord)) >> 16; // first 16 bits are for num_splits - int num_splits_dynamic = reinterpret_cast(num_splits_dynamic_u); - if (num_splits_dynamic <= 1) { - sink_val = params.mainloop.ptr_Sink[bidh]; + if constexpr (Split) { + if constexpr (Varlen) { + uint32_t num_splits_dynamic_u = reinterpret_cast(get<3>(block_coord)) >> 16; // first 16 bits are for num_splits + int num_splits_dynamic = reinterpret_cast(num_splits_dynamic_u); + if (num_splits_dynamic <= 1) { + sink_val = params.mainloop.ptr_Sink[bidh]; + printf("Sink val: %f, Split: %d, Varlen:%d, num_splits_dynamic: %d\n", sink_val, int(Split), int(Varlen), num_splits_dynamic); + } } } else { sink_val = params.mainloop.ptr_Sink[bidh]; From 14d3d19ef18b9be15f640f2bf63c29bc81b76864 Mon Sep 17 00:00:00 2001 From: aoxy Date: Wed, 21 Jan 2026 20:13:54 +0800 Subject: [PATCH 11/23] Remove log. --- hopper/flash_fwd_kernel_sm90.h | 1 - 1 file changed, 1 deletion(-) diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index fd2e2e923ba..29eaefb72c1 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -415,7 +415,6 @@ class FlashAttnFwdSm90 { int num_splits_dynamic = reinterpret_cast(num_splits_dynamic_u); if (num_splits_dynamic <= 1) { sink_val = params.mainloop.ptr_Sink[bidh]; - printf("Sink val: %f, Split: %d, Varlen:%d, num_splits_dynamic: %d\n", sink_val, int(Split), int(Varlen), num_splits_dynamic); } } } else { From 63a193180a4bd838bee7f95590051323922b6416 Mon Sep 17 00:00:00 2001 From: aoxy Date: Thu, 22 Jan 2026 10:04:58 +0800 Subject: [PATCH 12/23] Update scripts. --- reinstall.sh | 4 ++-- warp_reinstall.sh | 6 +++++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/reinstall.sh b/reinstall.sh index 56ac3d1a510..1014f90a96e 100755 --- a/reinstall.sh +++ b/reinstall.sh @@ -3,10 +3,10 @@ set -e export FLASH_ATTENTION_DISABLE_SM80="TRUE" export FLASH_ATTENTION_DISABLE_FP8="TRUE" -export FLASH_ATTENTION_DISABLE_SOFTCAP="TRUE" +# export FLASH_ATTENTION_DISABLE_SOFTCAP="TRUE" # export FLASH_ATTENTION_DISABLE_LOCAL="TRUE" # export FLASH_ATTENTION_DISABLE_BACKWARD="TRUE" -export FLASH_ATTENTION_DISABLE_APPENDKV="TRUE" +# export FLASH_ATTENTION_DISABLE_APPENDKV="TRUE" export FLASH_ATTENTION_FORCE_BUILD="TRUE" diff --git a/warp_reinstall.sh b/warp_reinstall.sh index 2c1a511309c..14bce9fa855 100755 --- a/warp_reinstall.sh +++ b/warp_reinstall.sh @@ -1,3 +1,7 @@ time ./reinstall.sh > my_install.log 2>&1 grep -i "error" my_install.log -tail -n 10 my_install.log \ No newline at end of file +tail -n 10 my_install.log + +cd hopper +python -m pytest -q -s test_flash_attn.py::test_flash_attn_output > test_all_output3 +python -m pytest -q -s test_flash_attn.py::test_flash_attn_varlen_output > test_all_varlen3 \ No newline at end of file From 4f31c274fbc0686f86d923138e61e0d616040011 Mon Sep 17 00:00:00 2001 From: aoxy Date: Thu, 22 Jan 2026 10:34:22 +0800 Subject: [PATCH 13/23] Fix. --- hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp index 1001aea2a99..417b245308f 100644 --- a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp @@ -882,14 +882,13 @@ struct CollectiveMainloopBwdSm90 { Tensor dS = make_tensor(tdPrdP.data(), scores.layout()); #pragma unroll for (int mi = 0; mi < size<0>(dS); ++mi) { - float dsink_val_cols = 0.0f; float const dP_sum_cur = [&] { if constexpr (!ShuffledPsum) return tLSErdPsum(mi); else return __shfl_sync(0xffffffff, tLSErdPsum(mi / 8), (mi % 8) * 4 + (thread_idx % 4)); }(); #pragma unroll for (int ni = 0; ni < size<1>(dS); ++ni) { - if constexpr (Has_sink) { dsink_val_cols += scores(mi, ni) * dS(mi, ni); } + if constexpr (Has_sink) { dsink_val += scores(mi, ni) * dS(mi, ni); } dS(mi, ni) = scores(mi, ni) * (dS(mi, ni) - dP_sum_cur); if constexpr (Has_softcap) { dS(mi, ni) *= dtanh(mi, ni); } } @@ -898,7 +897,7 @@ struct CollectiveMainloopBwdSm90 { if constexpr (!ShuffleLSE) return tLSErLSE(mi); else return __shfl_sync(0xffffffff, tLSErLSE(mi / 8), (mi % 8) * 4 + (thread_idx % 4)); }(); - dsink_val += dsink_val_cols * exp2f(sink_val * float(M_LOG2E) - lse_scaled); + dsink_val *= exp2f(sink_val * float(M_LOG2E) - lse_scaled); } } From 1c0dd9c619930841bc9dc02ec7b8c7ee1054a1fb Mon Sep 17 00:00:00 2001 From: Qubitium-ModelCloud Date: Wed, 21 Jan 2026 18:36:22 +0800 Subject: [PATCH 14/23] reduce chance of build oom (#2079) --- setup.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 3b1426ccddb..e41026a7a66 100644 --- a/setup.py +++ b/setup.py @@ -64,6 +64,7 @@ FORCE_CXX11_ABI = os.getenv("FLASH_ATTENTION_FORCE_CXX11_ABI", "FALSE") == "TRUE" USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" SKIP_CK_BUILD = os.getenv("FLASH_ATTENTION_SKIP_CK_BUILD", "TRUE") == "TRUE" if USE_TRITON_ROCM else False +NVCC_THREADS = os.getenv("NVCC_THREADS") or "4" @functools.lru_cache(maxsize=None) def cuda_archs() -> str: @@ -121,8 +122,7 @@ def check_if_rocm_home_none(global_option: str) -> None: def append_nvcc_threads(nvcc_extra_args): - nvcc_threads = os.getenv("NVCC_THREADS") or "2" - return nvcc_extra_args + ["--threads", nvcc_threads] + return nvcc_extra_args + ["--threads", NVCC_THREADS] def rename_cpp_to_cu(cpp_files): @@ -508,15 +508,23 @@ def __init__(self, *args, **kwargs) -> None: if not os.environ.get("MAX_JOBS"): import psutil + nvcc_threads = max(1, int(NVCC_THREADS)) + # calculate the maximum allowed NUM_JOBS based on cores max_num_jobs_cores = max(1, os.cpu_count() // 2) # calculate the maximum allowed NUM_JOBS based on free memory free_memory_gb = psutil.virtual_memory().available / (1024 ** 3) # free memory in GB - max_num_jobs_memory = int(free_memory_gb / 9) # each JOB peak memory cost is ~8-9GB when threads = 4 + # Assume worst-case peak observed memory usage of ~5GB per NVCC thread. + # Limit: peak_threads = max_jobs * nvcc_threads and peak_threads * 5GB <= free_memory. + max_num_jobs_memory = max(1, int(free_memory_gb / (5 * nvcc_threads))) # pick lower value of jobs based on cores vs memory metric to minimize oom and swap usage during compilation max_jobs = max(1, min(max_num_jobs_cores, max_num_jobs_memory)) + print( + f"Auto set MAX_JOBS to `{max_jobs}`, NVCC_THREADS to `{nvcc_threads}`. " + "If you see memory pressure, please use a lower `MAX_JOBS=N` or `NVCC_THREADS=N` value." + ) os.environ["MAX_JOBS"] = str(max_jobs) super().__init__(*args, **kwargs) From 9ff82f4c252216d8816c1bad205a62dc78862bfe Mon Sep 17 00:00:00 2001 From: aoxy Date: Wed, 28 Jan 2026 17:29:18 +0800 Subject: [PATCH 15/23] Fix bf16. --- hopper/flash_api.cpp | 17 ++++++++++++----- hopper/flash_attn_interface.py | 6 +----- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp index 4def6028d28..8d9fc0820bc 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -1047,8 +1047,7 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq } if (learnable_sink_.has_value()) { - at::Tensor learnable_sink = learnable_sink_.value(); - TORCH_CHECK(learnable_sink.dtype() == torch::kFloat32, "Learnable sink must have dtype fp32"); + at::Tensor learnable_sink = learnable_sink_.value().to(torch::kFloat32); CHECK_DEVICE(learnable_sink); CHECK_CONTIGUOUS(learnable_sink); TORCH_CHECK(learnable_sink.stride(-1) == 1, "Learnable sink tensor must have contiguous last dimension"); CHECK_SHAPE(learnable_sink, num_heads); @@ -1488,12 +1487,11 @@ std::vector mha_bwd( if (learnable_sink_.has_value()) { - at::Tensor learnable_sink = learnable_sink_.value(); - TORCH_CHECK(learnable_sink.dtype() == torch::kFloat32, "Learnable sink must have dtype fp32"); + at::Tensor learnable_sink = learnable_sink_.value().to(torch::kFloat32); CHECK_DEVICE(learnable_sink); CHECK_CONTIGUOUS(learnable_sink); TORCH_CHECK(learnable_sink.stride(-1) == 1, "Learnable sink tensor must have contiguous last dimension"); CHECK_SHAPE(learnable_sink, num_heads); - if (dsink_.has_value()) { + if (dsink_.has_value() && dsink_.value().dtype() == torch::kFloat32) { dsink = dsink_.value(); TORCH_CHECK(dsink.dtype() == torch::kFloat32, "dsink must have dtype fp32"); CHECK_DEVICE(dsink); CHECK_CONTIGUOUS(dsink); @@ -1543,6 +1541,15 @@ std::vector mha_bwd( softmax_d.zero_(); } + if (learnable_sink_.has_value()) { + if (!dsink_.has_value()) { + dsink = dsink.to(learnable_sink_.value().dtype()); + } else if (dsink_.value().dtype() != torch::kFloat32) { + dsink = dsink.to(dsink_.value().dtype()); + dsink_.value().copy_(dsink); + } + } + return { dq, dk, dv, softmax_d, dsink, softmax_lse_log2, dq_accum, dk_accum, dv_accum }; } diff --git a/hopper/flash_attn_interface.py b/hopper/flash_attn_interface.py index 678bcd85625..f1e76e8e4cb 100644 --- a/hopper/flash_attn_interface.py +++ b/hopper/flash_attn_interface.py @@ -62,7 +62,6 @@ def _flash_attn_forward( ] rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)] seqlens_rotary = maybe_contiguous(seqlens_rotary) - learnable_sink = None if learnable_sink is None else learnable_sink.float() out, softmax_lse, *rest = flash_attn_3_cuda.fwd( q, k, @@ -130,9 +129,6 @@ def _flash_attn_backward( ): # dq, dk, dv are allocated by us so they should already be contiguous dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] - dsink_dtype = None if dsink is None else dsink.dtype - learnable_sink = None if learnable_sink is None else learnable_sink.float() - dsink = None if dsink is None else dsink.float() dq, dk, dv, softmax_d, dsink, *rest = flash_attn_3_cuda.bwd( dout, q, @@ -159,7 +155,7 @@ def _flash_attn_backward( learnable_sink, dsink, ) - return dq, dk, dv, softmax_d, (None if dsink is None else dsink.to(dsink_dtype)) + return dq, dk, dv, softmax_d, dsink class FlashAttnQKVPackedFunc(torch.autograd.Function): From 0fc53d84fbf6807c154015bb2b9364b9458c60e0 Mon Sep 17 00:00:00 2001 From: aoxy Date: Thu, 29 Jan 2026 14:17:31 +0800 Subject: [PATCH 16/23] Fix bugs. --- hopper/flash_api.cpp | 1 - hopper/flash_attn_interface.py | 6 +++--- hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp | 5 +++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp index 8d9fc0820bc..55ec2cd2971 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -1493,7 +1493,6 @@ std::vector mha_bwd( CHECK_SHAPE(learnable_sink, num_heads); if (dsink_.has_value() && dsink_.value().dtype() == torch::kFloat32) { dsink = dsink_.value(); - TORCH_CHECK(dsink.dtype() == torch::kFloat32, "dsink must have dtype fp32"); CHECK_DEVICE(dsink); CHECK_CONTIGUOUS(dsink); TORCH_CHECK(dsink.stride(-1) == 1, "dsink tensor must have contiguous last dimension"); CHECK_SHAPE(dsink, num_heads); diff --git a/hopper/flash_attn_interface.py b/hopper/flash_attn_interface.py index f1e76e8e4cb..86a18b16810 100644 --- a/hopper/flash_attn_interface.py +++ b/hopper/flash_attn_interface.py @@ -256,7 +256,7 @@ def backward(ctx, dout, *args): dsink, ) dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None, None, dsink, None + return dqkv, None, None, None, None, None, None, None, None, None, None, None, dsink class FlashAttnFunc(torch.autograd.Function): @@ -348,7 +348,7 @@ def backward(ctx, dout, *args): dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink, None + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink class FlashAttnVarlenFunc(torch.autograd.Function): @@ -455,7 +455,7 @@ def backward(ctx, dout, *args): dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink, None + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink def flash_attn_qkvpacked_func( diff --git a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp index 417b245308f..1001aea2a99 100644 --- a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp @@ -882,13 +882,14 @@ struct CollectiveMainloopBwdSm90 { Tensor dS = make_tensor(tdPrdP.data(), scores.layout()); #pragma unroll for (int mi = 0; mi < size<0>(dS); ++mi) { + float dsink_val_cols = 0.0f; float const dP_sum_cur = [&] { if constexpr (!ShuffledPsum) return tLSErdPsum(mi); else return __shfl_sync(0xffffffff, tLSErdPsum(mi / 8), (mi % 8) * 4 + (thread_idx % 4)); }(); #pragma unroll for (int ni = 0; ni < size<1>(dS); ++ni) { - if constexpr (Has_sink) { dsink_val += scores(mi, ni) * dS(mi, ni); } + if constexpr (Has_sink) { dsink_val_cols += scores(mi, ni) * dS(mi, ni); } dS(mi, ni) = scores(mi, ni) * (dS(mi, ni) - dP_sum_cur); if constexpr (Has_softcap) { dS(mi, ni) *= dtanh(mi, ni); } } @@ -897,7 +898,7 @@ struct CollectiveMainloopBwdSm90 { if constexpr (!ShuffleLSE) return tLSErLSE(mi); else return __shfl_sync(0xffffffff, tLSErLSE(mi / 8), (mi % 8) * 4 + (thread_idx % 4)); }(); - dsink_val *= exp2f(sink_val * float(M_LOG2E) - lse_scaled); + dsink_val += dsink_val_cols * exp2f(sink_val * float(M_LOG2E) - lse_scaled); } } From a17af5cba2f9f80d73f16299f44dbfb1ed6006e1 Mon Sep 17 00:00:00 2001 From: aoxy Date: Thu, 29 Jan 2026 16:19:19 +0800 Subject: [PATCH 17/23] Fix bugs. --- hopper/flash_api.cpp | 7 ++++--- hopper/test_flash_attn.py | 10 +++++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp index 55ec2cd2971..52d0a5e6658 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -1046,8 +1046,9 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq params.kv_batch_idx = reinterpret_cast(kv_batch_idx.data_ptr()); } + at::Tensor learnable_sink; if (learnable_sink_.has_value()) { - at::Tensor learnable_sink = learnable_sink_.value().to(torch::kFloat32); + learnable_sink = learnable_sink_.value().to(torch::kFloat32); CHECK_DEVICE(learnable_sink); CHECK_CONTIGUOUS(learnable_sink); TORCH_CHECK(learnable_sink.stride(-1) == 1, "Learnable sink tensor must have contiguous last dimension"); CHECK_SHAPE(learnable_sink, num_heads); @@ -1485,9 +1486,9 @@ std::vector mha_bwd( params.dv = head_size_v; params.dv_rounded = head_size_v_rounded; - + at::Tensor learnable_sink; if (learnable_sink_.has_value()) { - at::Tensor learnable_sink = learnable_sink_.value().to(torch::kFloat32); + learnable_sink = learnable_sink_.value().to(torch::kFloat32); CHECK_DEVICE(learnable_sink); CHECK_CONTIGUOUS(learnable_sink); TORCH_CHECK(learnable_sink.stride(-1) == 1, "Learnable sink tensor must have contiguous last dimension"); CHECK_SHAPE(learnable_sink, num_heads); diff --git a/hopper/test_flash_attn.py b/hopper/test_flash_attn.py index ceeec938339..050a36b6d47 100644 --- a/hopper/test_flash_attn.py +++ b/hopper/test_flash_attn.py @@ -1,6 +1,7 @@ import os import math import itertools +import random import pytest import torch @@ -136,7 +137,8 @@ def test_flash_attn_output( else: qv_ref = None if has_learnable_sink: - learnable_sink_ref = torch.randn(nheads, device=device, dtype=torch.float32, requires_grad=True) + sink_dtype = random.choice([torch.float32, torch.float16, torch.bfloat16]) + learnable_sink_ref = torch.randn(nheads, device=device, dtype=sink_dtype, requires_grad=True) learnable_sink = learnable_sink_ref.detach().requires_grad_() else: learnable_sink_ref, learnable_sink = None, None @@ -384,7 +386,8 @@ def test_flash_attn_varlen_output( else: qv_ref = None if has_learnable_sink: - learnable_sink_ref = torch.randn(nheads, device=device, dtype=torch.float32, requires_grad=True) + sink_dtype = random.choice([torch.float32, torch.float16, torch.bfloat16]) + learnable_sink_ref = torch.randn(nheads, device=device, dtype=sink_dtype, requires_grad=True) learnable_sink = learnable_sink_ref.detach().requires_grad_() else: learnable_sink_ref, learnable_sink = None, None @@ -720,7 +723,8 @@ def test_flash_attn_kvcache( else: qv = None if has_learnable_sink: - learnable_sink_ref = torch.randn(nheads, device=device, dtype=torch.float32, requires_grad=True) + sink_dtype = random.choice([torch.float32, torch.float16, torch.bfloat16]) + learnable_sink_ref = torch.randn(nheads, device=device, dtype=sink_dtype, requires_grad=True) learnable_sink = learnable_sink_ref.detach().requires_grad_() else: learnable_sink_ref, learnable_sink = None, None From 91624e5e8047a633fd66c135a9fabb7aba444071 Mon Sep 17 00:00:00 2001 From: aoxy Date: Wed, 6 May 2026 21:46:32 +0800 Subject: [PATCH 18/23] Add use_fa4_sink option. Co-authored-by: Copilot --- hopper/flash.h | 1 + hopper/flash_api.cpp | 4 +++- hopper/flash_attn_interface.py | 25 ++++++++++++++++++++---- hopper/flash_fwd_kernel_sm90.h | 2 +- hopper/flash_fwd_launch_template.h | 2 +- hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp | 4 +++- hopper/softmax.h | 19 +++++++++++------- hopper/test_flash_attn.py | 3 +++ 8 files changed, 45 insertions(+), 15 deletions(-) diff --git a/hopper/flash.h b/hopper/flash.h index 30c1b9db311..a4703bfb14d 100644 --- a/hopper/flash.h +++ b/hopper/flash.h @@ -160,6 +160,7 @@ struct Flash_fwd_params : public Qkv_params { int arch; int num_sm; void *__restrict__ learnable_sink_ptr; + bool use_fa4_sink; }; //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp index 52d0a5e6658..468894b9b66 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -681,7 +681,8 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq int num_splits, std::optional pack_gqa_, int const sm_margin, - std::optional learnable_sink_ + std::optional learnable_sink_, + bool use_fa4_sink ) { auto dprops = at::cuda::getCurrentDeviceProperties(); @@ -1054,6 +1055,7 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq CHECK_SHAPE(learnable_sink, num_heads); params.learnable_sink_ptr = learnable_sink.data_ptr(); params.pack_gqa = false; // Disable pack_gqa if learnable sink is provided + params.use_fa4_sink = use_fa4_sink; } else { params.learnable_sink_ptr = nullptr; } diff --git a/hopper/flash_attn_interface.py b/hopper/flash_attn_interface.py index 86a18b16810..9c46e5739bb 100644 --- a/hopper/flash_attn_interface.py +++ b/hopper/flash_attn_interface.py @@ -50,7 +50,8 @@ def _flash_attn_forward( num_splits=1, pack_gqa=None, sm_margin=0, - learnable_sink=None): + learnable_sink=None, + use_fa4_sink=False): q, k, k_new, v_new = [maybe_contiguous(x) for x in (q, k, k_new, v_new)] v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v cu_seqlens_q, cu_seqlens_k, cu_seqlens_k_new = [ @@ -98,6 +99,7 @@ def _flash_attn_forward( pack_gqa, sm_margin, learnable_sink, + use_fa4_sink, ) return out, softmax_lse, *rest @@ -173,6 +175,7 @@ def forward( num_heads_q=None, sm_margin=0, learnable_sink=None, + use_fa4_sink=False, ): if softmax_scale is None: softmax_scale = qkv.shape[-1] ** (-0.5) @@ -204,6 +207,8 @@ def forward( attention_chunk=attention_chunk, softcap=softcap, sm_margin=sm_margin, + learnable_sink=learnable_sink, + use_fa4_sink=use_fa4_sink, ) # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) ctx.save_for_backward(q, k, v, out, softmax_lse, learnable_sink) @@ -256,7 +261,7 @@ def backward(ctx, dout, *args): dsink, ) dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension - return dqkv, None, None, None, None, None, None, None, None, None, None, None, dsink + return dqkv, None, None, None, None, None, None, None, None, None, None, None, dsink, None class FlashAttnFunc(torch.autograd.Function): @@ -279,6 +284,7 @@ def forward( deterministic=False, sm_margin=0, learnable_sink=None, + use_fa4_sink=False, ): if softmax_scale is None: softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) @@ -305,6 +311,7 @@ def forward( pack_gqa=pack_gqa, sm_margin=sm_margin, learnable_sink=learnable_sink, + use_fa4_sink=use_fa4_sink, ) # ctx.save_for_backward(q, k, v, out_padded, softmax_lse) ctx.save_for_backward(q, k, v, out, softmax_lse, learnable_sink) @@ -348,7 +355,7 @@ def backward(ctx, dout, *args): dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink, None class FlashAttnVarlenFunc(torch.autograd.Function): @@ -377,6 +384,7 @@ def forward( deterministic=False, sm_margin=0, learnable_sink=None, + use_fa4_sink=False, ): if softmax_scale is None: softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (-0.5) @@ -407,6 +415,7 @@ def forward( pack_gqa=pack_gqa, sm_margin=sm_margin, learnable_sink=learnable_sink, + use_fa4_sink=use_fa4_sink, ) # ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink) @@ -455,7 +464,7 @@ def backward(ctx, dout, *args): dq = dq[..., : q.shape[-1]] # We could have padded the head dimension dk = dk[..., : k.shape[-1]] dv = dv[..., : v.shape[-1]] - return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink, None def flash_attn_qkvpacked_func( @@ -470,6 +479,7 @@ def flash_attn_qkvpacked_func( num_heads_q=None, sm_margin=0, learnable_sink=None, + use_fa4_sink=False, ): """dropout_p should be set to 0.0 during evaluation If Q, K, V are already stacked into 1 tensor, this function will be faster than @@ -517,6 +527,7 @@ def flash_attn_qkvpacked_func( num_heads_q, sm_margin, learnable_sink, + use_fa4_sink, ) @@ -536,6 +547,7 @@ def flash_attn_func( deterministic=False, sm_margin=0, learnable_sink=None, + use_fa4_sink=False, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads @@ -598,6 +610,7 @@ def flash_attn_func( deterministic, sm_margin, learnable_sink, + use_fa4_sink, ) @@ -623,6 +636,7 @@ def flash_attn_varlen_func( deterministic=False, sm_margin=0, learnable_sink=None, + use_fa4_sink=False, ): return FlashAttnVarlenFunc.apply( q, @@ -646,6 +660,7 @@ def flash_attn_varlen_func( deterministic, sm_margin, learnable_sink, + use_fa4_sink, ) @@ -685,6 +700,7 @@ def flash_attn_with_kvcache( sm_margin=0, # Can be tuned if some SMs are used for communication return_softmax_lse=False, learnable_sink=None, + use_fa4_sink=False, ): """ If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from @@ -813,6 +829,7 @@ def flash_attn_with_kvcache( pack_gqa=pack_gqa, sm_margin=sm_margin, learnable_sink=learnable_sink, + use_fa4_sink=use_fa4_sink, ) # return (out, softmax_lse) if return_softmax_lse else out return (out, softmax_lse, *rest) if return_softmax_lse else out diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index 29eaefb72c1..a071ee10654 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -427,7 +427,7 @@ class FlashAttnFwdSm90 { float const k_descale = params.mainloop.ptr_k_descale == nullptr ? 1.0f : params.mainloop.ptr_k_descale[bidb * get<0>(params.mainloop.stride_k_descale) + bidh_kv * get<1>(params.mainloop.stride_k_descale)]; softmax_scale_log2 *= q_descale * k_descale; } - flash::Softmax softmax(softmax_scale_log2, sink_val); + flash::Softmax softmax(softmax_scale_log2, sink_val, params.mainloop.use_fa4_sink); // Attention output (GEMM-II) accumulator. Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_MNK_PV{})); bool tile_valid; diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index 43b91192126..d14091c8b24 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -127,7 +127,7 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { params.cu_seqlens_q, params.cu_seqlens_k, params.cu_seqlens_knew, params.seqused_q, params.seqused_k, params.leftpad_k, params.seqlens_rotary, - reinterpret_cast(params.learnable_sink_ptr) + reinterpret_cast(params.learnable_sink_ptr), params.use_fa4_sink }; typename CollectiveEpilogue::Arguments epilogue_args { static_cast(params.o_ptr), diff --git a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp index 984bab16242..fea9809eabb 100644 --- a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp @@ -398,6 +398,7 @@ struct CollectiveMainloopFwdSm90 { int const* const leftpad_k = nullptr; int const* const seqlens_rotary = nullptr; float const* const ptr_Sink = nullptr; + bool const use_fa4_sink = false; }; // Device side kernel params @@ -456,6 +457,7 @@ struct CollectiveMainloopFwdSm90 { int const* const leftpad_k = nullptr; int const *const seqlens_rotary = nullptr; float const* const ptr_Sink = nullptr; + bool const use_fa4_sink = false; }; static Params @@ -567,7 +569,7 @@ struct CollectiveMainloopFwdSm90 { !Split ? 1 : args.num_splits, args.kv_batch_idx, args.cu_seqlens_q, args.cu_seqlens_k, args.cu_seqlens_k_new, - args.seqused_q, args.seqused_k, args.leftpad_k, args.seqlens_rotary, args.ptr_Sink}; + args.seqused_q, args.seqused_k, args.leftpad_k, args.seqlens_rotary, args.ptr_Sink, args.use_fa4_sink}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance diff --git a/hopper/softmax.h b/hopper/softmax.h index 6ee01362c62..697fe5f62f7 100644 --- a/hopper/softmax.h +++ b/hopper/softmax.h @@ -97,9 +97,10 @@ struct Softmax { float const softmax_scale_log2; float const softmax_scale; float const sink_val; + bool const use_fa4_sink; - CUTLASS_DEVICE Softmax(float const softmax_scale_log2_, float const sink_val_ = -INFINITY) - : softmax_scale_log2(softmax_scale_log2_), softmax_scale(softmax_scale_log2_ * float(M_LN2)), sink_val(sink_val_) {}; + CUTLASS_DEVICE Softmax(float const softmax_scale_log2_, float const sink_val_ = -INFINITY, bool const use_fa4_sink_ = false) + : softmax_scale_log2(softmax_scale_log2_), softmax_scale(softmax_scale_log2_ * float(M_LN2)), sink_val(sink_val_), use_fa4_sink(use_fa4_sink_) {}; template __forceinline__ __device__ TensorT max_get_scale(Tensor0 &acc_s) { @@ -109,10 +110,14 @@ struct Softmax { TensorT scores_scale; if constexpr (Is_first) { if constexpr (Has_sink) { - const float sink_scaled = sink_val / softmax_scale; - #pragma unroll - for (int mi = 0; mi < size(row_max); ++mi) { row_max(mi) = sink_scaled; } - flash::template reduce_max(scores, row_max); + if (use_fa4_sink) { + flash::template reduce_max(scores, row_max); + } else { + const float sink_scaled = sink_val / softmax_scale; + #pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { row_max(mi) = sink_scaled; } + flash::template reduce_max(scores, row_max); + } } else { flash::template reduce_max(scores, row_max); } @@ -153,7 +158,7 @@ struct Softmax { float sum = row_sum(mi); if constexpr (Has_sink) { static constexpr float max_offset = float(Max_offset); - const float max_scaled = max(mi) == -INFINITY ? 0.f : (row_max(mi) * softmax_scale_log2) - max_offset; + const float max_scaled = row_max(mi) == -INFINITY ? 0.f : (row_max(mi) * softmax_scale_log2) - max_offset; sum += exp2f(sink_val * float(M_LOG2E) - max_scaled); } float inv_sum = (sum == 0.f || sum != sum) ? 0.f : 1.f / sum; diff --git a/hopper/test_flash_attn.py b/hopper/test_flash_attn.py index 050a36b6d47..9f7b75533fc 100644 --- a/hopper/test_flash_attn.py +++ b/hopper/test_flash_attn.py @@ -216,6 +216,7 @@ def test_flash_attn_output( pack_gqa=pack_gqa, num_splits=num_splits, learnable_sink=learnable_sink, + use_fa4_sink=True, ) print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") @@ -512,6 +513,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): pack_gqa=pack_gqa, num_splits=num_splits, learnable_sink=learnable_sink, + use_fa4_sink=True, ) out = output_pad_fn(out_unpad) if query_unused_mask is not None: @@ -952,6 +954,7 @@ def test_flash_attn_kvcache( num_splits=num_splits, return_softmax_lse=True, learnable_sink=learnable_sink, + use_fa4_sink=True, ) if varlen_q: out = output_pad_fn(out) From b6cb59549d09e91f45ba510afcd8515835607d71 Mon Sep 17 00:00:00 2001 From: aoxy Date: Thu, 7 May 2026 20:24:09 +0800 Subject: [PATCH 19/23] Add test. --- .../test_flash_attn_varlen_sink_dump_load.py | 513 ++++++++++++++++++ 1 file changed, 513 insertions(+) create mode 100644 hopper/test_flash_attn_varlen_sink_dump_load.py diff --git a/hopper/test_flash_attn_varlen_sink_dump_load.py b/hopper/test_flash_attn_varlen_sink_dump_load.py new file mode 100644 index 00000000000..4e3139d8011 --- /dev/null +++ b/hopper/test_flash_attn_varlen_sink_dump_load.py @@ -0,0 +1,513 @@ +#!/root/.envs/swa/bin/python +import argparse +from pathlib import Path +from typing import Any + +import pytest +import torch + +from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_fa3 +from test_util import generate_qkv + +try: + from flash_attn.cute.interface import flash_attn_varlen_func as flash_attn_varlen_func_cute + + _CUTE_IMPORT_ERROR = None +except Exception as exc: # pragma: no cover - only hit when cute isn't installed + flash_attn_varlen_func_cute = None + _CUTE_IMPORT_ERROR = exc + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _log(msg: str) -> None: + print(f"[varlen-sink-test] {msg}", flush=True) + + +def _build_case(seed: int = 20260507): + """Build a single reproducible varlen case used for dump/load replay.""" + _log(f"Building case with seed={seed}") + torch.random.manual_seed(seed) + device = "cuda" + dtype = torch.bfloat16 + + batch_size = 4 + seqlen_q = 512 + seqlen_k = 512 + nheads = 24 + nheads_k = 2 + headdim = 256 + + _log( + "Case config: " + f"batch={batch_size}, seqlen_q={seqlen_q}, seqlen_k={seqlen_k}, " + f"nheads={nheads}, nheads_k={nheads_k}, headdim={headdim}, dtype={dtype}" + ) + + q = torch.randn(batch_size, seqlen_q, nheads, headdim, device=device, dtype=dtype) + k = torch.randn(batch_size, seqlen_k, nheads_k, headdim, device=device, dtype=dtype) + v = torch.randn(batch_size, seqlen_k, nheads_k, headdim, device=device, dtype=dtype) + + ( + q_unpad, + k_unpad, + v_unpad, + _qv_unpad, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + _q, + _k, + _v, + _qv, + _output_pad_fn, + _dq_pad_fn, + _dk_pad_fn, + ) = generate_qkv( + q, + k, + v, + query_padding_mask=None, + key_padding_mask=None, + qv=None, + kvpacked=False, + ) + + q_unpad = q_unpad.detach().to(dtype) + k_unpad = k_unpad.detach().to(dtype) + v_unpad = v_unpad.detach().to(dtype) + + g_unpad = torch.randn( + q_unpad.shape[0], q_unpad.shape[1], v_unpad.shape[-1], device=device, dtype=dtype + ) + sink = torch.randn(nheads, device=device, dtype=dtype) + + _log( + "Built tensors: " + f"q_unpad={tuple(q_unpad.shape)}, k_unpad={tuple(k_unpad.shape)}, " + f"v_unpad={tuple(v_unpad.shape)}, g_unpad={tuple(g_unpad.shape)}, sink={tuple(sink.shape)}" + ) + + return { + "q_unpad": q_unpad.cpu(), + "k_unpad": k_unpad.cpu(), + "v_unpad": v_unpad.cpu(), + "cu_seqlens_q": cu_seqlens_q.cpu(), + "cu_seqlens_k": cu_seqlens_k.cpu(), + "seqused_q": None if seqused_q is None else seqused_q.cpu(), + "seqused_k": None if seqused_k is None else seqused_k.cpu(), + "max_seqlen_q": int(max_seqlen_q), + "max_seqlen_k": int(max_seqlen_k), + "learnable_sink": sink.cpu(), + "g_unpad": g_unpad.cpu(), + "causal": True, + "window_size": (-1, -1), + "softcap": 0.0, + "deterministic": True, + "num_splits": 1, + "pack_gqa": False, + } + + +def _to_cuda_tensor(x, requires_grad: bool = False): + if x is None: + return None + y = x.to(device="cuda") + if requires_grad: + y = y.detach().requires_grad_(True) + return y + + +def _normalize_window_size_for_cute(window_size): + left, right = window_size + return (None if left < 0 else left, None if right < 0 else right) + + +def _empty_grad_like(q: torch.Tensor) -> torch.Tensor: + return torch.empty(0, device=q.device, dtype=q.dtype) + + +def _run_case_fa3(case, use_fa4_sink: bool, enable_sink: bool): + effective_use_fa4_sink = use_fa4_sink and enable_sink + if use_fa4_sink and not enable_sink: + _log("Running FA3 case: use_fa4_sink=True ignored because sink is disabled") + else: + _log( + f"Running FA3 case: use_fa4_sink={effective_use_fa4_sink}, " + f"enable_sink={enable_sink}" + ) + + q = _to_cuda_tensor(case["q_unpad"], requires_grad=True) + k = _to_cuda_tensor(case["k_unpad"], requires_grad=True) + v = _to_cuda_tensor(case["v_unpad"], requires_grad=True) + sink = _to_cuda_tensor(case["learnable_sink"], requires_grad=True) if enable_sink else None + g_unpad = _to_cuda_tensor(case["g_unpad"], requires_grad=False) + + out, lse = flash_attn_varlen_func_fa3( + q, + k, + v, + _to_cuda_tensor(case["cu_seqlens_q"]), + _to_cuda_tensor(case["cu_seqlens_k"]), + case["max_seqlen_q"], + case["max_seqlen_k"], + seqused_q=_to_cuda_tensor(case["seqused_q"]), + seqused_k=_to_cuda_tensor(case["seqused_k"]), + causal=case["causal"], + window_size=tuple(case["window_size"]), + softcap=case["softcap"], + deterministic=case["deterministic"], + num_splits=case["num_splits"], + pack_gqa=case["pack_gqa"], + learnable_sink=sink, + use_fa4_sink=effective_use_fa4_sink, + ) + + dq, dk, dv = torch.autograd.grad(out, (q, k, v), g_unpad, retain_graph=enable_sink) + if enable_sink: + (dsink,) = torch.autograd.grad(out, sink, g_unpad) + else: + dsink = _empty_grad_like(q) + + result = { + "out": out.detach().cpu(), + "lse": lse.detach().cpu(), + "dq": dq.detach().cpu(), + "dk": dk.detach().cpu(), + "dv": dv.detach().cpu(), + "dsink": dsink.detach().cpu(), + } + _log(f"FA3 run done: use_fa4_sink={effective_use_fa4_sink}, enable_sink={enable_sink}") + return result + + +def _run_case_cute(case, use_fa4_sink: bool, enable_sink: bool): + if flash_attn_varlen_func_cute is None: + raise RuntimeError(f"flash_attn.cute import failed: {_CUTE_IMPORT_ERROR}") + + if use_fa4_sink: + _log("Running CuTe case: use_fa4_sink=True requested, but CuTe has no such switch; ignoring") + else: + _log("Running CuTe case: use_fa4_sink=False") + _log(f"CuTe run config: enable_sink={enable_sink}") + + q = _to_cuda_tensor(case["q_unpad"], requires_grad=True) + k = _to_cuda_tensor(case["k_unpad"], requires_grad=True) + v = _to_cuda_tensor(case["v_unpad"], requires_grad=True) + sink = _to_cuda_tensor(case["learnable_sink"], requires_grad=True) if enable_sink else None + g_unpad = _to_cuda_tensor(case["g_unpad"], requires_grad=False) + + out, lse = flash_attn_varlen_func_cute( + q=q, + k=k, + v=v, + cu_seqlens_q=_to_cuda_tensor(case["cu_seqlens_q"]), + cu_seqlens_k=_to_cuda_tensor(case["cu_seqlens_k"]), + max_seqlen_q=case["max_seqlen_q"], + max_seqlen_k=case["max_seqlen_k"], + seqused_q=_to_cuda_tensor(case["seqused_q"]), + seqused_k=_to_cuda_tensor(case["seqused_k"]), + causal=case["causal"], + window_size=_normalize_window_size_for_cute(case["window_size"]), + learnable_sink=sink, + softcap=case["softcap"], + num_splits=case["num_splits"], + pack_gqa=case["pack_gqa"], + deterministic=case["deterministic"], + return_lse=True, + ) + dq, dk, dv = torch.autograd.grad(out, (q, k, v), g_unpad, retain_graph=enable_sink) + if enable_sink: + # cute not support sink gradient, force disable to avoid autograd error + # (dsink,) = torch.autograd.grad(out, sink, g_unpad) + dsink = torch.zeros_like(sink) + else: + dsink = _empty_grad_like(q) + + + result = { + "out": out.detach().cpu(), + "lse": lse.detach().cpu(), + "dq": dq.detach().cpu(), + "dk": dk.detach().cpu(), + "dv": dv.detach().cpu(), + "dsink": dsink.detach().cpu(), + } + _log(f"CuTe run done: enable_sink={enable_sink}") + return result + + +def _run_case(case, use_fa4_sink: bool, backend: str, enable_sink: bool): + if backend == "fa3": + return _run_case_fa3(case, use_fa4_sink=use_fa4_sink, enable_sink=enable_sink) + if backend == "cute": + return _run_case_cute(case, use_fa4_sink=use_fa4_sink, enable_sink=enable_sink) + raise ValueError(f"Unknown backend: {backend}") + + +def _assert_finite(tag: str, result): + _log(f"[CHECK] finite: {tag}") + for name, t in result.items(): + finite = torch.isfinite(t) + if not finite.all(): + bad = (~finite).sum().item() + raise AssertionError(f"{tag}.{name} contains NaN/Inf (bad_count={bad})") + _log(f" [OK] {tag}.{name} finite, shape={tuple(t.shape)}, dtype={t.dtype}") + + +def _assert_close(name: str, a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float): + if a.numel() == 0 and b.numel() == 0: + _log(f"[CHECK] close: {name} | both tensors empty, skip numeric check") + return + if a.numel() == 0 or b.numel() == 0: + raise AssertionError(f"{name} mismatch: one tensor is empty and the other is not") + + a_f = a.float() + b_f = b.float() + diff = (a_f - b_f).abs().max().item() + scale = max(a_f.abs().max().item(), b_f.abs().max().item(), 1.0) + limit = atol + rtol * scale + _log( + f"[CHECK] close: {name} | max_diff={diff:.6e}, limit={limit:.6e} " + f"(atol={atol:.2e}, rtol={rtol:.2e}, scale={scale:.6e})" + ) + if diff > limit: + raise AssertionError(f"{name} max diff {diff:.6e} > limit {limit:.6e}") + + +def _compare_results(tag: str, a, b): + _log(f"[CHECK] compare result group: {tag}") + out_equal = torch.equal(a["out"], b["out"]) + lse_equal = torch.equal(a["lse"], b["lse"]) + _log("=" * 88) + _log(f"[BITWISE][{tag}] out torch.equal: {'PASS' if out_equal else 'FAIL'}") + _log(f"[BITWISE][{tag}] lse torch.equal: {'PASS' if lse_equal else 'FAIL'}") + _log("=" * 88) + + tol = { + "out": (5e-3, 5e-2), + "lse": (1e-2, 5e-2), + "dq": (1e-2, 8e-2), + "dk": (1e-2, 8e-2), + "dv": (1e-2, 8e-2), + # "dsink": (1e-2, 8e-2), + } + for key, (atol, rtol) in tol.items(): + _assert_close(f"{tag}.{key}", a[key], b[key], atol=atol, rtol=rtol) + _log(f"[OK] compare result group passed: {tag}") + + +def dump_case(path: Path, backend: str = "fa3", enable_sink: bool = True): + _log(f"Dump start -> {path} (backend={backend}, enable_sink={enable_sink})") + case = _build_case() + result_off = _run_case(case, use_fa4_sink=False, backend=backend, enable_sink=enable_sink) + result_on = _run_case(case, use_fa4_sink=True, backend=backend, enable_sink=enable_sink) + + _assert_finite("dump.fa4_off", result_off) + _assert_finite("dump.fa4_on", result_on) + + payload = { + "backend": backend, + "enable_sink": enable_sink, + "case": case, + "result_fa4_off": result_off, + "result_fa4_on": result_on, + } + path.parent.mkdir(parents=True, exist_ok=True) + torch.save(payload, path) + _log(f"Dump finished -> {path}") + + +def load_and_check( + path: Path, + backend: str = "fa3", + enable_sink: bool = True, + allow_backend_mismatch: bool = True, +): + _log( + f"Load start <- {path} (backend={backend}, enable_sink={enable_sink}, " + f"allow_backend_mismatch={allow_backend_mismatch})" + ) + payload = torch.load(path, map_location="cpu") + case = payload["case"] + + dumped_backend = payload.get("backend", "fa3") + if dumped_backend != backend: + if allow_backend_mismatch: + _log( + f"[CHECK] backend mismatch allowed: dump backend={dumped_backend}, " + f"replay backend={backend}" + ) + else: + raise AssertionError( + f"Backend mismatch: dump was produced with backend={dumped_backend}, " + f"but check requested backend={backend}" + ) + + dumped_enable_sink = payload.get("enable_sink", True) + if dumped_enable_sink != enable_sink: + raise AssertionError( + f"Sink enable mismatch: dump was produced with enable_sink={dumped_enable_sink}, " + f"but check requested enable_sink={enable_sink}" + ) + + replay_off = _run_case(case, use_fa4_sink=False, backend=backend, enable_sink=enable_sink) + replay_on = _run_case(case, use_fa4_sink=True, backend=backend, enable_sink=enable_sink) + + _assert_finite("replay.fa4_off", replay_off) + _assert_finite("replay.fa4_on", replay_on) + + base_tag = ( + f"replay_vs_dump.{backend}_vs_{dumped_backend}" + if dumped_backend != backend + else "replay_vs_dump" + ) + _compare_results(f"{base_tag}.fa4_off", replay_off, payload["result_fa4_off"]) + _compare_results(f"{base_tag}.fa4_on", replay_on, payload["result_fa4_on"]) + _compare_results(f"replay.{backend}.fa4_on_vs_off", replay_on, replay_off) + _log("Load+check finished") + + +def dump_and_check( + path: Path, + backend: str = "fa3", + enable_sink: bool = True, + check_backend: str | None = None, + allow_backend_mismatch: bool = True, +): + _log("Dump+check pipeline start") + dump_case(path, backend=backend, enable_sink=enable_sink) + replay_backend = backend if check_backend is None else check_backend + load_and_check( + path, + backend=replay_backend, + enable_sink=enable_sink, + allow_backend_mismatch=allow_backend_mismatch, + ) + _log("Dump+check pipeline finished") + + +def compare_fa3_vs_cute(enable_sink: bool = True): + if flash_attn_varlen_func_cute is None: + raise RuntimeError(f"flash_attn.cute import failed: {_CUTE_IMPORT_ERROR}") + + _log(f"FA3 vs CuTe compare start (enable_sink={enable_sink})") + case = _build_case() + + fa3_off = _run_case(case, use_fa4_sink=False, backend="fa3", enable_sink=enable_sink) + cute_off = _run_case(case, use_fa4_sink=False, backend="cute", enable_sink=enable_sink) + _assert_finite("compare.fa3.fa4_off", fa3_off) + _assert_finite("compare.cute.fa4_off", cute_off) + _compare_results("fa3_vs_cute.fa4_off", fa3_off, cute_off) + + fa3_on = _run_case(case, use_fa4_sink=True, backend="fa3", enable_sink=enable_sink) + cute_on = _run_case(case, use_fa4_sink=True, backend="cute", enable_sink=enable_sink) + _assert_finite("compare.fa3.fa4_on", fa3_on) + _assert_finite("compare.cute.fa4_on", cute_on) + _compare_results("fa3_vs_cute.fa4_on", fa3_on, cute_on) + + _compare_results("fa3.fa4_on_vs_off", fa3_on, fa3_off) + _compare_results("cute.fa4_on_vs_off", cute_on, cute_off) + _log("FA3 vs CuTe compare finished") + + +def test_flash_attn_varlen_sink_dump_load_compare(tmp_path: Path): + dump_path = tmp_path / "flash_attn_varlen_sink_case.pt" + dump_and_check(dump_path, backend="fa3", enable_sink=True) + + +@pytest.mark.skipif(flash_attn_varlen_func_cute is None, reason="CuTe backend is not available") +def test_flash_attn_varlen_sink_dump_load_compare_cute(tmp_path: Path): + dump_path = tmp_path / "flash_attn_varlen_sink_case_cute.pt" + dump_and_check(dump_path, backend="cute", enable_sink=True) + + +def test_flash_attn_varlen_sink_dump_load_compare_no_sink(tmp_path: Path): + dump_path = tmp_path / "flash_attn_varlen_no_sink_case.pt" + dump_and_check(dump_path, backend="fa3", enable_sink=False) + + +def _parse_args(): + parser = argparse.ArgumentParser(description="Varlen FlashAttention sink precision dump/load test") + parser.add_argument( + "--mode", + choices=["dump", "check", "dump_and_check", "compare_fa3_cute"], + default="dump_and_check", + help="dump only, check only, dump then check, or compare FA3 vs CuTe directly", + ) + parser.add_argument( + "--backend", + choices=["fa3", "cute"], + default="fa3", + help="attention backend to run (used by dump/check/dump_and_check)", + ) + parser.add_argument( + "--check-backend", + choices=["fa3", "cute"], + default=None, + help="backend used in check stage for dump_and_check; defaults to --backend", + ) + parser.add_argument( + "--strict-backend-match", + action="store_true", + help="fail check when dump backend and replay backend are different", + ) + parser.add_argument( + "--enable-sink", + action=argparse.BooleanOptionalAction, + default=True, + help="enable or disable learnable_sink", + ) + parser.add_argument( + "--path", + type=Path, + default=Path(__file__).resolve().parent / "varlen_sink_precision_case.pt", + help="path for saved tensors/results", + ) + return parser.parse_args() + + +def main(): + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required to run this test") + + args = _parse_args() + _log( + f"Running mode={args.mode}, backend={args.backend}, check_backend={args.check_backend}, " + f"strict_backend_match={args.strict_backend_match}, enable_sink={args.enable_sink}, " + f"path={args.path}" + ) + + if args.mode == "dump": + dump_case(args.path, backend=args.backend, enable_sink=args.enable_sink) + elif args.mode == "check": + load_and_check( + args.path, + backend=args.backend, + enable_sink=args.enable_sink, + allow_backend_mismatch=not args.strict_backend_match, + ) + elif args.mode == "dump_and_check": + dump_and_check( + args.path, + backend=args.backend, + enable_sink=args.enable_sink, + check_backend=args.check_backend, + allow_backend_mismatch=not args.strict_backend_match, + ) + else: + compare_fa3_vs_cute(enable_sink=args.enable_sink) + + +if __name__ == "__main__": + main() + +# usage examples +# python test_flash_attn_varlen_sink_dump_load.py --backend fa3 --mode dump_and_check --enable-sink --path ./with_sink.pt +# python test_flash_attn_varlen_sink_dump_load.py --backend fa3 --mode dump_and_check --no-enable-sink --path ./no_sink.pt +# python test_flash_attn_varlen_sink_dump_load.py --backend cute --mode check --enable-sink --path ./with_sink.pt +# python test_flash_attn_varlen_sink_dump_load.py --backend cute --mode check --no-enable-sink --path ./no_sink.pt From 05f9cfe160c170c2bde35104124cdc02491eb3d7 Mon Sep 17 00:00:00 2001 From: aoxy Date: Fri, 8 May 2026 11:24:12 +0800 Subject: [PATCH 20/23] Update test. --- .../test_flash_attn_varlen_sink_dump_load.py | 187 +++++------------- 1 file changed, 49 insertions(+), 138 deletions(-) diff --git a/hopper/test_flash_attn_varlen_sink_dump_load.py b/hopper/test_flash_attn_varlen_sink_dump_load.py index 4e3139d8011..292716544b3 100644 --- a/hopper/test_flash_attn_varlen_sink_dump_load.py +++ b/hopper/test_flash_attn_varlen_sink_dump_load.py @@ -3,22 +3,13 @@ from pathlib import Path from typing import Any -import pytest import torch from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_fa3 from test_util import generate_qkv -try: - from flash_attn.cute.interface import flash_attn_varlen_func as flash_attn_varlen_func_cute - _CUTE_IMPORT_ERROR = None -except Exception as exc: # pragma: no cover - only hit when cute isn't installed - flash_attn_varlen_func_cute = None - _CUTE_IMPORT_ERROR = exc - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +from sgl_kernel.flash_attn import flash_attn_varlen_func as flash_attn_varlen_func_sgl def _log(msg: str) -> None: @@ -39,11 +30,7 @@ def _build_case(seed: int = 20260507): nheads_k = 2 headdim = 256 - _log( - "Case config: " - f"batch={batch_size}, seqlen_q={seqlen_q}, seqlen_k={seqlen_k}, " - f"nheads={nheads}, nheads_k={nheads_k}, headdim={headdim}, dtype={dtype}" - ) + _log("Case config: " f"batch={batch_size}, seqlen_q={seqlen_q}, seqlen_k={seqlen_k}, " f"nheads={nheads}, nheads_k={nheads_k}, headdim={headdim}, dtype={dtype}") q = torch.randn(batch_size, seqlen_q, nheads, headdim, device=device, dtype=dtype) k = torch.randn(batch_size, seqlen_k, nheads_k, headdim, device=device, dtype=dtype) @@ -81,16 +68,10 @@ def _build_case(seed: int = 20260507): k_unpad = k_unpad.detach().to(dtype) v_unpad = v_unpad.detach().to(dtype) - g_unpad = torch.randn( - q_unpad.shape[0], q_unpad.shape[1], v_unpad.shape[-1], device=device, dtype=dtype - ) + g_unpad = torch.randn(q_unpad.shape[0], q_unpad.shape[1], v_unpad.shape[-1], device=device, dtype=dtype) sink = torch.randn(nheads, device=device, dtype=dtype) - _log( - "Built tensors: " - f"q_unpad={tuple(q_unpad.shape)}, k_unpad={tuple(k_unpad.shape)}, " - f"v_unpad={tuple(v_unpad.shape)}, g_unpad={tuple(g_unpad.shape)}, sink={tuple(sink.shape)}" - ) + _log("Built tensors: " f"q_unpad={tuple(q_unpad.shape)}, k_unpad={tuple(k_unpad.shape)}, " f"v_unpad={tuple(v_unpad.shape)}, g_unpad={tuple(g_unpad.shape)}, sink={tuple(sink.shape)}") return { "q_unpad": q_unpad.cpu(), @@ -122,30 +103,17 @@ def _to_cuda_tensor(x, requires_grad: bool = False): return y -def _normalize_window_size_for_cute(window_size): - left, right = window_size - return (None if left < 0 else left, None if right < 0 else right) - - -def _empty_grad_like(q: torch.Tensor) -> torch.Tensor: - return torch.empty(0, device=q.device, dtype=q.dtype) - - def _run_case_fa3(case, use_fa4_sink: bool, enable_sink: bool): effective_use_fa4_sink = use_fa4_sink and enable_sink if use_fa4_sink and not enable_sink: _log("Running FA3 case: use_fa4_sink=True ignored because sink is disabled") else: - _log( - f"Running FA3 case: use_fa4_sink={effective_use_fa4_sink}, " - f"enable_sink={enable_sink}" - ) + _log(f"Running FA3 case: use_fa4_sink={effective_use_fa4_sink}, " f"enable_sink={enable_sink}") q = _to_cuda_tensor(case["q_unpad"], requires_grad=True) k = _to_cuda_tensor(case["k_unpad"], requires_grad=True) v = _to_cuda_tensor(case["v_unpad"], requires_grad=True) sink = _to_cuda_tensor(case["learnable_sink"], requires_grad=True) if enable_sink else None - g_unpad = _to_cuda_tensor(case["g_unpad"], requires_grad=False) out, lse = flash_attn_varlen_func_fa3( q, @@ -167,85 +135,55 @@ def _run_case_fa3(case, use_fa4_sink: bool, enable_sink: bool): use_fa4_sink=effective_use_fa4_sink, ) - dq, dk, dv = torch.autograd.grad(out, (q, k, v), g_unpad, retain_graph=enable_sink) - if enable_sink: - (dsink,) = torch.autograd.grad(out, sink, g_unpad) - else: - dsink = _empty_grad_like(q) - result = { "out": out.detach().cpu(), "lse": lse.detach().cpu(), - "dq": dq.detach().cpu(), - "dk": dk.detach().cpu(), - "dv": dv.detach().cpu(), - "dsink": dsink.detach().cpu(), } _log(f"FA3 run done: use_fa4_sink={effective_use_fa4_sink}, enable_sink={enable_sink}") return result -def _run_case_cute(case, use_fa4_sink: bool, enable_sink: bool): - if flash_attn_varlen_func_cute is None: - raise RuntimeError(f"flash_attn.cute import failed: {_CUTE_IMPORT_ERROR}") - - if use_fa4_sink: - _log("Running CuTe case: use_fa4_sink=True requested, but CuTe has no such switch; ignoring") - else: - _log("Running CuTe case: use_fa4_sink=False") - _log(f"CuTe run config: enable_sink={enable_sink}") +def _run_case_sgl(case, use_fa4_sink: bool, enable_sink: bool): + _log("Running SGL kernel version of flash_attn_varlen_func...") q = _to_cuda_tensor(case["q_unpad"], requires_grad=True) k = _to_cuda_tensor(case["k_unpad"], requires_grad=True) v = _to_cuda_tensor(case["v_unpad"], requires_grad=True) sink = _to_cuda_tensor(case["learnable_sink"], requires_grad=True) if enable_sink else None - g_unpad = _to_cuda_tensor(case["g_unpad"], requires_grad=False) - - out, lse = flash_attn_varlen_func_cute( - q=q, - k=k, - v=v, - cu_seqlens_q=_to_cuda_tensor(case["cu_seqlens_q"]), - cu_seqlens_k=_to_cuda_tensor(case["cu_seqlens_k"]), - max_seqlen_q=case["max_seqlen_q"], - max_seqlen_k=case["max_seqlen_k"], + + out, lse, *rest = flash_attn_varlen_func_sgl( + q, + k, + v, + _to_cuda_tensor(case["cu_seqlens_q"]), + _to_cuda_tensor(case["cu_seqlens_k"]), + case["max_seqlen_q"], + case["max_seqlen_k"], seqused_q=_to_cuda_tensor(case["seqused_q"]), seqused_k=_to_cuda_tensor(case["seqused_k"]), causal=case["causal"], - window_size=_normalize_window_size_for_cute(case["window_size"]), - learnable_sink=sink, + window_size=tuple(case["window_size"]), softcap=case["softcap"], + # deterministic=case["deterministic"], num_splits=case["num_splits"], pack_gqa=case["pack_gqa"], - deterministic=case["deterministic"], - return_lse=True, + return_softmax_lse=True, + sinks=sink, ) - dq, dk, dv = torch.autograd.grad(out, (q, k, v), g_unpad, retain_graph=enable_sink) - if enable_sink: - # cute not support sink gradient, force disable to avoid autograd error - # (dsink,) = torch.autograd.grad(out, sink, g_unpad) - dsink = torch.zeros_like(sink) - else: - dsink = _empty_grad_like(q) - result = { "out": out.detach().cpu(), "lse": lse.detach().cpu(), - "dq": dq.detach().cpu(), - "dk": dk.detach().cpu(), - "dv": dv.detach().cpu(), - "dsink": dsink.detach().cpu(), } - _log(f"CuTe run done: enable_sink={enable_sink}") + _log(f"SGL kernel run done, enable_sink={enable_sink}") return result def _run_case(case, use_fa4_sink: bool, backend: str, enable_sink: bool): if backend == "fa3": return _run_case_fa3(case, use_fa4_sink=use_fa4_sink, enable_sink=enable_sink) - if backend == "cute": - return _run_case_cute(case, use_fa4_sink=use_fa4_sink, enable_sink=enable_sink) + if backend == "sgl": + return _run_case_sgl(case, use_fa4_sink=use_fa4_sink, enable_sink=enable_sink) raise ValueError(f"Unknown backend: {backend}") @@ -271,10 +209,7 @@ def _assert_close(name: str, a: torch.Tensor, b: torch.Tensor, atol: float, rtol diff = (a_f - b_f).abs().max().item() scale = max(a_f.abs().max().item(), b_f.abs().max().item(), 1.0) limit = atol + rtol * scale - _log( - f"[CHECK] close: {name} | max_diff={diff:.6e}, limit={limit:.6e} " - f"(atol={atol:.2e}, rtol={rtol:.2e}, scale={scale:.6e})" - ) + _log(f"[CHECK] close: {name} | max_diff={diff:.6e}, limit={limit:.6e} " f"(atol={atol:.2e}, rtol={rtol:.2e}, scale={scale:.6e})") if diff > limit: raise AssertionError(f"{name} max diff {diff:.6e} > limit {limit:.6e}") @@ -291,10 +226,6 @@ def _compare_results(tag: str, a, b): tol = { "out": (5e-3, 5e-2), "lse": (1e-2, 5e-2), - "dq": (1e-2, 8e-2), - "dk": (1e-2, 8e-2), - "dv": (1e-2, 8e-2), - # "dsink": (1e-2, 8e-2), } for key, (atol, rtol) in tol.items(): _assert_close(f"{tag}.{key}", a[key], b[key], atol=atol, rtol=rtol) @@ -328,32 +259,20 @@ def load_and_check( enable_sink: bool = True, allow_backend_mismatch: bool = True, ): - _log( - f"Load start <- {path} (backend={backend}, enable_sink={enable_sink}, " - f"allow_backend_mismatch={allow_backend_mismatch})" - ) + _log(f"Load start <- {path} (backend={backend}, enable_sink={enable_sink}, " f"allow_backend_mismatch={allow_backend_mismatch})") payload = torch.load(path, map_location="cpu") case = payload["case"] dumped_backend = payload.get("backend", "fa3") if dumped_backend != backend: if allow_backend_mismatch: - _log( - f"[CHECK] backend mismatch allowed: dump backend={dumped_backend}, " - f"replay backend={backend}" - ) + _log(f"[CHECK] backend mismatch allowed: dump backend={dumped_backend}, " f"replay backend={backend}") else: - raise AssertionError( - f"Backend mismatch: dump was produced with backend={dumped_backend}, " - f"but check requested backend={backend}" - ) + raise AssertionError(f"Backend mismatch: dump was produced with backend={dumped_backend}, " f"but check requested backend={backend}") dumped_enable_sink = payload.get("enable_sink", True) if dumped_enable_sink != enable_sink: - raise AssertionError( - f"Sink enable mismatch: dump was produced with enable_sink={dumped_enable_sink}, " - f"but check requested enable_sink={enable_sink}" - ) + raise AssertionError(f"Sink enable mismatch: dump was produced with enable_sink={dumped_enable_sink}, " f"but check requested enable_sink={enable_sink}") replay_off = _run_case(case, use_fa4_sink=False, backend=backend, enable_sink=enable_sink) replay_on = _run_case(case, use_fa4_sink=True, backend=backend, enable_sink=enable_sink) @@ -361,11 +280,7 @@ def load_and_check( _assert_finite("replay.fa4_off", replay_off) _assert_finite("replay.fa4_on", replay_on) - base_tag = ( - f"replay_vs_dump.{backend}_vs_{dumped_backend}" - if dumped_backend != backend - else "replay_vs_dump" - ) + base_tag = f"replay_vs_dump.{backend}_vs_{dumped_backend}" if dumped_backend != backend else "replay_vs_dump" _compare_results(f"{base_tag}.fa4_off", replay_off, payload["result_fa4_off"]) _compare_results(f"{base_tag}.fa4_on", replay_on, payload["result_fa4_on"]) _compare_results(f"replay.{backend}.fa4_on_vs_off", replay_on, replay_off) @@ -391,28 +306,25 @@ def dump_and_check( _log("Dump+check pipeline finished") -def compare_fa3_vs_cute(enable_sink: bool = True): - if flash_attn_varlen_func_cute is None: - raise RuntimeError(f"flash_attn.cute import failed: {_CUTE_IMPORT_ERROR}") - - _log(f"FA3 vs CuTe compare start (enable_sink={enable_sink})") +def compare_fa3_vs_sgl(enable_sink: bool = True): + _log(f"FA3 vs sgl compare start (enable_sink={enable_sink})") case = _build_case() fa3_off = _run_case(case, use_fa4_sink=False, backend="fa3", enable_sink=enable_sink) - cute_off = _run_case(case, use_fa4_sink=False, backend="cute", enable_sink=enable_sink) + sgl_off = _run_case(case, use_fa4_sink=False, backend="sgl", enable_sink=enable_sink) _assert_finite("compare.fa3.fa4_off", fa3_off) - _assert_finite("compare.cute.fa4_off", cute_off) - _compare_results("fa3_vs_cute.fa4_off", fa3_off, cute_off) + _assert_finite("compare.sgl.fa4_off", sgl_off) + _compare_results("fa3_vs_sgl.fa4_off", fa3_off, sgl_off) fa3_on = _run_case(case, use_fa4_sink=True, backend="fa3", enable_sink=enable_sink) - cute_on = _run_case(case, use_fa4_sink=True, backend="cute", enable_sink=enable_sink) + sgl_on = _run_case(case, use_fa4_sink=True, backend="sgl", enable_sink=enable_sink) _assert_finite("compare.fa3.fa4_on", fa3_on) - _assert_finite("compare.cute.fa4_on", cute_on) - _compare_results("fa3_vs_cute.fa4_on", fa3_on, cute_on) + _assert_finite("compare.sgl.fa4_on", sgl_on) + _compare_results("fa3_vs_sgl.fa4_on", fa3_on, sgl_on) _compare_results("fa3.fa4_on_vs_off", fa3_on, fa3_off) - _compare_results("cute.fa4_on_vs_off", cute_on, cute_off) - _log("FA3 vs CuTe compare finished") + _compare_results("sgl.fa4_on_vs_off", sgl_on, sgl_off) + _log("FA3 vs sgl compare finished") def test_flash_attn_varlen_sink_dump_load_compare(tmp_path: Path): @@ -420,10 +332,9 @@ def test_flash_attn_varlen_sink_dump_load_compare(tmp_path: Path): dump_and_check(dump_path, backend="fa3", enable_sink=True) -@pytest.mark.skipif(flash_attn_varlen_func_cute is None, reason="CuTe backend is not available") -def test_flash_attn_varlen_sink_dump_load_compare_cute(tmp_path: Path): - dump_path = tmp_path / "flash_attn_varlen_sink_case_cute.pt" - dump_and_check(dump_path, backend="cute", enable_sink=True) +def test_flash_attn_varlen_sink_dump_load_compare_sgl(tmp_path: Path): + dump_path = tmp_path / "flash_attn_varlen_sink_case_sgl.pt" + dump_and_check(dump_path, backend="sgl", enable_sink=True) def test_flash_attn_varlen_sink_dump_load_compare_no_sink(tmp_path: Path): @@ -435,19 +346,19 @@ def _parse_args(): parser = argparse.ArgumentParser(description="Varlen FlashAttention sink precision dump/load test") parser.add_argument( "--mode", - choices=["dump", "check", "dump_and_check", "compare_fa3_cute"], + choices=["dump", "check", "dump_and_check", "compare_fa3_sgl"], default="dump_and_check", - help="dump only, check only, dump then check, or compare FA3 vs CuTe directly", + help="dump only, check only, dump then check, or compare FA3 vs sgl directly", ) parser.add_argument( "--backend", - choices=["fa3", "cute"], + choices=["fa3", "sgl"], default="fa3", help="attention backend to run (used by dump/check/dump_and_check)", ) parser.add_argument( "--check-backend", - choices=["fa3", "cute"], + choices=["fa3", "sgl"], default=None, help="backend used in check stage for dump_and_check; defaults to --backend", ) @@ -500,7 +411,7 @@ def main(): allow_backend_mismatch=not args.strict_backend_match, ) else: - compare_fa3_vs_cute(enable_sink=args.enable_sink) + compare_fa3_vs_sgl(enable_sink=args.enable_sink) if __name__ == "__main__": @@ -509,5 +420,5 @@ def main(): # usage examples # python test_flash_attn_varlen_sink_dump_load.py --backend fa3 --mode dump_and_check --enable-sink --path ./with_sink.pt # python test_flash_attn_varlen_sink_dump_load.py --backend fa3 --mode dump_and_check --no-enable-sink --path ./no_sink.pt -# python test_flash_attn_varlen_sink_dump_load.py --backend cute --mode check --enable-sink --path ./with_sink.pt -# python test_flash_attn_varlen_sink_dump_load.py --backend cute --mode check --no-enable-sink --path ./no_sink.pt +# python test_flash_attn_varlen_sink_dump_load.py --backend fa3 --mode check --enable-sink --path ./with_sink.pt +# python test_flash_attn_varlen_sink_dump_load.py --backend sgl --mode check --no-enable-sink --path ./no_sink.pt From 0aa05b6a05631213827789c8e9f461bb6f2ecf10 Mon Sep 17 00:00:00 2001 From: aoxy Date: Fri, 22 May 2026 22:57:06 +0800 Subject: [PATCH 21/23] Use_fa4_sink to template. --- hopper/flash_fwd_kernel_sm90.h | 3 ++- hopper/flash_fwd_launch_template.h | 10 ++++++---- hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp | 7 +++---- hopper/softmax.h | 13 +++++++------ 4 files changed, 18 insertions(+), 15 deletions(-) diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index a071ee10654..5f9c8334914 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -50,6 +50,7 @@ class FlashAttnFwdSm90 { static constexpr bool SameHeadDim = CollectiveMainloop::SameHeadDim; static constexpr bool LargeHeadDimV = CollectiveMainloop::LargeHeadDimV; static constexpr bool Has_sink = CollectiveMainloop::Has_sink; + static constexpr bool Use_fa4_sink = CollectiveMainloop::Use_fa4_sink; static_assert(CollectiveMainloop::LargeHeadDimV == CollectiveEpilogue::LargeHeadDimV); using SeqlenInfo_t = typename CollectiveMainloop::SeqlenInfo_t; @@ -427,7 +428,7 @@ class FlashAttnFwdSm90 { float const k_descale = params.mainloop.ptr_k_descale == nullptr ? 1.0f : params.mainloop.ptr_k_descale[bidb * get<0>(params.mainloop.stride_k_descale) + bidh_kv * get<1>(params.mainloop.stride_k_descale)]; softmax_scale_log2 *= q_descale * k_descale; } - flash::Softmax softmax(softmax_scale_log2, sink_val, params.mainloop.use_fa4_sink); + flash::Softmax softmax(softmax_scale_log2, sink_val); // Attention output (GEMM-II) accumulator. Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_MNK_PV{})); bool tile_valid; diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index d14091c8b24..079213cb7ee 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -26,7 +26,7 @@ using namespace cute; template + bool PackGQA, bool Split, bool V_colmajor, bool Has_sink, bool Use_fa4_sink> void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { static_assert(!(Is_causal && Is_local), "Causal and Local cannot be enabled at the same time"); static_assert(!(AppendKV && V_colmajor), "AppendKV and V_colmajor cannot be enabled at the same time"); @@ -51,7 +51,7 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { using ClusterShape = cute::Shape, _1, _1>; using CollectiveMainloop = std::conditional_t< Arch >= 90, - flash::CollectiveMainloopFwdSm90, + flash::CollectiveMainloopFwdSm90, flash::CollectiveMainloopFwdSm80 >; using CollectiveEpilogue = flash::CollectiveEpilogueFwd; @@ -127,7 +127,7 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { params.cu_seqlens_q, params.cu_seqlens_k, params.cu_seqlens_knew, params.seqused_q, params.seqused_k, params.leftpad_k, params.seqlens_rotary, - reinterpret_cast(params.learnable_sink_ptr), params.use_fa4_sink + reinterpret_cast(params.learnable_sink_ptr) }; typename CollectiveEpilogue::Arguments epilogue_args { static_cast(params.o_ptr), @@ -213,10 +213,12 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { APPENDKV_SWITCH(params.knew_ptr, AppendKV, [&] { SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { static constexpr bool PackGQA_Sink = PackGQA && !Has_sink; + BOOL_SWITCH(Has_sink && params.use_fa4_sink, Use_fa4_sink, [&] { // Only use Cluster if number of tiles along seqlen_q is even and not varlen CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; - run_flash_fwd(params, stream); + run_flash_fwd(params, stream); + }); }); }); }); diff --git a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp index fea9809eabb..7ce0ccefd3e 100644 --- a/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp +++ b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp @@ -30,7 +30,7 @@ using namespace cute; template + bool MmaPV_is_RS, bool IntraWGOverlap, bool PackGQA_, bool Split_, bool V_colmajor_, bool Has_sink_, bool Use_fa4_sink_=false> struct CollectiveMainloopFwdSm90 { static constexpr int kStages = Stages; @@ -53,6 +53,7 @@ struct CollectiveMainloopFwdSm90 { static constexpr bool Split = Split_; static constexpr bool V_colmajor = V_colmajor_; static constexpr bool Has_sink = Has_sink_; + static constexpr bool Use_fa4_sink = Use_fa4_sink_; static constexpr bool Transpose_V = Is_FP8 && !V_colmajor; static constexpr bool Use_TMA_Q = !PackGQA; static constexpr bool Use_TMA_KV = !PagedKVNonTMA; @@ -398,7 +399,6 @@ struct CollectiveMainloopFwdSm90 { int const* const leftpad_k = nullptr; int const* const seqlens_rotary = nullptr; float const* const ptr_Sink = nullptr; - bool const use_fa4_sink = false; }; // Device side kernel params @@ -457,7 +457,6 @@ struct CollectiveMainloopFwdSm90 { int const* const leftpad_k = nullptr; int const *const seqlens_rotary = nullptr; float const* const ptr_Sink = nullptr; - bool const use_fa4_sink = false; }; static Params @@ -569,7 +568,7 @@ struct CollectiveMainloopFwdSm90 { !Split ? 1 : args.num_splits, args.kv_batch_idx, args.cu_seqlens_q, args.cu_seqlens_k, args.cu_seqlens_k_new, - args.seqused_q, args.seqused_k, args.leftpad_k, args.seqlens_rotary, args.ptr_Sink, args.use_fa4_sink}; + args.seqused_q, args.seqused_k, args.leftpad_k, args.seqlens_rotary, args.ptr_Sink}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance diff --git a/hopper/softmax.h b/hopper/softmax.h index 697fe5f62f7..1e26f644d82 100644 --- a/hopper/softmax.h +++ b/hopper/softmax.h @@ -89,7 +89,7 @@ __forceinline__ __device__ void scale_apply_exp2(Tensor &tenso //////////////////////////////////////////////////////////////////////////////////////////////////// -template +template struct Softmax { using TensorT = decltype(make_tensor(Shape>{})); @@ -97,10 +97,9 @@ struct Softmax { float const softmax_scale_log2; float const softmax_scale; float const sink_val; - bool const use_fa4_sink; - CUTLASS_DEVICE Softmax(float const softmax_scale_log2_, float const sink_val_ = -INFINITY, bool const use_fa4_sink_ = false) - : softmax_scale_log2(softmax_scale_log2_), softmax_scale(softmax_scale_log2_ * float(M_LN2)), sink_val(sink_val_), use_fa4_sink(use_fa4_sink_) {}; + CUTLASS_DEVICE Softmax(float const softmax_scale_log2_, float const sink_val_ = -INFINITY) + : softmax_scale_log2(softmax_scale_log2_), softmax_scale(softmax_scale_log2_ * float(M_LN2)), sink_val(sink_val_) {}; template __forceinline__ __device__ TensorT max_get_scale(Tensor0 &acc_s) { @@ -110,7 +109,7 @@ struct Softmax { TensorT scores_scale; if constexpr (Is_first) { if constexpr (Has_sink) { - if (use_fa4_sink) { + if constexpr (Use_fa4_sink) { flash::template reduce_max(scores, row_max); } else { const float sink_scaled = sink_val / softmax_scale; @@ -158,7 +157,9 @@ struct Softmax { float sum = row_sum(mi); if constexpr (Has_sink) { static constexpr float max_offset = float(Max_offset); - const float max_scaled = row_max(mi) == -INFINITY ? 0.f : (row_max(mi) * softmax_scale_log2) - max_offset; + const float max_scaled = Use_fa4_sink + ? (row_max(mi) == -INFINITY ? 0.f : (row_max(mi) * softmax_scale_log2) - max_offset) + : (row_max(mi) * softmax_scale_log2) - max_offset; sum += exp2f(sink_val * float(M_LOG2E) - max_scaled); } float inv_sum = (sum == 0.f || sum != sum) ? 0.f : 1.f / sum; From 58ecfeffc41f4c7e453a18eb8354e2e851d36c32 Mon Sep 17 00:00:00 2001 From: aoxy Date: Sat, 23 May 2026 18:20:27 +0800 Subject: [PATCH 22/23] fix v1. --- hopper/flash_fwd_launch_template.h | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index 079213cb7ee..e30bde49933 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -213,13 +213,21 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { APPENDKV_SWITCH(params.knew_ptr, AppendKV, [&] { SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { static constexpr bool PackGQA_Sink = PackGQA && !Has_sink; - BOOL_SWITCH(Has_sink && params.use_fa4_sink, Use_fa4_sink, [&] { - // Only use Cluster if number of tiles along seqlen_q is even and not varlen - CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { - static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; - run_flash_fwd(params, stream); - }); - }); + // Only instantiate Use_fa4_sink=true/false when Has_sink=true. + // When Has_sink=false, Use_fa4_sink is always false (no extra kernel). + if constexpr (Has_sink) { + BOOL_SWITCH(params.use_fa4_sink, Use_fa4_sink, [&] { + CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { + static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; + run_flash_fwd(params, stream); + }); + }); + } else { + CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA_Sink ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] { + static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; + run_flash_fwd(params, stream); + }); + } }); }); }); From 37d3734fcc2d5d0dbac889eb1d9c40a519e4fbca Mon Sep 17 00:00:00 2001 From: aoxy Date: Sun, 24 May 2026 12:10:47 +0800 Subject: [PATCH 23/23] Add test. --- hopper/test_flash_attn.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/hopper/test_flash_attn.py b/hopper/test_flash_attn.py index 9f7b75533fc..484c4227a53 100644 --- a/hopper/test_flash_attn.py +++ b/hopper/test_flash_attn.py @@ -104,11 +104,14 @@ ) # @pytest.mark.parametrize('seqlen_q,seqlen_k', [(128, 128)]) @pytest.mark.parametrize("has_learnable_sink", [False] + ([True] if not DISABLE_SINK else [])) +@pytest.mark.parametrize("use_fa4_sink", [False] + ([True] if not DISABLE_SINK else [])) def test_flash_attn_output( - seqlen_q, seqlen_k, d, causal, local, softcap, V_colmajor, deterministic, has_qv, mha_type, dtype, has_learnable_sink + seqlen_q, seqlen_k, d, causal, local, softcap, V_colmajor, deterministic, has_qv, mha_type, dtype, has_learnable_sink, use_fa4_sink ): if V_colmajor and (seqlen_k % 16 != 0 or dtype != torch.float8_e4m3fn): pytest.skip("V_colmajor requires seqlen_k to be a multiple of 16 and dtype to be float8_e4m3fn") + if not has_learnable_sink and use_fa4_sink: + pytest.skip("use_fa4_sink requires has_learnable_sink to be True") device = "cuda" # set seed torch.random.manual_seed(0) @@ -216,7 +219,7 @@ def test_flash_attn_output( pack_gqa=pack_gqa, num_splits=num_splits, learnable_sink=learnable_sink, - use_fa4_sink=True, + use_fa4_sink=use_fa4_sink, ) print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") @@ -356,9 +359,12 @@ def test_flash_attn_output( ], ) @pytest.mark.parametrize("has_learnable_sink", [False] + ([True] if not DISABLE_SINK else [])) +@pytest.mark.parametrize("use_fa4_sink", [False] + ([True] if not DISABLE_SINK else [])) def test_flash_attn_varlen_output( - seqlen_q, seqlen_k, d, add_unused_qkv, causal, local, softcap, deterministic, has_qv, mha_type, dtype, has_learnable_sink + seqlen_q, seqlen_k, d, add_unused_qkv, causal, local, softcap, deterministic, has_qv, mha_type, dtype, has_learnable_sink, use_fa4_sink ): + if not has_learnable_sink and use_fa4_sink: + pytest.skip("use_fa4_sink requires has_learnable_sink to be True") device = "cuda" # set seed torch.random.manual_seed(seqlen_q + seqlen_k + d + int(causal) * 2 + int(local)) @@ -513,7 +519,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): pack_gqa=pack_gqa, num_splits=num_splits, learnable_sink=learnable_sink, - use_fa4_sink=True, + use_fa4_sink=use_fa4_sink, ) out = output_pad_fn(out_unpad) if query_unused_mask is not None: