Port sliding-window attention follow-up fixes

[umangyadav]

Note to reviewers :

Review commit by commit.

Summary

• preserve source mask semantics while reconciling sliding-window and KV-cache sequence operands, clip bounds, and nested broadcasts in the TosaToRock matcher
• reconstruct tuning-key problems with a full-cache current_seq_len default and document the inclusive position / W + 1 window semantics
• add active F16, BF16, F32, and I8 boundary coverage using rocMLIR-supported attention tuning configurations

Commit breakdown

1. Reconcile sliding-window clipping with the KV-cache matcher

   • Bug: clip bounds were detected separately from recovery of the underlying currentSeqLen, so the matcher could detach a clip from its operand or fold negative signed bounds into unsigned Rock masking.
   • Why: returning the unclipped input and validated bounds together gives both masks one consistent reconstruction path, while rejecting negative bounds in the callers prevents failed clip recognition from being mistaken for an absent clamp.

2. Fix sliding-window sequence-length reconstruction

   • Bug: MIGraphX can broadcast currentSeqLen through multiple multiply-by-one layers, while the matcher peeled only the outer layer.
   • Why: resolving all broadcast-only layers reaches the original block argument and lets lowering rebuild the correct head shape.

3. Preserve sliding-window-only mask semantics

   • Bug: folding a lower sliding-window mask without a separate KV-cache upper mask introduced masking for keys after currentSeqLen that did not exist in the source IR.
   • Why: requiring an independent KV-cache mask preserves semantics; otherwise the lower-bound select remains explicit in the elementwise region.

4. Default sliding-window tuning to full cache

   • Bug: current_seq_len is runtime-only and omitted from tuning keys, but rocmlir-gen required it when reconstructing a problem with sliding_window_size.
   • Why: synthesizing seq_len_k - 1 per group represents a full cache and makes serialized tuning problems reproducible without changing tuning identity.

5. Clarify attention window semantics

   • Bug: documentation described currentSeqLen like a length and slidingWindowSize like a count, leaving the inclusive endpoint ambiguous.
   • Why: defining currentSeqLen as the last valid zero-based position and W as look-back distance makes the effective W + 1 range explicit.

6. Add active sliding-window boundary coverage

   • Bug: existing cases did not activate the lower bound across an N-block boundary, so N-loop start regressions could escape coverage.
   • Why: random-data tests for F16, BF16, F32, and I8 place the lower bound at 95 inside the [64, 128) block and use supported attn:v2 configurations to reach GPU validation.

Behavioral note

A sliding-window lower mask without a separate KV-cache upper mask still folds into rock.attention, but the lower-bound select remains in the QK elementwise region instead of becoming a slidingWindowSize attribute. This conservative path preserves source semantics, though it may be less optimized for such models.

Compatibility and scope

This ports the applicable review follow-ups from rocmlirTriton PR #356. rocMLIR already samples valid sliding-window KV-cache configurations in its attention sweep tooling, so the equivalent rocmlirTriton-specific sweep changes are not ported.

Test plan

• Incremental build
• clang-format, clang-tidy, flake8, and YAPF
• Focused sliding-window lit tests: 7 passed
• Python performance tests: 136 passed
• F16/BF16/F32/I8 boundary GPU tests: 4 passed
• ninja -C build check-rock-e2e: 636 passed, 8 unsupported
• ninja -C build check-rocmlir: 1,652 passed, 35 unsupported, 4 expected failures

[Copilot]

Pull request overview

Ports sliding-window attention fixes across matching, tuning reconstruction, documentation, and boundary testing.

Changes:

• Reconciles KV-cache and sliding-window masks while preserving clip semantics.
• Defaults omitted runtime sequence positions to full-cache values.
• Adds inclusive-window documentation and multi-type GPU boundary coverage.

Reviewed changes

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

Show a summary per file

