Recombine fully masked split-KV attention partials on the host

[umangyadav]

Summary

With split-KV (split_kv > 1), the attention kernel writes one partial output O_s and one log-sum-exp L_s per split, and the host wrapper generated by rocmlir-gen merges them into the final output. This PR makes that host merge match the kernel from #534:

• Causal split-KV now works when seq_len_q is larger than the key tile. rocmlir-gen used to exit with "Causal masking + split-KV is not supported with sequenceLengthQ > nPerBlock".
• Valid-split counts are computed per query row under causal and prefix-causal masking. Before, every causal row got a count of 1 (hence the limit above), and every prefix-causal row got row 0's count, which could drop splits holding later rows' keys when seq_len_q > 1.
• The merge is NaN-safe for rows where every split is empty, and runs in f32 for f16/bf16 storage.

Most of this ports host-side changes from ROCm/rocMLIR#2371 and ROCm/rocMLIR#2448 (details under "Relation to rocMLIR").

Technical details

All code changes are in mlir/tools/rocmlir-gen/rocmlir-gen.cpp.

Per-row valid-split counts (computeValidSplitKV)

Under causal or prefix-causal masking, the number of splits holding a row's visible keys depends on the row, so the host emits one count per (batch-head, query row). Non-causal configs keep one count per batch-head. For row q:

• rowLast = clamp(q + prefix_offset, 0, lastValid), where prefix_offset is 0 for plain causal and lastValid is last_valid_kv_index (default seq_len_k - 1)
• keysPerSplit = nPerBlock * ceil(ceil((rowLast + 1) / nPerBlock) / split_kv)
• count = ceil((rowLast + 1) / keysPerSplit)

nPerBlock is GEMM0's N tile, which is the key tile in this lowering: the Triton attention lowering keeps Q * K^T, so keys are GEMM0's N axis. rocMLIR computes the transposed product and uses MPerBlock instead.

The counts are conservative. The kernel partitions each M block's keys using the block's last row, while the host uses each row on its own, so the host can include extra empty splits but never drops a split that holds the row's data. The extra splits are harmless because #534 makes the kernel write zero output and -inf LSE for empty splits.

Per-row mask layout (createMaskSplitKV)

The counts are laid out as [batchHeads, 1, 1, 1] (per batch-head) or [batchHeads, 1, seq_len_q, 1] (per row) and compared against a range over the split axis. Split s is valid for a row when s < count. Invalid splits get LSE -inf and contribute zero to the weighted sum.

NaN-safe merge (computeFinalAttentionStage)

m   = max(max_s L_s, lowest finite)
w_s = exp(L_s - m)
O   = sum_s(w_s * O_s) / max(sum_s w_s, 1)

If a row has at least one finite L_s, both clamps are no-ops. If every split is -inf, all weights are 0 and the output row is 0, matching the kernel's result for a fully masked row. Without the clamps, -inf - (-inf) turned the whole row into NaN.

f32 merge for f16/bf16

The partial outputs and LSE are cast to f32 before the merge, and only the final output is cast back to the storage type.

Causal limitation removed

The old causal count assumed every row's visible keys fit in split 0, which only holds while seq_len_q <= nPerBlock, hence the exit. With per-row counts, the exit is removed.

Relation to rocMLIR

• Ported from [AIROCMLIR-707] Fix split-kv attention masking and sweep RMS for attention configs rocMLIR#2371: per-row valid-split counts, the per-row mask layout, the f32 merge for f16/bf16, and the attention-splitkv-host-perrow-mask.mlir and attention-splitkv-host-f32-combine.mlir tests. The kernel and sweep-script changes in [AIROCMLIR-707] Fix split-kv attention masking and sweep RMS for attention configs rocMLIR#2371 are outside this PR; the Triton lowering's kernel handling of empty splits is in Handle fully masked attention rows in kernel lowering #534.
• Ported from [Attention] Handle fully masked reference rows rocMLIR#2448: the two clamps in the merge.
• Specific to rocmlirTriton: the key tile is GEMM0's NPerBlock rather than MPerBlock, and the causal seq_len_q > nPerBlock exit is removed. rocMLIR develop still has that exit; removing it relies on the empty-split guarantees Handle fully masked attention rows in kernel lowering #534 adds to the Triton lowering.
• No rocMLIR back-port needed: the ported code is already in rocMLIR develop, and removing the exit depends on Handle fully masked attention rows in kernel lowering #534's Triton-lowering changes.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 95.23810% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	95.24%	0 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop     #535   +/-   ##
==========================================
  Coverage           ?   84.72%           
==========================================
  Files              ?      208           
  Lines              ?    36300           
  Branches           ?     8179           
==========================================
  Hits               ?    30753           
  Misses             ?     3323           
  Partials           ?     2224           

