perf(rocm): MFMA batch-invariant GEMM and chunked Triton attention for strict R/R

[Flink-ddd]

Stacked on #394 (perf/rocm-fixed-paged-bitwise). Draft: strict R/R is bitwise consistent, the whole training side beats native, and the end-to-end step is within 4.8% of native on a healthy node; rollout per-token throughput is still 1.75x behind, so this is not yet a net end-to-end win. Numbers and the open items are below.

What changes

1. MFMA batch-invariant deterministic GEMM (rlkernel.det_gemm.triton_mfma_rocm.v1)

The gfx942 strict GEMM from #325 used scalar-FMA leaves plus a BF16 K-tree (no matrix cores), which cost 3-30x versus hipBLASLt and dominated every R/R phase. rl_engine/kernels/ops/triton/matmul/mfma_gemm.py replaces it on gfx942:

• v_mfma_f32_16x16x16_bf16 with matrix_instr_nonkdim=16 and kpack=2 pinned (both change the in-tile K order, so they are part of the contract).
• K consumed in ascending BLOCK_K=64 tiles inside fixed CHUNK_K=1024 chunks; chunk partials combined ascending in FP32; one BF16 rounding.
• Every output element depends only on its row/column, so tile shape, warp count, num_stages, weight layout ([N,K] view or [K,N] prepared), row count M, and the split-K decode schedule are bitwise identical — tests/test_rocm_mfma_gemm.py asserts this across nine configs, four Qwen3-8B TP4 shapes, sub-batches and strided rows.
• RL_KERNEL_DET_GEMM_BACKEND=auto|triton_tree|triton_mfma; auto selects MFMA on gfx942, the previous tree contract stays selectable.

Reuse survey (asked in the task): vLLM's VLLM_BATCH_INVARIANT registers its Triton matmul overrides only for CUDA SM80/XPU (ROCm gets nothing at the aten level and the ROCm attention backends do not claim supports_batch_invariance), and AITER's Triton GEMMs autotune BLOCK_K and use split-K, so neither is reusable as-is; only the fixed-config idea carries over.

Microbench (MI300X, per rank TP4 shapes, median):

shape	hipBLASLt	old K-tree	MFMA (this PR)
gate_up M=4096 (4096x6144)	337 us	15494 us	482 us
down M=4096 (3072x4096)	189 us	6806 us	242 us
qkv decode M=8 (4096x1536)	10 us	82 us	26 us
gate_up decode M=8	19 us	82 us	23 us

2. Chunked-KV Triton flash attention (rlkernel.rocm.triton_chunked_flash_attention.v1)

rl_engine/kernels/ops/triton/attention/chunked_flash_attn.py serves Megatron full-sequence forward, vLLM prefill, prefix-cached extend and single-token paged decode from one arithmetic path: 64-key blocks inside 512-token KV chunks, each chunk's online softmax from an empty state, chunks merged ascending with the exact FA2 rescale, fully masked blocks/chunks leave a row's state untouched, FP contraction disabled. A monolithic (per query tile) and a split (per sequence/kv-head/chunk plus merge) schedule are bit-identical, which is what makes decode agree with training. Mixed decode/extend/prefill vLLM batches go through one launch. tests/test_rocm_triton_chunked_attention.py covers prefill-vs-decode-vs-extend identity, batch composition, mixed batches, the unmasked fast path and the core binding. Opt in with RL_KERNEL_ROCM_ATTENTION_BACKEND=triton; the #394 CK fixed-M128 path remains the default.

3. Harness

• examples/vime_rocm_attention_ablation/run_pr377_workload.py (parameterized P/P or R/R arm of the PR377 workload) and summarize_pr377_runs.py (the consistency/performance tables below).
• launch_arm.sh/run.py: strip inherited proxy variables (an operator-shell proxy turned every >30 s generation into a 502 retry loop), pin GPU_ARCHS for AITER JIT builds, forward the two new knobs.
• Validator accepts the two Triton contract ids; benchmarks/profile_rocm_rollout_decode.py + summarize_rollout_trace.py for kernel-level decode profiles.

