[Hopper] Support sink attention (learnable sink)

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)

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

hopper/epilogue_bwd.hpp

@@ -19,7 +19,7 @@ namespace flash {
 using namespace cute;
 
 template <class TileShape_MNK_, class Element_, class ArchTag_,
-          int NumEpilogueThreads_, bool Varlen_, bool dKV_swapAB_, int AtomLayoutKdKV=1>
+          int NumEpilogueThreads_, bool Varlen_, bool dKV_swapAB_, int AtomLayoutKdKV=1, bool Has_sink=false>
 struct CollectiveEpilogueBwd {
 
     using TileShape_MNK = TileShape_MNK_;
@@ -115,6 +115,7 @@ struct CollectiveEpilogueBwd {
         int* dv_semaphore;
         int const* cu_seqlens;
         int const* seqused;
+        float* dsink_ptr;
     };
 
     // Device side kernel params
@@ -128,6 +129,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
@@ -149,7 +151,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
@@ -169,7 +171,8 @@ struct CollectiveEpilogueBwd {
           SharedStorage& shared_storage,
           TiledMma tiled_mma,
           int thread_idx,
-          cute::tuple<int32_t, int32_t, int32_t> const& block_coord
+          cute::tuple<int32_t, int32_t, int32_t> const& block_coord,
+          float dsink_val=0.0f
           ) {
 
         auto [n_block, bidh, bidb] = block_coord;
@@ -268,6 +271,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<float> sum_op;
+            dsink_val = Allreduce<cutlass::NumThreadsPerWarp>::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
@@ -319,7 +330,7 @@ struct CollectiveEpilogueBwd {
 };
 
 template <class TileShape_MNK_, class ElementAccum, class ArchTag_,
-          int NumEpilogueThreads_, bool Varlen_, bool Deterministic>
+          int NumEpilogueThreads_, bool Varlen_, bool Deterministic, bool Has_sink=false>
 struct CollectiveEpilogueBwdGQA {
 
     using TileShape_MNK = TileShape_MNK_;
@@ -376,6 +387,7 @@ struct CollectiveEpilogueBwdGQA {
         int* dv_semaphore;
         int const* cu_seqlens;
         int const* seqused;
+        float* dsink_ptr;
     };
 
     // Device side kernel params
@@ -392,6 +404,7 @@ struct CollectiveEpilogueBwdGQA {
         int const num_batch;
         int const* cu_seqlens = nullptr;
         int const* seqused = nullptr;
+        float* dsink_ptr;
     };
 
     static Params
@@ -403,7 +416,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.num_batch, args.cu_seqlens, args.seqused};
+                args.num_batch, args.cu_seqlens, args.seqused, args.dsink_ptr};
     }
 
     /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
@@ -419,7 +432,8 @@ struct CollectiveEpilogueBwdGQA {
           SharedStorage& shared_storage,
           TiledMma tiled_mma,
           int thread_idx,
-          cute::tuple<int32_t, int32_t, int32_t> const& block_coord
+          cute::tuple<int32_t, int32_t, int32_t> const& block_coord,
+          float dsink_val=0.0f
           ) {
 
         auto [n_block, bidh, bidb] = block_coord;
@@ -510,6 +524,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<float> sum_op;
+            dsink_val = Allreduce<cutlass::NumThreadsPerWarp>::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);
         }

hopper/epilogue_fwd.hpp

@@ -411,7 +411,8 @@ struct CollectiveEpilogueFwd {
     store_zero(
          Params const& params,
          int thread_idx,
-         cute::tuple<int32_t, int32_t, int32_t, int32_t> const& block_coord
+         cute::tuple<int32_t, int32_t, int32_t, int32_t> 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;
@@ -438,13 +439,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;
                 }
             }
         }

hopper/flash.h

@@ -165,6 +165,7 @@ struct Flash_fwd_params : public Qkv_params {
 
     int arch;
     int num_sm;
+    void *__restrict__ learnable_sink_ptr;
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -211,6 +212,7 @@ struct Flash_bwd_params : public Flash_fwd_params {
 
     bool deterministic;
     index_t dq_accum_split_stride;
+    void *__restrict__ dsink_ptr;
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////

hopper/flash_api.cpp

@@ -704,7 +704,8 @@ mha_fwd(at::Tensor q,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seql
         std::optional<at::Tensor> scheduler_metadata_,  // (b + 1)
         int64_t num_splits,
         std::optional<bool> pack_gqa_,
-        int64_t sm_margin
+        int64_t sm_margin,
+        std::optional<const at::Tensor> learnable_sink_
         ) {
 
     auto dprops = at::cuda::getCurrentDeviceProperties();
@@ -1089,6 +1090,18 @@ mha_fwd(at::Tensor q,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seql
         params.kv_batch_idx = reinterpret_cast<int *>(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
+    } else {
+        params.learnable_sink_ptr = nullptr;
+    }
+
     at::Tensor out_accum, softmax_lse_accum;
     auto outaccum_type = at::ScalarType::Float;
     if (params.num_splits > 1) {
@@ -1286,7 +1299,9 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_bwd(
     int64_t window_size_right,
     double softcap,
     bool deterministic,
-    int64_t sm_margin
+    int64_t sm_margin,
+    std::optional<const at::Tensor> learnable_sink_,
+    std::optional<at::Tensor> dsink_
 ) {
 
     #ifdef FLASHATTENTION_DISABLE_BACKWARD
@@ -1473,7 +1488,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> 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));
@@ -1530,6 +1545,28 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> 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<int>();
     // Will be zero'ed out in the backward preprocess kernel
@@ -1564,6 +1601,13 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_bwd(
         softmax_d.zero_();
     }
 
+    if (learnable_sink_.has_value()) {
+        if (dsink_.has_value() && dsink_.value().dtype() != torch::kFloat32) {
+            dsink = dsink.to(dsink_.value().dtype());
+            dsink_.value().copy_(dsink);
+        }
+    }
+
     return { softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum };
 }