File	Description
mlir/utils/performance/perfRunner.py	Documents tuning reconstruction defaults.
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	Adds defaults, validation, and clarified options.
mlir/test/rocmlir-gen/problem-key.mlir	Tests tuning-key reconstruction.
mlir/test/rocmlir-gen/options.mlir	Updates option validation tests.
mlir/test/rocmlir-gen/attention-sliding-window.mlir	Tests full-cache defaults.
mlir/test/e2e/PrAttentionI8.toml	Adds I8 boundary coverage.
mlir/test/e2e/PrAttentionF32.toml	Adds F32 boundary coverage.
mlir/test/e2e/PrAttentionF16.toml	Adds F16 boundary coverage.
mlir/test/e2e/PrAttentionBF16.toml	Adds BF16 boundary coverage.
mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir	Verifies lower-only masks remain explicit.
mlir/lib/Conversion/TosaToRock/TosaToRock.cpp	Reworks mask, broadcast, and clip matching.
mlir/include/mlir/Dialect/Rock/IR/RockOps.td	Documents inclusive window semantics.

Suppressed comments (1)

mlir/lib/Conversion/TosaToRock/TosaToRock.cpp:2393

• The sliding-window path has the same invalid-clip fallthrough: a negative-bound clip makes tryClipPattern fail, but getValueSkipping below still skips minimum/maximum via seqLenSkip. If the KV-cache mask uses the same clip, both paths report the same unclipped block argument and empty bounds, allowing the masks to fold into unsigned Rock masking with changed semantics. Treat a recognized-but-unsupported clip as a match failure instead of continuing.

    auto maybeClip = tryClipPattern(seqLenCandidate);
    if (succeeded(maybeClip)) {
      seqLenCandidate = maybeClip->input;
      clipMin = maybeClip->clipMin;
      clipMax = maybeClip->clipMax;
    }

---

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

---

Scope

Ports six sliding-window attention follow-up fixes: reconciling clip-bound detection with the KV-cache matcher in TosaToRock.cpp, peeling nested currentSeqLen broadcasts, requiring an independent KV-cache mask before folding a sliding-window lower bound, defaulting current_seq_len to the full-cache position in rocmlir-gen, clarifying the inclusive [max(0, P-W), P] / W + 1 window semantics in RockOps.td, and adding F16/BF16/F32/I8 E2E boundary coverage.

Findings

One Major and four Minor. The Major (mlir/lib/Conversion/TosaToRock/TosaToRock.cpp:2253) is that the new negative-clip-bound rejection does not actually reject anything: because both callers treat tryClipPattern failure as "no clip present" and then resolve through seqLenSkip, which itself contains tosa::MaximumOp/tosa::MinimumOp, the clamp is skipped over and the mask still folds with the clamp silently discarded. The Minors are a duplicated broadcast-peel loop (TosaToRock.cpp:2282), a dead sequenceLengthK > 0 && sub-condition (mlir/tools/rocmlir-gen/rocmlir-gen.cpp:1365), and two comments that describe the wrong adjacent code (rocmlir-gen.cpp:1346, mlir/utils/performance/perfRunner.py:1751).

None of these are regressions relative to develop; the Major is a latent gap in newly added defensive code, so the verdict is COMMENT rather than REQUEST_CHANGES.

Notes

• Spot-checked the RockOps.td doc change against the matcher: greater(col, seqLen) and greater(seqLen - W, col) both fold to -inf, so the attended range really is the inclusive [max(0, P-W), P], i.e. W + 1 positions. The wording is correct.
• Requiring an independent KV-cache mask (TosaToRock.cpp:2683) also means a prefix-causal mask that co-occurs with a sliding-window mask but no KV-cache mask now fails to fold at all, rather than falling back to the old seqLen = slidingWindowSeqLen path. Returning failure() is the safe choice here since the peel loop has already advanced the input, so this looks intentional — just noting the coverage consequence.
• The new E2E configs check out: 159 - 64 = 95, which lands inside the [64, 128) N block for NPerBlock = 64, and the attn:v2 perf configs match ones already used elsewhere in the same TOML files.
• tosa-to-rock-attention-sliding-window-neg.mlir covers mismatched seq-len and mismatched (both non-negative) clip bounds, but has no case with a negative clip bound — see the Major finding.

CI status

No checks are in the fail or cancel bucket. Python performance script tests and Python format/lint checks pass; Jenkins, C/C++ premerge checks, Build and Test, MIGraphX, Code coverage, and the parameter sweeps are still pending at review time.

[rocmlir-pr-reviewer]

