[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 14 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 |
3 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: