[Attention] Handle fully masked rows

[umangyadav]

Reference

Port of ROCm/rocMLIR#2448.

Motivation

attentionSweeps.py can fail because causal sliding-window configurations can contain fully masked attention rows. In the investigation documented by the rocMLIR PR, 113 configurations in a 1,000-case MFMA GPU-vs-CPU sweep were reported as failures, and every one contained at least one fully masked row.

For example, a decode configuration with seq_len_q=1, current_seq_len=2, sliding_window_size=1, and causal masking can produce a row whose scores are all -inf. Softmax is not defined for this row: in the direct formula every numerator is exp(-inf) = 0 and the denominator is also zero, while the numerically stable formula computes

m = max(-inf, ..., -inf) = -inf

followed by

x_i - m = -inf - (-inf) = NaN.

A fully masked query has no eligible value vectors, so its contribution is defined as the zero vector. rocmlirTriton's GPU online-softmax implementation now replaces exp2(-inf - (-inf)) contributions with zero, keeps the running sum defined while both row maxima are -inf, and selects zero when the final exponential sum is zero. These guards prevent a masked row from poisoning later arithmetic with NaN.

The TOSA-based CPU reference used the stable softmax formula without equivalent degenerate-row handling. It therefore produced NaN instead of zero. The split-KV host combine had the same all--inf reduction problem. After all paths consistently produce zero, an entirely masked output can make the GPU and reference tensors identically zero; that also exposes the verifier's normalized-RMS 0/0 edge case addressed here.

Summary

• Define fully masked GPU and CPU-reference softmax rows as zero instead of allowing NaN from -inf - (-inf) and division by zero.
• Keep split-KV host recombination defined when all partial LSE values are -inf.
• Preserve the original reduction values used for LSE, so a fully masked row retains the expected -inf LSE.
• Make exact all-zero GPU/reference matches report zero normalized RMS and add regression coverage for all paths.

Port notes

rocMLIR already had its GPU guards in its buffer-based attention lowering. rocmlirTriton uses the newer GridwiseAttnToBlockwise lowering, so this port adds the equivalent score, running-sum, and final-output guards to that implementation.

Mathematical rationale

For attention scores x_i, let

• m = max_i(x_i)
• L be the lowest finite value of the compute type
• m_safe = max(m, L)
• e_i = exp(x_i - m_safe)
• A = sum_i(e_i)
• p_i = e_i / max(A, 1)

For every ordinary row containing a finite score, m >= L, so m_safe = m. At least one score equals m, hence one exponential is exp(0) = 1 and A >= 1. Both guards are therefore identities and leave the existing softmax unchanged.

For a fully masked row, every x_i = -inf, so m = -inf, m_safe = L, every e_i = exp(-inf - L) = 0, and A = 0. The guarded denominator is one and every p_i = 0. The same argument applies when combining split-KV partial LSE values. The unguarded max and sum remain the inputs to LSE, preserving log(0) + (-inf) = -inf for fully masked rows.

The verifier computes normalized RMS as

E = sqrt(sum_i((gpu_i - ref_i)^2)) / (M * sqrt(N)),

where M = max_i(max(abs(gpu_i), abs(ref_i))). Exact all-zero tensors have zero squared-error sum and M = 0, making the literal expression 0/0. Defining E = 0 when the numerator is exactly zero preserves the fundamental property that identical tensors have zero error. The guard checks the numerator rather than M: nonzero finite mismatches necessarily have M > 0, while non-finite mismatches continue through the existing failure path.

Test plan

• Build rocmlir-gen, rocmlir-opt, rocmlir-driver, and conv-validation-wrappers.
• Run 11 focused generator, CPU-reference, split-KV, verifier, and GPU-lowering lit tests.
• Run 8 broader attention GPU-lowering and E2E tests.
• Verify the fully masked causal sliding-window GPU-vs-CPU E2E case reports [1 1 1].

Made with Cursor

[Copilot]

Pull request overview

Ports fully masked attention-row handling from ROCm/rocMLIR#2448.

Changes:

