Skip to content

[Attention] Handle fully masked reference rows - #2448

Open
umangyadav wants to merge 4 commits into
developfrom
fix/fully-masked-attention
Open

[Attention] Handle fully masked reference rows#2448
umangyadav wants to merge 4 commits into
developfrom
fix/fully-masked-attention

Conversation

@umangyadav

@umangyadav umangyadav commented Aug 12, 2026

Copy link
Copy Markdown
Member

Motivation

attentionSweeps.py had been failing because causal sliding-window configurations can contain fully masked attention rows. In the investigated 1,000-case MFMA GPU-vs-CPU sweep, 113 configurations 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 rocMLIR's GPU attention kernel deliberately represents its contribution as the zero vector rather than as a probability distribution. Its online-softmax implementation explicitly replaces exp2(-inf - (-inf)) contributions with zero and selects zero for the final output when the accumulated 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 where the GPU produced zero, causing the attention sweep verification failures. The split-KV host combine had the same all--inf reduction problem. After both paths consistently produce zero, an entirely masked output can make the GPU and reference tensors identically zero; that also exposed the verifier's normalized-RMS 0/0 edge case addressed by this PR.

Summary

  • Define fully masked CPU-reference softmax rows and split-KV recombination as zero instead of allowing NaN from -inf - (-inf) and division by zero.
  • 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 both behaviors.

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. This avoids both undefined operations while giving the zero contribution produced by the GPU path. 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 produce a nonzero or NaN numerator and continue through the existing failure path.

Test plan

  • Run focused rocmlir-gen and verifier lit tests (13 passed).
  • Run 100 sampled MFMA attention configurations with GPU-vs-CPU verification (100 passed).
  • Verify fully masked f32, f16, bf16, and i8 cases, plus mixed valid/masked rows.
  • Re-run an exact previously failing f16 split_kv=64 sweep configuration (446208/446208 elements matched, RMS 0, [1 1 1]).

Made with Cursor

Keep CPU and split-KV softmax normalization defined when every logit is masked, and report zero normalized RMS for exact all-zero matches.

Co-authored-by: Cursor <cursoragent@cursor.com>
@umangyadav
umangyadav requested a review from causten as a code owner August 12, 2026 18:30
@umangyadav
umangyadav requested a balanced review from Copilot August 12, 2026 18:33
@umangyadav umangyadav self-assigned this Aug 12, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Handles fully masked attention rows without producing NaNs while preserving -inf LSE values.

Changes:

  • Adds guarded softmax and split-KV normalization.
  • Handles exact all-zero normalized-RMS comparisons.
  • Adds and updates regression tests.

Reviewed changes

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

Show a summary per file
File Description
mlir/tools/rocmlir-gen/rocmlir-gen.cpp Implements safe normalization for masked rows.
mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp Prevents normalized-RMS 0/0.
mlir/test/rocmlir-driver/verify-all-zero.mlir Tests all-zero verification.
mlir/test/rocmlir-gen/attention-sliding-window.mlir Tests fully masked softmax and LSE generation.
mlir/test/rocmlir-gen/attention-splitkv-host-f32-combine.mlir Tests safe split-KV recombination.
mlir/test/rocmlir-gen/attention-kernel.mlir Updates normalization checks.
mlir/test/rocmlir-gen/attention-kernel-gqa.mlir Updates GQA checks.
mlir/test/rocmlir-gen/attention-kernel-gqa-kvcache.mlir Updates GQA 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-attention checks.

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

@umangyadav umangyadav added the claude-review Trigger automated PR review by claude[bot]; auto-removed after the run. label Aug 12, 2026
@@ -0,0 +1,36 @@
// RUN: rocmlir-driver --host-pipeline=runner %s \

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This new test actually executes, but it has no feature gate. mlir/test/CMakeLists.txt:124-137 only adds mlir-runner, mlir_runner_utils, and conv-validation-wrappers to ROCMLIR_TEST_DEPENDS when MLIR_ENABLE_ROCM_RUNNER is set, and mlir/test/rocmlir-driver/lit.local.cfg has no directory-level guard (unlike mlir/test/e2e/lit.local.cfg, which sets config.unsupported). So in a build configured without the ROCm runner this test lands in the always-on check-rocmlir suite and fails on a missing tool / missing shared library rather than being skipped. Please add // REQUIRES: rocm-runner as the first line of the file (the feature is registered at mlir/test/lit.site.cfg.py.in:44-46). Note this is a pure host-side verifier test, so rocm-runner is the correct gate — no GPU or e2e gate is needed. Closest checklist bullet is the Minor "Lit test missing // RUN: line, -verify-diagnostics, or FileCheck prefix coverage", raised to Major here because the consequence is a broken test suite in a supported build configuration.

