[perf][rocm]: restore full Graph and cache strict Attention schedules

[inaniloquentee]

Context

This is the post-merge ROCm performance follow-up to #390, based on the VIME scheduling work merged by #377. It keeps the strict Qwen3 Dense arithmetic contract and changes only RL-Kernel runtime/adapter code.

Changes

• Restore vLLM FULL_AND_PIECEWISE Graph execution on ROCm by keeping strict IPC collective generations in device-visible state.
• Namespace vLLM AOT caches by the RL-Kernel operator route so stale P/P or R/R compiled graphs cannot be reused silently.
• Cache immutable CP2 Attention position transport, sort order, and inverse order across Megatron layers and recompute.

The final diff against test changes two RL-Kernel files:

• rl_engine/integrations/vllm_runtime.py
• rl_engine/kernels/ops/rocm/attention/strict_runtime.py

No Megatron-LM, vLLM, or VIME source is modified. CUDA dispatch is unchanged.

Matched MI300X configuration

Item	Value
Commit	1ba8448
Hardware	1 node, 8x AMD Instinct MI300X (gfx942)
Model / dtype	Qwen3-8B / BF16
Training topology	TP4 / CP2, 8 GPUs
Rollout topology	2 engines x TP4, all 8 GPUs
Router	round-robin
Batch	rollout batch 2, global batch 2, 1 sample per prompt
Sequence limits	response length 32, max 256 tokens/GPU
Seeds	training 1234, rollout 42
Placement	colocated; rollout offload enabled; training offload disabled
vLLM Graph	FULL_AND_PIECEWISE, capture max 32, route-specific AOT cache
Attention	ROCM_AITER_FA; direct paged CK; Split-KV disabled
Operator matrix	Attention P/P vs R/R; FFN and Logp remain R/R in both arms
Fallback	none

Round 1 includes process-local compile/warmup effects. Steady values below are the mean of rounds 2 and 3.

Performance

Metric	P/P	R/R	R/R vs P/P
Artifact	pr390-full-graph-route-cache-3round-p-p-v132	pr390-position-plan-cache-3round-r-r-v135	-
Rollout rounds	2.877867 / 1.090509 / 1.051628 s	2.511761 / 1.140868 / 1.025610 s	-
Steady rollout	1.071069 s	1.083239 s	+0.012171 s (+1.14%)
Steady Logprob	0.416055 s	0.627624 s	+0.211569 s (+50.85%)
Steady actor train	1.443130 s	1.878836 s	+0.435705 s (+30.19%)
Steady train compute	2.292586 s	2.939528 s	+0.646942 s (+28.22%)
Steady weight update	1.157943 s	1.139518 s	-0.018425 s (-1.59%)
Step rounds	23.276181 / 6.917838 / 6.267942 s	20.989418 / 7.439283 / 7.297853 s	-
Steady step	6.592890 s	7.368568 s	+0.775677 s (+11.77%)

The remaining steady gap is training-side rather than rollout-side. Relative to the immediately preceding strict run, the position-plan cache reduces actor train from 1.964626 s to 1.878836 s (-4.37%) and train compute from 3.069056 s to 2.939528 s (-4.22%). End-to-end wait varied by about 0.24 s between runs, so the optimization claim is limited to the directly measured compute phases.

Exactness

Metric	P/P	R/R
Compared Logprob elements	192	192
Mismatch count	97	0
Maximum absolute difference	1.036945	0
Train-rollout Logprob absolute difference	0.031487	0
torch.equal	false	true
Frozen sources unchanged during run	true	true

The final R/R run completed all three steps with no validation errors and no arithmetic fallback.

[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: f317fef4-656e-46f1-8856-df77cc49f474

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

[maxiaosong1124]

LGTM!