• Prevents NaNs in GPU, CPU-reference, and split-KV softmax paths.
• Preserves -inf LSE values for fully masked rows.
• Handles exact all-zero verifier results and adds regression coverage.

Reviewed changes

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

Show a summary per file

File	Description
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	Adds guarded CPU and split-KV normalization.
mlir/lib/Dialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp	Guards GPU online softmax and final output.
mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp	Avoids normalized-RMS 0/0.
mlir/test/rocmlir-gen/attention-splitkv-host-combine.mlir	Tests split-KV guards.
mlir/test/rocmlir-gen/attention-sliding-window.mlir	Tests fully masked CPU-reference lowering.
mlir/test/rocmlir-gen/attention-kernel.mlir	Updates softmax checks.
mlir/test/rocmlir-gen/attention-kernel-gqa.mlir	Updates GQA checks.
mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir	Updates KV-cache checks.
mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse.mlir	Updates LSE checks.
mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache-lse-splitkv.mlir	Updates split-KV LSE checks.
mlir/test/rocmlir-gen/attention-kernel-f16.mlir	Updates f16 checks.
mlir/test/rocmlir-gen/attention-kernel-causal.mlir	Updates causal checks.
mlir/test/rocmlir-driver/verify-all-zero.mlir	Tests all-zero verification.
mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir	Adds fully masked E2E coverage.
mlir/test/Dialect/Rock/toblockwise_attention_lowering.mlir	Checks GPU softmax guards.
mlir/test/Dialect/Rock/toblockwise_attention_lowering_nPerBlockG1.mlir	Checks chunked-output guards.

---

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

[rocmlir-pr-reviewer]

Verdict: COMMENT  ·  Findings: 3 (0 Critical, 1 Major, 2 Minor)

---

Scope

Port of the rocMLIR fully-masked-attention fix. Three layers change: (1) the GPU GridwiseAttnToBlockwise lowering gains -inf - (-inf) guards in expSubstractMaxFromGemm0, updateRowSum, and scaleFinalOutput; (2) the rocmlir-gen TOSA CPU reference and the split-KV host combine clamp the softmax max to lowest-finite and the exp-sum denominator to >= 1, while keeping the unclamped values for LSE; (3) mcpuVerify defines normalized RMS as zero when the squared-error sum is exactly zero. Plus a new E2E fixture, a new verifier lit test, a new split-KV host-combine lit test, and CHECK refreshes across the attention generator tests.

Findings

• Major: mlir/test/rocmlir-driver/verify-all-zero.mlir:27 declares and calls mcpuVerifyFloat with one operand more than the wrapper actually takes.
• Minor: mlir/lib/Dialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp:236 puts a 4-op per-element guard in the innermost softmax loop where the cheaper row-vector clamp the CPU reference already uses would do.
• Minor: mlir/lib/Dialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp:202 assert lacks a descriptive message.

Notes

Math spot-checks hold. For any row with a finite score the row max exceeds lowest-finite and the exp-sum is at least exp(0) = 1, so both tosa.maximum clamps in rocmlir-gen.cpp (computeFinalAttentionStage, createCpuAttentionKernelWithMlir) are identities and existing results are unchanged. LSE still reads the unclamped qkMaxs / expsSums, so a fully masked row keeps log(0) + (-inf) = -inf on both sides. The sumDiffSq == 0.0 guard in conv-validation-wrappers.cpp:287 is the right predicate: a NaN or non-finite mismatch does not compare equal to zero and still falls through to the division, and any nonzero finite mismatch implies maxMag > 0.

The two overloads of the new createFloatSplatTensor helper are both used, so no dead code. FileCheck variables introduced in the refreshed lowering tests (zero32, zero64, zeroRow, negInf) are all bound by existing CHECK-DAG lines.

rocMLIR back-port check: the PR touches shared paths (mlir/lib/Dialect/Rock/, mlir/tools/rocmlir-gen/), and the description satisfies option (a) by linking the parallel rocMLIR PR and adding port notes explaining the lowering divergence. No finding.

CI status

No failed or cancelled checks. Jenkins, Build and Test, and MIGraphX were still pending at review time; the review check is this pipeline itself.

[rocmlir-pr-reviewer]

Verdict: APPROVE -- submitted as COMMENT (automated reviews are advisory)  ·  New findings: 0 (0 Critical, 0 Major, 0 Minor)

---

Scope

Defines fully masked attention rows (causal + sliding window) as producing zero instead of NaN. Adds -inf/zero guards to the GPU lowering in mlir/lib/Dialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp (score exponent, running-sum rescale, final normalize), tosa.maximum clamps in the TOSA CPU reference and split-KV host combine in mlir/tools/rocmlir-gen/rocmlir-gen.cpp, and a 0/0 guard for normalized RMS in mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp. 11 test files updated or added.

Findings

No blocking issues found. Both defects raised in the previous review are fixed at head 236637970: the verify-all-zero.mlir call/declaration now matches the 7-parameter mcpuVerifyFloat ABI, and the type-equality assert in GridwiseAttnToBlockwise.cpp now carries a descriptive message. The remaining Minor (guard cost / GPU-vs-CPU strategy asymmetry) was explicitly deferred by the author to a separate ticket covering both rocMLIR and rocmlirTriton, so it is closed here rather than re-raised.

Notes

Spot-checks performed on this revision:

• Ordering in the CPU reference is correct — the LSE block consumes the unclamped qkMaxs and expsSums, so a fully masked row still yields log(0) + (-inf) = -inf rather than log(1) + lowestFinite. The new SAFE prefix in attention-sliding-window.mlir pins exactly this (tosa.add of tosa.cast %[[MAX]], not the clamped value).
• Both clamps are mathematical identities on ordinary rows: any row with a finite score has m >= lowestFinite and contributes exp(0) = 1, so the sum is at least one.
• The verifier guard keys off the numerator rather than maxMag; any nonzero finite mismatch implies maxMag > 0, and non-finite diffs still propagate to a NaN RMS and fail. The pre-existing valNum == gpuNum early-equality branch keeps -inf vs -inf LSE elements passing.
• All new FileCheck variables in the two toblockwise_attention_lowering*.mlir files are bound by CHECK-DAG lines above their uses, so -enable-var-scope is satisfied; each new RUN line declares exactly one prefix and uses it, so --allow-unused-prefixes=false is satisfied.
• mlir/test/fusion/pr-e2e/lit.local.cfg gates the whole directory, so the new E2E test needs no REQUIRES: line or CMake registration.

rocMLIR back-port: the PR touches shared paths (GridwiseAttnToBlockwise.cpp, rocmlir-gen.cpp, conv-validation-wrappers.cpp), and the description links the parallel rocMLIR PR plus a "Port notes" section explaining the lowering difference. Option (a) is satisfied. No external/ or *-patches/ changes, so the subtree-bump and [EXTERNAL] commit rules do not apply.

CI status

No checks in the fail or cancel buckets. py-checks and detect pass; notify-fork-pr is skipped. The auto-review check itself is in progress, which is expected.

[Copilot]

Pull request overview

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

[rocmlir-pr-reviewer]

Verdict: APPROVE -- submitted as COMMENT (automated reviews are advisory)  ·  New findings: 0 (0 Critical, 0 Major, 0 Minor)

---

Scope

Port of rocMLIR#2448: defines fully masked attention rows (all scores -inf) as producing a zero output vector instead of NaN, on both the GPU lowering (GridwiseAttnToBlockwise.cpp) and the TOSA CPU reference / split-KV host combine (rocmlir-gen.cpp), and makes the RMS verifier report zero error for exactly-matching all-zero tensors (conv-validation-wrappers.cpp). 13 test files updated or added, including a new GPU-vs-CPU E2E case with direct and split-KV pipelines.

Findings

No blocking issues found. All three findings from the previous review round are resolved on this revision:

• The mcpuVerifyFloat ABI mismatch in mlir/test/rocmlir-driver/verify-all-zero.mlir is fixed — the call and the func.func private declaration now both use the wrapper's real 7-operand signature, and the unused arith.constant false is gone.
• The message-less assert in selectZeroIfBothNegInf now carries "operands must share a type".
• The GPU/CPU asymmetry note was answered by the author, who is filing it as a separate codegen-improvement ticket for both rocMLIR and rocmlirTriton. Nothing further needed here.

