feat(attention): add strict bitwise ROCm path

[inaniloquentee]

Latest status — 25 Aug 2026

• Single card: bitwise identical. out, lse, dQ, dK, dV all 0 mismatch across repeat, train-vs-rollout, and batch composition (B=2/4 x S=256/512/2048).
• TP: bitwise identical, verified by single-card head-sharding only. TP=2/4/8 x S=512..4096 give 0 on all five tensors. This is simulated TP (one process, sharded heads); attention TP performs no cross-rank reduction, so the arithmetic is the same, but it has not been re-run on real multi-GPU since the per-KV-group change.
• Not yet measured: CPU and H100.
• Benchmark numbers below are stale. They predate the per-KV-group launch rule, which trades roughly 3x forward time for the TP invariance above. Re-measurement on MI300X is in progress; the distributed latency sweep and the PR-325-style charts are not done yet.

Summary

This PR adds the ROCm strict Attention path and makes it reachable from Vime through explicit, contract-aware dispatch. It reuses qualified high-performance vendor operators where their reduction behavior is controllable, and keeps RL-Kernel implementations where vendor paths did not pass the expanded bitwise probes.

Scope note: #235 lists "ROCm / domestic accelerator backend" under Non-goals, and its PR1–PR8 stack is CUDA-only by design. This PR does not implement a #235 deliverable; it extends the #235 attention contract to ROCm — same AttentionContract, same attention-domain LSE, same fixed global_block_index merge order, same fail-closed policy — so a ROCm deployment is held to the definition #235 established on CUDA.

Area	Production path	Strict-consistency control
Attention core	AITER/CK dense MHA via aiter.ops.mha.mha_fwd/mha_bwd	Same core on training and rollout; deterministic backward; dense non-Split-K API; num_splits=1 provenance
Dispatch	KernelRegistry.get_attention_op on AttentionContract	Registered only when aiter.ops.mha loads; explicit request never degrades to another backend
Vime boundary	rl_engine.integrations.vime.attention.attention_provider	Structural request in, (out, lse) + provenance out; RL-Kernel never imports Vime
Batch layout	One logical row materialized per call	Raw AITER is batch-composition invariant for most shapes but not all (measured below); the per-row rule is what makes it unconditional
Q/K RMSNorm	Transformer Engine through #230	Same TE functional forward/backward on both sides; bitwise admission probe; fail closed
RoPE	rlkernel.rocm.deterministic_rope	Same post-RoPE boundary and position semantics
QKV / O projection	RL-Kernel deterministic GEMM	TE GEMM is not selected because expanded batch-composition probes drifted
CP communication	Self-owned rccl_ag_rs	Q/K/V/positions AG, platform core, Out/LSE RS; autograd enabled
CP schedule	rlkernel.attention.strict_ring_state.v1	Fixed logical order; communication/compute decoupled; overlap disabled
TP degree	Invariant	Every launch is pinned to one batch row and one KV group, so a head shard's result does not depend on the batch size or TP degree that produced it
Fallback	None	Missing or mismatched vendor controls fail closed

The self-owned deterministic Attention core remains a reference/debug oracle, not the ROCm production arithmetic path. Strict equality is defined within the same ROCm hardware and runtime stack; this PR makes no cross-vendor bitwise claim.

Changes

