diff --git a/3rdparty/composable_kernel b/3rdparty/composable_kernel index 1aa93ef551a..f5573f56d9d 160000 --- a/3rdparty/composable_kernel +++ b/3rdparty/composable_kernel @@ -1 +1 @@ -Subproject commit 1aa93ef551a31405aef5c8c14e869241ba96639d +Subproject commit f5573f56d9d4981def16f575ddb14535b93bb9bb diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py index 77f548846f4..c7270d4902f 100644 --- a/aiter/ops/mha.py +++ b/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, 0 means no sink 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, @@ -2446,7 +2459,7 @@ def flash_attn_varlen_func( softmax_scale=None, 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, 0 means no sink bias=None, alibi_slopes=None, deterministic=False, @@ -2756,7 +2769,7 @@ def flash_attn_fp8_pertensor_func( k_descale, v_descale, causal=False, - window_size=(-1, -1), # -1 means infinite context window + window_size=(-1, -1, 0), # -1 means infinite context window, 0 means no sink softmax_scale=None, ): if softmax_scale is 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, diff --git a/csrc/cpp_itfs/mha_fwd_generate.py b/csrc/cpp_itfs/mha_fwd_generate.py index 24dc5b9f709..0889d9d070e 100644 --- a/csrc/cpp_itfs/mha_fwd_generate.py +++ b/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 = false, 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); }""" diff --git a/csrc/include/mha_common.h b/csrc/include/mha_common.h index 211fde00e71..e101008c33c 100644 --- a/csrc/include/mha_common.h +++ b/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); diff --git a/csrc/include/mha_fwd.h b/csrc/include/mha_fwd.h index 84e84e386ad..9f7794638c7 100644 --- a/csrc/include/mha_fwd.h +++ b/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, diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index 615ef0a150b..cfebff343af 100644 --- a/csrc/include/rocm_ops.hpp +++ b/csrc/include/rocm_ops.hpp @@ -743,6 +743,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, \ @@ -903,6 +904,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, \ diff --git a/csrc/include/torch/mha_fwd.h b/csrc/include/torch/mha_fwd.h index 48fdfc2b8f3..94b5cf056db 100644 --- a/csrc/include/torch/mha_fwd.h +++ b/csrc/include/torch/mha_fwd.h @@ -13,6 +13,7 @@ std::vector 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 cu_seqlens_q, diff --git a/csrc/include/torch/mha_varlen_fwd.h b/csrc/include/torch/mha_varlen_fwd.h index 3b7ab8dea8a..cbdabcfc85d 100644 --- a/csrc/include/torch/mha_varlen_fwd.h +++ b/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 out, // [total_q, hq, d] diff --git a/csrc/py_itfs_ck/mha_fwd_kernels.cu b/csrc/py_itfs_ck/mha_fwd_kernels.cu index 1ab440790d6..b56ae67164f 100644 --- a/csrc/py_itfs_ck/mha_fwd_kernels.cu +++ b/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(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 cu_seqlens_q_, @@ -233,7 +235,7 @@ 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) { @@ -241,9 +243,10 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d] } 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 { diff --git a/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu b/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu index 7f18a24bf93..4f47e73f6b8 100644 --- a/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu +++ b/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu @@ -171,6 +171,7 @@ mha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, batch_stride_o, mask.left, mask.right, + mask.sink, static_cast(mask.type), min_seqlen_q, p_dropout, @@ -344,6 +345,7 @@ mha_varlen_fwd( bool is_causal, int window_size_left, int window_size_right, + int sink_size, bool return_softmax_lse, bool return_dropout_randval, std::optional out_, // [total_q, hq, d] @@ -456,7 +458,7 @@ mha_varlen_fwd( 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, max_seqlen_q, max_seqlen_k); // casual } else if (window_size_left == -1 && window_size_right == -1) { @@ -464,10 +466,10 @@ mha_varlen_fwd( } 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, max_seqlen_q, max_seqlen_k); // local } - + bool has_sink = mask.sink > 0; CHECK_SHAPE(q, total_q, num_heads, head_size_q); if (!paged_KV) { const int total_k = k.size(0); @@ -594,7 +596,8 @@ mha_varlen_fwd( true, //is_group_mode mask.type, bias_type, - has_lse); + has_lse, + has_sink); TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd_splitkv"); } else @@ -633,7 +636,6 @@ mha_varlen_fwd( drop_seed_offset, const_cast&>(cu_seqlens_q_padded_), const_cast&>(cu_seqlens_k_padded_)); - float t = aiter::mha_fwd(args, stream_config, dtype_str, @@ -643,6 +645,7 @@ mha_varlen_fwd( has_lse, qscale_type, false, // use_ext_asm + has_sink, 1); // how_v3_bf16_cvt TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd"); } diff --git a/csrc/py_itfs_cu/asm_mha_fwd.cu b/csrc/py_itfs_cu/asm_mha_fwd.cu index 5d67e7d64e1..d4c901fd2a7 100644 --- a/csrc/py_itfs_cu/asm_mha_fwd.cu +++ b/csrc/py_itfs_cu/asm_mha_fwd.cu @@ -133,6 +133,7 @@ mha_fwd_args get_asm_fmha_fwd_args(bool has_lse, batch_stride_o, mask.left, mask.right, + 0, // sink_size static_cast(mask.type), 0, // min_seqlen_q p_dropout, diff --git a/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu b/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu index 8447c40e017..29dcbc285c7 100644 --- a/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu +++ b/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu @@ -142,6 +142,7 @@ mha_fwd_args get_asm_mha_varlen_fwd_args(bool has_lse, batch_stride_o, mask.left, mask.right, + 0, // sink_size static_cast(mask.type), min_seqlen_q, p_dropout,