Flag	Coverage Δ
gfx950	84.61% <95.24%> (?)
mfma	84.50% <95.24%> (?)
navi4x	84.53% <95.24%> (?)

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

Files with missing lines	Coverage Δ
...ialect/Rock/Transforms/GridwiseAttnToBlockwise.cpp	86.29% <ø> (ø)
mlir/tools/rocmlir-gen/rocmlir-gen.cpp	87.76% <95.24%> (ø)

🚀 New features to boost your workflow:

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

[Copilot]

Copilot review overview

🟢 Approval recommended

The recombination logic is consistent with kernel partitioning and includes targeted regression coverage.

Review effort: Balanced
Findings: None

What changed in this PR

Updates host-side split-KV attention recombination to correctly handle fully masked rows and causal multi-tile queries.

Changes:

• Adds per-row split validity masks for causal attention.
• Performs narrow-type recombination in f32 with safe normalization.
• Adds IR and end-to-end regression coverage.

File	Description
mlir/​tools/​rocmlir-gen/​rocmlir-gen.cpp	Implements robust host recombination and per-row masks.
mlir/​test/​rocmlir-gen/​attention-splitkv-host-perrow-mask.mlir	Tests causal per-row validity layout.
mlir/​test/​rocmlir-gen/​attention-splitkv-host-f32-combine.mlir	Tests f32 recombination for f16/bf16 storage.
mlir/​test/​rocmlir-gen/​attention-splitkv-host-combine.mlir	Tests safe all-masked normalization.
mlir/​test/​fusion/​pr-e2e/​attention/​rock-attention-fully-masked.mlir	Adds split-KV end-to-end coverage.

---

💡 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, 0 Major, 3 Minor)

---

Scope

Host-side (rocmlir-gen) split-KV attention recombination. The PR removes the causalMasking && sequenceLengthQ > nPerBlock bail-out, extends computeValidSplitKV() to emit per-(batch-head, query-row) validity counts under causal / prefix-causal masking, teaches createMaskSplitKV() the [batch, 1, seqQ, 1] constant layout, guards the LSE re-normalization against an all--inf split column (clamp the reduce_max to the lowest finite value, clamp the exp-sum to 1.0), and runs the combine in f32 for f16/bf16 storage. Three new mlir/test/rocmlir-gen/ FileCheck tests plus four split-KV RUN lines in rock-attention-fully-masked.mlir.

Findings

Three Minor findings, none blocking:

• mlir/tools/rocmlir-gen/rocmlir-gen.cpp:3566 — the f32-combine guard is a two-type allow-list rather than a width test, so other sub-f32 float storage types silently combine at storage precision.
• mlir/tools/rocmlir-gen/rocmlir-gen.cpp:1471 — the numPerBlock / itersPerBlock / numValidKV computation is duplicated verbatim in both branches of computeValidSplitKV().
• mlir/test/rocmlir-gen/attention-splitkv-host-perrow-mask.mlir:30 — the per-row validity counts, the thing this PR actually computes, are wildcarded in the tosa.const CHECK.

Notes

• I worked through the safety argument in the new computeValidSplitKV() comment and it holds. For a row q inside an M block whose last row is q_last >= q, itersPerBlock_host <= itersPerBlock_gpu, so numValidKV_host = ceil((q+1)/itersPerBlock_host) >= ceil((q+1)/itersPerBlock_gpu), i.e. the host can only over-count. Over-counted splits carry the kernel's defined 0 / -inf contribution from the stacked kernel PR and drop out of the weighted sum. numValidKV is also bounded by splitKV and is always >= 1, so the new assert(nPerBlock > 0) is the only precondition that needed stating.
• Both new clamps are no-ops outside the degenerate case: reduce_max is only -inf when every split LSE is -inf, and sum exp(lse - max) is >= 1 whenever any split is finite. Good.
• Decode-path behaviour is preserved: for seqLenQ = 1, q = 0, the per-row branch reproduces the old effectiveLastValidKVIndex = 0 result exactly.
• License headers on all three new .mlir files match the first-party rocmlirTriton template.
• rocMLIR back-port: mlir/tools/rocmlir-gen/ is on the shared-with-rocMLIR path list. The description carries "no rocMLIR backport necessary", which satisfies option (b), but a one-line reason (e.g. that it depends on the Triton-only kernel guard in the stacked PR) would make the note auditable later.
• The PR description notes this is stacked and should be retargeted to develop after the base merges; the diff here is correctly limited to the single split-KV commit.

CI status

No failed or cancelled checks. Jenkins (Build and Test, MIGraphX, Code coverage) is still pending at the reviewed SHA, so the new E2E RUN lines are unverified as of this review.

[rocmlir-pr-reviewer]

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

---

Scope

Removes the rocmlir-gen restriction that rejected causal masking combined with split-KV when sequenceLengthQ > nPerBlock, and reworks the host-side split-KV recombination to match. computeValidSplitKV() now emits per-(batch-head, query-row) validity counts under causal / prefix-causal masking, createMaskSplitKV() grows a [batch, 1, seqQ, 1] constant layout to broadcast them, and computeFinalAttentionStage() upcasts f16/bf16 partials to f32 for the combine while clamping the LSE max to the lowest finite value and the exp-sum to 1 so an all-masked reduction stays defined. Three new rocmlir-gen lit tests plus four new split-KV RUN lines in rock-attention-fully-masked.mlir cover it.

Findings

Two Minor findings, both about defensive coverage rather than correctness:

• mlir/test/rocmlir-gen/attention-splitkv-host-perrow-mask.mlir:18 — the newly-enabled seqLenQ > nPerBlock causal configuration has no rocmlir-gen lit coverage.
• mlir/tools/rocmlir-gen/rocmlir-gen.cpp:3514 — the per-row constant layout hardcodes the split axis as dim 1 while index is a parameter.

All three findings from the previous review round are addressed on this revision: the duplicated numPerBlock/itersPerBlock/numValidKV chain is now the single numValidSplits lambda at rocmlir-gen.cpp:1449, the per-row tosa.const CHECKs pin dense<1> instead of wildcarding values, and the new STRADDLE RUN line adds a case where the counts genuinely differ across rows ([[2], [2], [3], [3]], which I re-derived by hand and confirms the min(q + prefixOffset, lastValidIndex) chain). The narrow-storage-type discussion is closed by the author's rationale that rock.attention restricts result/LSE to F32/F16/BF16.

Notes

• I checked the conservativeness claim in the new computeValidSplitKV() comment. Because the per-row itersPerBlock is derived from a lastValid no larger than the block's, its denominator is never larger than the GPU's, so the host count is always at least the number of splits holding real data for that row — the comment's "cannot exclude a split containing valid row data" direction holds. The count also can never exceed splitKV.
• tosa.maximum(sumExpNorm, 1.0) is a no-op whenever at least one split is valid, since the max-normalized term contributes exp(0) = 1; it only bites in the all-masked case. The subsequent applyMask(outExp, ..., 0.0f) clears the 0 * -inf NaN that the -inf result mask would otherwise produce.
• mlir/tools/rocmlir-gen/ is on the checklist's shared-with-rocMLIR path list; the PR description's "no rocMLIR backport necessary" note satisfies option (b), though a one-line reason would make it auditable.

CI status

No failing checks. Jenkins (Build and Test, MIGraphX, Code coverage) is still pending at the reviewed SHA, so the new E2E RUN lines in rock-attention-fully-masked.mlir — including MULTIBLOCK-SPLIT, the only coverage of the newly-enabled regime — have not reported yet.

[rocmlir-pr-reviewer]

This header comment still describes only the original configuration ("causal, seqLenQ = 4, splitKV = 4" plus the f16 f32-combine check), but the file now carries four RUN lines: CHECK, PREFIX, STRADDLE (prefix_offset=62, where counts genuinely differ across rows), and the new MULTITILE (seq_len_q 96 against nPerBlockG0=32). MULTITILE is the regime the removed causal exit used to reject, so it is the PR's headline behavior change and a reader scanning the comment would not know it is covered here. Please extend the comment to name the four configurations and what each pins. Checklist Minor: "Lit test missing // RUN: line, -verify-diagnostics, or FileCheck prefix coverage."

[rocmlir-pr-reviewer]

Resolved -- addressed in this revision.

[rocmlir-pr-reviewer]

The four new split-KV RUN lines each pin a single unanchored [1 1 1], while the pre-existing DIRECT / MIXED-DIRECT / F16-DIRECT / PREFIX-DIRECT lines each pin two ([1 1 1] followed by a -NEXT: [1 1 1]). Because a non-NEXT FileCheck directive scans forward to the first match, a single unanchored [1 1 1] can be satisfied by a different verification result than the intended one, so a regression in the split-KV output verification could be masked by another passing result line. MULTIBLOCK-SPLIT here is the weakest of the four since it pins nothing else. Please either pin the full expected set of result lines (with -NEXT anchoring, as the DIRECT prefixes do) or add a short comment stating why only one result is checked for the split-KV path. Checklist Minor: "Lit test missing // RUN: line, -verify-diagnostics, or FileCheck prefix coverage."

[rocmlir-pr-reviewer]

Resolved -- addressed in this revision.

[rocmlir-pr-reviewer]

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

---

Scope

Host-side recombination of split-KV attention partials in mlir/tools/rocmlir-gen/rocmlir-gen.cpp: per-row valid-split counts under causal / prefix-causal masking (computeValidSplitKV), a per-row mask layout (createMaskSplitKV), a NaN-safe merge that runs in f32 for f16/bf16 storage (computeFinalAttentionStage), and removal of the seq_len_q > nPerBlock causal exit. Test changes add three new rocmlir-gen lit tests and four split-KV RUN lines to the fully-masked E2E test.

Findings

Two Minor, both in tests:

• mlir/test/rocmlir-gen/attention-splitkv-host-perrow-mask.mlir:13 — header comment still describes only the original single configuration although the file now has four RUN lines.
• mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir:52 — new split-KV RUN lines each pin a single unanchored result line, unlike the pre-existing DIRECT lines which pin two.

No Critical or Major findings remain, hence COMMENT.

Notes

• Both previously-flagged items are addressed in commit de63ba40: the new MULTITILE RUN line exercises seq_len_q (96) greater than the key tile (nPerBlockG0=32), and createMaskSplitKV now asserts the split axis is dim 1 on the per-row path.
• Hand-checked the expected counts in the new tests against computeValidSplitKV: MULTITILE's 32x1 then 64x2 and STRADDLE's 2,2,3,3 are both correct, and the {{...}} regex brackets balance.
• The conservativeness argument in the new computeValidSplitKV comment holds: using each row's own last index rather than the block's last row can only raise the count, so the host never drops a split that holds valid data.
• Both createFloatSplatTensor overloads (const APFloat & and float) exist, so the two new call sites resolve.
• rocMLIR back-port: mlir/tools/rocmlir-gen/ is a shared path, and the PR description's "No rocMLIR back-port needed" note satisfies the requirement.
• License headers and trailing newlines are correct on all three new .mlir files.

CI status

No checks in the fail or cancel bucket. The only non-passing entry is this auto-review job itself (review, in progress), which is expected. Note that no build or test job ran on this SHA, so there is no compile signal for the new C++.

[rocmlir-pr-reviewer]

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

---

Scope

Host-side split-KV attention recombination in mlir/tools/rocmlir-gen/rocmlir-gen.cpp: per-(batch-head, query-row) valid-split counts under causal / prefix-causal masking, a NaN-safe merge (clamped LSE max and weight sum), an f32 merge for f16/bf16 storage, and removal of the sequenceLengthQ > nPerBlock && causalMasking exit. Four LIT/E2E tests cover the new behaviour.

Findings

No remaining issues. All seven previously raised findings are addressed in this revision: the duplicated split-count arithmetic is now a single numValidSplits lambda, createMaskSplitKV asserts the split axis is dim 1, the per-row constants are pinned exactly (STRADDLE [2],[2],[3],[3]; MULTITILE 32x1 + 64x2), and both test-header comments now document the configurations and the single-result split-KV checks.

Notes

I re-derived the per-row counts for every FileCheck expectation and they match: MULTITILE (splitKV=2, nPerBlock=32, seqLenQ=96) yields 1 for rows 0-31 and 2 for rows 32-95; STRADDLE (prefixOffset=62) yields [2,2,3,3]. The conservativeness argument in the computeValidSplitKV comment holds: the host itersPerBlock is never larger than the kernel's, so the count can only over-approximate, and count <= splitKV always. Over-counted splits rely on the empty-split zero-output / -inf-LSE guarantee from the base branch, which the PR description calls out.

rocMLIR back-port: mlir/tools/rocmlir-gen/ is a shared path, and the description's "Relation to rocMLIR" section supplies the required justification. No vendored-subtree (external/) paths are touched.

CI status

No failing or cancelled checks.

[dhernandez0]

The counts are conservative. The kernel partitions each M block's keys using the block's last row, while the host uses each row on its own, so the host can include extra empty splits but never drops a split that holds the row's data. The extra splits are harmless because #534 makes the kernel write zero output and -inf LSE for empty splits.

Note that at some point, we might want to have real early exit. Where if a workgroup has not work to do, it exits without even writing to memory, which is more efficient. If that is implemented, the CPU version needs to be updated as well. Because we are assuming memory is initialized.

It might be worth adding a note/comment or a TODO.

Also, the original implementation of rocmlir-gen (before this PR) assumed that data was not initialized, because originally that's how early exit was done. But we found some bugs that made us revert that to always init memory.

[dhernandez0]

the reason for this is that split-kv chooses where each split happens dynamically, based on seqLen and tile size. That would be hard to replicate on the CPU (but possible).

If the split doesn't match, LSE/output of attention won't match. So, we decided to just match the output after the "second kernel", then it should match.

[dhernandez0]

can you check if this is explained somewhere? if not it might be worth adding as a comment.

[dhernandez0]

nit: if we assume memory is initialized. I think we can simplify this even more. Just do the reduction over all splits?

[dhernandez0]

this reminds me, we should add a Release compilation to CI. Otherwise, we could merge PRs without [[maybe_unused]] here.

[dhernandez0]

nit: floating point < f32? just in case attn is extended to smaller data types in the future