Notes

Spot-checks that passed on this revision:

• The guards are identities on ordinary rows: any row with a finite score has m >= lowestFinite, so the max clamp is a no-op, and at least one exponential is exp2(0) = 1, so the sum clamp is a no-op. Existing numerics are unchanged.
• createConstantFloatOp is called with APFloat::opOK for -inf and 0.0; both convert exactly for f16/f32/f64, so the internal status assert holds.
• createAttentionRowStateCorrections consumes the guarded maxRowDiffExp, so the accumulator rescale stays consistent with the guarded row sum. createFirstGemmNegInfPadding means fully padded M rows also benefit.
• The CPU reference deliberately keeps the unclamped max and sum as LSE inputs, preserving -inf LSE for a fully masked row — matching the PR description.
• FileCheck variables referenced by the new CHECK lines (zero32, zero64, zeroRow, negInf) are all defined by CHECK-DAG in both lowering tests.
• rocMLIR back-port check: satisfied. The PR touches shared paths (mlir/lib/Dialect/Rock/, mlir/tools/rocmlir-gen/, mlir/lib/ExecutionEngine/) and the description links the parallel rocMLIR PR plus a "Port notes" section explaining the intentional divergence in the GPU lowering.
• No external/triton or external/llvm-project paths are touched, so the [EXTERNAL] commit-prefix and subtree-bump rules do not apply.

CI status

No failing or cancelled checks. Jenkins, ml-ci-internal.amd.com, Build and Test, and MIGraphX are still pending at review time — the GPU E2E claim in the test plan is worth confirming against the completed Jenkins run before merge.

[dhernandez0]

I worry this adds extra instructions for every attention kernel to cover a corner case. Can you try changing the initialization of initial value to -FLT_MAX (APFloat::getLargest(semantics, /*Negative=*/true)) instead of -inf? hopefully we will be able to get rid of some run-time checks then?

Also, if it's possible disable the checks (enable if slidingWindowSize > 0:)

[umangyadav]

I worry this adds extra instructions for every attention kernel to cover a corner case. Can you try changing the initialization of initial value to -FLT_MAX (APFloat::getLargest(semantics, /*Negative=*/true)) instead of -inf? hopefully we will be able to get rid of some run-time checks then?

Also, if it's possible disable the checks (enable if slidingWindowSize > 0:)

This PR is just port of missing codegen logic to handle fully masked rows from rocMLIR to rocmlirTriton and it is also adapting CPU reference to handle fully masked rows so that it doesn't lead to false failures.

I can do a follow up PR to improve codegen and explore different strategies there. it may require additional perf benchmarking and can lead to larger change.

[dhernandez0]

I worry this adds extra instructions for every attention kernel to cover a corner case. Can you try changing the initialization of initial value to -FLT_MAX (APFloat::getLargest(semantics, /*Negative=*/true)) instead of -inf? hopefully we will be able to get rid of some run-time checks then?
Also, if it's possible disable the checks (enable if slidingWindowSize > 0:)

This PR is just port of missing codegen logic to handle fully masked rows from rocMLIR to rocmlirTriton and it is also adapting CPU reference to handle fully masked rows so that it doesn't lead to false failures.

I can do a follow up PR to improve codegen and explore different strategies there. it may require additional perf benchmarking and can lead to larger change.

I understand this is a bugfix, I'm just saying it can be implemented more efficiently. I'm not sure if the extra instructions are irrelevant in practice or not. But at least, can you gate the checks if slidingWindowSize > 0? or get some perf numbers of this change to see what the impact is.

[justinrosner]

Is it possible to update this test to use a combination of valid and fully masked rows?

[justinrosner]

It might be good to have an additional data type as well (e.g., f16) just for completeness.

[justinrosner]

It looks like we are currently adding this logic in all cases, even when there is no masking/padding. Could we tighten the use of when these extra instructions are added?

[dhernandez0]

can you add mlir/test/rocmlir-gen/attention-splitkv-host-perrow-mask.mlir as well?

[umangyadav]

Superseded by #534