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.py b/hopper/benchmark_attn.py index 33e5d282716..9b96e071b3b 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 @@ -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 = {} @@ -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=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 @@ -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/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/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 6d9b5f4f596..1bce5e07ff4 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,14 @@ struct CollectiveEpilogueBwd { gmem_tiled_copy_dKV, tdKVrdK, tdKVgdK, tdKVcdKV, tdKVpdK, std::min(seqlen_info.seqlen - n_block * kBlockN, kBlockN) ); } + + 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); + } + } } CUTLASS_DEVICE void @@ -318,7 +329,7 @@ struct CollectiveEpilogueBwd { }; template + int NumEpilogueThreads_, bool Varlen_, bool Deterministic, bool Has_sink=false> struct CollectiveEpilogueBwdGQA { using TileShape_MNK = TileShape_MNK_; @@ -374,6 +385,7 @@ struct CollectiveEpilogueBwdGQA { int* dv_semaphore; int const* cu_seqlens; int const* seqused; + float* dsink_ptr; }; // Device side kernel params @@ -389,6 +401,7 @@ struct CollectiveEpilogueBwdGQA { int* dv_semaphore; int const* cu_seqlens = nullptr; int const* seqused = nullptr; + float* dsink_ptr; }; static Params @@ -400,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.cu_seqlens, args.seqused, args.dsink_ptr}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -416,7 +429,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 +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::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); } 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..a4703bfb14d 100644 --- a/hopper/flash.h +++ b/hopper/flash.h @@ -159,6 +159,8 @@ struct Flash_fwd_params : public Qkv_params { int arch; int num_sm; + void *__restrict__ learnable_sink_ptr; + bool use_fa4_sink; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -205,6 +207,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..468894b9b66 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -680,7 +680,9 @@ 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_, + bool use_fa4_sink ) { auto dprops = at::cuda::getCurrentDeviceProperties(); @@ -1045,6 +1047,19 @@ 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()) { + 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); + 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; + } + at::Tensor out_accum, softmax_lse_accum; auto outaccum_type = at::ScalarType::Float; if (params.num_splits > 1) { @@ -1233,7 +1248,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 +1431,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 +1488,28 @@ 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()) { + 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() && dsink_.value().dtype() == torch::kFloat32) { + dsink = dsink_.value(); + 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 +1543,16 @@ std::vector mha_bwd( softmax_d.zero_(); } - return { dq, dk, dv, softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum }; + 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 }; } std::vector diff --git a/hopper/flash_attn_interface.py b/hopper/flash_attn_interface.py index 06782fa409b..9c46e5739bb 100644 --- a/hopper/flash_attn_interface.py +++ b/hopper/flash_attn_interface.py @@ -49,7 +49,9 @@ def _flash_attn_forward( scheduler_metadata=None, num_splits=1, pack_gqa=None, - sm_margin=0): + sm_margin=0, + 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 = [ @@ -96,6 +98,8 @@ def _flash_attn_forward( num_splits, pack_gqa, sm_margin, + learnable_sink, + use_fa4_sink, ) return out, softmax_lse, *rest @@ -122,10 +126,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 +154,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 +174,8 @@ def forward( deterministic=False, 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) @@ -197,9 +207,11 @@ 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) + 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 +225,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 +237,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 +257,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, None, None, None, None, None, None, None, None, None, None, None, dsink, None class FlashAttnFunc(torch.autograd.Function): @@ -268,6 +283,8 @@ def forward( pack_gqa=None, 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) @@ -293,9 +310,11 @@ def forward( num_splits=num_splits, 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) + 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 +326,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 +349,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, None, None, None, None, None, None, None, None, None, None, None, None, None, dsink, None class FlashAttnVarlenFunc(torch.autograd.Function): @@ -361,6 +383,8 @@ def forward( pack_gqa=None, 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) @@ -390,9 +414,11 @@ def forward( num_splits=num_splits, 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) + 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 +432,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 +458,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, 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( @@ -449,6 +478,8 @@ def flash_attn_qkvpacked_func( deterministic=False, 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 @@ -495,6 +526,8 @@ def flash_attn_qkvpacked_func( deterministic, num_heads_q, sm_margin, + learnable_sink, + use_fa4_sink, ) @@ -513,6 +546,8 @@ def flash_attn_func( pack_gqa=None, 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 @@ -574,6 +609,8 @@ def flash_attn_func( pack_gqa, deterministic, sm_margin, + learnable_sink, + use_fa4_sink, ) @@ -598,6 +635,8 @@ def flash_attn_varlen_func( pack_gqa=None, deterministic=False, sm_margin=0, + learnable_sink=None, + use_fa4_sink=False, ): return FlashAttnVarlenFunc.apply( q, @@ -620,6 +659,8 @@ def flash_attn_varlen_func( pack_gqa, deterministic, sm_margin, + learnable_sink, + use_fa4_sink, ) @@ -658,6 +699,8 @@ 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, + 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 @@ -785,6 +828,8 @@ def flash_attn_with_kvcache( num_splits=num_splits, 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_bwd_kernel_sm90.h b/hopper/flash_bwd_kernel_sm90.h index b93a0219161..613715dad2a 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,7 @@ class FlashAttnBwdSm90 { // static constexpr uint32_t LoadRegisterRequirement = 40; // static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 232 : 152; + // Kernel level shared memory storage struct SharedStorage { struct TensorStorage : cute::aligned_struct<128> { @@ -198,6 +200,7 @@ class FlashAttnBwdSm90 { CollectiveMainloop mainloop; CollectiveEpilogue epilogue; + // 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(); @@ -261,16 +264,16 @@ 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); } - } epilogue.store_tail(); } diff --git a/hopper/flash_bwd_launch_template.h b/hopper/flash_bwd_launch_template.h index b6e8810b25f..ec0198b0468 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,10 +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_combine_kernel.h b/hopper/flash_fwd_combine_kernel.h index a22e05969d9..c6acfa3d6fa 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_cur); + 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..728d92002a3 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, [&] { - 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; + 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 (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 47b3817cd28..5f9c8334914 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -49,6 +49,8 @@ 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 constexpr bool Use_fa4_sink = CollectiveMainloop::Use_fa4_sink; static_assert(CollectiveMainloop::LargeHeadDimV == CollectiveEpilogue::LargeHeadDimV); using SeqlenInfo_t = typename CollectiveMainloop::SeqlenInfo_t; @@ -113,7 +115,6 @@ class FlashAttnFwdSm90 { alignas(16) typename CollectiveMainloop::MainloopPipelineKVNew::SharedStorage pipeline_v_new; alignas(16) typename TileScheduler::SharedStorage smem_scheduler; } pipelines; - }; static constexpr int SharedStorageSize = sizeof(SharedStorage); @@ -406,14 +407,28 @@ 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 = -INFINITY; + if constexpr (Has_sink) { + 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]; + } + } + } else { + sink_val = 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); + 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 +460,7 @@ 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); + 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..e30bde49933 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; @@ -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), @@ -210,10 +211,23 @@ 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, [&] { - // 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); + SINK_SWITCH(params.learnable_sink_ptr != nullptr, Has_sink, [&] { + static constexpr bool PackGQA_Sink = PackGQA && !Has_sink; + // 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); + }); + } }); }); }); diff --git a/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_bwd_sm90_tma_gmma_ws.hpp index ec34e20eca1..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; @@ -320,6 +321,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 +352,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 +409,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 +686,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 +794,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 +882,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/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp b/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp index ba699f17105..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; @@ -52,6 +52,8 @@ 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 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; @@ -396,6 +398,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 +456,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 +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.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/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/softmax.h b/hopper/softmax.h index 8fcdb6bd070..1e26f644d82 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,18 @@ 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) { + if constexpr (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); + } cute::fill(scores_scale, 1.f); } else { Tensor scores_max_prev = make_fragment_like(row_max); @@ -141,6 +155,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 = 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; 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 +169,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_flash_attn.py b/hopper/test_flash_attn.py index 80d4dc0c15c..484c4227a53 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 @@ -37,6 +38,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,11 +103,15 @@ ], ) # @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 + 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) @@ -133,6 +139,12 @@ 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: + 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 # 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 +167,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 +185,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 +217,9 @@ 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, + 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()}") @@ -243,7 +259,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 +272,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 +293,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,9 +358,13 @@ def test_flash_attn_output( (2048, 2048), ], ) +@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 + 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)) @@ -360,6 +392,12 @@ 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: + 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 # 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 +464,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 +482,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 +516,10 @@ 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, + use_fa4_sink=use_fa4_sink, ) out = output_pad_fn(out_unpad) if query_unused_mask is not None: @@ -523,7 +567,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 +593,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 +615,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 +678,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 +696,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 +730,12 @@ 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: + 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 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 +879,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 +894,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 +958,9 @@ 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, + use_fa4_sink=True, ) if varlen_q: out = output_pad_fn(out) 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..292716544b3 --- /dev/null +++ b/hopper/test_flash_attn_varlen_sink_dump_load.py @@ -0,0 +1,424 @@ +#!/root/.envs/swa/bin/python +import argparse +from pathlib import Path +from typing import Any + +import torch + +from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_fa3 +from test_util import generate_qkv + + +from sgl_kernel.flash_attn import flash_attn_varlen_func as flash_attn_varlen_func_sgl + + +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 _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 + + 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, + ) + + result = { + "out": out.detach().cpu(), + "lse": lse.detach().cpu(), + } + _log(f"FA3 run done: use_fa4_sink={effective_use_fa4_sink}, enable_sink={enable_sink}") + return result + + +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 + + 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=tuple(case["window_size"]), + softcap=case["softcap"], + # deterministic=case["deterministic"], + num_splits=case["num_splits"], + pack_gqa=case["pack_gqa"], + return_softmax_lse=True, + sinks=sink, + ) + + result = { + "out": out.detach().cpu(), + "lse": lse.detach().cpu(), + } + _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 == "sgl": + return _run_case_sgl(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), + } + 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_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) + 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.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) + 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.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("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): + dump_path = tmp_path / "flash_attn_varlen_sink_case.pt" + dump_and_check(dump_path, backend="fa3", 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): + 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_sgl"], + default="dump_and_check", + help="dump only, check only, dump then check, or compare FA3 vs sgl directly", + ) + parser.add_argument( + "--backend", + choices=["fa3", "sgl"], + default="fa3", + help="attention backend to run (used by dump/check/dump_and_check)", + ) + parser.add_argument( + "--check-backend", + choices=["fa3", "sgl"], + 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_sgl(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 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 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 new file mode 100755 index 00000000000..1014f90a96e --- /dev/null +++ b/reinstall.sh @@ -0,0 +1,23 @@ +#!/bin/bash +set -e + +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" + +export FLASH_ATTENTION_FORCE_BUILD="TRUE" + +# export FLASH_ATTENTION_DISABLE_HDIM64="TRUE" +# export FLASH_ATTENTION_DISABLE_HDIM96="TRUE" +# export FLASH_ATTENTION_DISABLE_HDIM128="TRUE" +# export FLASH_ATTENTION_DISABLE_HDIM192="TRUE" + +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/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) diff --git a/warp_reinstall.sh b/warp_reinstall.sh new file mode 100755 index 00000000000..14bce9fa855 --- /dev/null +++ b/warp_reinstall.sh @@ -0,0 +1,7 @@ +time ./reinstall.sh > my_install.log 2>&1 +grep -i "error" my_install.log +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