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

examples/vime_rocm_attention_ablation/run_full_pp_200.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+
+from examples.vime_rocm_attention_ablation.run import (
+    MatrixConfig,
+    _prepare_run_dir,
+    build_arm_environment,
+    frozen_input_manifest,
+    public_arm_environment,
+)
+from examples.vime_rocm_attention_ablation.validate_artifacts import (
+    CASE_IMPLEMENTATIONS,
+    validate_arm,
+    write_report,
+)
+
+
+RUN_DIR = Path("/app/model/vime-runs/pr393-200round-p-p")
+
+
+def main() -> int:
+    root = Path("/workspace/RL-Kernel-pr390")
+    config = MatrixConfig(
+        vime_root=Path("/workspace/vime"),
+        rl_kernel_root=root,
+        megatron_root=Path("/workspace/Megatron-LM-vime"),
+        model_root=Path("/app/model/Qwen3-8B"),
+        reference_checkpoint=Path("/app/model/Qwen3-8B_torch_dist"),
+        prompt_data=Path("/app/model/dapo-math-17k/dapo-math-17k.jsonl"),
+        run_dir=RUN_DIR,
+        launcher=root / "examples/vime_rocm_attention_ablation/launch_arm.sh",
+        num_rollout=200,
+    )
+    config.validate(require_paths=True)
+    _prepare_run_dir(RUN_DIR)
+    frozen_before = frozen_input_manifest(config)
+    write_report(RUN_DIR / "frozen-inputs.before.json", frozen_before)
+
+    case_id = "P/P"
+    arm_dir = RUN_DIR / "arms/p-p"
+    for directory in (
+        arm_dir / "readbacks",
+        arm_dir / "dump",
+        arm_dir / "checkpoint",
+        arm_dir / "mismatch_sidecars",
+    ):
+        directory.mkdir(parents=True, exist_ok=False)
+    environment = build_arm_environment(config, case_id, arm_dir, arm_index=0)
+    environment["VLLM_GPU_MEMORY_UTILIZATION"] = "0.38"
+    launch = {
+        "schema_version": "rlkernel.vime_rocm_attention_arm_launch.v1",
+        "case_id": case_id,
+        "expected_implementations": CASE_IMPLEMENTATIONS[case_id],
+        "frozen_input_fingerprint": frozen_before["fingerprint"],
+        "command": ["bash", str(config.launcher.resolve())],
+        "environment": public_arm_environment(environment),
+        "started_at": datetime.now(timezone.utc).isoformat(),
+    }
+    write_report(arm_dir / "launch.json", launch)
+    with (arm_dir / "launcher.log").open("w", encoding="utf-8") as log_handle:
+        process = subprocess.run(
+            ["bash", str(config.launcher.resolve())],
+            cwd=config.rl_kernel_root,
+            env=environment,
+            stdout=log_handle,
+            stderr=subprocess.STDOUT,
+            check=False,
+        )
+    report = validate_arm(arm_dir, case_id, launcher_returncode=process.returncode)
+    write_report(arm_dir / "validation.json", report)
+    frozen_after = frozen_input_manifest(config)
+    write_report(RUN_DIR / "frozen-inputs.after.json", frozen_after)
+    summary = {
+        "run_dir": str(RUN_DIR),
+        "num_rollout": 200,
+        "launcher_returncode": process.returncode,
+        "passed": report["passed"],
+        "errors": report["errors"],
+        "metrics": report["metrics"],
+        "frozen_sources_match": frozen_before["fingerprint"]
+        == frozen_after["fingerprint"],
+    }
+    write_report(RUN_DIR / "single-arm-summary.json", summary)
+    print(json.dumps(summary, indent=2, sort_keys=True), flush=True)
+    return 0 if report["passed"] else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/vime_rocm_attention_ablation/run_full_rr_200.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, "/workspace/RL-Kernel-pr390")
+from examples.vime_rocm_attention_ablation import run_full_rr_single_arm_v90 as base
+
+
+if __name__ == "__main__":
+    matrix_config = base.MatrixConfig
+    base.MatrixConfig = lambda **kwargs: matrix_config(num_rollout=200, **kwargs)
+    base.RUN_DIR = Path("/app/model/vime-runs/pr393-200round-r-r")
+    sys.exit(base.main())

rl_engine/integrations/vllm_runtime.py