Full analysis (where the time went, microbenchmarks, one-round tables, node caveat): benchmarks/results/pr396_rocm_mi300x/report.md.

Results (1 round, PR377 workload: Qwen3-8B, actor TP4/CP2, 2 x TP4 rollout engines, 8 samples, 7168 response, 4096 tokens/GPU, seeds 1234)

Healthy node after host GPU reset (vLLM memory utilization 0.38) — primary numbers

Config	Mismatch Count	Max |dlogp|	torch.equal
P/P native	23294 / 42042	2.342051	false
R/R strict (MFMA GEMM + Triton attention)	0 / 28652	0	true

Metric	P/P native	R/R (Triton attn)	R/R vs P/P
rollout time	56.31 s	67.17 s	19.3% slower
effective tokens/GPU/s	93.32	53.32	42.9% lower
update weights	2.59 s	1.18 s	54.6% faster
log probs	8.71 s	7.19 s	17.4% faster
actor train	14.92 s	10.67 s	28.5% faster
train time	24.31 s	18.36 s	24.5% faster
actor train tok/s	2886.6	2779.7	3.7% lower
end-to-end step	83.69 s	87.71 s	4.8% slower

Mean sampled response length was 5255 tokens (P/P) vs 3581 (R/R), so the per-token rollout throughput (93 vs 53 tokens/GPU/s, 1.75x) is the honest measure of the remaining gap, not the 19% rollout-time difference. Against the #394 R/R baseline on a healthy node (round 0: 39 tokens/GPU/s, log probs 25.5 s, actor train 50.1 s, step 183.0 s) this branch is 1.36x faster in rollout throughput and 4.7x faster in actor train.

Same degraded node, same conditions for all three arms (vLLM memory utilization 0.30; see caveat)

Config	Mismatch Count	Max |dlogp|	torch.equal
P/P native	20305 / 37034	2.840120	false
R/R strict (MFMA GEMM + CK attention)	0 / 33320	0	true
R/R strict (MFMA GEMM + Triton attention)	0 / 28652	0	true

Metric	P/P native	R/R (CK attn)	R/R (Triton attn)	Triton R/R vs P/P
rollout time	172.13 s	336.20 s	312.16 s	81.4% slower
effective tokens/GPU/s	26.89	12.39	11.47	57.3% lower
update weights	7.39 s	3.53 s	8.35 s	13.0% slower
log probs	31.32 s	44.33 s	37.63 s	20.1% slower
actor train	63.25 s	54.56 s	48.15 s	23.9% faster
train time	95.58 s	99.45 s	86.38 s	9.6% faster
actor train tok/s	601.6	629.3	616.2	2.4% higher
end-to-end step	276.24 s	441.30 s	409.79 s	48.3% slower

Against the #394 R/R baseline on the same workload (round 0: rollout 104.9 s, log probs 25.5 s, actor train 50.1 s, step 183.0 s, measured when the node was healthy), the earlier healthy-node R/R run of this branch gave actor train 11.6 s, log probs 21.0 s, step 169.1 s with 0/33320 mismatches.

Caveat for the degraded-node table. The node had degraded during this work (a native P/P round went from 69 s to 276 s uniformly across phases while single-GPU microbenchmarks stayed unchanged; the driver held GPU contexts of killed vLLM workers). After a host-side GPU reset the primary table above was measured with 8-GPU copy bandwidth 3.8-3.9 TB/s and TP4 all-reduce 1 MB at 43-53 us. The degraded-node relative numbers are same-condition, the absolute ones are inflated for every arm.

Open items before this is a net win

• Rollout decode is still ~1.75x native per token; the kernel-level decode profile (benchmarks/profile_rocm_rollout_decode.py) on the healthy node is the next step to attribute it (GEMM is within 1.3x of hipBLASLt, the attention core is faster than CK).
• Triton attention is faster than CK on the training side and in this rollout, but the CK path stays default until the 30-round run confirms it.
• CP training still all-gathers Q/K/V and computes the full sequence on each CP rank; the per-row-invariant kernel allows computing only the local zigzag rows (halves attention compute at CP2).

[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: 9481cdc6-9697-4cb6-a093-4bf91f52b095

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