feat(attention): add strict bitwise ROCm path

[inaniloquentee]

Adds the strict ROCm Attention path and makes it reachable from Vime through contract-aware dispatch, plus a Triton core that is bitwise-identical to the native reference core.

Scope: #235 lists ROCm under Non-goals and its PR1–PR8 stack is CUDA-only. This does not implement a #235 deliverable; it extends the #235 attention contract to ROCm — same AttentionContract, same attention-domain LSE, same global_block_index merge order, same fail-closed policy.

Operator-only throughout: seeded Q/K/V, no checkpoint or serving engine. 8×MI300X (gfx942), torch 2.12.0+rocm7.14.0, HIP 7.14.60850, Triton 3.7.0. Qwen3-8B shapes (Hq=32, Hkv=8, D=128), BF16 causal unless noted.

What is claimed, and what is not

Scope	Claim
Batch composition	bitwise — B>1 rejected at the core
Padding	bitwise — key_padding_mask rejected
TP degree	bitwise — one KV group per launch
CP degree	bitwise — RCCL used as transport only
Triton core vs native reference core	bitwise
ROCm production core vs reference core	not claimed — different kernels
CUDA vs ROCm production core	not claimed — FA4 CuTe vs AITER CK

Bitwise equality is asserted within one ROCm hardware and runtime stack. The self-owned deterministic core stays a reference oracle, not the production arithmetic path.

Algorithmic arrangements

The core does not reimplement attention. It removes each source of arithmetic variation from the vendor kernel and records what was removed, so a mismatch is a contract violation rather than a debugging session.

#	Arrangement	Mechanism
1	Split-KV structurally impossible	AITER has no num_splits; bind the dense non-split API (split_kv_control="dense_non_split_api"). Non-DISABLED SplitKVSpec raises.
2	Batch composition cannot vary	_validate_inputs rejects q.size(0) != 1. The batched launch does not exist, so it cannot disagree.
3	Padding never enters a reduction	key_padding_mask rejected; each unpadded logical row is materialized.
4	TP-degree independence	One launch per (batch row, KV group), so a shard's result never depends on the launch head count.
5	RCCL is transport, never reduction	all_gather + root-owned scatter; reduce_scatter gathers then evaluates a fixed balanced rank tree locally.
6	Vendor kernel fingerprinted	sha256 of the aiter.ops.mha source file in provenance — a version number does not pin Python-level dispatch.
7	Fail closed	Missing AITER entry point, missing extension, or a dispatch that resolves elsewhere all raise. No substitution.
8	Reference core shared with CUDA	deterministic_attention.cu hipified to .hip; the new Triton core is bit-identical to it.

Arrangements 2 and 4 are load-bearing rather than defensive, because raw AITER fails both only at some shapes — the failure mode most likely to reach production unnoticed:

Failure	Shapes measured	Worst drift
Batch composition	fails at B=2/S=512, B=4/S=256, B=4/S=512; clean at B=2/4 × S=128, 1024–4096	1.5625e-02
TP degree	fails at 5 of 12 (S, TP) points; clean at S=512 all TP, and S=4096/TP=2	7.8125e-03

Changes

• Contract-aware dispatch: KernelRegistry.register_attention_backend / get_attention_op(contract, requested_backend=...), candidate list separate from the legacy attn priority maps and empty by default, so a WS2 caller can never be served by an SDPA-shaped wrapper without attention-domain LSE.
• Register aiter.rocm.ck_dense_mha only when aiter.ops.mha genuinely imports; an explicit request fails rather than silently becoming the PyTorch path.
• rl_engine/integrations/vime/attention.py: runtime provider taking and returning structural objects, so RL-Kernel never imports Vime.
• Derive and validate query_position_ids / key_position_ids so a training full-sequence call and a rollout chunk provably describe the same logical tokens.
• AttentionContract.cross_rank_fingerprint() for cross-rank preflight; requested_backend="auto" rejected under CP>1 without it.
• New: rl_engine/kernels/ops/triton/attention/deterministic_attn.py — Triton core bitwise-identical to _C.deterministic_attention_*, so the reference arithmetic is exercisable without the vendor kernel. tests/test_triton_deterministic_attention.py, 71 tests.
• New: benchmarks/benchmark_ws2_rocm_attention.py + benchmarks/results/ws2_rocm_mi300x/ — measurement matrix and charts, reusing the PR feat(ffn): add deterministic distributed Triton FFN for ROCm #325 / feat(logprob): add deterministic ROCm vocab-parallel path #328 helpers and figure style.
• Fix: csrc/ops.cpp — the merge from feat/rocm-deterministic-collectives dropped an #if !defined(USE_ROCM) around the prefix_shared_attention registration but kept its #endif (13 #endif vs 12 #if). ROCm builds failed with #endif without #if.

