Skip to content

fix(rocke): bound the THD token base in the SDPA backward (AICK-3016) - #15045

Draft
afagaj wants to merge 17 commits into
users/jafagani/rocke-bwd-cpp-stackfrom
users/jafagani/rocke-bwd-thd-tok0-bound
Draft

fix(rocke): bound the THD token base in the SDPA backward (AICK-3016)#15045
afagaj wants to merge 17 commits into
users/jafagani/rocke-bwd-cpp-stackfrom
users/jafagani/rocke-bwd-thd-tok0-bound

Conversation

@afagaj

@afagaj afagaj commented Nov 4, 2026

Copy link
Copy Markdown
Contributor

JIRA ID : AICK-3016

Summary: A bad THD offset (cu_seqlens_q) no longer makes the SDPA backward read or write outside its tensors; the sequence is skipped instead.

Stacked on #14365; retarget to develop after it merges.

Motivation

Raised by @amd-khushbu in the review of #14365. In THD mode the backward kernels take the first token of each sequence, tok0 = cu_seqlens_q[b] (or OFF_Q[b] * mult / div), and use it as a memory offset for the Q / O / dO / LSE reads and the dQ convert store. Nothing bounds it: the host check RAGGED_RANGE is never called on the launch path (registry.py only runs shape checks, and only with tensor_meta), and a host value check would need a device-to-host copy on every launch. block_table entries, by contrast, already have a device guard. This existed before #14365; that PR only mirrored it into C++.

