feat(rocm): enable bitwise Qwen3 Dense VIME alignment

[inaniloquentee]

Summary

• extend the ROCm attention ablation from [feat][rocm]: add end-to-end Vime rollout-training attention ablation matrix #385 with the optimized strict VIME path from Add VIME Qwen3-8B TP4/CP2 consistency experiment and results #377
• enable the Qwen3 Dense R/R route for attention, packed SwiGLU FFN, tensor-parallel collectives, and vocab-parallel log-prob
• fix the ROCm AITER paged-KV layout ambiguity when TP4 leaves exactly two local KV heads
• add reproducible launch, provenance, artifact validation, and topology regression coverage

Topology

• 8x AMD MI300X, colocated execution
• Megatron training: TP4 / CP2
• rollout: two vLLM engines, TP4 per engine
• router: round-robin across engines; no TP-rank-level routing change is required

Performance

The runtime keeps the PR377 scheduling path. The KV fix uses vLLM's native zero-copy transpose/split views. It adds no tensor copy, device synchronization, kernel launch, or TP communication. Layer diagnostics are environment-gated and install no probes in the default path.

Validation

R/R MI300X validation run v22:

• torch_equal=true
• element_count=64
• mismatch_count=0
• max_abs_diff=0.0
• train_rollout_logprob_abs_diff=0.0
• frozen input fingerprints match before and after the run
• all attention, FFN, and log-prob readback provenance checks passed

Tests:

• pytest -q tests/test_framework_runtime_adapters.py: 31 passed
• pytest -q tests/test_vime_rocm_attention_topology.py tests/test_rocm_packed_ffn.py: 11 passed
• git diff --check: passed

This PR intentionally prioritizes the strict R/R arm; the full four-arm ablation matrix was not rerun.

Hot-cache R/R vs production Attention

This is an Attention-only isolated comparison, not a fully native-stack comparison. P/P uses production Attention for both Megatron training and vLLM rollout; R/R uses RL-Kernel strict Attention for both. FFN and linear Logp remain on the same RL-Kernel R/R path in both arms.

Configuration:

• source: 872780c75b3f783e2842aaae3ea1e86bffc94f3d (tracked sources clean and frozen fingerprints matched)
• hardware: 8x AMD Instinct MI300X VF (gfx942)
• software: PyTorch 2.12.0+rocm7.14.0a20260608, HIP 7.14.60850, vLLM 0.26.1rc1.dev543+g11d76fc72
• model/dtype: Qwen3-8B, bf16
• Megatron: TP4 / CP2 / PP1, 8 GPUs
• rollout: 2 vLLM engines, TP4 per engine, round_robin, eager mode, custom all-reduce disabled
• placement: colocated; offload_train=false, offload_rollout=true
• workload: one GRPO step, rollout_batch_size=2, samples_per_prompt=1, global_batch_size=2
• limits: max_response_length=32, max_tokens_per_gpu=256
• seeds: training 1234, rollout 42
• routes held fixed: FFN=R/R and Logp=R/R

Single-step Vime timer results after JIT caches were populated:

Metric	P/P production Attention	R/R strict bitwise	R/R vs P/P
rollout_time	4.997880 s	14.354721 s	+187.22%
effective_tokens_per_gpu_per_sec	1.600679	0.557308	-65.18%
update_weights_time	2.659139 s	1.200663 s	-54.85%
log_probs_time	9.181175 s	4.888454 s	-46.76%
actor_train_time	20.772190 s	10.542294 s	-49.25%
train_time	31.878997 s	15.936777 s	-50.01%
actor_train_tok_per_s	20.363765	40.124095	+97.04%
step_time	40.020603 s	31.899600 s	-20.29%

The strict path costs 2.872x rollout latency and 65.18% rollout throughput in this small workload, while the strict training path is faster; the measured end-to-end step time is 20.29% lower. This is a single-step directional comparison rather than a long-run throughput benchmark.

Strict R/R bitwise result (fresh independent sidecar validation):

• passed=true, require_exact=true, torch_equal=true
• sample_count=2, element_count=64
• mismatch_count=0, max_abs_diff=0.0
• train_rollout_logprob_abs_diff=0.0
• mismatch_kl=0.0, mismatch_k3_kl=0.0
• route readbacks proved RL-Kernel Attention/FFN/Logp on both Megatron and vLLM with no fallback

For contrast, P/P production Attention was not bitwise aligned: torch_equal=false, mismatch_count=37/64, max_abs_diff=0.2500038146972656, and mean absolute train/rollout Logp difference 0.016262680292129517.

Run artifacts on the AMD host:

• P/P: /app/model/vime-runs/rocm-dense-attention-pp-8gpu-tp4-v31
• R/R: /app/model/vime-runs/rocm-dense-attention-rr-8gpu-tp4-v32

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

Run ID: e2298057-ff90-402c-9a02-2d738fd8ba7e

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!