[Attention] Handle fully masked reference rows - #2448
Conversation
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>
There was a problem hiding this comment.
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.
| @@ -0,0 +1,36 @@ | |||
| // RUN: rocmlir-driver --host-pipeline=runner %s \ | |||
There was a problem hiding this comment.
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.
| // RUN: rocmlir-driver --host-pipeline=runner %s \ | |
| // REQUIRES: rocm-runner |
There was a problem hiding this comment.
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]] |
There was a problem hiding this comment.
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".
There was a problem hiding this comment.
Resolved -- addressed in this revision.
| static Value createFloatSplatTensor(OpBuilder builder, Location loc, | ||
| RankedTensorType type, | ||
| const APFloat &value) { | ||
| assert(isa<FloatType>(type.getElementType())); |
There was a problem hiding this comment.
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.
| assert(isa<FloatType>(type.getElementType())); | |
| assert(isa<FloatType>(type.getElementType()) && "expected a float element type"); |
There was a problem hiding this comment.
Resolved -- addressed in this revision.
There was a problem hiding this comment.
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-runnergate, so it runs in builds wheremlir-runner/libconv-validation-wrapperswere never built.mlir/test/rocmlir-gen/attention-kernel.mlir:29(Minor) — FileCheck operand matching weakened to%{{.*}}rather than checking the newly insertedtosa.maximum; same pattern in five sibling tests.mlir/tools/rocmlir-gen/rocmlir-gen.cpp:2985(Minor) —assertwithout 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), andmaximum(sum, 1)is an identity whenever the row max is attained (one term is exactlyexp(0) == 1, and adding non-negative terms is monotone in round-to-nearest). Ordinary rows are therefore bit-identical to before. Feeding the unguardedqkMaxs/expsSumsinto the LSE computation (rocmlir-gen.cpp:4596-4609) correctly preserves-infLSE for masked rows. - Guarding on
sumDiffSq == 0.0rather thanmaxMag == 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 = 0in the new test maps toPrintOption::Off, and the hand-written@mcpuVerifyFloatdeclaration 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:
createFloatSplatTensortakesOpBuilderby value, matching the surroundingapplyMask;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.
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>
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>
There was a problem hiding this comment.
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:1now carries// REQUIRES: rocm-runner, matching the feature registered atmlir/test/lit.site.cfg.py.in:44-46and theMLIR_ENABLE_ROCM_RUNNER-gated deps atmlir/test/CMakeLists.txt:124-137.- The
tosa.sub/tosa.reciprocaldef-use chains are re-asserted in all sixattention-kernel-*.mlirfiles 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 containingexp(max - max) = 1sums to at least1.0in any IEEE format, so real GPU/reference divergence is not masked.tosa.maximumlowers toarith.maximumf, which propagatesNaN, so a genuinely poisoned row still fails verification rather than being silently rescued. - The
sumDiffSq == 0.0guard is checked on the numerator, as the comment claims: aNaNmismatch leavessumDiffSqasNaN(which is!= 0.0), so it flows into the existing failure path andRMS_passis 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=runnerimplementation) runsconvert-linalg-to-affine-loopsbefore the arith/memref/func-to-LLVM conversions, and the default (non-bare-pointer) func lowering explodesmemref<?xf32>into the descriptor fields thatmcpuVerifyFloatexpects. Without the fix, the test would print[0 1 1], so it is real regression coverage. - Follow-up, not blocking:
cast<FloatType>(softmaxType)atrocmlir-gen.cpp:4583(andcast<FloatType>(computeType)at:3372) will abort if-softmax_dtypeis set toi8/i32, whichtypeFromStringaccepts. Before this change the same input produced a TOSA verifier diagnostic ontosa.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-genFileCheck) 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 1would 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.
justinrosner
left a comment
There was a problem hiding this comment.
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 Report❌ Patch coverage is 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
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Motivation
attentionSweeps.pyhad 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 isexp(-inf) = 0and the denominator is also zero, while the numerically stable formula computesm = max(-inf, ..., -inf) = -inffollowed 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 withNaN.The TOSA-based CPU reference used the stable softmax formula without equivalent degenerate-row handling. It therefore produced
NaNwhere the GPU produced zero, causing the attention sweep verification failures. The split-KV host combine had the same all--infreduction 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-RMS0/0edge case addressed by this PR.Summary
NaNfrom-inf - (-inf)and division by zero.-infLSE.Mathematical rationale
For attention scores
x_i, letm = max_i(x_i)Lbe the lowest finite value of the compute typem_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, som_safe = m. At least one score equalsm, hence one exponential isexp(0) = 1andA >= 1. Both guards are therefore identities and leave the existing softmax unchanged.For a fully masked row, every
x_i = -inf, som = -inf,m_safe = L, everye_i = exp(-inf - L) = 0, andA = 0. The guarded denominator is one and everyp_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, preservinglog(0) + (-inf) = -inffor 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 andM = 0, making the literal expression0/0. DefiningE = 0when the numerator is exactly zero preserves the fundamental property that identical tensors have zero error. The guard checks the numerator rather thanM: nonzero finite mismatches necessarily haveM > 0, while non-finite mismatches produce a nonzero orNaNnumerator and continue through the existing failure path.Test plan
split_kv=64sweep configuration (446208/446208 elements matched, RMS 0,[1 1 1]).Made with Cursor