@@ -78,6 +78,12 @@
     "rl_kernel::rocm_det_gemm_linear_all_reduce_inference",
     "rl_kernel::qwen3_ffn_packed_tp_inference_rocm",
 )
+_ROCM_FULL_GRAPH_CACHE_NAMESPACE = "rl_kernel_rocm_full_graph_v1"
+_ROCM_GRAPH_ROUTE_ENVIRONMENT = (
+    "RL_KERNEL_ATTENTION_CASE",
+    "RL_KERNEL_FFN_CASE",
+    "RL_KERNEL_LOGP_CASE",
+)
 
 
 @dataclass
@@ -892,7 +898,7 @@ def apply(value: torch.Tensor | None) -> torch.Tensor | None:
 
 
 def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None:
-    """Keep stateful ROCm TP reductions outside replayed HIP graph segments."""
+    """Keep device-sequenced TP reductions inside replayed accelerator graphs."""
 
     if vllm_config is None:
         from vllm.config import get_current_vllm_config_or_none
@@ -912,17 +918,29 @@ def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None:
         ]
         return
 
+    from vllm import envs as vllm_envs
     from vllm.config import CUDAGraphMode
 
-    compilation.cudagraph_mode = CUDAGraphMode.PIECEWISE
-    # The IPC transport owns peer-visible sequence and staging state that is
-    # advanced once per invocation. Replaying the transport inside a captured
-    # graph can reuse capture-time state and stale peer payloads. Keep just
-    # these opaque reductions eager while vLLM captures the pure compute around
-    # them as piecewise HIP graphs.
-    for op in _ROCM_STATEFUL_GRAPH_SPLITTING_OPS:
-        if op not in splitting_ops:
-            splitting_ops.append(op)
+    route_key = "_".join(
+        re.sub(r"[^a-z0-9]+", "-", os.getenv(name, "unset").lower()).strip("-")
+        for name in _ROCM_GRAPH_ROUTE_ENVIRONMENT
+    )
+    cache_namespace = f"{_ROCM_FULL_GRAPH_CACHE_NAMESPACE}_{route_key}"
+    cache_root = os.path.normpath(os.fspath(vllm_envs.VLLM_CACHE_ROOT))
+    if os.path.basename(cache_root) != cache_namespace:
+        # vLLM's AOT key cannot see implementations behind torch custom ops.
+        # Keep its normal config/code/compiler hashing under an RL-Kernel ABI
+        # namespace so an older custom-op artifact cannot be replayed silently.
+        os.environ["VLLM_CACHE_ROOT"] = os.path.join(
+            cache_root, cache_namespace
+        )
+    compilation.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
+    # ROCm IPC generations are allocated and consumed on device. Replayed
+    # reductions therefore advance their generation instead of reusing the
+    # capture-time payload, so these ops can remain in the full HIP graph.
+    splitting_ops[:] = [
+        op for op in splitting_ops if op not in _ROCM_STATEFUL_GRAPH_SPLITTING_OPS
+    ]
 
 
 def _patch_rocm_weight_cache_refresh() -> None:

rl_engine/kernels/ops/rocm/attention/strict_runtime.py

@@ -52,7 +52,6 @@
 # The sequence reorder and the position validation are platform-neutral tensor
 # bookkeeping. They are bound from the CUDA runtime rather than reimplemented
 # so the two runtimes cannot drift into two different global orderings.
-_sort_by_position = StrictCUDAAttentionRuntime._sort_by_position
 _gather_sequence = StrictCUDAAttentionRuntime._gather_sequence
 _validate_local_positions = StrictCUDAAttentionRuntime._validate_local_positions
 _validate_global_positions = StrictCUDAAttentionRuntime._validate_global_positions
