fix(rocm): stabilize direct paged CK attention for bitwise R/R

[inaniloquentee]

This follows merged #393 on test and makes direct paged ROCm attention usable with a fixed arithmetic schedule for strict training/rollout agreement. AITER's shape-dependent M64/M128 dispatch produced different BF16 results between full prefill and decode; incomplete tail page tables also reproduced an HSA kernel fault. The opt-in RL-Kernel CK entrypoint fixes these cases while consuming the vLLM KV cache directly.

Changes

• Instantiate fixed M64 or M128 CK templates in RL-Kernel using installed CK headers.
• Preserve the real maximum query length instead of launching artificial empty query blocks.
• Disable implicit scalar FMA contraction in this extension to preserve the softmax rescaling order; explicit MFMA instructions remain enabled.
• Pad page tables to complete 128-token KV tiles and sanitize unused columns of active requests and inactive graph rows.
• Retain 64-bit global addressing when either K/V view spans more than 2 GB.
• Correct decode request-to-query mapping and one-query-per-row metadata.
• Share prepared attention metadata across layers without dense KV materialization on the direct path.
• Add RL_KERNEL_ROCM_FIXED_PAGED_TILE=128; tile 64 is also supported, while the default 0 retains the existing AITER entrypoint.
• Include the tile and optional KV scheduling bound in the HIP Graph cache identity.
• Forward the settings through Ray and report the actual selected entrypoint.
• Cache immutable ROCm logp masks and TP shard starts, and skip identity masking for Vime's all-active requests.
• Make the gfx942 direct-root deterministic GEMM reduction accept decode row count as a runtime value, avoiding repeated compilation for dynamic decode batches.
• Fuse adjacent deterministic GEMM reduction levels during gfx942 no-grad execution while preserving the canonical BF16 rounding boundaries.
• Match deterministic GEMM leaf tiles to Qwen3-8B TP4 small-M decode shapes and fuse paired leaves with their first reduction.
• Keep leaf fusion disabled for the short-K attention output projection where it measured negative.
• Rename the lazy CK source to fixed_paged_ck.hip so its ROCm/HIP implementation is explicit.
• Add fixed-paged correctness probes, identical-input benchmarks, and PR377-workload runners.

The specialization currently supports BF16, head dimension 128, page size 16, and no dropout, softcap, bias, or sliding window. It requires installed aiter_meta CK headers and the ROCm C++ toolchain. The extension must be warmed before HIP Graph capture.

Changes are confined to RL-Kernel and its examples. CUDA implementations and installed Megatron, vLLM, Vime, and AITER sources are unchanged.

Validation

R/R validation run:

/app/model/vime-runs/pr394-nograd-gemm-reduction-rr-3round-v165

Configuration	Value
Hardware	8 × AMD Instinct MI300X VF, gfx942
Model	Qwen3-8B, BF16
Actor topology	TP4 / CP2 / PP1
Rollout topology	2 engines × TP4, round-robin router
Batch	rollout batch 1; 8 samples per prompt; global batch 8
Length limits	response 7168; max training tokens/GPU 4096; vLLM model length 40960
Seeds	training 1234; rollout 1234; second rollout engine 1235
Execution	FULL_AND_PIECEWISE HIP Graph; capture max 32
Placement	colocated; rollout offload enabled; training offload disabled
Routes	strict RL-Kernel attention, FFN, and logp
Candidate settings	fixed tile 128; paged KV scheduling bound 8192

The direct rollout path reported:

• rl_kernel_fixed_paged_ck_m128
• dense_kv_materialized=false
• no HSA kernel fault
• mismatch_count=0
• max_abs_diff=0
• torch.equal=true

The emitted comparison sidecars also passed exact comparison. This validates strict bitwise agreement for the tested workload and shapes; it is not a guarantee for every possible shape.

Full-native P/P vs strict R/R

The previous P/P number published in this PR was an attention-only ablation and was not a valid all-native baseline. The comparison below replaces it with a full-native P/P run under the same PR377 workload and framework-level training/rollout-consistency mode.

P/P:

/app/model/vime-runs/pr394-full-native-pp-vime-tis-3round-v167

R/R:

/app/model/vime-runs/pr394-nograd-gemm-reduction-rr-3round-v165

P/P uses native attention, FFN, and logp. R/R uses strict RL-Kernel attention, FFN, and logp.

Round	Response tokens P/P / R/R	P/P step	R/R step	P/P rollout	R/R rollout
0	~33.0k / ~32.9k	69.022 s	182.984 s	42.139 s	104.867 s
1	~42.7k / ~40.7k	83.916 s	216.023 s	74.860 s	129.021 s
2	57,344 / 57,344	76.658 s	295.243 s	67.297 s	196.746 s

Round 2 has exactly equal generated-token work. Relative to P/P, R/R is:

Phase	P/P	R/R	R/R delta
Rollout	67.297 s	196.746 s	+129.449 s
Actor train	—	—	+73.662 s
Logp	—	—	+14.422 s
End-to-end step	76.658 s	295.243 s	+218.585 s

The difference is a real implementation-performance gap rather than a test bug. Trace-visible copy_, to, and item operations account for only about 0.36–0.61 seconds and cannot explain the rollout delta.

The primary R/R bottleneck is the strict deterministic GEMM/FFN implementation, followed by fixed-schedule attention. P/P uses --use-rollout-logprobs, while R/R uses the strict linear-logp provider, so the logp measurements are not an identical internal path. This can explain part of the logp difference but not the rollout or actor-train gaps.

The fixed-paged attention path already avoids dense KV materialization. Therefore, the remaining large rollout gap is mainly arithmetic and scheduling cost, not KV-cache communication or materialization.

Operator benchmarks

Identical-input HIP Graph microbenchmarks compared the R/R AITER dynamic dispatcher with fixed M128 using seven-group medians and 200 launches per group.

Input / KV layout	AITER dynamic	Fixed M128	Change
B4, Q1, KV7168; >2 GB interleaved pool	0.286207 ms	0.272859 ms	4.7% faster
B4, Q128, KV7168; >2 GB interleaved pool	0.290410 ms	0.282509 ms	2.7% faster
B1, Q4096, KV4096; small interleaved pool	0.177159 ms	0.181009 ms	2.2% slower
B4, Q1, KV7168; small interleaved pool	0.233388 ms	0.268397 ms	15.0% slower

The candidate remains opt-in. It restores strict agreement and safe tail-page access and benefits the tested large-cache decode shapes, but it is not uniformly faster across all attention shapes.

Additional checks

Fixed M64 and M128 each passed 804 bit-pattern comparisons covering:

• full and suffix execution
• causal/LSE vs decode
• HIP Graph vs eager
• KV lengths 513, 4096, 4171, and 7168
• batch sizes 1 and 4
• full, tail128, tail17, and tail1 queries
• separate, interleaved, and greater-than-2-GB KV layouts
• reversed physical pages
• graph batches shrinking 4 → 3 → 2 → 1

Backward finite-gradient and real-adapter padding probes also passed. These tests cover the reproduced HSA fault but do not claim to cover every possible HSA failure.

git diff --check, changed Python syntax checks, launcher bash -n, and the applicable ROCm deterministic-GEMM tests passed.

Reproduction

python -m examples.vime_rocm_attention_ablation.run_full_rr_fixed_paged_ck \
  --tile 128 \
  --run-dir /app/model/vime-runs/fixed-paged-ck-rr

[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: Advanced

Run ID: 48445bc4-5292-46b3-825b-603af445c647

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.}

[Flink-ddd]

LGTM