From e88cec46b8844f582a2f8f847fbefb11ba22dbb2 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 8 Sep 2026 09:46:22 +0000 Subject: [PATCH 1/4] perf(rocm): restore strict full graph rollout --- rl_engine/integrations/vllm_runtime.py | 28 +++++++++++++++++--------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index e43d9032..34487454 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -78,6 +78,7 @@ "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" @dataclass @@ -892,7 +893,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 +913,24 @@ 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) + cache_root = os.path.normpath(os.fspath(vllm_envs.VLLM_CACHE_ROOT)) + if os.path.basename(cache_root) != _ROCM_FULL_GRAPH_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, _ROCM_FULL_GRAPH_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: From 293813a0af4e59ae70897fc6af84cac8817e41ce Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 8 Sep 2026 10:06:54 +0000 Subject: [PATCH 2/4] fix(rocm): isolate graph cache by operator route --- rl_engine/integrations/vllm_runtime.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index 34487454..e032ea01 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -79,6 +79,11 @@ "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 @@ -916,13 +921,18 @@ def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None: from vllm import envs as vllm_envs from vllm.config import CUDAGraphMode + 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) != _ROCM_FULL_GRAPH_CACHE_NAMESPACE: + 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, _ROCM_FULL_GRAPH_CACHE_NAMESPACE + cache_root, cache_namespace ) compilation.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE # ROCm IPC generations are allocated and consumed on device. Replayed From 1ba8448671f8f011b8592bee68e91b3fce135398 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 8 Sep 2026 10:43:59 +0000 Subject: [PATCH 3/4] perf(rocm): cache training attention position plans --- .../ops/rocm/attention/strict_runtime.py | 112 +++++++++++++++--- 1 file changed, 98 insertions(+), 14 deletions(-) diff --git a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py index 53f1de30..e6693b0e 100644 --- a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py +++ b/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,19 +259,12 @@ 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 @@ -200,13 +272,26 @@ def forward_with_lse( 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( From 463cc4da05026e3d8600f55f09c1dc8ebccc3be3 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 8 Sep 2026 21:59:30 +0800 Subject: [PATCH 4/4] examples(rocm): add 200-round PP and RR runners --- .../run_full_pp_200.py | 94 +++++++++++++++++++ .../run_full_rr_200.py | 14 +++ 2 files changed, 108 insertions(+) create mode 100644 examples/vime_rocm_attention_ablation/run_full_pp_200.py create mode 100644 examples/vime_rocm_attention_ablation/run_full_rr_200.py diff --git a/examples/vime_rocm_attention_ablation/run_full_pp_200.py b/examples/vime_rocm_attention_ablation/run_full_pp_200.py new file mode 100644 index 00000000..14ff1253 --- /dev/null +++ b/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()) diff --git a/examples/vime_rocm_attention_ablation/run_full_rr_200.py b/examples/vime_rocm_attention_ablation/run_full_rr_200.py new file mode 100644 index 00000000..0e5680ed --- /dev/null +++ b/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())