• Add contract-aware attention dispatch: KernelRegistry.register_attention_backend and get_attention_op(contract, requested_backend=...). The candidate list is separate from the legacy attn/attention priority maps and starts empty on every platform, so a WS2 caller can never be served by an SDPA-shaped wrapper that does not export attention-domain LSE, and the strict core can never be picked up by a legacy get_op caller.
• Register aiter.rocm.ck_dense_mha only when aiter.ops.mha genuinely imports and exposes mha_fwd/mha_bwd. An explicit request for the native backend fails immediately if the vendor stack is missing; it will not silently turn into the PyTorch reference path.
• Add rl_engine/integrations/vime/attention.py: a runtime attention provider for Vime's Megatron backend. It accepts and returns structural objects, so RL-Kernel never imports Vime, and signals native fallback through an attention_provider_unavailable marker rather than an imported exception type.
• Materialize each logical batch row independently. Raw AITER mha_fwd is batch-composition sensitive in BF16 at some shapes (S=256 B=4, S=512 B=2/B=4; up to 1.5625e-02) while invariant at others, so batching would otherwise change the bits for a subset of shapes only — the failure mode most likely to reach production unnoticed.
• Derive and validate position identity (query_position_ids / key_position_ids) so a training-side full-sequence call and a rollout-side chunk provably describe the same logical tokens.
• Add validate_cross_config_alignment() and bind the TP/CP degree. A head shard computed under TP=4 is not bit-identical to the same shard under TP=8 at some shapes (measured below), because the AITER reduction order depends on the launch head count. Making the path TP-degree invariant is possible — executing one KV group per launch removes the dependence — but costs 2.7–3.9x forward time, so the degree is bound and a mismatch fails closed instead.
• Add AttentionContract.cross_rank_fingerprint() — a rank-independent identity for cross-rank preflight, mirroring LogprobContract.cross_rank_fingerprint. requested_backend="auto" is rejected under CP>1 without it.
• Add benchmarks/benchmark_rocm_attention.py plus the MI300X artifacts under benchmarks/results/pr319_rocm_mi300x/.

ROCm Validation

Environment	Value
GPU	8 x AMD Instinct MI300X (gfx942:sramecc+:xnack-)
PyTorch / HIP	2.12.0+rocm7.14.0a20260608 / 7.14.60850
AITER source	aiter.ops.mha, source fingerprint recorded in every report
TE	2.14.0.dev0+e6ede467
Workload	BF16, Qwen3-8B dense heads Hq=32, Hkv=8, D=128

Gate	Coverage	Result
Attention contracts	tests/test_attention_correctness.py + tests/test_flashinfer_pr7_attention.py	156 passed, 41 skipped
CP + comparison harness	tests/test_cp_attention.py + tests/test_attention_comparison.py + rl_engine/tests/test_dispatch.py	83 passed
Vime provider + dispatch	tests/test_vime_attention_provider.py + tests/test_attention_dispatch.py	45 passed
Strict RCCL AG+RS	2 ranks: TP=1, CP=2	2/2 ranks; Out/LSE/dQ/dK/dV bitwise
Strict RCCL AG+RS	4 ranks: TP=2, CP=2	4/4 ranks; Out/LSE/dQ/dK/dV bitwise
Strict RCCL AG+RS	8 ranks: TP=2, CP=2, 2 replicas	8/8 ranks; Out/LSE/dQ/dK/dV bitwise
Train vs rollout	Same contract, both roles	Out/LSE max abs 0
Batch composition	B=2/B=4 x S=256/512/2048 vs the same rows one at a time	Out/LSE max abs 0
Backward determinism	Repeated dQ/dK/dV	max abs 0
Fail-closed	CP>1, decode, dropout, sliding window, soft-cap, ALiBi, FP32, TP-degree mismatch, non-contiguous positions	All refused

Every strict report records fallback=false, native_attention_arithmetic=true, actual_backend=aiter.rocm.ck_dense_mha, deterministic backward, disabled Split-KV, and communication_backend=rccl_ag_rs.

Benchmark results

Operator-only benchmark — seeded Q/K/V, no checkpoint, tokenizer, or serving engine. BF16, causal prefill, Qwen3-8B dense head layout.

Platform	Hardware	Runtime
mi300x	AMD Instinct MI300X ×8	torch 2.12.0+rocm7.14.0a20260608, HIP 7.14.60850

Single device (BF16, causal prefill)

