[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 12 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 |
| } | ||
| #ifndef USE_ROCM | ||
| // include softmax_offset if provided | ||
| if (softmax_offset_buf != nullptr) { |
There was a problem hiding this comment.
🟡 The sink slot is written under a different predicate than every reader uses.
Here the pack appends the sink slot when softmax_offset_buf != nullptr, but all three readers key off softmax_type:
common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu:1198—if (softmax_type != NVTE_VANILLA_SOFTMAX)common/fused_attn_rocm/fused_attn_ck.cpp:1236— samecommon/fused_attn_rocm/fused_attn.cpp:510— same (ctx_next_id++)
Note the bias slot two lines above is symmetric (bias_type on both sides); sink is the odd one out, and this now affects CUDA too since the positional fix at line 310 makes softmax_offset_buf reach this branch for the first time.
Today it is benign in both directions: for VANILLA_SOFTMAX JAX passes a shape-(0,) buffer (cpp_extensions/attention.py:3501) whose untyped_data() is typically non-null, so the writer appends a slot the readers skip — harmless. But if XLA ever hands back a null pointer for a real (1,h,1,1) offset, the writer omits the slot while the reader still consumes tensors[i], which is a preallocated-but-unset entry → dptr == nullptr → has_sink silently false → attention computed without the sink, no error raised.
Cheapest fix that keeps both sides on one predicate: pass softmax_type down (both call sites at :310 and :596 have it in scope) and gate on softmax_type != NVTE_VANILLA_SOFTMAX, matching the bias_type check.
There was a problem hiding this comment.
Fixed — the writer and the readers are now on one predicate.
PrepareFusedAttnForwardAuxTensors takes an NVTE_Softmax_Type softmax_type (right after bias_type, so it reads symmetrically with the bias slot) and the sink slot is appended under softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX instead of softmax_offset_buf != nullptr. It is threaded through PrepareFusedAttnBackwardAuxTensors and forwarded on both branches of the #ifndef USE_ROCM split (the real softmax_type in both — only bias_type/backend are dummied on CUDA), plus both call sites (FusedAttnForwardImpl, FusedAttnBackwardImpl), where it was already in scope.
This keeps CUDA correct rather than just ROCm: fused_attn_f16_arbitrary_seqlen.cu:1198 and fused_attn/fused_attn.cpp:709 use the same predicate, so for VANILLA_SOFTMAX the writer no longer appends a slot the reader skips, and for a real offset the slot can no longer be dropped because of a null pointer.
This approach is also consistent with other code comments that suggest using aux as the preferred mechanism for retrieving tensor dimensions in the future.
| // ASM v3 bwd does not compute the sink gradient; force the CK tile path so | ||
| // has_sink users still get a correct (if slower) d_sink, mirroring the fwd | ||
| // ASM v3 sink guard above. | ||
| ck_args.uses_bwd_v3 = nvte_ck_uses_bwd_v3 && !has_sink; |
There was a problem hiding this comment.
🟡 This && !has_sink is deliberately not mirrored into the workspace-size query at line 777 (ws_size_args.uses_bwd_v3 = nvte_ck_uses_bwd_v3;), and that asymmetry is load-bearing — please add a comment so it doesn't get "fixed" later.
Why it matters: has_sink is not the same in the sizing pass and the execution pass. In JAX, GetFusedAttnBackwardWorkspaceSizes (jax/csrc/extensions/attention.cpp:489) creates aux_input_tensors and never populates it, so nvte_fused_attn_bwd reads a preallocated-but-unset slot, devPtrSoftmaxOffset is nullptr, and has_sink is false during sizing while it is true during execution. (PyTorch differs — it passes the populated pack to both calls, so has_sink is true in both.)
Current code is safe: line 777 keeps the raw env value in both passes, so the two passes agree, and ck_attn_bwd_workspace_size returns max(v2_bytes, v3_bytes) — i.e. it already reserves at least what the tile path needs. Propagating && !has_sink to line 777 would make sizing return v2_bytes only while JAX's sizing pass still computes has_sink == false → the two passes disagree → under-allocated dq_acc, the exact failure mode the cu_seqlens_sizing_sentinel comment at jax/csrc/extensions/attention.cpp:505 documents.
Also worth updating: ck_fused_attn/src/ck_fused_attn_bwd.cpp:480 still claims "Gating mirrors ck_attn_bwd's use_asm_v3", which is no longer literally true once has_sink enters the picture.
| return false; | ||
| } | ||
|
|
||
| // joint filter based on sliding window and attn_mask |
There was a problem hiding this comment.
🟢 With the softmax_type filter deleted here, the softmax_type parameter of is_ck_backend_supported is now unused, so there is no CK-side guard on sink at all — every config the CK backend accepts will also be accepted with a sink.
Two things worth confirming:
- Is sink actually correct across the whole accepted surface (dropout, THD/group mode, MQA/GQA
dk/dvexpansion, deterministic bwd)? The bwd accumulatesd_sinkviaatomicAddper head, so group mode and the GQA reduction path are the ones I'd want covered. Because there's no guard, an unsupported combination now silently produces a wrongd_sinkinstead of falling back to unfused.tests/jax/test_fused_attn.pydoes parametrizeLEARNABLE_SOFTMAXacross those axes, so if CI is green over the full matrix that answers it — just calling it out explicitly. - If no narrower guard is needed, drop the now-dead
softmax_typeparameter (and its argument at the call site) rather than leaving it silently unread.
There was a problem hiding this comment.
Took option 2 — the dead parameter is gone (definition in fused_attn_ck.cpp, declaration in fused_attn_ck.h, argument at the nvte_get_fused_attn_backend call site in fused_attn_rocm/fused_attn.cpp).
On your first question: no narrower guard is needed, and the full matrix is the evidence. The complete LEARNABLE_SOFTMAX sweep in tests/jax/test_fused_attn.py — 486 tests spanning dropout, THD/group mode, MQA/GQA, ragged layouts and deterministic bwd, i.e. exactly the atomicAdd-per-head d_sink paths you flagged — passes with zero failures. So sink is correct across the surface CK accepts, and there is nothing left for a CK-side sink filter to reject.
is_aotriton_backend_supported keeps its softmax_type parameter: it still genuinely rejects sink (fused_attn_aotriton.cpp:73), so there is no signature-parity reason to touch it.
| const size_t kv_max_seqlen, DType dtype, | ||
| [[maybe_unused]]NVTE_Bias_Type bias_type, | ||
| [[maybe_unused]]NVTE_Fused_Attn_Backend backend, | ||
| void *softmax_buf, | ||
| void *rng_state_buf, void *bias_buf, |
There was a problem hiding this comment.
🟢 Formatting — this signature won't survive the repo's hooks. [[maybe_unused]] needs a space before the type, and lines 138/142 (plus 595) have trailing whitespace, which the trailing-whitespace hook in .pre-commit-config.yaml covers for .cpp/.h. The wrapping also isn't what clang-format -style=file produces for this file.
| const size_t kv_max_seqlen, DType dtype, | |
| [[maybe_unused]]NVTE_Bias_Type bias_type, | |
| [[maybe_unused]]NVTE_Fused_Attn_Backend backend, | |
| void *softmax_buf, | |
| void *rng_state_buf, void *bias_buf, | |
| const size_t kv_max_seqlen, DType dtype, | |
| [[maybe_unused]] NVTE_Bias_Type bias_type, | |
| [[maybe_unused]] NVTE_Fused_Attn_Backend backend, | |
| void *softmax_buf, void *rng_state_buf, void *bias_buf, |
Same trailing-whitespace nit in three other spots in this PR: ck_fused_attn/src/ck_fused_attn_fwd.cpp (// denominator). , which also has a //(aiter's ... missing the space after //), fused_attn_rocm/fused_attn_ck.cpp (blank line after ck_args.how_v3_bf16_cvt), and tests/jax/test_fused_attn.py (blank line after optional_dgrad_idx = 3, which black will rewrite).
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: