[BUG][CK][MHA] Fix for MHA with softmax-sink

[shurale-nkn]

Motivation

Fix the sink shape contract in the MHA backward bindings: [batch, num_heads]
→ [num_heads]. The sink is a learnable per-head softmax offset — one scalar per
query head, shared across the batch — so a batch dimension was never meaningful.
The declared contract was inconsistent with the kernel and with itself:

• CK reads it per head, sink_ptr[i_nhead], and accumulates the gradient
  atomically over all batches into a [num_heads] d_sink. A [H] gradient can
  only belong to an [H] parameter — the output side already encoded the right
  semantics.
• Our own forward path already requires exactly [H]: aiter/ops/mha.py:1285.
• The Triton path, which is the one production path where sink is actually live
  today, also takes [H]: aiter/ops/triton/attention/mha.py:139-141.
  So [H] is not a choice, it is the only shape consistent with every other
  surface. Backward was the sole outlier.

Changed

• csrc/py_itfs_ck/mha_bwd_kernels.cu, csrc/py_itfs_ck/mha_varlen_bwd_kernels.cu
  — TORCH_CHECK now requires 1-D [num_heads]; comments updated.
• csrc/include/mha_bwd.h, csrc/include/torch/mha_bwd.h,
  csrc/include/torch/mha_varlen_bwd.h — doc comments.
• op_tests/test_mha.py, op_tests/test_mha_varlen.py — see Tests below.

Consumer impact: no regression

The three consumers of this interface:

Consumer	Path to CK backward sink	Shape it passes	Impact
TransformerEngine	aiter::mha_bwd_args C++ API directly (ck_fused_attn_bwd.cpp:534-535)	(1, H, 1, 1), i.e. H contiguous fp32 (jax/cpp_extensions/attention.py:452)	Contract now matches what TE already passes
PyTorch	sets sink_ptr/d_sink_ptr to nullptr (SFINAE, mha_bwd_ck.hip:19-33)	—	Unaffected; generates _nsink instances only
Primus-Turbo	none — hard dispatch sends sink to Triton, never CK (attention_aiter_impl.py:10-12,146,265)	[H_q] on the Triton path	Unaffected

TE is the only consumer that actually feeds sink to CK backward, and it already
supplies an H-element buffer. Under the old [B, H] claim combined with CK's
batch-major indexing, the kernel would index past the end of that buffer for
batch > 0. This change removes that, so for TE the fix is strictly corrective.
Primus-Turbo gates on sink is not None and routes to Triton unconditionally, so
sink never reaches CK there. That Triton path takes [H], which independently
corroborates that [H] is the intended contract rather than a convenience.
Note the check is a runtime TORCH_CHECK, not compile-time: any caller still
passing [B, H] is rejected at the binding boundary with a clear message rather
than silently computing wrong numbers.

Tests

test_mha.py / test_mha_varlen.py sink tests were producing inputs no forward
pass can reach: sink was passed to backward but not to forward, so the LSE did
not include the sink term and P_sink = exp(sink - lse) could exceed 1. Now:

• sink is passed to forward as sink_ptr, so the LSE is the one the kernel
  actually produces and the reference formula is exact rather than approximate;
• the sink draw range changed (see inline comment at the call site);
• tolerances tightened from rtol=0.02, atol=0.5 to rtol=1e-3, atol=1e-3,
  which is roughly a 450x margin over the observed max error.
  All sink tests pass: 18 in test_mha.py, 4 in test_mha_varlen.py.

Technical Details

This is PR in CK
ROCm/rocm-libraries#10519

this is PR in TransformerEngine
ROCm/TransformerEngine#678

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[shurale-nkn]

so buffer magnitude change .

1. Updated the buffer data range. With the previous distribution
   sink = torch.empty(batch, nhead, device=device, dtype=torch.float32).uniform_( 30.0, 60.0 )
   , the BWD output values were extremely small in magnitude, making the validation ineffective since even incorrect implementations could still pass due to the negligible numerical differences. The range has been changed to reflect realistic input data, providing more meaningful validation coverage.
   Below are examples of the output values produced with different input tensor ranges.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4612 --add-label <label>