Validation

Gate	Result
Attention contracts, CP harness, provider + dispatch	284 passed
Triton bitwise core	71 passed
Train vs rollout, backward determinism	max abs 0
Fail-closed (decode, dropout, sliding window, soft-cap, ALiBi, FP32, 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, communication_backend=rccl_ag_rs.

Results

Measured after the per-KV-group rule landed. Full data: benchmarks/results/ws2_rocm_mi300x/.

Triton core vs native reference core — 0 mismatched elements

dtype	S	out	lse	dQ	dK	dV
bf16	512 / 1024 / 2048 / 4096	0	0	0	0	0
fp16	512 / 1024 / 2048 / 4096	0	0	—	—	—

Three things had to be reproduced rather than re-derived: the sequential FMA chain (contraction index is the loop, head dim is the vector), the 256-lane partial + stride-halving softmax fold, and expf/logf. The last is the trap — every Triton exp/log intrinsic lowers to a bare v_exp_f32, ~1 ULP from the expf the C++ kernel calls, which alone broke parity on ~14% of elements. The helpers re-emit hipcc's two-term argument reduction around that same instruction, with an inline-asm barrier to stop LLVM refolding it into an FMA. Verified over 4M+ inputs including subnormals, ±inf and NaN.

TP-degree invariance

raw_launch is one launch for all heads; per_kv_group is what the provider runs.

S	TP	raw_launch out max-abs	per_kv_group
512	2 / 4 / 8	0	0
1024	2 / 4 / 8	7.8125e-03	0
2048	2	0	0
2048	4 / 8	3.9e-03 / 1.9e-03	0
4096	2 / 4	0	0
4096	8	3.9e-03	0

Cost of the per-KV-group schedule — corrects the earlier ~3x estimate

S	sdpa (ms)	raw_launch (ms)	per_kv_group (ms)	vs raw	vs sdpa
512	0.0712	0.2579	1.7759	6.89x	24.95x
1024	0.1302	0.2513	1.9985	7.95x	15.35x
2048	0.2802	0.2917	1.7507	6.00x	6.25x
4096	0.7046	0.5682	2.0472	3.60x	2.91x

The real multiplier is 4.11–7.31x, not ~3x, worst at short sequence where per-launch overhead dominates. Still worth buying, given the invariance table above, but the bill should be stated correctly.

Single device (BF16, B=1)

strict-aiter is the core called once for all heads — not the production schedule; see the table above. sdpa is a speed baseline only, with no accuracy comparison mixed in.

S	Platform	Path	Fwd (ms)	Fwd+bwd (ms)	Fwd peak MiB	Fwd+bwd peak MiB	out max-abs vs FP64
512	mi300x	sdpa	0.0785	0.2836	12.1	32.2	8.20e-03
512	mi300x	pytorch-native	0.1981	0.9766	44.2	76.3	1.39e-02
512	mi300x	strict-aiter	0.2475	0.6328	14.1	288.2	2.47e-02
512	mi300x	reference-native	1.0180	3.1217	36.1	78.1	7.74e-03
512	mi300x	triton-bitwise	1.3627	4.8656	36.1	78.1	7.74e-03
512	cpu	sdpa	31.5270	53.7609	0.0	0.0	8.61e-03
512	cpu	pytorch-native	15.3921	24.7810	32.7	46.4	1.95e-02
2048	mi300x	sdpa	0.2875	1.0837	48.3	128.8	7.99e-03
2048	mi300x	pytorch-native	1.0848	2.3871	564.0	1076.0	1.80e-02
2048	mi300x	strict-aiter	0.2938	1.9570	56.3	4224.8	2.03e-02
2048	mi300x	reference-native	12.7428	47.8157	528.2	1080.5	7.80e-03
2048	mi300x	triton-bitwise	19.4210	76.7789	528.3	1080.5	7.80e-03
2048	cpu	sdpa	304.6115	533.5067	50.5	31.7	9.31e-03
2048	cpu	pytorch-native	216.1569	422.9041	543.6	776.1	1.79e-02
4096	mi300x	sdpa	0.6965	3.2163	96.5	257.5	9.60e-03
4096	mi300x	pytorch-native	3.9513	9.3608	2160.0	4208.0	1.40e-02
4096	mi300x	strict-aiter	0.5569	5.7263	112.5	16641.5	2.14e-02
4096	mi300x	reference-native	49.3536	173.3246	2080.5	4209.0	7.81e-03
4096	mi300x	triton-bitwise	105.4049	306.5966	2080.5	4209.0	7.81e-03

S=1024 and the FP16 sweep are in the report; same shape of result. The host run covers S<=2048 and only the two paths that exist there; its peak memory is an RSS high-water delta from /proc, not an allocator statistic, so it is an approximation and not directly comparable to the device figures.

Two readings. The deterministic backward is the binding constraint, and it is worse than materializing the whole score matrix: at S=4096 strict-aiter peaks at 16.6 GiB while reference-native — which writes a full FP32 [B, Hq, Sq, Skv] buffer — peaks at 4.2 GiB. This belongs to AITER's mha_bwd, not to the integration: toggling only the deterministic flag on raw AITER gives 1188 → 4329 → 16753 MiB at S=1024/2048/4096 (O(S²)) against 180 → 265 → 433 MiB with it off, and raw AITER at S=4096 peaks within 65 MiB of the provider. And the deterministic cores are the most accurate of the four against an FP64 oracle (7.8e-03 vs 9.6e-03 SDPA, 2.1e-02 AITER) — determinism is not bought with accuracy here.

Distributed CP (RCCL AG/RS)

All-gather Q/K/V and position ids over the CP group, strict core once on the full sequence, reduce-scatter (out, lse) back to this rank's query range. Acceptance is bitwise against CP=1. S=4096.

Topology	World	TP	CP	Rep	local Hq/Hkv	Fwd (ms)	p95	MiB/rank	out / lse / repeat
tp1_cp2	2	1	2	1	32/8	1.8352	1.8726	160.5	bitwise
tp2_cp2	4	2	2	1	16/4	1.2114	1.3138	80.3	bitwise
tp1_cp4	4	1	4	1	32/8	1.3973	1.4363	160.5	bitwise
tp2_cp2_x2	8	2	2	2	16/4	1.2297	1.2948	80.3	bitwise
tp2_cp4	8	2	4	1	16/4	1.2594	1.3294	80.3	bitwise
tp1_cp8	8	1	8	1	32/8	1.4112	2.2029	160.5	bitwise

0 mismatched elements summed across every rank. Requires the native extension built for the platform; otherwise it fails closed on the ROCm deterministic RoPE operator rather than substituting another. The separate multi-rank acceptance run of the full strict shared core (all five tensors, RoPE included) is under benchmarks/results/pr319_rocm_mi300x/distributed/; the stale pre-per-KV-group performance artifacts that used to sit beside it have been removed.

Charts

Every measured cell is 0; n/m marks the FP16 gradient columns, which are not measured.

Blue is one launch for all heads, orange is one per KV group; the flat orange line at the 1e-12 floor is bitwise. reference-native and triton-bitwise allocate identical buffers, so their memory curves coincide and the later-drawn series hides the earlier.

Known limits

Stated rather than worked around, because a silent approximation would be worse than a refusal.

• Backward memory is O(S²) — ~16.6 GiB for one sequence at S=4096. Long-context training on this path is not viable until AITER offers a deterministic backward with bounded workspace; forward-only / rollout is unaffected.
• TP invariance costs 4.1–7.3x forward time. The single-launch alternative is faster and not invariant.
• CP>1 and decode are refused by the provider. Doing the cross-rank merge inside it would add a second merge order; decode needs KV-cache identity metadata the dense core does not materialize. CP goes through the transport path above.
• The Triton core's bitwise claim does not extend to CUDA. The nvcc expf/logf reductions are not ported and no CUDA device was available to verify them, so the op raises on CUDA unless the caller passes require_bitwise_libm=False. A test pins this on both platforms.
• The Triton core is a parity core, not a FlashAttention replacement — it materializes the full FP32 score matrix.
• _C does not register deterministic_attention_forward_fp32. The .cu defines it but the pybind registration is missing, so the native forward_fp32 raises AttributeError. Pre-existing, untouched.
• A stale comment contradicts the shipped TP policy. vime/attention.py still says RL-Kernel "binds the degree rather than paying ~3x forward time", from before the per-KV-group loop; the provenance dict below it correctly reports tp_degree_invariant: True.
• H100 not measured yet. The benchmark supports it (strict-fa4 replaces strict-aiter, the CP transport dispatches to CUDA IPC, and the Triton core is built with the bitwise opt-out since the nvcc libm sequence is unported); the run itself is pending. No cross-vendor claim.
• The host run covers S<=2048. S=4096 on CPU was dropped after a first attempt was killed; the host column is absolute-latency context, not a headline.

Reproduce

PYTORCH_ROCM_ARCH=gfx942 python setup.py build_ext --inplace

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 \
  tests/test_triton_deterministic_attention.py

python benchmarks/benchmark_ws2_rocm_attention.py \
  --seq-lens 512,1024,2048,4096 --dtypes bf16,fp16 \
  --warmup 5 --samples 20 --training-samples 10 \
  --output-dir benchmarks/results/ws2_rocm_mi300x

[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: 3c2674e4-e976-4ae3-b404-79b191f36cfb

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

---

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