Dead store: seqLenCandidate is initialized from *maybeNonOne here and then unconditionally overwritten four lines below by seqLenCandidate = peelBroadcasts(seqLenCandidate);, so the initializer never reaches a read. Fold the two into one statement -- drop this line and write Value seqLenCandidate = peelBroadcasts(*maybeNonOne); at line 2287, keeping the explanatory comment directly above it. This also matches how trySlidingWindowPattern already does it at line 2384. (Checklist: Major / dead-code bullet -- flagged Minor because it is a single redundant assignment with no behavioral effect.)

[rocmlir-pr-reviewer]

Verdict: COMMENT  ·  New findings: 1 (0 Critical, 0 Major, 1 Minor)

---

Scope

Re-review of the sliding-window attention follow-up port at head 9a80cd9. Six commits touch the TOSA→Rock attention matcher (TosaToRock.cpp), the rock.attention op documentation (RockOps.td), rocmlir-gen tuning-key reconstruction, perfRunner.py config identity, and lit/E2E test coverage.

Findings

All five findings from the previous review are addressed at this revision:

• Tri-state clip handling: tryClipPattern now returns FailureOr<ClipResult> and both callers reject negative bounds with a hard failure() rather than silently dropping them; covered by the new sliding_window_kvcache_negative_clip case.
• Duplicated broadcast-peel loop: factored into peelBroadcasts and used from both tryKVCachePattern and trySlidingWindowPattern.
• Dead sub-condition sequenceLengthK > 0 && removed; the required-arg loop already rejects non-positive values.
• The KV-cache-position rationale comment moved onto the currentSeqLen.empty() default, with a separate rationale for the negativity check.
• The perfRunner.py comment no longer mislabels sliding_window_size as runtime-only.

One new Minor issue: a dead store at mlir/lib/Conversion/TosaToRock/TosaToRock.cpp:2283.

Notes

Spot-checks that came back clean:

• The W + 1 window documentation in RockOps.td matches the greater(P - W, col) lowering and the slidingWindowSize > seq_len_k rejection in rocmlir-gen.
• currentSeqLen and slidingWindowSize print before qk = elementwise in the op's assemblyFormat, so the new CHECK-NOT assertions in the negative lit tests genuinely constrain the attribute region.
• The default-current_seq_len lit test is self-consistent: -g 2 with seq_len_k 64 yields memref<2xi32> and arith.constant 63 : i32.
• Hardcoded -perf_config attn:v2:... in the E2E TOML files is an established pattern in these suites, not new.

The four-line negative-clip guard is now duplicated verbatim in both callers. That is below the threshold where a helper pays for itself, so no change requested.

Behavioral note worth recording: getSeqLenMask returning failure() for a sliding-window-only mask also forgoes folding an otherwise-recognizable prefixOffset from the same select chain. This is conservative and semantically correct, but the fusion-coverage cost is slightly wider than the sliding-window mask alone.

Pre-existing issues (out of scope)

tryClipPattern still conflates "no clip present" with "a clip is present but not in the recognized min(max(x, lo), hi) shape". Because seqLenSkip includes tosa::MaximumOp and tosa::MinimumOp, an unrecognized one-sided clamp such as tosa.minimum(seqLen, cap) is skipped through rather than rejected. The pre-PR code had the identical fall-through, so this is not introduced here — noting it only as a possible follow-up.

CI status

No checks in /tmp/pr/checks.json are in the fail or cancel bucket. "C/C++ premerge checks", "Python format and lint checks", and "Python performance script tests" are still in progress at review time.

[codecov]

Codecov Report

❌ Patch coverage is 71.79487% with 22 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mlir/lib/Conversion/TosaToRock/TosaToRock.cpp	70.67%	10 Missing and 12 partials ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2450      +/-   ##
===========================================
+ Coverage    82.57%   83.55%   +0.99%     
===========================================
  Files          120      121       +1     
  Lines        42852    43177     +325     
  Branches      7110     7185      +75     
===========================================
+ Hits         35381    36076     +695     
+ Misses        4815     4513     -302     
+ Partials      2656     2588      -68     

Flag	Coverage Δ
gfx120x	83.48% <71.79%> (+0.95%)	⬆️
gfx950	83.36% <71.79%> (+1.02%)	⬆️

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	87.07% <100.00%> (+0.52%)	⬆️
mlir/lib/Conversion/TosaToRock/TosaToRock.cpp	79.27% <70.67%> (+0.86%)	⬆️

... and 38 files with indirect coverage changes

🚀 New features to boost your workflow:

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