@@ -148,6 +147,86 @@ def __init__(
         self._page_bounds_epoch_owner = object()
         self._page_bounds_validation: _PageBoundsValidation | None = None
         self._causal_prefill_position_cache: tuple[torch.device, int, torch.Tensor] | None = None
+        self._position_plan_cache: dict[
+            tuple[int, int, int, bool],
+            tuple[
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor,
+            ],
+        ] = {}
+
+    def _position_plan(
+        self,
+        query_position_ids: torch.Tensor,
+        key_position_ids: torch.Tensor,
+        *,
+        plan: AttentionCPCommunicationPlan | None,
+        cp_world_size: int,
+        causal: bool,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        """Cache immutable CP position transport and reorder indices across layers."""
+
+        key = (id(query_position_ids), id(key_position_ids), cp_world_size, causal)
+        cached = self._position_plan_cache.get(key)
+        if (
+            cached is not None
+            and cached[0] is query_position_ids
+            and cached[1] is key_position_ids
+        ):
+            return cached[2:]
+
+        if cp_world_size == 1:
+            global_query_positions = query_position_ids
+            global_key_positions = key_position_ids
+        else:
+            if plan is None:
+                raise RuntimeError("CP Attention position plan requires a communication plan")
+            global_query_positions, global_key_positions = (
+                self._communication.all_gather_position_ids(
+                    query_position_ids,
+                    key_position_ids,
+                    plan,
+                )
+            )
+        query_sort = torch.argsort(global_query_positions, dim=1)
+        if global_key_positions is global_query_positions:
+            key_sort = query_sort
+        else:
+            key_sort = torch.argsort(global_key_positions, dim=1)
+        query_positions_sorted = torch.gather(
+            global_query_positions, 1, query_sort
+        )
+        key_positions_sorted = torch.gather(global_key_positions, 1, key_sort)
+        _validate_global_positions(
+            query_positions_sorted,
+            key_positions_sorted,
+            causal,
+        )
+        inverse_query_sort = torch.argsort(query_sort, dim=1)
+        value = (
+            query_position_ids,
+            key_position_ids,
+            query_positions_sorted,
+            key_positions_sorted,
+            query_sort,
+            key_sort,
+            inverse_query_sort,
+        )
+        if len(self._position_plan_cache) >= 16:
+            self._position_plan_cache.pop(next(iter(self._position_plan_cache)))
+        self._position_plan_cache[key] = value
+        return value[2:]
 
     def new_page_bounds_epoch(self) -> object:
         """Issue a proof scope while its page table and lengths remain immutable.
@@ -180,33 +259,39 @@ def forward_with_lse(
         plan = None
         if cp_world_size == 1:
             global_q, global_k, global_v = q, k, v
-            global_q_positions = query_position_ids
-            global_k_positions = key_position_ids
             communication_backend = "none"
             self.communication_executed = False
         else:
             plan = self._communication_plan(contract, q.size(2), k.size(2))
             global_q = self._communication.all_gather_query(q, plan)
             global_k, global_v = self._communication.all_gather_kv(k, v, plan)
-            global_q_positions, global_k_positions = self._communication.all_gather_position_ids(
-                query_position_ids,
-                key_position_ids,
-                plan,
-            )
             communication_backend = self.communication_backend_id
             self.communication_executed = True
 
         if positions_are_sorted:
             if cp_world_size != 1:
                 raise RuntimeError("pre-sorted Attention positions are supported only at CP=1")
             q_sorted, k_sorted, v_sorted = global_q, global_k, global_v
-            q_positions_sorted, k_positions_sorted = global_q_positions, global_k_positions
+            q_positions_sorted, k_positions_sorted = query_position_ids, key_position_ids
             q_sort = None
+            inverse_q_sort = None
         else:
-            q_sorted, q_positions_sorted, q_sort = _sort_by_position(global_q, global_q_positions)
-            k_sorted, k_positions_sorted, k_sort = _sort_by_position(global_k, global_k_positions)
+            (
+                q_positions_sorted,
+                k_positions_sorted,
+                q_sort,
+                k_sort,
+                inverse_q_sort,
+            ) = self._position_plan(
+                query_position_ids,
+                key_position_ids,
+                plan=plan,
+                cp_world_size=cp_world_size,
+                causal=causal,
+            )
+            q_sorted = _gather_sequence(global_q, q_sort)
+            k_sorted = _gather_sequence(global_k, k_sort)
             v_sorted = _gather_sequence(global_v, k_sort)
-            _validate_global_positions(q_positions_sorted, k_positions_sorted, causal)
 
         paged_schedule = bool(getattr(self._core, "supports_paged_schedule", False))
         if paged_schedule:
@@ -238,9 +323,8 @@ def forward_with_lse(
             )
 
         if cp_world_size > 1:
-            if q_sort is None:
+            if q_sort is None or inverse_q_sort is None:
                 raise RuntimeError("CP Attention requires a framework position reorder")
-            inverse_q_sort = torch.argsort(q_sort, dim=1)
             out_rank_packed = _gather_sequence(out_sorted, inverse_q_sort)
             lse_rank_packed = _gather_sequence(lse_sorted, inverse_q_sort)
             shard = self._communication.reduce_scatter_strict_result(