Change

  • decode_thd (Python _attention_bwd_addr.py and its C++ twin) takes an optional tok_rows. Under ws_seg = 0 the workspace row count is also the storage token bound (total_q in the unified family, max_total_q / max_total_kv in the general one). A sequence with tok0 < 0 or tok0 + len > rows is decoded as empty at token 0 (len = 0, tok0 = 0), so none of its rows is read or written. Two compares and two selects per CTA, CTA-uniform. No tok_rows, or ws_seg != 0: no guard op.
  • Wired into every THD decode whose kernargs already carry the bound, so the kernel ABI does not change: unified main, general main, split dQ (q side), prep, convert and slice reduce, and the work-list weights (they must match the main kernel's lengths).
  • New helper thd_token_guard / rocke_attn_bwd_thd_token_guard; host-evaluated unit tests in test_attention_bwd_addr.py.
  • Design docs: attention_bwd.md 12.7 and attention_bwd_unified.md section 8 describe the new device behaviour.
  • Goldens re-blessed (THD configs only): differential llvm_gfx_all.json (ll at llvm20/22/23/24: unified main 236, main 15, aux 13, dq 6 configs per flavor; plus ir_canonical), and library/tests/golden/attention_bwd_ir_sha256.json. Batched and dense kernels are unchanged.

Not covered (no bound reaches the device without an ABI change): the general family with slot rows (ws_seg = 1, no max_total_* given) and the kv side of the split dQ kernel (its ABI has no ws_rows_kv). There the offsets stay trusted, as documented in 12.7.

Testing

CPU only (no GPU run).

  • Byte identity, check_byte_identity.py --only attention_bwd,attention_dense_bwd: GREEN, 20/20 families, at llvm20 and llvm22.
  • ctest --test-dir build -j4: 100% passed, 50/50, including rocke_native_golden_ll (C++ emitters against the re-blessed hashes).
  • run_diff.py --python-only --check-golden (ll at llvm20/22/23/24, and ir --canonical) for the backward families: GOLDEN OK.
  • pytest on test_attention_bwd_ir_golden.py, test_attention_bwd_addr.py, test_attention_bwd_host_checks.py, test_attention_bwd_paged_guard.py, test_attention_bwd_shared_helpers.py: 207 passed, 23 skipped.
  • Negative check: with the guard call disabled, test_thd_token_guard[0] and [runtime0] fail (a sequence past the bound and a negative start keep their raw tok0); restored, all pass.
  • pre-commit run --from-ref <base> --to-ref HEAD: clean.

🤖 Generated with Claude Code

afagaj and others added 14 commits November 4, 2026 01:43
@afagaj @claude
test(rocke): C++/Python lowerer byte-identity gate for attention back…
e19e0d0 
…ward (AICK-2762)

Python builds each attention backward kernel (general, dense, unified
atomic/direct; gfx942 and gfx950; shipped defaults and every tuned-table
entry), both engines lower it, and the .ll must match byte for byte at
every golden LLVM flavor.

ws_bounds='buffer' on the general kernel is also pinned (one table key per
arch): it is the only knob value that emits tile.buffer_load.

The test needs the rocke_engine binding and skips without it. It is in no
CTest list, so it is a local gate and never runs in CI. CI's check of the
C++ side is the ctest rocke_native_golden_ll (the CI-check commit).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@afagaj @claude
test(rocke): C++ HIP lowering ctest for tile.buffer_load (AICK-2762)
07aa78e 
The attention backward kernels emit tile.buffer_load with
ws_bounds='buffer'. Its C++ HIP lowering had no check that runs in CI:
test_buffer_load_scalar.py compares both engines, but skips its C++ side
without the rocke_engine binding, which the CI test tree does not build.

- New ctest rocke_buffer_load_hip lowers the op through the C++ HIP route
  for f16, bf16, f32 and i32 on gfx942 and gfx950. It requires the two
  full lines of Python's _op_tile_buffer_load template, spelled with the
  kernel's own value names: the raw_buffer_load_b16 / _b32 call with its
  operands in order (rsrc, voffset, soffset, 0), and the memcpy into the
  result. A swapped or dropped operand fails it. It is in the quick and
  ffm-quick tiers of ROCKE_ENGINE_test_categories_external.
- The JSON importer and recipe VM comments still said tile.buffer_load is
  aliased to a dtype-suffixed opcode. It has its own dtype-generic opcode,
  so the exact name resolves it. Comment-only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twins of the backward aux kernels and addr/worklist/…
e8c9cc1 
…caps helpers (AICK-2762)

First layer of the C++ port of the attention backward builders. Adds
C++ twins of attention_bwd_aux (the prep, convert, reduce, page zero and
page convert kernels) and the helpers they use: _attention_bwd_addr,
_attention_bwd_worklist and the caps header. addr and worklist are
partial here; later layers fill them in.

Each twin has a parity pair under library/tests/parity: a Python emit
script and a C driver build the same configs, and the .ll must match
byte for byte. The Python modules' docstrings now point at their twins.

attention_fwd_ext's LOG2E is defined locally in attention_bwd_aux.cpp:
that forward helper has no C twin.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the attention band helpers (AICK-2762)
0658b66 
Adds the shared platform band helper (rocke_band_*,
helper_rocke.helpers.attention_band.h) and the backward band twin
(_attention_bwd_band), with a parity pair.

All backward twins use one rocke_attn_bwd_ / ROCKE_ATTN_BWD_ prefix,
and the addr twin takes I32_FAR from the band helper, as Python does.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the backward spec, caps, tiles and LDS plann…
66af609 
…er (AICK-2762)

Ports the main backward kernel spec (instance_attention_bwd.h,
attention_bwd.cpp), the caps checks, the tile geometry and the LDS
planner. Parity pairs: layout, lds, spec and spec_extra.

The spec pair has a config for every LAYOUT_SPILLING_COMBOS row
(configs 255-264 cover the 10 rows the others miss), so the parity
checks prove C++ refuses each row. ROCKE_ATTN_BWD_NONE (INT32_MIN), the
None of an Optional field, is documented as reserved: it always reads
as None, never as a value.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twins of the backward fragment and block-GEMM helper…
742810e 
…s (AICK-2762)

Ports _attention_bwd_frag and _attention_bwd_gemm, with a parity pair.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twins of the backward counted, valu, sinks and paged…
615cc27 
… helpers (AICK-2762)

Ports four small backward helpers, each with a parity pair, and the
addr pieces they need. The Python docstrings point at the twins.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
refactor(rocke): remove duplicated attention backward C++ helpers (AI…
011b7a9 
…CK-2762)

The parallel ports left private copies of shared helpers. frag now
calls valu's cast_c_frag and lds's DmaSlabLayout, and attention_bwd.cpp
builds COUNTED_ISSUES, COUNTED_VMCNT_MAX, BF16_CASTS, BWD_EXP2S and
STATS_VECS from the counted and valu exports.

The shared DmaSlabLayout now computes in int64_t, as the removed private
copy did, so a large slab offset no longer wraps; its row_bytes, size
and offset C entry points return int64_t. Every narrowing to int (the
lane-map coordinates, the xlane index, align_up, the LDS plan's nbytes
and access offsets) is range-checked and raises ROCKE_ERR_VALUE instead
of wrapping.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the backward stage scheduler (AICK-2762)
89a7fbc 
Ports _attention_bwd_sched (the stage-scheduler tables) and the first
part of _attention_bwd_ring. Parity pairs: sched, sched_ir and
sched_params.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the backward tile-step body and prefetch rin…
4ebd732 
…g (AICK-2762)

Ports _attention_bwd_body and the rest of _attention_bwd_ring, plus the
addr, counted and gemm pieces they need. Unknown arch raises, as in
Python. Parity pair: body.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the backward main-kernel builder (AICK-2762)
9bdaffa 
Ports the general attention backward main kernel builder
(attention_bwd_main.cpp) and the rest of the addr and worklist
helpers. The main parity pair builds every phase-1 general case.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the split dQ kernel builder (AICK-2762)
f27968c 
Ports the split dQ kernel (attention_bwd_dq.cpp) and its body
(_attention_bwd_dqbody). Parity pair: dq.

Refusals: validate refuses the plain-tuple tables DQ_SPILLING_COMBOS and
DQ_SCHED_SPILLING as Python does, with Python's messages and in Python's
order. DQ_RESOURCE_REFUSED is keyed by spec_digest and stays Python, so
a caller that launches the kernel must consult Python's
validate_attn_bwd_dq_spec first. The dq pair has one refused config per
table row, and test_attention_bwd_dq_refusal_rows.py (in
rocke_library_attention_cpu_pytest) fails if Python adds a row without
one.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the dense backward kernels (AICK-2762)
8f0d239 
Ports the dense backward spec, main and dQ kernels and the dense shell
(_attention_dense_bwd_shell). The main and dQ builders gain _kv_tile
hooks so the dense shell can reuse them. Dense body_knobs names are
refused as Python does. Parity pairs: dense_main and dense_dq.

A C spec is a plain struct, so a caller can edit it after init or skip
init. The dense and dense-dQ validators therefore re-run init before
anything else, so such a spec is refused with init's own message.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@afagaj @claude
feat(rocke): C++ twin of the unified paged backward main kernel (AICK…
acf86fd 
…-2762)

Ports the unified backward main kernel (attention_bwd_unified.cpp). It
shares the _kv_tile hooks with the dense kernels through
instance_attention_bwd_main_internal.h. Parity pair: unified_main.

Its validate() re-runs post_init on a copy first, as main and dQ do
through resolved(). New ctest rocke_attention_bwd_spec_reinit checks
that the unified, dense and dense-dQ validators refuse a field edited
after init, and a spec never initialised, with init's message.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@github-actions github-actions Bot added documentation project: hip-kernel-provider rocke-lib: attention Touches rocKE library attention/FMHA kernels labels Nov 4, 2026
@afagaj afagaj mentioned this pull request Nov 4, 2026
Open
@github-actions github-actions Bot added the rocke: platform Touches rocKE platform label Nov 4, 2026
afagaj and others added 2 commits November 4, 2026 09:01
@afagaj @claude
test(rocke): CI check of the attention backward C++ twins against the…
46ceded 
… golden (AICK-2762)

CI never compared the backward C++ twins with Python. check_byte_identity.py
builds the engine from source, so it cannot run in the installed test tree;
test_attention_bwd_ir_cpp_parity.py is in no CTest list; and the
differential golden had no entries for these families.

- Bless the 20 attention_bwd / attention_dense_bwd parity families into
  golden/llvm_gfx_all.json: ll at llvm20, llvm22, llvm23 and llvm24
  (llvm24 is its own section, with the same hashes as llvm23), 2,646
  configs each, and canonical IR (15 families, 1,846 configs; the other 5
  emit no IR). rocke_differential_golden_ll now fails when a backward
  Python builder changes.
- tests/cmake/NativeGoldenTests.cmake builds those 20 C++ emitters with the
  engine and installs them as test binaries. A listed family with no
  _emit.c is a FATAL_ERROR; a platform-only copy (no library tree) skips
  with a STATUS message. Warnings do not fail these test drivers, as in
  run_diff.py, which builds them with no warning flags.
- New check_native_golden.py runs each emitter over its configs and
  compares the .ll with the same golden hashes. It enforces the emitters'
  exit codes: 0 must be a blessed config with the blessed .ll; 1 an
  unblessed one with no output; 2 + "unknown config" the end of the
  family; anything else fails. When every index up to the family's cap
  answers, it asks for index cap and fails if that config exists. A missing emitter, a family with
  no golden or a crash fails, never skips. test_check_native_golden.py
  covers each case with a scripted emitter.
- run_diff.py: attention_bwd_body, attention_bwd_spec and
  attention_bwd_unified_main get a 512-config cap; they have 256 or more
  configs, and the default cap is 256.
- New ctest rocke_native_golden_ll (build tree and installed tree, quick
  and ffm-quick tiers) runs it at llvm20, llvm22, llvm23 and llvm24.
- New test_attention_bwd_layout_spill_rows.py (in
  rocke_library_attention_cpu_pytest): every LAYOUT_SPILLING_COMBOS row
  has a spec parity config refused at that row, present in the C table.
  attention_bwd_spec_emit.c is installed with the library tests (the
  install skips *.c otherwise) so this check also runs in CI.
- New ctest rocke_attention_bwd_dma_slab_int64: the DMA slab layout keeps
  64-bit offsets (expected values from Python).
- golden/README.md: for these families rocke_differential_golden_ll and
  rocke_native_golden_ll together are the parity verdict in CI.

Python == golden and C++ == golden, so C++ == Python in CI. A tampered
hash, a removed family, an extra config and a missing binary each fail
the new ctest.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@afagaj @claude
docs(rocke): describe the attention backward C++ twins (AICK-2762)
d83c2d6 
The design doc, the dispatcher notes and the barrier-fence parity files
still said the C++ engine has no backward builder (the "C++-engine
waiver"). The twins now exist and CI checks them, so:

- attention_bwd.md: 0.3 replaces the waiver with what is ported (every
  kernel builder and its helpers), the three checks that hold the
  dual-engine rule (check_byte_identity.py, the ctest
  rocke_native_golden_ll, test_attention_bwd_ir_cpp_parity.py) and what
  is not ported (routing, plans, tuning, and the RESOURCE_REFUSED /
  DQ_RESOURCE_REFUSED tables, which a future C++ caller must consult in
  Python first). The file map lists the twins and their checks; 0.4,
  8.4 and the section 11 gate table follow.
- library/dispatch/AGENTS.md: the dispatcher is Python-only; the C++
  engine has the kernel builders, not this routing.
- barrier_fence_emit.py, test_barrier_fence_parity_emit.py: there is
  still no barrier_fence_emit.c, but because the twins' own parity
  pairs compare every fenced barrier, not because no twin exists.

Docs and comments only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@afagaj
afagaj force-pushed the users/jafagani/rocke-bwd-cpp-stack branch from 3c32658 to d83c2d6 Compare November 4, 2026 15:01
@afagaj @claude
fix(rocke): bound the THD token base in the SDPA backward (AICK-3016)
67c10ea 
The THD decode used tok0 = OFF[b] (cu_seqlens_q in the unified family)
as a memory offset with no check, so a bad offset read Q/O/dO and wrote
dQ outside the tensors. decode_thd now takes tok_rows: under ws_seg = 0
the workspace row count (total_q, or max_total_q / max_total_kv) also
bounds the storage tokens, and a sequence with tok0 < 0 or
tok0 + len > rows decodes as empty at token 0. No ABI change; the guard
is wired into every THD decode whose kernargs carry the bound (unified
main, general main, split dQ q side, prep, convert / reduce, work list),
mirrored in the C++ twins. Backward IR goldens re-blessed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@afagaj
afagaj force-pushed the users/jafagani/rocke-bwd-thd-tok0-bound branch from a27a466 to 67c10ea Compare November 4, 2026 15:29
@afagaj
afagaj force-pushed the users/jafagani/rocke-bwd-cpp-stack branch 2 times, most recently from bd91264 to 94ea1c1 Compare November 7, 2026 03:43

This branch has not been deployed

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

Reviewers

No reviews

Assignees

No one assigned

Labels

documentation organization: ROCm project: hip-kernel-provider rocke: platform Touches rocKE platform rocke-lib: attention Touches rocKE library attention/FMHA kernels

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@afagaj