Fix causal+sliding window mask alignment for autoregressive decode

[juliantorr-es]

Summary

Fix for issue #158. The CK (Composable Kernel) FlashAttention wrapper uses bottom-right mask 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).

Root Cause

In csrc/flash_attn_ck/mha_fwd.cpp (and the other three CK wrapper files), the mask identifier string uses "b:" (bottom-right alignment):

std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + "0";

When seqlen_q == 1 (decode step) and window_size_left >= 0 (sliding window), bottom-right alignment computes:

y_tmp = seqlen_q - seqlen_k  → negative for decode
y = 1 + left_size + y_tmp   → can become negative for long prompts
x_start = -y + i_y + 1      → large positive → masks early KV positions

Example with window=128, seqlen_q=1, seqlen_k=256: x_start = 127, masking the first 127 KV tokens from the single query. Every decode step sees a truncated window, producing degenerate repetitive output.

Fix

Changed "b:" to "t:" (top-left alignment) in all 8 mask construction sites across 4 files:

File	Lines
csrc/flash_attn_ck/mha_fwd.cpp	2 lines (causal + local branches)
csrc/flash_attn_ck/mha_bwd.cpp	2 lines
csrc/flash_attn_ck/mha_varlen_fwd.cpp	2 lines
csrc/flash_attn_ck/mha_varlen_bwd.cpp	2 lines

This matches the HuggingFace causal attention convention (top-left aligned).

Verification

With top-left alignment during decode (window=128, seqlen_q=1, seqlen_k=256):

x_tmp = 0
y_tmp = 0  
y = 1 + 128 + 0 = 129
x_start = -129 + 0 + 1 = -128  → no early masking
x_end = min(0 + 1, 256) = 1    → attend to position 0 only

Correct: query at accumulated position N sees KV range [N-128, N].

Related Issues

• Closes [Issue]: Inferencing GPT-OSS 20B with FlashAttentionV2 in ROCm7.0 on MI300X results in unusable & degenerate repetitive output #158 (GPT-OSS 20B degenerate output on MI300X with FlashAttentionV2)
• Affects any model using sliding window attention on CK-backed ROCm flash-attention

[rocking5566]

Thanks for the fix! tridao is a read-only mirror we re-sync from Dao-AILab upstream, so PRs merged there get overwritten — please don't target it directly.

Two options:

• Retarget to ck_improve_main — lands in the ROCm fork right away. We batch PRs from it back to upstream, but the timing of that isn't fixed.
• If you're in a hurry, open the PR directly against Dao-AILab upstream — once it merges there, we can sync it into the fork at any time.

Either works; pick whichever gets you unblocked fastest.