Handle fully masked attention rows in kernel lowering by umangyadav · Pull Request #534 · ROCm/rocmlirTriton

umangyadav · 2026-09-23T13:49:44Z

Stack (merge bottom-up):

[NFC] Apply clang-format to GridwiseAttnToBlockwise.cpp #537 — [NFC] clang-format of GridwiseAttnToBlockwise.cpp (merged)

Handle fully masked attention rows in kernel lowering #534 — kernel empty-row guards (base develop)

Recombine fully masked split-KV attention partials on the host #535 — host split-KV recombination (base Handle fully masked attention rows in kernel lowering #534)

Supersedes #416.

Summary

Define a fully masked attention row (every score is -inf) as a zero output row with LSE -inf instead of NaN, in both the kernel lowering and the rocmlir-gen CPU reference.
Emit the kernel guard only when a fully masked row is possible (see below). Other attention kernels are unchanged.
Rows that were finite before produce exactly the same values.

Cases that can produce fully masked rows

The kernel emits the guard only for these cases:

Causal + sliding window. With L = lastValidKVIndex, W = slidingWindowLookBack and prefix offset offset (0 for plain causal), query i keeps the keys max(0, L - W) <= j <= min(i + offset, L). That range is empty when i + offset < L - W. For example, with L = 2 and W = 1, query 0 may only see key 0, but the window starts at key 1 (the case in rock-attention-fully-masked.mlir).
Split-KV (splitKV > 1). Each split sees only a slice of the keys, and a slice can hold no valid key, for example when all of its keys are past lastValidKVIndex.
Pre-softmax fusion that could produce -inf, such as a fused bias or mask containing -inf, or scaling by a constant larger than 1 in magnitude, which can overflow a finite score (2 * -FLT_MAX = -inf). Only arith.extf and multiplication by constants with |c| <= 1 are treated as safe.

These cases never produce a fully masked row, so their kernels are unchanged:

Causal, prefix-causal and KV-cache masking without a sliding window never mask key 0, which is in the first tile the kernel processes.
A KV-cache sliding window without causal gives every query the same non-empty key range max(0, L - W) .. L.
Padded query rows are fully masked, but they are dropped when the output is stored, so their values are never visible.

Mathematical rationale

For the scores x_i of one query row, let

m = max_i(x_i)
L be the lowest finite value of the softmax type (-FLT_MAX for f32)
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 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, so every p_i = 0 and the output row is zero. LSE uses the unguarded sum A, so it stays -inf because log(0) = -inf.

The kernel computes the same thing tile by tile: its running max now starts at L instead of -inf, so a fully masked tile adds 0 instead of NaN.

rocMLIR back-port

Not needed. GridwiseAttnToBlockwise.cpp has no rocMLIR counterpart (rocMLIR lowers attention in GridwiseGemmToBlockwise.cpp), and the rocmlir-gen CPU-reference change only matters alongside this kernel guard.

Test plan

Rock dialect lowering tests (toblockwise_attention_lowering*.mlir, prefix-causal, sliding-window, KV-cache, pre-softmax)
Direct-path rock-attention-fully-masked.mlir (f32 and f16, fully masked and mixed rows, prefix-causal)
rocmlir-gen attention-kernel FileCheck tests
fusion/pr-e2e/attention and generated PrAttention E2E tests

Locally on gfx950: 840/840 Rock-dialect and rocmlir-gen lit tests, 74/74 fusion/pr-e2e/attention tests and 1042/1042 PrAttention E2E tests pass.

Made with Cursor

Pure clang-format run over the existing file. No functional change; the diff is whitespace and comment rewrapping only. Split out of the fully masked attention row work so that change reviews as functional-only. Co-authored-by: Cursor <cursoragent@cursor.com>

Gate a finite running-max sentinel and denominator clamp on compile-time empty-row risk, look through arith.extf scaling, and keep the CPU softmax reference defined for all-masked rows. Co-authored-by: Cursor <cursoragent@cursor.com>

Copilot

Copilot review overview

🟡 Changes recommended

The critical multiplication-overflow case must remain conservatively guarded.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1

Open (1)

Finite multiplier can overflow scores and produce NaN softmax · New

What changed in this PR

Adds handling for fully masked attention rows in CPU references and Rock kernel lowering.

Changes:

