[feat][rocm]: add end-to-end Vime rollout-training attention ablation matrix

.github/workflows/ci.yml

@@ -89,6 +89,8 @@ jobs:
             tests/test_attention_cross_config_binding.py \
             tests/test_attention_preprocess.py \
             tests/test_attention_projection.py \
+            tests/test_framework_runtime_adapters.py \
+            tests/test_rocm_e2e_ablation.py \
             tests/test_cp_attention.py \
             tests/test_cp_attention_transformer_engine.py
 

examples/vime_qwen3_8b_rocm_ablation/README.md

@@ -0,0 +1,93 @@
+# Vime Qwen3-8B ROCm Attention ablation
+
+This is the ROCm end-to-end counterpart of PR230's production/RL-Kernel
+operator matrix. It launches the real Vime orchestration once per Attention
+cell and requires runtime evidence from both sides:
+
+| Case | Megatron training | vLLM rollout |
+|---|---|---|
+| `P/P` | framework-native | framework-native |
+| `P/R` | framework-native | RL-Kernel AITER/CK |
+| `R/P` | RL-Kernel AITER/CK | framework-native |
+| `R/R` | RL-Kernel AITER/CK | RL-Kernel AITER/CK |
+
+FFN and Logp remain fixed at `P/P`, so only the Attention implementation
+changes. Each cell starts in a fresh process and inherits the same model,
+checkpoint, prompt data, seeds, token limits, and one-rollout pre-update state.
+
+This is not an operator microbenchmark. The subprocess must run both vLLM
+rollout and Megatron training. A return code of zero is insufficient: the
+runner fails the cell if either framework emitted no executed Attention
+readback, selected the wrong P/R route, reported fallback, or failed to prove
+the strict ROCm runtime on an R side.
+
+## Required host state
+
+- A ROCm PyTorch build with visible AMD GPUs.
+- AITER with `aiter.ops.mha.mha_fwd` and `mha_bwd` available.
+- Vime, Megatron-LM, vLLM and RL-Kernel importable by every Ray worker.
+- A frozen Qwen3 model/checkpoint and prompt file.
+- A Vime launcher that honors `RL_KERNEL_ABLATION_OUTPUT_DIR` for case-local
+  output, so one cell cannot update the input checkpoint used by the next.
+- Megatron startup wired through
+  `rl_engine.integrations.megatron_runtime.initialize_from_environment`, so
+  the training worker installs the selected P/R plan and emits its readback.
+
+The executable run requires these immutable input variables:
+
+```bash
+export MODEL_ROOT=/models/Qwen3-8B
+export TORCH_DIST_ROOT=/models/Qwen3-8B_torch_dist
+export VIME_CKPT=/checkpoints/qwen3-8b-pre-update
+export PROMPT_DATA=/data/dapo-math-17k.jsonl
+export NUM_ROLLOUT=1
+export TRAIN_SEED=1234
+export ROLLOUT_SEED=42
+```
+
+The parent launcher must propagate the P/R and readback variables into Ray
+workers. `rocm_python_entrypoint.sh` is provided for launchers that replace
+their Python executable. Point `RL_KERNEL_REAL_PYTHON` at the real interpreter
+and configure Vime to invoke this wrapper.
+
+## Review the launch contract
+
+Without `--run`, the runner writes only a review summary and does not require a
+ROCm host:
+
+```bash
+python examples/vime_qwen3_8b_rocm_ablation/run.py \
+  --output-dir /tmp/rocm-attention-ablation \
+  -- bash /path/to/vime/scripts/run-qwen3-8B-rocm.sh
+```
+
+## Execute the full matrix
+
+```bash
+python examples/vime_qwen3_8b_rocm_ablation/run.py \
+  --run \
+  --output-dir /tmp/rocm-attention-ablation \
+  -- bash /path/to/vime/scripts/run-qwen3-8B-rocm.sh
+```
+
+Use `--case R/R` (repeatable) to run a subset while debugging. The final
+acceptance run should execute all four cells.
+
+## Evidence and pass boundary
+
+Each case directory contains the combined process log and the unmodified JSON
+readbacks emitted by `FrameworkOperatorIntegration`. The aggregate is a human-
+readable `summary.md`; no generated result JSON is checked into the repository.
+
+For each R side, accepted evidence includes:
+
+- semantic backend `rlkernel.attention.deterministic.v1`;
+- `runtime_platform=rocm`;
+- actual runtime `rlkernel.rocm.attention.aiter_ck_ag_rs.v1`;
+- AITER/CK fixed no-Split-KV schedule;
+- no native, PyTorch-reference, Triton, or fallback execution.
+
+The rollout route consumes vLLM's paged cache, reconstructs logical KV order,
+and invokes the same strict AITER/CK core as the training route. The training
+route binds the Megatron CP process group to the RCCL AG/RS transport and
+preserves the explicit global position order used by the PR230 contract.

