Add sink_size parameter for mha_fwd and varlen_mha_fwd api

aiter/ops/mha.py

@@ -21,6 +21,7 @@ def cmdGenFunc_mha_fwd(
     is_causal: bool,
     window_size_left: int,
     window_size_right: int,
+    sink_size: int,
     return_softmax_lse: bool,
     return_dropout_randval: bool,
     cu_seqlens_q: Optional[torch.Tensor] = None,
@@ -164,6 +165,7 @@ def gen_mha_fwd_fake_tensors(
     is_causal: bool,
     window_size_left: int,
     window_size_right: int,
+    sink_size: int,
     return_softmax_lse: bool,
     return_dropout_randval: bool,
     cu_seqlens_q: Optional[torch.Tensor] = None,
@@ -196,6 +198,7 @@ def mha_fwd(
     is_causal: bool,
     window_size_left: int,
     window_size_right: int,
+    sink_size: int,
     return_softmax_lse: bool,
     return_dropout_randval: bool,
     cu_seqlens_q: Optional[torch.Tensor] = None,
@@ -270,6 +273,7 @@ def cmdGenFunc_mha_varlen_fwd(
     is_causal: bool,
     window_size_left: int,
     window_size_right: int,
+    sink_size: int,
     return_softmax_lse: bool,
     return_dropout_randval: bool,
     out: Optional[torch.Tensor] = None,
@@ -435,6 +439,7 @@ def gen_mha_varlen_fwd_fake_tensor(
     is_causal: bool,
     window_size_left: int,
     window_size_right: int,
+    sink_size: int,
     return_softmax_lse: bool,
     return_dropout_randval: bool,
     out: Optional[torch.Tensor] = None,
@@ -502,6 +507,7 @@ def mha_varlen_fwd(
     is_causal: bool,
     window_size_left: int,
     window_size_right: int,
+    sink_size: int,
     return_softmax_lse: bool,
     return_dropout_randval: bool,
     out: Optional[torch.Tensor] = None,
@@ -1211,6 +1217,7 @@ def _flash_attn_forward(
     causal: bool,
     window_size_left: int,
     window_size_right: int,
+    sink_size: int,
     bias: Optional[torch.Tensor],
     alibi_slopes: Optional[torch.Tensor],
     q_descale: Optional[torch.Tensor],
@@ -1288,6 +1295,7 @@ def _validate_cu(name: str, x: Optional[torch.Tensor]):
             causal,
             window_size_left,
             window_size_right,
+            sink_size,
             return_lse,
             return_softmax,
             cu_seqlens_q,
@@ -1706,6 +1714,7 @@ def forward(
             causal=causal,
             window_size_left=int(window_size[0]),
             window_size_right=int(window_size[1]),
+            sink_size=int(window_size[2]) if len(window_size) == 3 else 0,
             bias=bias,
             alibi_slopes=alibi_slopes,
             q_descale=None,
@@ -1824,7 +1833,7 @@ def flash_attn_func(
     dropout_p=0.0,
     softmax_scale=None,
     causal=False,
-    window_size=(-1, -1),  # -1 means infinite context window
+    window_size=(-1, -1, 0),  # -1 means infinite context window
     bias=None,
     alibi_slopes=None,
     deterministic=True,
@@ -1923,6 +1932,7 @@ def _flash_attn_varlen_forward(
     logits_soft_cap: float = 0.0,
     window_size_left: int = -1,
     window_size_right: int = -1,
+    sink_size: int = 0,
     bias: Optional[torch.Tensor] = None,
     alibi_slopes: Optional[torch.Tensor] = None,
     q_descale: Optional[torch.Tensor] = None,
@@ -1943,6 +1953,7 @@ def _flash_attn_varlen_forward(
     # mask
     window_size_left = -1 if window_size_left >= max_seqlen_k else window_size_left
     window_size_right = -1 if window_size_right >= max_seqlen_k else window_size_right
+    sink_size = 0 if sink_size >= max_seqlen_k else sink_size
     mask = causal == True and window_size_left == -1  # causal mask
     nmask = (
         causal == False and window_size_left == -1 and window_size_right == -1
@@ -2026,6 +2037,7 @@ def _validate(name: str, t: torch.Tensor):
             causal,
             window_size_left,
             window_size_right,
+            sink_size,
             return_lse,
             return_softmax,
             out=out,
@@ -2293,6 +2305,7 @@ def forward(
             logits_soft_cap=logits_soft_cap,
             window_size_left=window_size[0],
             window_size_right=window_size[1],
+            sink_size=window_size[2] if len(window_size) > 2 else 0,
             bias=bias,
             alibi_slopes=alibi_slopes,
             q_descale=None,
@@ -2777,6 +2790,7 @@ def flash_attn_fp8_pertensor_func(
         causal=causal,
         window_size_left=int(window_size[0]),
         window_size_right=int(window_size[1]),
+        sink_size=int(window_size[2]) if len(window_size) == 3 else 0,
         bias=None,
         alibi_slopes=None,
         q_descale=q_descale,
@@ -2803,7 +2817,7 @@ def flash_attn_varlen_fp8_pertensor_func(
     min_seqlen_q=0,
     logits_soft_cap=0.0,
     causal=False,
-    window_size=(-1, -1),  # -1 means infinite context window
+    window_size=(-1, -1, 0),  # -1 means infinite context window
     softmax_scale=None,
 ):
     if softmax_scale is None:
@@ -2832,6 +2846,7 @@ def flash_attn_varlen_fp8_pertensor_func(
         logits_soft_cap=logits_soft_cap,
         window_size_left=int(window_size[0]),
         window_size_right=int(window_size[1]),
+        sink_size=int(window_size[2]) if len(window_size) == 3 else 0,
         bias=None,
         alibi_slopes=None,
         q_descale=q_descale,

csrc/cpp_itfs/mha_fwd_generate.py

@@ -39,6 +39,7 @@
                                   bool has_dropout,
                                   quant_scale_enum qscale_type,
                                   bool use_ext_asm,
+                                  bool has_sink,
                                   int how_v3_bf16_cvt = 1,
                                   bool skip_min_seqlen_q = false)
 {{
@@ -54,7 +55,8 @@
                           qscale_type,
                           use_ext_asm,
                           how_v3_bf16_cvt,
-                          skip_min_seqlen_q);
+                          skip_min_seqlen_q,
+                          has_sink);
 }}
 
 mha_fwd_splitkv_traits get_mha_fwd_splitkv_traits(int head_size_q,
@@ -64,7 +66,8 @@
                                                   bool has_logits_soft_cap,
                                                   mask_enum mask_type,
                                                   bias_enum bias_type,
-                                                  bool has_lse)
+                                                  bool has_lse,
+                                                  bool has_sink)
 {{
     return mha_fwd_splitkv_traits(head_size_q,
                                   head_size_v,
@@ -73,7 +76,8 @@
                                   has_logits_soft_cap,
                                   mask_type,
                                   bias_type,
-                                  has_lse);
+                                  has_lse,
+                                  has_sink);
 }}
 {F_dispatch}
 
@@ -91,6 +95,7 @@
               bool has_lse,
               quant_scale_enum qscale_type,
               bool use_ext_asm,
+              bool has_sink,
               int how_v3_bf16_cvt,
               const void* seqstart_q_padding_ptr,
               const void* seqstart_k_padding_ptr,
@@ -110,6 +115,7 @@
                                      has_dropout,
                                      qscale_type,
                                      use_ext_asm,
+                                     has_sink,
                                      how_v3_bf16_cvt,
                                      args.min_seqlen_q != 0);
     float t = -1;
@@ -124,7 +130,8 @@
                       bool is_group_mode,
                       mask_enum mask_type,
                       bias_enum bias_type,
-                      bool has_lse)
+                      bool has_lse,
+                      bool has_sink)
 {
     int head_size_q = args.hdim_q;
     int head_size_v = args.hdim_v;
@@ -135,7 +142,8 @@
                                              args.logits_soft_cap > 0.f,
                                              mask_type,
                                              bias_type,
-                                             has_lse);
+                                             has_lse,
+                                             has_sink);
     return fmha_fwd_splitkv(traits, args, stream_config);
 }"""
 
@@ -161,8 +169,10 @@
                                      bias_type,
                                      has_lse,
                                      has_dropout,
+                                     false,
                                      quant_scale_enum::no_scale,
-                                     use_ext_asm);
+                                     use_ext_asm,
+                                     has_sink);
     return fmha_batch_prefill(traits, args, stream_config);
 }"""
 

csrc/include/mha_common.h

@@ -136,6 +136,7 @@ inline void print_fmha_fwd_args(ARG args)
     printf("batch_stride_o = %d\n", args.batch_stride_o);
     printf("window_size_left = %d\n", args.window_size_left);
     printf("window_size_right = %d\n", args.window_size_right);
+    printf("sink_size = %d\n", args.sink_size);
     printf("mask_type = %d\n", args.mask_type);
     printf("p_drop = %f\n", args.p_drop);
     printf("s_randval = %d\n", args.s_randval);

csrc/include/mha_fwd.h

@@ -23,7 +23,8 @@ struct mha_fwd_traits : public fmha_fwd_traits
                    quant_scale_enum qscale_type,
                    bool use_ext_asm,
                    int how_v3_bf16_cvt,
-                   bool skip_min_seqlen_q)
+                   bool skip_min_seqlen_q,
+                   bool has_sink)
         : fmha_fwd_traits{head_size_q,
                           head_size_v,
                           dtype,
@@ -35,7 +36,8 @@ struct mha_fwd_traits : public fmha_fwd_traits
                           has_lse,
                           has_dropout,
                           qscale_type,
-                          skip_min_seqlen_q},
+                          skip_min_seqlen_q,
+                          has_sink},
           use_ext_asm(use_ext_asm),
           how_v3_bf16_cvt(how_v3_bf16_cvt)
     {
@@ -53,7 +55,8 @@ struct mha_fwd_splitkv_traits : public fmha_fwd_splitkv_traits
                            bool has_logits_soft_cap,
                            mask_enum mask_type,
                            bias_enum bias_type,
-                           bool has_lse)
+                           bool has_lse,
+                           bool has_sink)
         : fmha_fwd_splitkv_traits{head_size_q,
                                   head_size_v,
                                   dtype,
@@ -63,7 +66,8 @@ struct mha_fwd_splitkv_traits : public fmha_fwd_splitkv_traits
                                   mask_type,
                                   bias_type,
                                   has_lse,
-                                  false} // do_fp8_static_quant
+                                  false, // do_fp8_static_quant
+                                  has_sink} 
     {
     }
 };