Adds finite max sentinels and denominator clamps.
Introduces compile-time empty-row risk detection.
Expands lowering, generator, and end-to-end tests.

File	Review
`mlir/tools/rocmlir-gen/rocmlir-gen.cpp`	Stabilizes CPU-reference softmax.
`mlir/test/rocmlir-gen/attention-sliding-window.mlir`	Tests masked-row reference IR.
`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 checks.
`mlir/test/rocmlir-gen/attention-kernel-f16.mlir`	Covers f16 guards.
`mlir/test/rocmlir-gen/attention-kernel-causal.mlir`	Updates causal checks.
`mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir`	Adds fully masked end-to-end coverage.
`mlir/test/Dialect/Rock/toblockwise_attention_lowering.mlir`	Tests risk classification and sentinels.
`mlir/test/Dialect/Rock/toblockwise_attention_lowering_nPerBlockG1.mlir`	Tests chunked normalization.
`mlir/test/Dialect/Rock/gridwise-attention-sliding-window.mlir`	Covers sliding-window guards.
`mlir/test/Dialect/Rock/gridwise-attention-prefix-causal.mlir`	Checks prefix-causal lowering.
`mlir/test/Dialect/Rock/gridwise-attention-kvcache-clamp.mlir`	Checks KV-cache lowering.
`mlir/test/Dialect/Rock/attn_to_blockwise_presoftmax.mlir`	Tests pre-softmax risk detection.
`mlir/lib/Dialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp`	Critical: finite constant multiplication can overflow to infinity, incorrectly bypassing guards and allowing NaN softmax results.

💡 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: 4 (0 Critical, 1 Major, 3 Minor)

Scope

Fixes the fully-masked-attention-row NaN in -rock-gridwise-attn-to-blockwise: when every score in a row is -inf, the running max stays -inf and maxRowDiff = -inf - (-inf) becomes NaN, poisoning exp2 and the accumulator. The PR gates two changes on a new compile-time predicate mayHaveFullyMaskedRows — a lowest-finite sentinel (APFloat::getLargest(..., Negative=true)) for the initial running max, and an arith.maxnumf(sumRow, 1.0) clamp on the normalization denominator only. The CPU reference in rocmlir-gen.cpp gets the matching tosa.maximum clamps, plus a new direct-path E2E test. 17 files: 1 transform, 1 tool, 14 lit tests, 1 new E2E test.

Findings

The numerics look right to me. max(sumRow, 1.0) is provably a no-op for any row with at least one unmasked score (that score contributes exp2(m - m) = 1), the LSE path correctly keeps the unclamped sumRow/maxRow so a fully masked row still reports -inf, and -FLT_MAX - m overflowing to -inf still gives the correct exp2 -> 0. isFiniteConstantScaleOf recursion is linear and acyclic.

One Major (process, not code): this PR touches two files on the rocMLIR shared-path list — mlir/lib/Dialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp and mlir/tools/rocmlir-gen/rocmlir-gen.cpp — and the description carries no back-port note. Plus three Minors on an undocumented invariant, a defensive gap, and a test-coverage regression. Details inline.

Notes

toblockwise_attention_lowering_nPerBlockG1.mlir mutates the pre-existing gridwise_attn_nperblockg1 onto the guarded path (adds causal + slidingWindowLookBack), so the detailed chunked-GEMM join/trans/reshape CHECKs no longer exercise the plain path; the new gridwise_attn_nperblockg1_plain is covered only by the weak PLAIN prefix. Worth considering duplicating the detailed block instead.
The file-global COUNT-COUNT-1: / COUNT-NOT: idiom in that same file is valid FileCheck but takes a moment to parse; a short comment explaining "exactly one guarded function in this file" would help.
New E2E test rock-attention-fully-masked.mlir carries the correct first-party license header ahead of the RUN lines.
No external/triton or external/llvm-project changes, so the [EXTERNAL] commit and patch-record rules do not apply here.

CI status

No failing or cancelled checks. py-checks and detect pass; copilot-pull-request-reviewer and this pipeline's own review check are still in progress. Note that no build or lit-test signal has reported yet, so the FileCheck updates are unverified by CI at review time.

codecov-commenter · 2026-09-24T00:44:30Z

Codecov Report

❌ Patch coverage is 89.70588% with 14 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp	86.60%	5 Missing and 8 partials ⚠️
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	97.44%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop     #534   +/-   ##
==========================================
  Coverage           ?   87.02%           
==========================================
  Files              ?      208           
  Lines              ?    39867           
  Branches           ?     8688           
==========================================
  Hits               ?    34693           
  Misses             ?     3082           
  Partials           ?     2092

Flag	Coverage Δ
gfx950	`84.57% <86.41%> (?)`
navi4x	`84.49% <86.41%> (?)`

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	`92.94% <97.44%> (ø)`
...ialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp	`89.56% <86.60%> (ø)`

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Treat compile-time scales as safe only when |c| <= 1, drop prePadG0M from the gate, and keep prefix-causal padding coverage without emitting the denominator clamp. Co-authored-by: Cursor <cursoragent@cursor.com>

Copilot

Copilot review overview

🟡 Changes recommended

The guard predicate overlooks overflow from the mandatory base-2 conversion, leaving some finite rows vulnerable to NaN.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1

Open (1)

Account for log2(e) conversion to prevent finite-row overflow and NaNs · New

Resolved since last review (1)

Finite multiplier can overflow scores and produce NaN softmax

rocmlir-pr-reviewer

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

Scope

Makes fully masked attention rows well-defined (output 0, LSE -inf) instead of NaN, in both GridwiseAttnToBlockwise.cpp (kernel lowering) and the rocmlir-gen CPU reference. The guard is two clamps — a finite row-max sentinel (-FLT_MAX/-65504) replacing the -inf accumulator init, and max(rowSum, 1) applied only to the output denominator so LSE keeps the raw sum. Both are provable no-ops for any row with at least one valid key. 16 test files updated/added, plus an NFC clang-format reflow of the pass.

Findings

No blocking issues found. All four findings from the previous review are fixed at 0c0c7b4a:

back-port justification is now in the PR description (GridwiseAttnToBlockwise.cpp has no rocMLIR counterpart; the rocmlir-gen change is only meaningful with it), and the author confirmed in-thread;
the prefixOffset >= 0 / unsigned-ugt invariant is documented on both Rock_AttentionOp and Rock_GridwiseAttentionOp in RockOps.td and restated in the mayHaveFullyMaskedRows comment;
preSoftmaxMayFullyMask now checks yield.getNumOperands() != 1 before indexing;
prePadG0M = 16 : index is restored in gridwise-attention-prefix-causal.mlir, with the unpadded case split into its own -split-input-file module so both shapes carry the CHECK-NOT.

Notes

The Copilot finding (a finite splat does not imply mulf stays finite) is also addressed: isNonOverflowingFloatScale now requires |scale| <= 1, and @gridwise_attn_overflowing_scale pins the -FLT_MAX case to the guarded path.
Spot-checked the two safety claims behind the narrowed predicate. Sliding-window-without-causal is safe because slidingWindowLowerBound = max(lastValidKVIndex - lookBack, 0) is always <= lastValidKVIndex and the N-loop start is advanced to the tile holding it, so the first processed tile always has a valid key. Dropping the prePadG0M clause is safe because rock.blockwise_reduce reduces per row along axis 1, so a padded row cannot contaminate a real one, and padded rows are not stored.
toblockwise_attention_lowering_nPerBlockG1.mlir uses a second COUNT prefix to assert exactly one arith.maxnumf file-wide — a good negative-coverage pattern worth reusing.
No external/ paths touched, so the [EXTERNAL] commit-subject rule does not apply. New test file carries the correct AMD/SPDX header.

CI status

No failed or cancelled checks. py-checks and detect pass; the two in-progress entries are the auto-review pipeline itself and the Copilot reviewer.

Copilot

Copilot review overview

🟢 Approval recommended

The guarded normalization is internally consistent and has broad lowering and end-to-end regression coverage.

Review effort: Balanced
Findings: None

Resolved since last review (1)

Account for log2(e) conversion to prevent finite-row overflow and NaNs

dhernandez0 · 2026-09-25T14:39:42Z

    If causal is enabled AND prefixOffset is not null, we implement prefix
-    causal masking where: mask = key_pos > (query_pos + offset).
+    causal masking where: mask = key_pos > (query_pos + offset). The offset is
+    expected to be non-negative; the comparison is unsigned (`ugt`), so key 0


do we have verifiers to check offset>=0?

It is runtime value

thanks! when we can write this assumption with llvm.assume offset>=0 I think triton uses it for dynamic shapes. It's out of scope but it'd be nice if you add a TODO.

dhernandez0 · 2026-09-25T14:45:11Z

+    bool mayHaveFullyMaskedRows =
+        op.getEnableSoftmax() &&
+        (preSoftmaxMayFullyMask(op) ||
+         (isCausal && slidingWindowLookBack > 0) || splitKV > 1);


can this happen with splitKV > 1 alone? I think if a workgroup is fully masked out, it should early exit. So, it never gets to compute anything?

also what about combinations of lastValidKVIndex + others?

can this happen with splitKV > 1 alone?

Good question. There is an early exit for split-KV, but it works per workgroup: it only skips a split when that split has no KV tiles at all (end <= start).
The issue is that masking is per row. So it is still possible that some rows in workgroup are full masked.

It can't happen just with splitKV > 1 alone. So there is possibility of improvement. But for now i think i can just add a TODO

dhernandez0 · 2026-09-25T14:46:08Z

+          rewriter, loc, blockMTensorType,
+          SplatElementsAttr::get(
+              blockMTensorType,
+              rewriter.getFloatAttr(elemTypeSoftmax, lowestFinite)));


nit: I wonder if we can just do this unconditionally. Does it have any perf impact?

We can do it unconditonally, it only adds a arith::maxnumf. It will have perf impact but not much i think. I find it better to leave it as a follow up investigation

dhernandez0 · 2026-09-25T14:47:30Z

+            createConstantFloatOp(rewriter, loc, blockMTensorType,
+                                  elemTypeSoftmax, 1.0, APFloat::opOK);
+        normalizationSum =
+            arith::MaxNumFOp::create(rewriter, loc, sumRow, oneFloat);


can you add a comment here? I don't understand this

also, we added no-NaN assumption recently. Does this PR break it? do we have NaNs at any point?

ah I see, we are doing exp(-big) so we get all zeros and the addition of all is still zero?

ah I see, we are doing exp(-big) so we get all zeros and the addition of all is still zero?

exactly, so it does max(exp(-big), 1) to avoid division by zero.

also, we added no-NaN assumption recently. Does this PR break it? do we have NaNs at any point?

With this PR change, it replaces previous NaNs with zeros. So it is improving on that assumption.

dhernandez0 · 2026-09-25T14:58:05Z

+    // through padded output/LSE views, so a NaN there is not observable and is
+    // not a reason to emit the guard. Arbitrary pre-softmax fusion is guarded
+    // unless its result is proven not to overflow finite QK scores.
+    bool mayHaveFullyMaskedRows =


This is out of scope for this PR. But we can add a TODO to add more conditions to early exit: There are cases that can't be known at compile time: preSoftmaxMayFullyMask and/or lastValidKVIndex. But in some cases (isCausal && slidingWindowLookBack) it might be known at compile time that nothing is to be computed?

yeah it can be candidate for early exit. I can add a TODO

…n-empty-row-guard

umangyadav requested review from causten, dhernandez0, justinrosner and pabloantoniom as code owners September 23, 2026 13:49

umangyadav marked this pull request as draft September 23, 2026 13:56

umangyadav and others added 2 commits September 23, 2026 22:38

umangyadav force-pushed the users/umayadav/attn-empty-row-guard branch from 0b72017 to 29c7a50 Compare September 23, 2026 22:43

umangyadav mentioned this pull request Sep 23, 2026

[NFC] Apply clang-format to GridwiseAttnToBlockwise.cpp #537

Merged

umangyadav changed the base branch from develop to users/umayadav/attn-format-nfc September 23, 2026 22:43

umangyadav mentioned this pull request Sep 23, 2026

Recombine fully masked split-KV attention partials on the host #535

Merged

umangyadav requested a balanced review from Copilot September 23, 2026 22:46

umangyadav self-assigned this Sep 23, 2026

umangyadav added the claude-review Request a Claude PR review label Sep 23, 2026

Copilot started reviewing on behalf of umangyadav September 23, 2026 22:47

View session

Copilot AI reviewed Sep 23, 2026

View reviewed changes

Comment thread mlir/lib/Dialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp Outdated