Suggested change
// RUN: rocmlir-driver --host-pipeline=runner %s \
// REQUIRES: rocm-runner

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolved -- addressed in this revision.

// CHECK_SCALE-DAG: %[[sqkTensorCast:.*]] = tosa.cast %[[sqkTensor]] : ([[squareShape]]) -> [[squareShape]]
// CHECK_SCALE-DAG: %[[sqkMaxs:.*]] = tosa.reduce_max %[[sqkTensorCast]] {{.*}} : ([[squareShape]]) -> [[reducedShape:tensor<.*>]]
// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %[[sqkMaxs]] : ([[squareShape]], [[reducedShape]]) -> [[squareShape]]
// CHECK_SCALE-DAG: %[[normilizedSqkTensor:.*]] = tosa.sub %[[sqkTensorCast]], %{{.*}} : ([[squareShape]], [[reducedShape]]) -> [[squareShape]]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Relaxing the second operand from %[[sqkMaxs]] to %{{.*}} (and likewise tosa.reciprocal %{{.*}} on line 32) drops the def-use check entirely: the test no longer verifies that the value subtracted from the scores derives from the tosa.reduce_max on line 28, only that some value of the right shape is subtracted. The same weakening appears in attention-kernel-causal.mlir, attention-kernel-f16.mlir, attention-kernel-gqa.mlir, attention-kernel-gqa-kvcache.mlir, attention-kernel-gqa-kvcache-lse.mlir, and attention-kernel-gqa-kvcache-lse-splitkv.mlir. Prefer the approach you already used in attention-splitkv-host-f32-combine.mlir and the new SAFE block of attention-sliding-window.mlir: insert a CHECK_SCALE-DAG line capturing the new clamp (e.g. %[[safeSqkMaxs:.*]] = tosa.maximum %[[sqkMaxs]], %{{.*}}) and keep the tosa.sub/tosa.reciprocal lines matching that capture. That keeps the chain asserted and also documents the new op in these tests. Checklist: Minor, "Lit test missing ... FileCheck prefix coverage".

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolved -- addressed in this revision.

