Enable softmax-sink FMHA kernels for TransformerEngine

[shurale-nkn]

Motivation

Required by ROCm/TransformerEngine#678 — [ROCm] Jax Add softmax sink (learnable
off-by-one) support for the ROCm/CK fused attention backend.

This PR is the TransformerEngine-side mirror of ROCm/rocm-libraries#10519 ([CK] [MHA] Fix for MHA with softmax-sink), delivered as a QoLA patch because the CK snapshot
vendored by our pinned AITER commit predates that fix.

Important

The patch is temporary. Once CK 0282f870b propagates into the AITER branch we build against,
3rdparty/QoLA/patches/aiter/0009-fmha-sink-fixes.patch can simply be deleted — no
replacement, no follow-up. Nothing else in TE depends on it.
On the AITER side the related change is PR #4612, but that one only touches verification —
the API and the runtime behaviour are unchanged, so it has no bearing on this patch.
Note that #10519 is deliberately larger than what we carry here: it also contains changes that
no TE call path can reach. See Not included, and why.

---

Shipped

/patches/aiter/0009-fmha-sink-fixes.patch — a single patch, code only, no comments
in the hunks. It supersedes the former 0008 / 0009 / 0010 and carries three changes:

File	Change	Effect
kernel/fmha_bwd_kernel.hpp	sink_ptr indexed by i_nhead instead of i_batch * nhead + i_nhead	Fixes d_softmax_offset and the q/k/v gradients for every batch index > 0
pipeline/block_fmha_pipeline_qr_ks_vs.hpp (bias, randval, K/V) and ..._async.hpp (bias, randval)	Sink-prefix window jump guarded on an actual sink phase instead of i_total_loops == 0	Fixes sink combined with a local (sliding-window) mask. The K/V check is live in QRKSVS, so forward output was garbage (rel. err ~1.0) for bias + SWA + sink
kernel/fmha_fwd_kernel.hpp	Sink logit normalized by GetSoftmaxScale() instead of raw kargs.scale_s	fp8 per-tensor quantization only; inert for NO_SCALE

TE reaches all three through softmax_offset (softmax_type != "vanilla"), which sets
has_sink with sink_size = 0.
The first two were already shipped as separate patches and are unchanged in behaviour — this PR
only merges them and strips the C++ comments they were injecting into the CK sources. The third is
new and pre-emptive: it is a verified no-op today, carried because its gate can open without
any TE source change. Receipt 700 (our own patch 0003) is what currently keeps
PERTENSOR instances out of the build, and it is rebased on every AITER bump.

---

Not included, and why

All of the following are present in #10519 and were deliberately left out of the QoLA patch,
because no TE call path reaches them.
Async K/V sink-prefix jump. Turns a dead branch (i_total_loops == 0 sits after
i_total_loops++) into one guarded by num_sink_loop > 0. TE pins sink_size = 0, so it never
fires either way — zero behavioural difference.
MX scale-window offsets. Not gated on MX, gated on sink_size > 0. The two hunks that look
unconditional (seqlen_k_start → kv_load_start on the scale-window origins) are arithmetic
no-ops when sink_seq_end == 0, because
kv_load_start = (sink_seq_end == 0 && seqlen_k_start > 0) ? seqlen_k_start : 0 makes the two
expressions equal in every branch. The rest sit inside the i_total_loops == num_sink_loop - 1
guard, which evaluates to == -1 when num_sink_loop == 0. The two async hunks additionally have
no target: that pipeline declares no scale windows and carries
static_assert(QScaleEnum != BlockAttentionQuantScaleEnum::MX). Since quantization is the wrong
gate here, this stays irrelevant even after fp8/mxfp land.
Host reference and tests. fmha_{fwd,bwd}_runner.hpp and test_fmha_{fwd,bwd}.cpp are
included only by CK's example_fmha_*.cpp and test_fmha_*.cpp, none of which QoLA builds — it
builds libmha_fwd / libmha_bwd from generated blobs plus cpp_itfs sources only.
Comment rewrites. Comment-only upstream churn. The merged patch carries no comments at all;
the rationale lives in the patch message instead, where it does not have to be rebased.

Technical Details

Test Plan

Test Result

Submission Checklist

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

[Micky774]

@shurale-nkn have you opened PR's to the upstream CK library (which is part of the rocm-libraries repo) for these patches yet? Ideally it would be reviewed there for correctness, and we'd accept the patch here off of that review.

[shurale-nkn]

@Micky774 ROCm/rocm-libraries#10519