[ROCm] Jax Add softmax sink (learnable off-by-one) support for the ROCm/CK fused attention backend - #678
[ROCm] Jax Add softmax sink (learnable off-by-one) support for the ROCm/CK fused attention backend#678shurale-nkn wants to merge 15 commits into
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Review summaryReviewed the sink-support diff (fwd + bwd for CK/AITER, JAX aux-tensor plumbing, one new test). Scope focus was the actual PR changes (merge-base → PR head), i.e. High-level verdict: approach is sound — sink is threaded through as a new optional aux tensor with dynamic slot indexing on both the fwd write and bwd read sides, the ASM v3 fallback is symmetric between fwd/bwd, and d_sink zeroing matches the CK atomicAdd contract. One likely CUDA regression and two comment/documentation nits — see inline. Highlights:
Copyright headers: OK — all 8 modified files carry AMD headers with end-year |
Re-review summaryRe-reviewed since the last pass (commit Prior findings — all addressed:
New findings (4 inline):
Verdict: approach is sound and the dynamic aux-slot indexing now matches the CUDA reference implementation's structure. Nothing blocking; the two 🟡 items are worth resolving before merge. Copyright headers: OK — all 8 modified source files carry AMD headers ending in |
|
Hey there @shurale-nkn, thanks for the contribution! While we're looking into it, could you provide us some context regarding why this feature enablement is wanted/needed? What use case are you trying to enable? Thanks! |
Hi @Micky774, I need sink support for train in MaxText. At the moment, TE is the only provider of fused attention for this framework, and without this PR, GPT-OSS will not work correctly. |
| bias_type, | ||
| attn_mask_type, | ||
| softmax_type, | ||
| dropout, |
There was a problem hiding this comment.
🟡 Dropping softmax_type from is_ck_backend_supported also flips the PyTorch ROCm path onto CK, which this PR neither mentions nor tests.
Chain: PyTorch's Python-level gating already permits FusedAttention for non-vanilla softmax, and the ROCm carve-outs are explicit — dot_product_attention/utils.py:1061 skips the thd/cuDNN-version disable under IS_HIP_EXTENSION, and :1479 skips the determinism disable the same way. Until this commit the only thing stopping it was the C++ layer: CK rejected sink here and AOTriton still does (fused_attn_aotriton.cpp:72), so nvte_get_fused_attn_backend returned NVTE_No_Backend and DPA silently fell back to UnfusedDotProductAttention. With the guard gone, CK is selected.
Concretely, tests/pytorch/attention/test_attention.py::test_dpa_softmax and ::test_dpa_softmax_thd (15 configs each, num_gqa_groups=8 + causal/padding/SWA (128,0)) now run against CK on ROCm — they aren't cuDNN-gated here because get_cudnn_version() returns (99, 0, 0) for HIP (pytorch/utils.py:698), and ci/pytorch.sh:89 runs the whole file at TEST_LEVEL 1. That's a meaningful surface: GQA dk/dv expansion plus THD, i.e. the atomicAdd-per-head d_sink paths, on a framework whose aux-pack plumbing this PR didn't touch.
The 486-test JAX sweep you cited answers the CK-kernel question, but not the PyTorch-binding one. Could you confirm the ROCm PyTorch attention job is green on this branch? If it isn't yet, gating on framework (or keeping a narrow CK-side guard until PyTorch is validated) would be safer than enabling both frameworks in one commit.
| size_t *workspace_size, | ||
| cudaStream_t stream){ | ||
|
|
||
| const bool has_sink = softmax_type != NVTE_VANILLA_SOFTMAX; |
There was a problem hiding this comment.
🟢 Good call dropping the devPtrSoftmaxOffset != nullptr term — that's what makes the earlier uses_bwd_v3 concern go away rather than just move. has_sink is now derived purely from softmax_type, which JAX passes identically to the sizing call (jax/csrc/extensions/attention.cpp:547) and the execution call (:683), so folding && !has_sink into nvte_ck_uses_bwd_v3 at :738 reaches ws_size_args.uses_bwd_v3 at :779 with the same value in both passes. ck_attn_bwd_workspace_size then returns max(v2_bytes, 0) == v2_bytes, which is exactly what the tile path needs. The ck_fused_attn_bwd.cpp:480 comment about mirroring use_asm_v3 is literally true again too.
One thing the old term was incidentally covering: has_sink is now true even when the offset pointer is null, so ck_args.sink_ptr (:644, :1035) and ck_args.d_sink_ptr (:1036) can reach CK as nullptr with has_sink == true, and the bwd cudaMemsetAsync(devPtrDSoftmaxOffset, ...) at :936 would memset a null pointer. Both sizing passes return before those points (:547, :897), so JAX is fine, and PyTorch's DPA always allocates the buffer for non-vanilla softmax (dot_product_attention.py:490-499). But pytorch/csrc/extensions/attention.cpp:231 and :308 explicitly tolerate softmax_type != VANILLA with SoftmaxOffset == nullopt, and that combination now reaches CK instead of being rejected — previously a silent no-op, now a null deref.
Since the pointer is a hard requirement once has_sink is set, an NVTE_CHECK just after the sizing early-return would turn that into a clear error instead of a GPU fault:
NVTE_CHECK(!has_sink || devPtrSoftmaxOffset != nullptr,
"softmax_offset is required for non-vanilla softmax");Minor while you're in here: removing the old has_sink line left a double blank at :484-485, and .clang-format sets MaxEmptyLinesToKeep: 1.
| void *softmax_buf, | ||
| void *rng_state_buf, void *bias_buf, | ||
| void *softmax_offset_buf = nullptr) { | ||
| #ifndef USE_ROCM |
There was a problem hiding this comment.
🟢 The sink slot is symmetric now, but the bias slot on this CUDA branch still isn't — and that only started mattering once a second optional tensor was appended behind it.
dummy_bias_type = NVTE_POST_SCALE_BIAS makes the writer always emit a bias slot at index 2. The CUDA bwd reader (common/fused_attn/fused_attn.cpp:701-711) walks the pack with a running index and consumes the bias slot only under the real bias_type:
if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) {
input_Bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]);
}
if (softmax_type != NVTE_VANILLA_SOFTMAX) {
input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]);
}So for NO_BIAS + non-vanilla softmax the writer produces [softmax, rng, bias(dummy), sink] while the reader skips bias and picks up tensors[2] — the dummy bias slot, carrying bias_buf and a {bias_batch, bias_heads, q, kv} shape — as input_SoftmaxOffset. Before sink existed the extra trailing slot was simply ignored, which is presumably why the dummy survived this long.
Not a regression from this PR (the base had the same layout via softmax_offset_buf != nullptr), and I can't run the CUDA path to confirm, so treat this as a question rather than a claim. But the fix looks like the one you already applied on the ROCm side — pass the real bias_type here too, since the reader is dynamic on both slots:
| #ifndef USE_ROCM | |
| #ifndef USE_ROCM |
…with dummy_bias_type replaced by bias_type at line 151. If that's deliberate scope-limiting for a ROCm fork PR, an upstream issue link in the description would be enough.
Re-review summaryRe-reviewed since the last pass (commit Prior findings — all addressed:
New findings (3 inline):
Also verified the new distributed-test collective accounting: Verdict: the aux-slot indexing is now consistent end to end and the v3-gating/workspace-sizing interaction is correct. Nothing blocking; the PyTorch-enablement question is the one worth answering before merge. Copyright headers: OK — all 9 files in scope carry AMD headers ending in |
## Motivation MHA softmax-sink-learnable support was needed in the Maxtext->TransformerEngine->Aiter framework, but testing revealed several bugs in the implementation. While I was fixing it, I found that tests is not sensitive to changes in the code. JIRA ID : AIJAX-322 ## Technical Details **Problem 1)** The external MHA interface contract for the backward (BWD) path uses the **_sink_** buffer with an incorrect shape [B, H]. This breaks the symmetry with the forward (FWD) path, where the corresponding _sink_ buffer has shape [H]. As a result, the backward computation produces incorrect sink gradients. The [H] shape is the only correct representation because this tensor corresponds to a trainable parameter with one shared value per attention head, not a unique value for every element in the batch. Therefore, introducing the batch dimension is semantically incorrect. In addition, [H] is the shape used by other frameworks. Because of this mismatch, framework integrations currently fall back to Triton implementations for Sink attention instead of using Composable Kernel (CK). _Fix:_ The first commit fixes the BWD. **Problem 2**) After applying the first commit, the existing test suite continues to pass even though a critical part of the kernel behavior has changed. Since the tests themselves were not modified, this indicates that they are effectively insensitive to the bug and unable to detect incorrect behavior. ``./bin/test_ck_tile_fmha_bwd_bf16 --gtest_filter='*Sink*'`` The root cause is the initialization range used by the tests. The generated input values are too large, which causes the Sink outputs to become extremely small and converge toward zero. As a consequence, the numerical differences introduced by the bug are masked, allowing the tests to pass regardless of whether the implementation is correct. Failed 3 out of 15, but should all 15 <img width="1088" height="228" alt="image" src="https://github.com/user-attachments/assets/23b38cc1-0f81-44da-8b46-0f2cf2eb7b9e" /> _Fix:_ The second commit adjusts the input initialization range to produce numerically meaningful outputs and improve test sensitivity. <img width="1154" height="521" alt="image" src="https://github.com/user-attachments/assets/678d5c1c-7e97-4d12-9117-6f4e814db8e0" /> _Fix:_ The 3-rd commit fixed test code, now all 15 passing. <img width="717" height="338" alt="image" src="https://github.com/user-attachments/assets/d20ad772-fb77-4cfc-9f7b-a9dbaa62e741" /> **Problem 3**) There is an issue with the Sliding Window Attention (SWA) window progression when used together with Sink attention. The bug is observable in external framework tests, but the current CK test suite does not contain any coverage capable of exposing this behavior. As a result, the implementation can regress without any CK tests failing. Test Coverage Improvement The fourth commit adds a set of test configurations specifically designed to exercise the affected SWA + Sink code path. These new tests reproduce the issue and fail with the current implementation, demonstrating the gap in the existing test coverage and providing a reliable way to validate the fix. <img width="1256" height="340" alt="image" src="https://github.com/user-attachments/assets/c9f050e1-a973-4aff-9d38-17ed21387a6f" /> Fix The fifth commit corrects the SWA window movement logic when Sink attention is enabled. With this change applied, the newly added tests pass, and the behavior matches the expectations observed in external framework implementations. ### **Related PRs** This is PR in CK #10519 this is PR in TransformerEngine ROCm/TransformerEngine#678 ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
[CK] [MHA] Fix for MHA with softmax-sink ## Motivation MHA softmax-sink-learnable support was needed in the Maxtext->TransformerEngine->Aiter framework, but testing revealed several bugs in the implementation. While I was fixing it, I found that tests is not sensitive to changes in the code. JIRA ID : AIJAX-322 ## Technical Details **Problem 1)** The external MHA interface contract for the backward (BWD) path uses the **_sink_** buffer with an incorrect shape [B, H]. This breaks the symmetry with the forward (FWD) path, where the corresponding _sink_ buffer has shape [H]. As a result, the backward computation produces incorrect sink gradients. The [H] shape is the only correct representation because this tensor corresponds to a trainable parameter with one shared value per attention head, not a unique value for every element in the batch. Therefore, introducing the batch dimension is semantically incorrect. In addition, [H] is the shape used by other frameworks. Because of this mismatch, framework integrations currently fall back to Triton implementations for Sink attention instead of using Composable Kernel (CK). _Fix:_ The first commit fixes the BWD. **Problem 2**) After applying the first commit, the existing test suite continues to pass even though a critical part of the kernel behavior has changed. Since the tests themselves were not modified, this indicates that they are effectively insensitive to the bug and unable to detect incorrect behavior. ``./bin/test_ck_tile_fmha_bwd_bf16 --gtest_filter='*Sink*'`` The root cause is the initialization range used by the tests. The generated input values are too large, which causes the Sink outputs to become extremely small and converge toward zero. As a consequence, the numerical differences introduced by the bug are masked, allowing the tests to pass regardless of whether the implementation is correct. Failed 3 out of 15, but should all 15 <img width="1088" height="228" alt="image" src="https://github.com/user-attachments/assets/23b38cc1-0f81-44da-8b46-0f2cf2eb7b9e" /> _Fix:_ The second commit adjusts the input initialization range to produce numerically meaningful outputs and improve test sensitivity. <img width="1154" height="521" alt="image" src="https://github.com/user-attachments/assets/678d5c1c-7e97-4d12-9117-6f4e814db8e0" /> _Fix:_ The 3-rd commit fixed test code, now all 15 passing. <img width="717" height="338" alt="image" src="https://github.com/user-attachments/assets/d20ad772-fb77-4cfc-9f7b-a9dbaa62e741" /> **Problem 3**) There is an issue with the Sliding Window Attention (SWA) window progression when used together with Sink attention. The bug is observable in external framework tests, but the current CK test suite does not contain any coverage capable of exposing this behavior. As a result, the implementation can regress without any CK tests failing. Test Coverage Improvement The fourth commit adds a set of test configurations specifically designed to exercise the affected SWA + Sink code path. These new tests reproduce the issue and fail with the current implementation, demonstrating the gap in the existing test coverage and providing a reliable way to validate the fix. <img width="1256" height="340" alt="image" src="https://github.com/user-attachments/assets/c9f050e1-a973-4aff-9d38-17ed21387a6f" /> Fix The fifth commit corrects the SWA window movement logic when Sink attention is enabled. With this change applied, the newly added tests pass, and the behavior matches the expectations observed in external framework implementations. ### **Related PRs** This is PR in CK ROCm/rocm-libraries#10519 this is PR in TransformerEngine ROCm/TransformerEngine#678 ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
4 participants
Description
Adds forward and backward support for softmax sink (NVTE_LEARNABLE_SOFTMAX) to the CK/AITER fused-attention backend on ROCm.
Fixes # (issue)
CK code in QoLa
as i_batch * nhead + i_nhead, but the buffer holds one value per head (shape
[nhead]), matching how the forward kernels already read it, so i_batch must not
factor into the offset.
Affected all bwd tests
18 tests fixed: POST_SCALE_BIAS-1HSS-{Mask,Seqlens,SegmentIDs}-SWA-DROP_0.0-<cfg>-LEARNABLE_SOFTMAX-<mask>Type of change
Changes
Please list the changes introduced in this PR:
Checklist: