Repository navigation
feat(rocke/gfx950): add causal_bottom_right to the dense attention kernel - #10999
Merged
Merged
Conversation
github-actions
Bot
added
project: hip-kernel-provider
rocke-lib: attention
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
dpeabody-amd
force-pushed
the
users/dpeabody/dense-causal-bottom-right
branch
from
August 19, 2026 23:05
81fb201 to
2678a23
Compare
AnaghaRaoAMD
requested changes
Aug 20, 2026
dpeabody-amd
force-pushed
the
users/dpeabody/dense-causal-bottom-right
branch
from
August 28, 2026 19:33
ad533bf to
88572fa
Compare
Contributor
Author
|
Rebased onto
Every build-derived number in the description described a kernel that no longer existed. I have not re-established on-hardware parity on this base, and I am not carrying the old Two things I found while re-verifying that were not review comments:
|
…rnel Align the causal diagonal to the bottom-right corner of the S_q x S_kv score matrix, so query i attends to keys j <= i + (seqlen_kv - seqlen_q) instead of j <= i. Chunked prefill needs this. vLLM and SGLang score a chunk of queries against a longer KV cache on every chunk, and those queries are the LAST seqlen_q positions of the sequence. Under top-left alignment query 0 would see only key 0: it builds, it runs, and it is silently wrong. The offset is a build-time constant, so it folds into three sites on the non-persistent path: the mask comparison, the KV-tile prune that bounds the loop, and the diagonal-only masking split. The third is an optimisation rather than a correctness requirement -- the diagonal only ever moves right, so tiles already below it stay below it -- but shifting it keeps below-diagonal tiles in the mask-free body instead of needlessly masking them. Off by default and free: with causal_bottom_right=False the offset is 0, no IR is emitted, and the compiled HSACO is byte-identical. Verified by hashing the emitted objects for causal, batched, non-causal, ragged, sliding-window, persistent, varlen, d64, long-context and paged variants against the unmodified kernel. Rejected rather than silently mishandled: persistent (that path derives the diagonal separately and needs its own offset), sliding_window (the window band would have to shift with it), seqlen_q > seqlen_kv, and offsets that are not a whole multiple of block_n (a partial-tile offset would cut into the pruned bound). Numeric parity against the CPU reference is green on MI355X, including a bottom-right case at S_q=256 / S_kv=512.
Two combinations were accepted and produced wrong numbers rather than an error. varlen. The diagonal offset is baked from seqlen_q/seqlen_kv, which under varlen are the PADDED MAXIMA -- the real per-sequence lengths only exist at runtime in cu_seqlens. Every sequence in the batch would get the same baked diagonal regardless of its own length, and the n_upper clamp keeps the result in bounds, so it is bad numerics rather than a fault. Bottom-right is specified per sequence for variable lengths, so it needs a runtime (seqlen_kv_b - seqlen_q_b); until that exists, reject the combination. gfx942. AttentionDenseSpec is SHARED with the gfx942 dense kernel, which never reads this field. Its supports_attention_dense() accepted such a spec, its builder emitted a top-left mask, and gfx942_kernel_name() inherited the `br` tag from the shared kernel_name() -- caching a wrong result under a symbol claiming to be right, and breaking that module's stated contract that supports() implies build() is correct. Reject it there alongside varlen/ragged/sliding_window, so dispatch falls through to another candidate instead. Also corrects the field comment. The chunked-prefill motivation overstated the reach: the aligned path requires seqlen_q % 256 == 0, `ragged` (the path that takes arbitrary lengths) is self-attention only so its offset is always 0, and varlen is now rejected. Only statically-shaped, tile-aligned specs reach a nonzero offset today; a general vLLM/SGLang chunk with an arbitrary KV length does not, and extending ragged/varlen is the follow-up that would get there. The gfx942 perturbation table and deferred-modes list are updated to match, keeping that suite's anti-staleness gates honest.
…rify it Nothing exercised this path, and the builder's own verify harness could not: it builds its oracle with is_causal=, which is TOP-LEFT. Against a bottom-right kernel that is a spurious FAIL whenever seqlen_q < seqlen_kv, and a vacuous PASS at equal lengths where the offset is 0. The TFLOPS formula was top-left-only too, under-reporting the work. Give the harness the shifted diagonal explicitly (ki <= qi + (Skv - Sq)) and the matching pair count, so --check means something on this path. Adds tests/test_attention_dense_bottom_right.py. The offset is a build-time constant, so most of what can go wrong is provable without a GPU and runs in the CPU lane: which specs are accepted, which are refused (including the two fixed in the previous commit), the symbol the kernel gets, and -- with the symbol normalised out of the IR so a rename cannot masquerade as a behaviour change -- that the default path is byte-identical, that bottom-right actually moves the emitted body, and that equal lengths collapse back to top-left. One gpu-marked test checks the numbers against the shifted-mask oracle and asserts they do NOT match the top-left one, which is what fails if the offset is dropped.
…ght kernel The flag had no producer: dispatch built the spec with causal=(mask_type != 0) and never set causal_bottom_right, so a request arriving with MaskType BOTTOM_RIGHT_CAUSAL (2) became causal=True with the field defaulted False -- a top-left mask for a bottom-right request, which is the exact failure the field was added to prevent. Map the ordinal through. Where the spec cannot serve it (unaligned lengths, or an explicit dense_persistent="on") construction raises and admits() already turns that into a decline, so the candidate falls through instead of returning a wrong answer. Under dense_persistent="auto", persistent is a heuristic about work size rather than something the caller asked for, and only the non-persistent grid implements the shifted diagonal -- so pick the grid that can serve the mask instead of declining. An explicit "on" is a caller request and is left to be rejected. gfx942 dispatch has the identical line and is deliberately NOT changed: that builder does not implement the mask, so wiring it there would produce a wrong result under a `br` symbol. It now refuses such specs at the supports layer.
…unreachable guard The validator required (seqlen_kv - seqlen_q) to be a whole number of block_n tiles, and the comment justified it as a property of the shape. It is not: it could never fire. seqlen_q must be a multiple of the 256-row query tile, seqlen_kv a multiple of block_n, and every legal block_n divides 256 -- so seqlen_q is a multiple of block_n as well and the difference always is too. The guard was dead code enforcing a restriction the surrounding rules already implied. The restriction existed only because n_upper floored. Computing the span as a ceil over (BLOCK_M - 1 + DIAG_OFF) instead removes the need for it: an offset that lands mid-tile still visits the tile it falls in. The qb term stays outside the ceil because qb*BLOCK_M is a multiple of block_n. diag_start keeps its floor, which is the conservative bound for the mask-free region and is correct for any offset. For every offset reachable today the two forms evaluate to the same build-time constant, so this changes no emitted code: the lowered IR is identical across causal, bottom-right at three length ratios, and the equal-length case. It matters for the ragged/varlen follow-up (AICK-2032), where lengths are arbitrary and no such alignment is available. Raised in review by Anagha Rao, who also supplied the ceil expression.
Chunked prefill at ARBITRARY lengths. Until now bottom-right required tile-aligned shapes -- seqlen_q a multiple of 256 and seqlen_kv a multiple of block_n -- which a real request is not, because the KV cache holds whatever it holds. The ragged path already handles non-multiple lengths by padding the boundary tiles on-chip, but it was restricted to self-attention, so its diagonal offset was always zero and bottom-right could not use it. Lift that restriction under causal_bottom_right only. Plain ragged cross-attention stays rejected: without a shifted diagonal there is nowhere to place the queries. No new kernel machinery was needed. The KV-tile bound is already a ceil, so an offset that lands mid-tile -- which is exactly what arbitrary lengths produce -- is covered, and ragged already ceils both the tile count and the query-block count. The interaction worth stating is the partial last KV tile. Ragged visits it and its out-of-range keys load as zeros; plain causal needs no key mask there because every query stops before the padding. Bottom-right shifts every query's reach to the right, so that argument has to be re-earned: the last REAL query sits at seqlen_q-1 and reaches seqlen_q-1 + (seqlen_kv - seqlen_q) = seqlen_kv-1, landing on the last real key and no further. Padded QUERY rows can reach past it, but they load zero and are dropped by the guarded store, so they cannot reach the output. Verified on gfx950 against a shifted-diagonal CPU reference at S_q/S_kv = 197/400 and 300/1234 -- the second spanning two query blocks with the second block partially padded. Green on both an MI355X and an MI350X. Full CPU lane shows no failure that develop does not already have. AICK-2032.
The golden IR hashes caught a claim that was not true. The feature was
described as free when off, and the compiled object is indeed byte-identical
-- but only because LLVM folds the extra instruction. The IR kept it:
develop: %lt = icmp slt i32 %mul, %sel
ours: %add = add nsw i32 %mul, 0 <-- emitted unconditionally
%lt = icmp slt i32 %add, %sel
diag_start added DIAG_TILES without the guard the mask site already had, so
every non-persistent causal kernel picked up an `add x, 0` and every hash on
that path moved. Hashing HSACO missed it; the golden pins IR, which is the
stricter contract and the right one.
Guard the add. The default path now emits IR identical to develop, verified
directly against the golden's own spec at 128/8 heads: 0 differing lines for
the default, lazy_rescale=False and ragged variants.
Also gate block_n to divisors of the 256-row query tile. The KV-tile bound
leaves the query-block term outside its ceil, which is exact only under that
condition; block_n=96 at seqlen_q=512 otherwise stops one tile short of the
tile holding the last query's diagonal and drops keys, silently. gfx942
already gates this identically. Placed in supports_attention_dense rather
than the shared spec, because the spec stays a permissive shape container
and each arch's supports_* decides what it can run -- moving it into
__post_init__ makes specs unconstructible that gfx942's contract grid builds
on purpose to prove they are rejected.
Pin two bottom-right cases in the golden, since the offset reaches codegen by
two routes: aligned lengths fold it to the top-left constant, ragged lengths
make it arbitrary and exercise the ceil. Adds 6 entries (2 cases x 3 llvm
flavors) and moves none, which is the property that matters.
The ragged relaxation was unreachable. dense_spec_for_request() still gated
`ragged` on seqlen_q == seqlen_kv, so a bottom-right request at arbitrary
lengths fell through to the aligned path and died on its 256-multiple rule:
197/ 400 ValueError: seqlen_q must be a multiple of 256, got 197
300/1234 ValueError: seqlen_q must be a multiple of 256, got 300
512/4097 ValueError: seqlen_kv must be a multiple of block_n=64, got 4097
100/8000 ValueError: seqlen_q must be a multiple of 256, got 100
Every shape the new tests declare supported. They passed only because they
build the spec by hand with ragged=True, which proves the kernel runs and
nothing about whether anything can ask it to.
Hoist bottom_right above the gate and widen it: ragged is self-attention only
because a short query block has no diagonal to sit on, and bottom-right is
precisely what supplies one. All four shapes now reach supports -> (True, '').
Chunked prefill is the case that needs it -- its KV cache is whatever length
it is.
Add a dispatch test over those four shapes so the gap cannot reopen, and one
over block_n on both sides of the new gate.
Correct two tests that were checking the wrong thing:
- the alignment-guard docstring stated the pre-PR contract as the global
rule; it holds on the non-ragged path, which is what this commit relaxes.
- the negative ragged case used causal=False, testing "not causal" rather
than "causal but top-left". causal_bottom_right is the axis that moved, so
it now pins causal=True with the flag off.
Drop the builder's two bottom-right branches: main() has no flag for the mask
and hard-codes seqlen_kv=seqlen_q, so neither could run. The GPU test covers
the same oracle at the same tolerance, plus a negative check against top-left.
Point the varlen comment at AICK-2049, where that work actually lives.
…ight Two things the rebase onto current develop exposed, plus the doc drift the ragged relaxation left behind. DISPATCH. bottom_right was read straight off mask_type, with no length check, so a BOTTOM_RIGHT_CAUSAL request at seqlen_q == seqlen_kv asked for a shifted diagonal whose offset is zero -- the same mask, the same emitted body as top-left. That cost three things: the spec refuses persistent=True, so an explicit dense_persistent=on DECLINED a request this kernel serves exactly; the auto downgrade dropped the large-Sq cohort off the ~970-TFLOPS grid-stride path for a diagonal that never moved; and kernel_name() minted a second `br` symbol byte-identical to the top-left kernel already in the launcher cache. Gate on sq != sk and all three stop, with the mask semantics untouched. The downgrade that remains is tracked in AICK-2145, so it reads as temporary rather than a permanent throughput cap. SINKS. develop gained attention sinks on this kernel in ecec2a7, after this branch forked. use_sinks is on the shared spec and dispatch now passes it through, so mask_type=2 with sinks is reachable and the causal_bottom_right validation says nothing about it. It composes: a sink is a per-query-head logit with no key position, so it seeds m/l the same way whichever keys the mask admits. Both features write the same softmax though, so the field comment now says the combination is pinned rather than argued -- the pin itself lands in the next commit. DOC DRIFT. c01c92f lifted the ragged self-attention restriction under this flag and four places still said otherwise, including the SCOPE paragraph a previous review round had already made me rewrite once: the field comment, the module docstring, the ragged field comment, and run_attention_dense_torch. All four now describe what the code does. Also corrects the padded-row justification added with the ragged relaxation. It said padded query rows load 0 and are dropped by the guarded store. The guarded store part is right; the zeroing is not. The ragged buffer resources bound the WHOLE tensor (B*Sq*Hq*D), not one batch element, so a padded row in batch b < B-1 reads batch b+1's real tokens. Only the last batch element pads with zeros. Correctness rests on the mask plus the qtok < seqlen_q store predicate; the buffer bound only prevents faults. A batch=2 case is added in the next commit to hold that down. Finally, asserts BLOCK_M % block_n == 0 at the top of build_attention_dense. supports_attention_dense rejects a block_n that breaks the KV-bound arithmetic, but nothing forces a caller through supports_* first and the failure mode is a short loop that drops keys silently.
…tch>1 ragged Coverage for the three things the previous commit changed or newly exposed, plus the golden entry that makes the sink interaction a gated contract rather than a claim. SINKS x BOTTOM-RIGHT. A new golden case, bottom_right_sinks_sq512. The sink seeds m/l before the KV loop and the shifted diagonal decides which keys that loop admits, so the two features meet in the same softmax; without a pinned hash a regression in their composition moves nothing. On CPU: acceptance, that kernel_name() carries BOTH `br` and `sinks` (or two different kernels share a launcher-cache entry), that all four combinations of the two flags emit four distinct bodies (so neither flag can be dropped when the other is set), and that the Sq == Skv collapse does not quietly depend on the sink being absent. On GPU: parity at 512/1024 aligned and 300/1234 ragged, against an oracle that puts the sink in the denominator as a zero-value column, plus the negative check against a top-left oracle. Re-blessing the golden also settles the zero-cost claim against the CURRENT base, which is the part the old evidence could not speak to. Three attention_dense/gfx950 cases added, none removed, and NO existing hash moved, in all three llvm flavors. THE GRID DOWNGRADE. Nothing pinned it. The existing dispatch test uses Sq=256, where work = nqb*Hq*B = 32 is below the 256-CTA threshold and persistent is off regardless -- it cannot observe the downgrade at all. The new test uses Sq=8192 (work = 1024) so the heuristic wants persistent, and asserts the grid is given up only where the diagonal actually moves: dropped at Sq != Skv, kept at Sq == Skv, and kept for a plain top-left request of the same size. Two more cover the explicit-instruction path -- 'on' declines when the diagonal moves, and is served when it does not, which is the capability the length gate restores. BATCH > 1 RAGGED. Every ragged case was batch=1, which is the one batch size where the padded-rows-load-zero story happens to be true. 300/1000 at batch=2 has padded query rows in batch 0 with a live neighbour behind them, so it exercises the mixed real/padded wave and the partial KV tile against the store predicate rather than against zeros. Also drops an unused torch import in the aligned GPU test.
The file went in with CRLF. `git ls-files --eol` reported `i/crlf` for it and `i/lf` for every other file this branch touches, and the original blob from 8b9864a carried 463 CR bytes, so this dates from the commit that added the test rather than from any later edit. It matters beyond tidiness: the tree is LF everywhere else, so the file is the only one whose committed bytes differ from what a Linux checkout of the rest of the repo looks like. Anything that reads bytes rather than lines sees it differently -- a patch generated on one platform will not apply on the other, and the pre-commit hooks run against a file the repo does not otherwise produce. Re-staged through `git add --renormalize`, which is why the diff reads 660 insertions and 660 deletions with no content change. The net diff of the branch is unaffected; only the committed bytes move.
dpeabody-amd
force-pushed
the
users/dpeabody/dense-causal-bottom-right
branch
from
August 28, 2026 19:59
88572fa to
30a9b6f
Compare
Contributor
Author
|
On-hardware parity is re-established on this base. I was not willing to carry the old numbers forward, and it turns out I did not have to.
Every case checks a shifted-diagonal oracle, and every case where the diagonal actually moves also asserts the result does not match the top-left oracle -- which is what fails if the offset is ever dropped. The sink cases use an oracle that puts the sink in the denominator as a zero-value column. That covers the two things the previous evidence could not speak to: the sink composition, which did not exist at the old merge base, and Two notes on how this was run, because it differs from how the earlier numbers were produced. It went through MI350X is queued and I will add it when it lands. Same ISA target, so it is family coverage rather than a distinct code path. |
Contributor
Author
|
MI350X is in as well: 9/9 passed on Slurm job IDs for the record: 67856654 (MI355X) and 67856686 (MI350X). Logs are under The description is updated. Nothing else outstanding from my side. |
1 similar comment
Contributor
Author
|
MI350X is in as well: 9/9 passed on Slurm job IDs for the record: 67856654 (MI355X) and 67856686 (MI350X). Logs are under The description is updated. Nothing else outstanding from my side. |
yraparti
approved these changes
Sep 3, 2026
yraparti
reviewed
Sep 3, 2026
yraparti
reviewed
Sep 3, 2026
yraparti
requested changes
Sep 3, 2026
AnaghaRaoAMD
approved these changes
Sep 4, 2026
yraparti
approved these changes
Sep 6, 2026
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
dlangbe
approved these changes
Sep 11, 2026
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Fix the gfx942 bottom-right dispatch mismatch and address the documented cleanup items.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds bottom-right causal masking to the gfx950 dense attention kernel, including ragged shapes, sink composition, dispatch wiring, and validation coverage.
Changes:
- Implements build-time diagonal offsets and tile-bound handling.
- Wires mask type 2 through gfx950 dispatch and rejects unsupported gfx942 usage.
- Adds CPU/GPU tests and LLVM IR golden cases.
File summaries
| File | Review summary |
|---|---|
dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py |
Adds representative bottom-right parity cases. |
dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json |
Records new LLVM golden hashes. |
dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py |
Covers unsupported gfx942 behavior. |
dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_bottom_right.py |
Adds functional and dispatch coverage; update the AOT verification reference to the reproducible JIT GPU tests. |
dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py |
Implements shifted masking; scope the byte-identity documentation to the unchanged top-left path and remove the internal Jira reference. |
dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py |
The gfx942 dispatch path can admit bottom-right requests while emitting a top-left mask; propagate or reject the mask ordinal and add coverage. |
dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx950.py |
Wires bottom-right dispatch; remove the internal Jira reference. |
Review details
Suppressed comments (3)
dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx950.py:98
- This adds an internal Jira identifier to a committed rocKE artifact. The repository's compliance guidance forbids Jira/Confluence/Perforce references in source and documentation; please describe the follow-up without the ticket ID (or use a public GitHub issue reference).
# persistent builder, tracked in AICK-2145.
dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py:64
- The claim here is too broad after adding the aligned bottom-right variant: with
ragged=Falseandcausal_bottom_right=True, the emitted body intentionally changes to shift the diagonal. Please scope the byte-identity statement to the unchanged top-left path (for example,causal_bottom_right=False) so this module documentation does not claim that the new kernel is identical to the old one.
block needs (see that field). The aligned path is emitted byte-identically when
``ragged=False`` (no TFLOPS impact).
dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py:391
- This adds an internal Jira identifier to a committed rocKE artifact. The repository's compliance guidance forbids Jira/Confluence/Perforce references in source and documentation; please describe the follow-up without the ticket ID (or use a public GitHub issue reference).
# tracked in AICK-2049.
- Files reviewed: 7/7 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Preserve unpublished changes and retain shape-specific cache identities for shifted bottom-right kernels.
dpeabody-amd
enabled auto-merge (squash)
September 24, 2026 03:19
Contributor
TheRock Submodule Bump ActivityNewest first
|
CongMa13
added a commit
that referenced
this pull request
Sep 29, 2026
Merging develop brought in #10999, which added four dense-attention cases: attention_dense/gfx950/bottom_right_bm128_sq512 attention_dense/gfx950/bottom_right_sinks_sq512 attention_dense/gfx950/bottom_right_sq512 attention_dense/gfx950/ragged_bottom_right_sq500 Their goldens were blessed before this branch added p10-p15 to the llvm23 datalayout. The four are new JSON keys, so the golden merges cleanly and no conflict flags the staleness -- but their llvm23 hashes still describe the short header, and rocke_golden_static fails on the merge result even though both sides pass on their own. Re-bless so they carry the fixed header. Verified the delta is exactly those four llvm23 entries, each +60 bytes, with llvm20 and llvm22 unchanged (0 of 150 cases) -- the same +60 every other llvm23 case took. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
JIRA ID : AICK-2014
Current merge-readiness status (
cd37e8f3774)Merged
developatedc1529384002b5d3566250a8f3d2f0ec23fab03and pushed the conflict resolution. GitHub now reportsMERGEABLE(no merge conflicts), but the PR is not yet cleared to merge for the reasons below.The resolution preserves the incoming gfx942 sliding-window admission and this PR's moving-bottom-right rejection, keeps the behavior-based dispatch coverage rather than restoring the deleted frozen capability table, and reconciles the representative-golden summary counts without re-blessing emitted hashes. The README clarifies that top-left and equal-length bottom-right gfx942 requests retain sliding-window support on both grids.
Fresh verification on the merged tree
The graph-API scope thread is resolved.
Remaining blockers / unverified lanes
test_attention_dense_gfx942_golden.pyfails only forattention_dense_gfx942/persist_swa_d128_bf16_w128, on llvm20 and llvm22. An isolated source-object comparison shows that the merged branch and pristinedevelopemit byte-identical IR for this case; the mismatch is inherited from the incoming sliding-window work, not introduced by bottom-right conflict resolution. The fixture has not been blindly re-blessed. Current full gfx942 numerical verification is needed before settling that pin.SKIPPED BY MATH-CI. The Multi-Arch workflow is in progress and its required summary has not yet reported.The earlier sections below retain their original commit-specific scope and results.
What
Adds bottom-right causal masking to the non-persistent gfx950 dense attention path. Query
iattends to keysj <= i + (seqlen_kv - seqlen_q). The flag defaults to off; equal Q/K lengths have zero offset.This serves a query block scored against a longer, already-populated KV cache without incorrectly applying the top-left diagonal.
API and supported scope
AttentionRequest.mask_typeaccepts the exportedAttentionMaskType, whose ordinals matchplan_utils::MaskType. Dense dispatch derivescausal_bottom_rightfrom the request rather than hardcoding it.use_sinksand multiple batches.seqlen_q > seqlen_kvremain rejected. Persistent and per-sequence varlen support remain tracked separately in AICK-2145 and AICK-2049.The hipDNN graph already represents bottom-right masks through its mask attributes. This PR implements the rocKE request/spec/kernel path, not a native attention pack or a new graph attribute. The descriptor-packaging example explicitly documents that a native rocKE attention pack is not present; native graph-to-rocKE execution is not claimed here.
Kernel and cache correctness
The compile-time offset reaches the causal comparison, the ceiling-based KV-tile bound, and the diagonal-only masking split. The tile bound relies on the enforced
block_m % block_n == 0invariant. Ragged batch coverage checks that padded query rows cannot overwrite the next batch.The refresh also fixes an integration regression introduced by combining bottom-right attention with develop's runtime-shape caching: different baked diagonal offsets were receiving the same cache key and kernel name. Moving bottom-right kernels now retain shape-specific identity, while the zero-offset runtime path remains unchanged. The kernel's shape-argument ABI is unchanged.
test_different_bottom_right_offsets_do_not_share_compiled_identityfailed before the fix and passes afterward. It exercises different emitted bodies for 256/512 and 512/1024 sequence lengths and requires distinct compiled identities.Integration
Merged
developat892310280375742482adc4fec36726029349c2aewhile preserving the previously unpublished branch work and existing merge history. The original kernel verification below is forb6de017d2f2fb2372e7126ad020238990c1b1416. The dispatch-only follow-up and its verification are recorded separately at the end.Test plan and results
library/tests/test_attention_builds.py -k dense: 25 tests passed on the ROCm toolchain. Four HSACO-budget subtests initially lackedllvm-objdump; all four passed when rerun with the ROCm LLVM tools onPATH(2 tests, 4 subtests).library/tests/test_attention_dense_gfx950_numeric.pyGPU lane on ROCm 7.2: 50 passed, no skips, including all 11 bottom-right cases. The cases cover both tile geometries, aligned/ragged lengths, sinks, and batch-two stores; the shifted oracle is also distinguished from top-left.library/tests/test_attention_dense_gfx942_numeric.pyGPU regression: 11 passed, no skips.llvm20,llvm22, andllvm23. Four bottom-right cases cover aligned, ragged, sinks, and the 128-row geometry. Regeneration changed only these four entries in each flavor.llvm20andllvm22.llvm20,llvm22, andllvm23.Numerical evidence comes from the checked-in tests through
run_attention_dense_torch(JIT). No external or cluster-only AOT suite is used as verification evidence, and no full-monorepo test result is claimed.Merge-readiness follow-up (
abf0ae7a3651bd2ae7cd5dba8ffe3f425b3b54d1)The readiness review found that the new
causal_bottom_rightcapability label excluded existing unified kernels even though their causal mask already uses the runtime KV/query length difference. The follow-up restores eligibility for generic unified 2D/3D, gfx942 dense-pipe, gfx950 D256 prefill, and D256 decode. Standalone gfx942 dense and gfx1250 WMMA remain rejected for a moving bottom-right diagonal.The blanket-decline assertions were replaced with routing regressions. All five affected regression cases failed before the correction. This follow-up changes only capability admission, tests, and documentation; no kernel emitter, ABI, compiled-cache policy, or golden data changes.
The GPU/build/byte-identity results above remain evidence for the original kernel commit; those lanes were not rerun for this dispatch-only follow-up. The graph-API scope thread is resolved.