examples/vime_qwen3_8b_rocm_ablation/rocm_python_entrypoint.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REAL_PYTHON="${RL_KERNEL_REAL_PYTHON:?RL_KERNEL_REAL_PYTHON must name the real Python executable}"
+RL_KERNEL_ROOT="${RL_KERNEL_ROOT:-$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)}"
+export RL_KERNEL_ROOT
+export PYTHONPATH="${RL_KERNEL_ROOT}:${PYTHONPATH:-}"
+
+if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then
+  : "${RL_KERNEL_ATTENTION_CASE:?the ablation runner must select an Attention case}"
+  : "${RL_KERNEL_FFN_CASE:?the ablation runner must freeze the FFN case}"
+  : "${RL_KERNEL_LOGP_CASE:?the ablation runner must freeze the Logp case}"
+  : "${RL_KERNEL_READBACK_DIR:?the ablation runner must provide a readback directory}"
+
+  export RL_KERNEL_VLLM_INTEGRATION=1
+  export RL_KERNEL_PLATFORM=rocm
+  export RL_KERNEL_ROCM_STRICT_ATTENTION=1
+  export RL_KERNEL_ROUTE_REPORT=1
+
+  exec "${REAL_PYTHON}" "$@" \
+    --seed "${TRAIN_SEED:-1234}" \
+    --rollout-seed "${ROLLOUT_SEED:-42}" \
+    --vllm-enable-deterministic-inference \
+    --vllm-attention-backend rocm_aiter_fa \
+    --vllm-disable-custom-all-reduce \
+    --deterministic-mode \
+    --accumulate-allreduce-grads-in-fp32
+fi
+
+exec "${REAL_PYTHON}" "$@"

examples/vime_qwen3_8b_rocm_ablation/run.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Launch the PR230 Attention P/R matrix through a real Vime ROCm job."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+from rl_engine.integrations.rocm_ablation import (
+    ROCM_ATTENTION_CASE_IDS,
+    run_rocm_attention_ablation,
+)
+
+
+def _case_id(value: str) -> str:
+    normalized = value.strip().upper()
+    if normalized not in ROCM_ATTENTION_CASE_IDS:
+        raise argparse.ArgumentTypeError(
+            f"case must be one of {', '.join(ROCM_ATTENTION_CASE_IDS)}"
+        )
+    return normalized
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=Path("runs/rocm-attention-ablation"),
+        help="case logs, framework readbacks, and summary location",
+    )
+    parser.add_argument(
+        "--case",
+        action="append",
+        type=_case_id,
+        dest="cases",
+        help="run only one matrix cell; repeat to select multiple cells",
+    )
+    parser.add_argument(
+        "--run",
+        action="store_true",
+        help="execute the orchestration command (default: review-only dry run)",
+    )
+    parser.add_argument(
+        "command",
+        nargs=argparse.REMAINDER,
+        help="Vime command after '--', for example: -- bash scripts/run-qwen3.sh",
+    )
+    args = parser.parse_args(argv)
+    if args.command[:1] == ["--"]:
+        args.command = args.command[1:]
+    if not args.command:
+        parser.error("a Vime orchestration command is required after '--'")
+    return args
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    results = run_rocm_attention_ablation(
+        args.command,
+        output_dir=args.output_dir.resolve(),
+        base_environment=os.environ,
+        case_ids=args.cases or ROCM_ATTENTION_CASE_IDS,
+        execute=args.run,
+    )
+    for result in results:
+        print(
+            f"[{result.status.upper()}] Attention={result.case_id} "
+            f"log={result.log_path} readbacks={result.readback_dir}"
+        )
+        for error in result.errors:
+            print(f"  - {error}")
+    print(f"summary={args.output_dir.resolve() / 'summary.md'}")
+    return 1 if any(result.status == "failed" for result in results) else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

rl_engine/integrations/__init__.py

@@ -14,6 +14,13 @@
 )
 from rl_engine.integrations.megatron import MegatronIntegration
 from rl_engine.integrations.megatron_runtime import install_megatron_integration
+from rl_engine.integrations.rocm_ablation import (
+    ROCM_ATTENTION_CASE_IDS,
+    RocmAblationCaseResult,
+    RocmAttentionAblationCase,
+    rocm_attention_ablation_matrix,
+    run_rocm_attention_ablation,
+)
 from rl_engine.integrations.vllm import VllmIntegration
 from rl_engine.integrations.vllm_runtime import configure_vllm_environment
 
@@ -22,11 +29,16 @@
     "IntegrationPlan",
     "MegatronIntegration",
     "OperatorAblationCase",
+    "ROCM_ATTENTION_CASE_IDS",
+    "RocmAblationCaseResult",
+    "RocmAttentionAblationCase",
     "VllmIntegration",
     "configure_integration_environment",
     "configure_vllm_environment",
     "install_megatron_integration",
     "integration_plan_from_environment",
     "operator_ablation_case",
     "operator_ablation_cases",
+    "rocm_attention_ablation_matrix",
+    "run_rocm_attention_ablation",
 ]