[HIP] [GFX950]OPUS MHA Support LSE

[shay-li77]

Motivation

The OPUS gfx950 bf16 forward kernels were inference-only: with no LSE output, return_lse=True — and therefore the whole autograd path, which asserts return_lse whenever is_grad — fell back to fmha_v3/CK and lost the OPUS speedup. This adds LSE to both OPUS forward kernels (D=128 and D_QK=192/D_V=128, batch + group/varlen) and opens the dispatch gates so those callers keep it.

Technical Details

• LSE is exact under the FAv4 selective max update. m_row is not the true row max (the running max is only refreshed past RESCALE_THRESHOLD), but l_row stays consistent with whatever base it holds and log-sum-exp is invariant to that base, so lse = ln2 * (m_row + log2(l_row)). Fully-masked rows (bottom-right causal, seqlen_q > seqlen_kv) get -inf, matching torch.logsumexp.
• Gated at runtime on ptr_lse != nullptr rather than a traits flag: a scalar branch, no extra instantiations, and callers that skip LSE skip the allocation and the store. Costs 2 VGPR on D=192 / 1 on D=128, occupancy unchanged.
• Output is fp32 natural log — batch [B, H, N], group [H, total_q]. In the API lse is an out-param like out; only return_lse changes the return type.
• Both dispatch gates dropped not return_lse (keeping not return_softmax), which also makes the autograd path reachable, i.e. the OPUS forward now feeds the v3/CK backward. Test side: 5 OPUS test functions → 3, with seqlen_q != seqlen_kv folded into the shared case lists.

Test Plan

• pytest op_tests/{test_mha,test_mha_varlen}.py -k opus. LSE is compared against an fp32 logsumexp reference, not attention_ref's — that one is cast back to the input dtype and at |lse| ~ 8 is coarser than the kernels' own error.
• Coverage: cross-attention both directions incl. fully-masked rows; causal head/tail merge engaged so the mirror Q block scans KV in reverse; KV padding with NaN-filled gaps; single / 17 / decode-shaped groups; autograd grads vs a torch fp32 reference. Every case asserts the output is bitwise-equal to the direct wrapper, so the checks cannot silently validate another backend (negative controls confirm they discriminate).
• Resources via -Rpass-analysis=kernel-resource-usage before/after; perf via run_perftest on the PR Add OPUS gfx950 bf16 fmha d192x128 kernel #4205 shapes (seqlen × batch = 16384, causal, bf16, nheads ∈ {16, 32}), LSE on vs off plus a rebuilt pre-LSE baseline.

Test Result

• Correctness: 84 passed (60 batch + 24 group). LSE max diff vs fp32 reference: D=192 ~1e-6; D=128 ~4e-3, which is that kernel's pre-existing bf16 Q pre-scale rather than the LSE — against an arithmetic-matched reference it is also 1.4e-6. -inf row patterns match exactly, no NaN. Grads within bf16 tolerance (causal: out 7.6e-3, dq 1.1e-2, dk 2.1e-2, dv 2.8e-2).
• Resources: D=192 248–250 → 250–252 VGPR, D=128 245/246 → 246/247; occupancy 2 waves/SIMD, VGPRs Spill: 0 throughout.
• Perf: LSE on vs off is 0.06%–0.20% mean across nheads ∈ {16, 32} in both modes (noise-dominated); the no-LSE path is within ±1% of a rebuilt pre-LSE baseline with no systematic direction.
• No regressions from opening the gates: the 320 test_flash_attn_output[…192-128…] cases that now reroute to OPUS give an identical 160/160 split with and without AITER_DISABLE_FMHA_OPUS=1 — those 160 failures are a pre-existing invalid parametrization (assert nheads % gqa_ratio == 0 → 6 % 8) that fails in the test's own setup on either backend.

Submission Checklist

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

[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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
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 4877 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[Copilot]

Pull request overview

This PR extends the gfx950 OPUS bf16 MHA forward path to optionally produce LSE (log-sum-exp) in fp32 (natural log) for both dense (batch) and group/varlen variants, plumbing the buffer from Python → pybind → host launcher → device kernels, and updating tests to validate LSE correctness and fully-masked-row behavior.

Changes:

• Add optional fp32 LSE output support to OPUS D=128 (batch) and D_QK=192/D_V=128 (batch + group/varlen) kernels and their host launchers/pybind bindings.
• Update Python wrappers/dispatch gating so OPUS can be selected when return_lse=True (still disallowing return_softmax / dropout mask).
• Expand OPUS test coverage to assert routing + output correctness and validate LSE vs a reference (including cross-attention and fully-masked rows).

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
op_tests/test_mha.py	Refactors OPUS dense tests to cover LSE + cross-attention cases and assert OPUS routing via flash_attn_func.
op_tests/test_mha_varlen.py	Adds OPUS varlen/group tests validating output + LSE and KV-padding safety (NaN gap checks).
csrc/py_itfs_cu/fmha_fwd_bf16_opus_kernels.cu	Plumbs optional lse tensor through the shared launcher into both OPUS variants with shape/stride validation.
csrc/include/torch/fmha_fwd_bf16_opus.h	Documents the new optional LSE output contract and adds the optional argument to the API.
csrc/include/rocm_ops.hpp	Exposes lse as an optional pybind argument (default nullopt).
csrc/include/fmha_fwd_hd192_v128_bf16_opus_kernel.hpp	Implements optional LSE store (fp32, natural log) and refactors S/P fragments to support it.
csrc/include/fmha_fwd_hd192_v128_bf16_opus_defs.h	Extends kargs with optional LSE pointer + strides for batch/group modes.
csrc/include/fmha_fwd_hd128_bf16_opus_kernel.hpp	Implements optional LSE store for D=128 kernel.
csrc/include/fmha_fwd_hd128_bf16_opus_defs.h	Extends D=128 kargs with optional LSE pointer + strides.
aiter/test_mha_common.py	Adds OPUS-specific fp32 LSE reference and checking helper.
aiter/ops/mha.py	Adds return_lse/lse plumbing in OPUS wrappers and enables OPUS dispatch when return_lse=True.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.