B	S	Backend	Fwd median (ms)	Fwd p95 (ms)	Fwd peak MiB	Fwd+bwd median (ms)	Fwd+bwd peak MiB
1	1024	native (SDPA)	0.1421	0.1495	36.1	0.4742	84.3
1	1024	triton	0.3042	0.3306	32.1	1.6989	52.2
1	1024	strict	0.4012	0.4156	40.1	1.0196	1108.5
1	2048	native (SDPA)	0.2976	0.3011	72.3	1.0757	168.8
1	2048	triton	0.6367	0.6474	64.2	2.3840	104.5
1	2048	strict	0.4271	0.4365	80.3	2.1179	4265.0
1	4096	native (SDPA)	0.7059	0.7387	144.5	3.2329	337.5
1	4096	triton	1.4373	1.4808	128.5	4.5723	209.0
1	4096	strict	0.7387	0.7481	160.5	5.8905	16722.0
2	2048	native (SDPA)	0.4518	0.4812	144.5	1.9088	337.5
2	2048	triton	0.9840	1.0009	128.5	2.6583	209.0
2	2048	strict	0.7842	0.8086	120.5	3.9364	4369.3
4	2048	native (SDPA)	0.8306	0.8653	289.0	3.5771	675.0
4	2048	triton	1.7539	1.7740	257.0	5.1061	418.0
4	2048	strict	1.3859	1.4039	226.1	7.6314	4610.3

Forward is competitive: at S=4096 the strict path is 0.7387 ms against 0.7059 ms for SDPA and 1.4373 ms for the Triton FlashAttention backend — roughly 1.9× faster than Triton and within 4% of SDPA.

What the strict contract costs

The forward is cheap; the cost is concentrated in the deterministic backward, and it belongs to AITER's mha_bwd rather than to this integration. Measured with raw AITER at B=1, toggling only the deterministic flag:

S	Bwd det=on (ms)	Bwd det=off (ms)	Time	det=on peak MiB	det=off peak MiB
1024	0.4433	0.1902	2.33×	1188.3	180.3
2048	1.3299	0.5682	2.34×	4328.5	264.5
4096	5.0607	2.0433	2.48×	16753.0	433.0

Deterministic backward peak memory scales as O(S²): every doubling of sequence length quadruples it (1.2 → 4.3 → 16.8 GiB). This is the binding constraint on the path today — see Known limits. (Peak figures include the already-resident forward tensors, so the absolute det=on column is the meaningful scaling signal rather than the ratio between columns.)

That the integration is not the source is confirmed directly: at S=4096 raw AITER peaks at 16657 MiB against the provider's 16722 MiB — a 65 MiB difference that is exactly the provider's transposed input copies.

Batch composition

Raw AITER mha_fwd is batch-composition invariant for most shapes and not for others. Measured on MI300X, BF16, Hq=32/Hkv=8/D=128, causal, batch vs the same rows submitted one at a time:

B	S	out max abs	lse max abs	Invariant
2	128	0	0	yes
2	256	0	0	yes
2	512	1.5625e-02	9.5367e-07	no
2	1024–4096	0	0	yes
4	128	0	0	yes
4	256	7.8125e-03	9.5367e-07	no
4	512	1.5625e-02	9.5367e-07	no
4	1024–4096	0	0	yes

Through the provider every shape above is bitwise identical (0), because each logical row is materialized on its own.

The shape-dependence is the point. A kernel that always broke would be caught on the first batched run; one that breaks only at S=256/S=512 passes casual testing and then silently loses bitwise consistency in production. The one-row-at-a-time rule is therefore load-bearing rather than defensive, and tests/test_vime_attention_provider.py::test_batch_composition_is_bitwise_invariant keeps S=512 in its parametrization deliberately.

TP-degree invariance

TP shards attention by head and performs no cross-rank reduction, so a rank computing its own head slice must match the corresponding slice of an unsharded run. Raw AITER does not provide that: its reduction order depends on the launch head count, and measured on MI300X the same head shard drifts by up to 7.8125e-03 between TP degrees at some shapes.

