From 98d39b86d124238c2ca2d3aae0f215c1433676ce Mon Sep 17 00:00:00 2001 From: Julian Torres Date: Mon, 22 Jun 2026 01:58:38 -0700 Subject: [PATCH] fix: use top-left mask alignment for causal+sliding window attention The CK wrapper uses bottom-right alignment ('b:' prefix) for causal and local attention masks. This produces correct coordinates during prefill (seqlen_q == seqlen_k) but breaks during autoregressive decode (seqlen_q=1, seqlen_k>1). With bottom-right alignment and a sliding window (e.g., left=128): y_tmp = seqlen_q - seqlen_k = 1 - 256 = -255 y = 1 + 128 + (-255) = -126 x_start = -(-126) + i_y + 1 = 127 // for single query token This masks the first ~127 KV positions from every decode step, producing degenerate repetitive output with models that use sliding window attention (e.g., GPT-OSS 20B with window=128 alternating every other layer). Fixes ROCm/flash-attention#158 --- csrc/flash_attn_ck/mha_bwd.cpp | 4 ++-- csrc/flash_attn_ck/mha_fwd.cpp | 4 ++-- csrc/flash_attn_ck/mha_varlen_bwd.cpp | 6 +++--- csrc/flash_attn_ck/mha_varlen_fwd.cpp | 6 +++--- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/csrc/flash_attn_ck/mha_bwd.cpp b/csrc/flash_attn_ck/mha_bwd.cpp index e038d504a25..805ceb922a3 100644 --- a/csrc/flash_attn_ck/mha_bwd.cpp +++ b/csrc/flash_attn_ck/mha_bwd.cpp @@ -279,7 +279,7 @@ mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num mask_info mask; if (is_causal) { - std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + "0"; + std::string mask_identify = "t:" + std::to_string(window_size_left) + "," + "0"; mask = mask_info::decode(mask_identify, seqlen_q, seqlen_k); // casual } else if (window_size_left == -1 && window_size_right == -1) { @@ -287,7 +287,7 @@ mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num } 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 = "t:" + std::to_string(window_size_left) + "," + std::to_string(window_size_right); mask = mask_info::decode(mask_identify, seqlen_q, seqlen_k); // local } diff --git a/csrc/flash_attn_ck/mha_fwd.cpp b/csrc/flash_attn_ck/mha_fwd.cpp index 57ffd57dd9d..47f4ad29fae 100644 --- a/csrc/flash_attn_ck/mha_fwd.cpp +++ b/csrc/flash_attn_ck/mha_fwd.cpp @@ -219,7 +219,7 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num 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 = "t:" + std::to_string(window_size_left) + "," + "0"; mask = mask_info::decode(mask_identify, seqlen_q, seqlen_k); // casual } else if (window_size_left == -1 && window_size_right == -1) { @@ -227,7 +227,7 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num } 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 = "t:" + std::to_string(window_size_left) + "," + std::to_string(window_size_right); mask = mask_info::decode(mask_identify, seqlen_q, seqlen_k); // local } diff --git a/csrc/flash_attn_ck/mha_varlen_bwd.cpp b/csrc/flash_attn_ck/mha_varlen_bwd.cpp index f0a1298f18e..1bb0c455353 100644 --- a/csrc/flash_attn_ck/mha_varlen_bwd.cpp +++ b/csrc/flash_attn_ck/mha_varlen_bwd.cpp @@ -292,15 +292,15 @@ mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads mask_info mask; if (is_causal) { - std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + "0"; - mask = mask_info::decode(mask_identify, max_seqlen_q, max_seqlen_k); // casual + std::string mask_identify = "t:" + std::to_string(window_size_left) + "," + "0"; + mask = mask_info::decode(mask_identify, max_seqlen_q, max_seqlen_k); // causal } else if (window_size_left == -1 && window_size_right == -1) { mask = mask_info::decode("0", max_seqlen_q, max_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 = "t:" + std::to_string(window_size_left) + "," + std::to_string(window_size_right); mask = mask_info::decode(mask_identify, max_seqlen_q, max_seqlen_k); // local } diff --git a/csrc/flash_attn_ck/mha_varlen_fwd.cpp b/csrc/flash_attn_ck/mha_varlen_fwd.cpp index f08ffb54970..8ec02b4b340 100644 --- a/csrc/flash_attn_ck/mha_varlen_fwd.cpp +++ b/csrc/flash_attn_ck/mha_varlen_fwd.cpp @@ -407,15 +407,15 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_si 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"; - mask = mask_info::decode(mask_identify, max_seqlen_q, max_seqlen_k); // casual + std::string mask_identify = "t:" + std::to_string(window_size_left) + "," + "0"; + mask = mask_info::decode(mask_identify, max_seqlen_q, max_seqlen_k); // causal } else if (window_size_left == -1 && window_size_right == -1) { mask = mask_info::decode("0", max_seqlen_q, max_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 = "t:" + std::to_string(window_size_left) + "," + std::to_string(window_size_right); mask = mask_info::decode(mask_identify, max_seqlen_q, max_seqlen_k); // local }