@@ -81,6 +85,7 @@ __attribute__((visibility("default"))) float mha_fwd(mha_fwd_args args,
                                                      bool has_lse,
                                                      quant_scale_enum qscale_type,
                                                      bool use_ext_asm,
+                                                     bool has_sink = false,
                                                      int how_v3_bf16_cvt                = 1,
                                                      const void* seqstart_q_padding_ptr = nullptr,
                                                      const void* seqstart_k_padding_ptr = nullptr,
@@ -93,7 +98,8 @@ mha_fwd_splitkv(mha_fwd_splitkv_args args,
                 bool is_group_mode,
                 mask_enum mask_type,
                 bias_enum bias_type,
-                bool has_lse);
+                bool has_lse,
+                bool has_sink = false);
 
 __attribute__((visibility("default"))) float
 mha_batch_prefill(mha_batch_prefill_args args,

csrc/include/rocm_ops.hpp

@@ -709,6 +709,7 @@ namespace py = pybind11;
           py::arg("is_causal"),                    \
           py::arg("window_size_left"),             \
           py::arg("window_size_right"),            \
+          py::arg("sink_size"),                    \
           py::arg("return_softmax_lse"),           \
           py::arg("return_dropout_randval"),       \
           py::arg("cu_seqlens_q")  = std::nullopt, \
@@ -869,6 +870,7 @@ namespace py = pybind11;
           py::arg("is_causal"),                          \
           py::arg("window_size_left"),                   \
           py::arg("window_size_right"),                  \
+          py::arg("sink_size"),                          \
           py::arg("return_softmax_lse"),                 \
           py::arg("return_dropout_randval"),             \
           py::arg("out")                 = std::nullopt, \

csrc/include/torch/mha_fwd.h

@@ -13,6 +13,7 @@ std::vector<at::Tensor> mha_fwd(at::Tensor& q,       // [b, sq, hq, d]
                                 bool is_causal,
                                 int window_size_left,
                                 int window_size_right,
+                                int sink_size,
                                 bool return_softmax_lse,
                                 bool return_dropout_randval,
                                 std::optional<at::Tensor> cu_seqlens_q,

csrc/include/torch/mha_varlen_fwd.h

@@ -21,6 +21,7 @@ mha_varlen_fwd(at::Tensor& q,                                 // [total_q, hq, d
                bool is_causal,
                int window_size_left,
                int window_size_right,
+               int sink_size,
                bool return_softmax_lse,
                bool return_dropout_randval,
                std::optional<at::Tensor> out,                // [total_q, hq, d]

csrc/py_itfs_ck/mha_fwd_kernels.cu

@@ -145,6 +145,7 @@ mha_fwd_args get_ck_fmha_fwd_args(bool has_lse,
                         batch_stride_o,
                         mask.left,
                         mask.right,
+                        mask.sink,
                         static_cast<ck_tile::index_t>(mask.type),
                         0, // min_seqlen_q
                         p_dropout,
@@ -161,6 +162,7 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d]
         bool is_causal,
         int window_size_left,
         int window_size_right,
+        int sink_size,
         bool return_softmax_lse,
         bool return_dropout_randval,
         std::optional<at::Tensor> cu_seqlens_q_,
@@ -233,17 +235,18 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d]
     if (is_causal) {
         // Causal is the special case where window_size_right == 0 and window_size_left < 0.
         window_size_right = 0;
-        std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + "0";
+        std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + "0" + "," + std::to_string(sink_size);
         mask = mask_info::decode(mask_identify, seqlen_q, seqlen_k); // casual
     }
     else if (window_size_left == -1 && window_size_right == -1) {
         mask = mask_info::decode("0", seqlen_q, seqlen_k); // no mask
     }
     else {
         // Local is the more general case where window_size_right >= 0 or window_size_left >= 0.
-        std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + std::to_string(window_size_right);
+        std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + std::to_string(window_size_right) + "," + std::to_string(sink_size);
         mask = mask_info::decode(mask_identify, seqlen_q, seqlen_k); // local
     }
+    bool has_sink = mask.sink > 0;
 
     TORCH_CHECK(!(bias_.has_value() && alibi_slopes_.has_value()), "cannot apply bias and alibi at the same time");
     bias_enum bias_type = bias_.has_value() ? bias_enum::elementwise_bias :
@@ -362,7 +365,8 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d]
                                  bias_type,
                                  has_lse,
                                  qscale_type,
-                                 false);
+                                 false,
+                                 has_sink);
         TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd");
     }
     else {