The provider removes the dependence by pinning every launch to one batch row and one KV group. Through the provider, all five tensors are bitwise equal at every TP degree measured:

S	TP	out	lse	dQ	dK	dV
512	2 / 4 / 8	0	0	0	0	0
1024	2 / 4 / 8	0	0	0	0	0
2048	2 / 4 / 8	0	0	0	0	0
4096	2 / 4 / 8	0	0	0	0	0

This is simulated TP: one process holding the full tensors and slicing heads per rank. Because attention TP has no collectives, the arithmetic a real rank performs is exactly this shard computation — but the real multi-GPU rerun after this change is still outstanding.

The cost is roughly 3x forward time. The cross-config contract therefore does not bind the TP degree: doing so would reject train/rollout comparisons that are in fact bitwise identical. validate_cross_config_alignment checks only what changes the arithmetic (dtype, head layout, causal, Split-KV, reduction spec).

Distributed (strict CP, RCCL AG+RS)

world	TP	CP	replicas	ranks passed	Out/LSE/dQ/dK/dV
2	1	2	1	2/2	bitwise
4	2	2	1	4/4	bitwise
8	2	2	2	8/8	bitwise

The strict CP path requires the native extension built for the active GPU platform; without it the run fails closed on the ROCm deterministic RoPE operator (ROCm deterministic RoPE is unavailable; rebuild rl_engine._C for ROCm) rather than substituting a different one. Build with PYTORCH_ROCM_ARCH=gfx942 python setup.py build_ext --inplace.

Artifacts: benchmarks/results/pr319_rocm_mi300x/distributed/.

Known limits

These are stated rather than worked around, because each one is a place where a silent approximation would be worse than a refusal:

• Backward memory is O(S²). At S=4096 the deterministic backward peaks at ~16.6 GiB for a single sequence. Long-context training on this path is not viable until AITER offers a deterministic backward with bounded workspace; the forward-only / rollout path is unaffected.
• CP>1 is refused by this provider. The strict ROCm core owns single-rank attention arithmetic. Performing the cross-rank (out, lse) merge inside the provider would introduce a second merge implementation with its own order, so it fails closed and points at the CP transport path.
• Decode is refused by this provider. It requires KV-cache identity metadata (cache_position, block table, prefix-cache key) that the dense core does not materialize.
• TP invariance costs ~3x forward time. Pinning each launch to one KV group is what makes the path TP-degree invariant; the single-launch alternative is faster but not invariant.
• The strict CP path needs the native extension built. It fails closed on the ROCm deterministic RoPE operator otherwise.
• No cross-vendor claim. Bitwise equality is asserted within one ROCm hardware and runtime stack only.

Reproduce

python -m pytest -q \
  tests/test_attention_correctness.py \
  tests/test_flashinfer_pr7_attention.py \
  tests/test_vime_attention_provider.py \
  tests/test_attention_dispatch.py

PYTORCH_ROCM_ARCH=gfx942 python setup.py build_ext --inplace

python benchmarks/benchmark_rocm_attention.py \
  --output benchmarks/results/pr319_rocm_mi300x/results.json

for n in 2 4 8; do
  torchrun --standalone --nproc-per-node=$n \
    scripts/ws2_p2p_nccl_attention_reference_check.py \
    --transport rccl_ag_rs --strict-shared-core \
    --output benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w$n.json
done

python scripts/ws2_pr7_flashinfer_attention_check.py \
  --strict --device cuda --no-dry-run \
  --mode decode --batch-size 2 --query-len 1 \
  --kv-seq-len 128 --page-size 16 \
  --output results/rocm-attention/strict-decode-b2-kv128.json

python scripts/ws2_p2p_nccl_attention_reference_check.py \
  --run-rocm-matrix --output-dir results/rocm-attention

Related to #235 and #294. Q/K RMSNorm wrapper integration is in #230. The Vime adapter boundary mirrors the selected-logprob provider in #328.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 75bf6b77-a0bd-493c-9278-76560eb721b6

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}