Comment thread mlir/tools/rocmlir-gen/rocmlir-gen.cpp Outdated
static Value createFloatSplatTensor(OpBuilder builder, Location loc,
RankedTensorType type,
const APFloat &value) {
assert(isa<FloatType>(type.getElementType()));

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This precondition assert carries no message, so a failure prints only the expression. Since the helper is now called from four sites with different element types, a message makes the failure self-explanatory. Checklist: Major, "Missing assert with descriptive message on non-trivial preconditions" (treated as Minor here because the assert itself is carried over from the old applyMask body). The same applies to assert(status == APFloat::opOK); at line 2996.

Suggested change
assert(isa<FloatType>(type.getElementType()));
assert(isa<FloatType>(type.getElementType()) && "expected a float element type");

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolved -- addressed in this revision.

@rocmlir-pr-reviewer rocmlir-pr-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


Scope

Makes the rocmlir-gen TOSA CPU reference for attention treat fully masked rows (all -inf scores) the same way the GPU kernel does — as a zero output row — by clamping the softmax normalization max to the lowest finite value and the exponent sum to 1, in both createCpuAttentionKernelWithMlir and the split-KV host combine (computeFinalAttentionStage). Also makes mcpuVerify report zero normalized RMS instead of 0/0 for exactly-matching all-zero tensors. One new lit test, ten updated FileCheck tests, one new createFloatSplatTensor helper factored out of applyMask.

Findings

  • mlir/test/rocmlir-driver/verify-all-zero.mlir:1 (Major) — new execution test has no // REQUIRES: rocm-runner gate, so it runs in builds where mlir-runner / libconv-validation-wrappers were never built.
  • mlir/test/rocmlir-gen/attention-kernel.mlir:29 (Minor) — FileCheck operand matching weakened to %{{.*}} rather than checking the newly inserted tosa.maximum; same pattern in five sibling tests.
  • mlir/tools/rocmlir-gen/rocmlir-gen.cpp:2985 (Minor) — assert without a descriptive message.

Notes

  • The numerics in the description check out. maximum(max, lowestFinite) is an identity for any row containing a finite score (every finite value is >= the lowest finite value), and maximum(sum, 1) is an identity whenever the row max is attained (one term is exactly exp(0) == 1, and adding non-negative terms is monotone in round-to-nearest). Ordinary rows are therefore bit-identical to before. Feeding the unguarded qkMaxs / expsSums into the LSE computation (rocmlir-gen.cpp:4596-4609) correctly preserves -inf LSE for masked rows.
  • Guarding on sumDiffSq == 0.0 rather than maxMag == 0.0 (conv-validation-wrappers.cpp:218) is the right choice: a NaN/inf mismatch yields a NaN or non-zero numerator and still flows through the existing failure path. %printDebug = 0 in the new test maps to PrintOption::Off, and the hand-written @mcpuVerifyFloat declaration matches what rocmlir-gen emits.
  • No E2E coverage was added for the scenario that actually regressed (a causal + sliding-window config with fully masked rows, verified GPU-vs-CPU). The description reports a manual 100-config sweep; consider registering one such config under mlir/test/e2e/ so the regression is caught in CI rather than by the sweep script.
  • Nit, no change required: createFloatSplatTensor takes OpBuilder by value, matching the surrounding applyMask; OpBuilder & would avoid the copy if the file is ever cleaned up.

CI status

No failing or cancelled checks at the reviewed SHA. C/C++ premerge checks is still in progress; the auto-review job's own review check being in progress is expected.

@rocmlir-pr-reviewer rocmlir-pr-reviewer Bot removed the claude-review Trigger automated PR review by claude[bot]; auto-removed after the run. label Aug 12, 2026
Preserve the safe-softmax def-use checks, gate the verifier test on runner support, and document helper preconditions.

Co-authored-by: Cursor <cursoragent@cursor.com>
@umangyadav umangyadav added the claude-review Trigger automated PR review by claude[bot]; auto-removed after the run. label Aug 12, 2026
Exercise direct attention/LSE and split-KV recombination when a causal sliding-window row has no eligible keys.

Co-authored-by: Cursor <cursoragent@cursor.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

@rocmlir-pr-reviewer rocmlir-pr-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


Scope

Fixes NaN from the TOSA-based CPU attention reference when a query row is fully masked (all scores -inf), which causal + sliding-window + short current_seq_len configurations can produce. mlir/tools/rocmlir-gen/rocmlir-gen.cpp now clamps the softmax normalization max to the lowest finite value of the compute type and the exponential sum to 1.0 (both in createCpuAttentionKernelWithMlir and in the split-KV computeFinalAttentionStage recombination), while feeding the unclamped max/sum into the LSE computation so a masked row still yields -inf. mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp defines normalized RMS as 0 when the squared-error sum is exactly zero, so two identical all-zero tensors no longer evaluate 0/0. Seven rocmlir-gen lit tests are updated and one host-only verifier lit test is added.

Findings

No blocking issues found. All three findings from the previous review are addressed:

  • mlir/test/rocmlir-driver/verify-all-zero.mlir:1 now carries // REQUIRES: rocm-runner, matching the feature registered at mlir/test/lit.site.cfg.py.in:44-46 and the MLIR_ENABLE_ROCM_RUNNER-gated deps at mlir/test/CMakeLists.txt:124-137.
  • The tosa.sub / tosa.reciprocal def-use chains are re-asserted in all six attention-kernel-*.mlir files via new %[[safeSqkMaxs]] / %[[safeExpsSums]] captures, rather than the %{{.*}} wildcards of the earlier revision.
  • Both asserts in createFloatSplatTensor (rocmlir-gen.cpp:2985, :2997) now carry descriptive messages.

Notes

  • Spot-checked the numerics. Both clamps are identities on ordinary rows: a finite row max makes maximum(max, lowest) a no-op, and a row containing exp(max - max) = 1 sums to at least 1.0 in any IEEE format, so real GPU/reference divergence is not masked. tosa.maximum lowers to arith.maximumf, which propagates NaN, so a genuinely poisoned row still fails verification rather than being silently rescued.
  • The sumDiffSq == 0.0 guard is checked on the numerator, as the comment claims: a NaN mismatch leaves sumDiffSq as NaN (which is != 0.0), so it flows into the existing failure path and RMS_pass is still 0. maxMag-based guarding would not have had that property.
  • Confirmed the new lit test is lowerable end to end: mhal::buildRunnerPipeline (the --host-pipeline=runner implementation) runs convert-linalg-to-affine-loops before the arith/memref/func-to-LLVM conversions, and the default (non-bare-pointer) func lowering explodes memref<?xf32> into the descriptor fields that mcpuVerifyFloat expects. Without the fix, the test would print [0 1 1], so it is real regression coverage.
  • Follow-up, not blocking: cast<FloatType>(softmaxType) at rocmlir-gen.cpp:4583 (and cast<FloatType>(computeType) at :3372) will abort if -softmax_dtype is set to i8/i32, which typeFromString accepts. Before this change the same input produced a TOSA verifier diagnostic on tosa.exp. Nonsense input to a developer tool either way, but an explicit early check would give a better message.
  • Follow-up, not blocking: the added coverage is IR-shape (rocmlir-gen FileCheck) plus a host-only verifier test. A numerical regression test pinning a fully-masked configuration such as -current_seq_len=2 -sliding_window_size=1 --causal -seq_len_q 1 would guard the actual GPU-vs-CPU comparison that originally broke; worth considering for the e2e suite in a separate change.

CI status

No failing or cancelled checks. C/C++ premerge checks, Python format and lint checks, and Python performance script tests were still in progress at review time; clang-format and the lit suites are the ones to watch before merge.

@rocmlir-pr-reviewer rocmlir-pr-reviewer Bot removed the claude-review Trigger automated PR review by claude[bot]; auto-removed after the run. label Aug 12, 2026

@justinrosner justinrosner left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left some review comments on the rocmlirTriton version of this PR, but some of the same things apply here:

  • rock-attention-fully-masked.mlir:
    • Can we have a test with a combination of valid and fully masked rows?
    • Also a test with an additional datatype (e.g., f16) so that we can be sure that this is working for more than just f32

* [Rock] Guard inactive blockwise reduction threads

The NR-small LDS tree reduction derives per-thread LDS coordinates from
tid without checking that the thread owns a reduction slice. When
blockSize is not divisible by the product of the non-reduction dims, the
surplus workitems compute coordinates that either alias a valid thread's
slice or run past the logical workspace. The shortfall grows when
rthreads is further clamped to divide the reduction dimension, which is
the shape the attention configs hit. Because no barrier separates the
reduction load loop from the result store, an idle workitem in one wave
can double-count a partial that a valid workitem in another wave has
already written.

Guard both threadwise loops on tid < rthreads * nrDimProd, keeping the
barriers workgroup-uniform so the whole workgroup still reaches them.
The cross-lane DPP, permlane and ds_swizzle fast paths already require
exact thread packing, so the guard is only emitted for the LDS fallback.

Covered by two lit sections: one where blockSize leaves the remainder
idle, and one where the rthreads clamp does. Both fail without the
guard. Negative EXACT checks assert that exact-packing configs stay
guard-free.

Co-authored-by: Cursor <cursoragent@cursor.com>

* [Rock] Strengthen inactive reduction test coverage

Clarify the clamped-thread hazard and protect another exact-packing fallback from unnecessary guards.

Co-authored-by: Cursor <cursoragent@cursor.com>

* [Rock] Clarify inactive reduction test comment

Document that the barrier-placement check remains valid independently of guard canonicalization.

Co-authored-by: Cursor <cursoragent@cursor.com>

* [Rock] Harden inactive reduction thread guard

Make the DPP factoring dependency explicit, avoid unused guard IR, and keep unrelated builder behavior unchanged.

Co-authored-by: Cursor <cursoragent@cursor.com>

* [Rock] Generalize inactive reduction test comments

Keep the regression documentation focused on reusable test behavior rather than the motivating attention configurations.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
@codecov

codecov Bot commented Aug 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.96970% with 3 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
mlir/tools/rocmlir-gen/rocmlir-gen.cpp 95.45% 0 Missing and 2 partials ⚠️
...lect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp 98.18% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@             Coverage Diff             @@
##           develop    #2448      +/-   ##
===========================================
+ Coverage    82.57%   83.65%   +1.09%     
===========================================
  Files          120      121       +1     
  Lines        42852    43243     +391     
  Branches      7110     7185      +75     
===========================================
+ Hits         35381    36174     +793     
+ Misses        4815     4502     -313     
+ Partials      2656     2567      -89     
Flag Coverage Δ
gfx120x 83.57% <96.97%> (+1.05%) ⬆️
gfx950 83.47% <96.97%> (+1.12%) ⬆️

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

Files with missing lines Coverage Δ
...lect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp 90.38% <98.18%> (+0.29%) ⬆️
mlir/tools/rocmlir-gen/rocmlir-gen.cpp 87.17% <95.45%> (+0.62%) ⬆️

... and 46 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants