diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d3e5f6e2..711b8df4 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/examples/vime_qwen3_8b_rocm_ablation/README.md b/examples/vime_qwen3_8b_rocm_ablation/README.md new file mode 100644 index 00000000..31674e3e --- /dev/null +++ b/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. diff --git a/examples/vime_qwen3_8b_rocm_ablation/rocm_python_entrypoint.sh b/examples/vime_qwen3_8b_rocm_ablation/rocm_python_entrypoint.sh new file mode 100755 index 00000000..c957274f --- /dev/null +++ b/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}" "$@" diff --git a/examples/vime_qwen3_8b_rocm_ablation/run.py b/examples/vime_qwen3_8b_rocm_ablation/run.py new file mode 100644 index 00000000..21f5daea --- /dev/null +++ b/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()) diff --git a/rl_engine/integrations/__init__.py b/rl_engine/integrations/__init__.py index 49b7193d..fc1a229e 100644 --- a/rl_engine/integrations/__init__.py +++ b/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,6 +29,9 @@ "IntegrationPlan", "MegatronIntegration", "OperatorAblationCase", + "ROCM_ATTENTION_CASE_IDS", + "RocmAblationCaseResult", + "RocmAttentionAblationCase", "VllmIntegration", "configure_integration_environment", "configure_vllm_environment", @@ -29,4 +39,6 @@ "integration_plan_from_environment", "operator_ablation_case", "operator_ablation_cases", + "rocm_attention_ablation_matrix", + "run_rocm_attention_ablation", ] diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py index b1b5f970..3f19816b 100644 --- a/rl_engine/integrations/framework_operators.py +++ b/rl_engine/integrations/framework_operators.py @@ -22,6 +22,8 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_FA4_SCHEDULE_ID, STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, AttentionContract, AttentionDType, AttentionMode, @@ -120,6 +122,30 @@ def _require_nvidia_cuda(tensor: torch.Tensor, module: str) -> None: raise RuntimeError(f"strict {module} R/R requires NVIDIA CUDA tensors") +def _require_attention_accelerator(tensor: torch.Tensor) -> str: + """Return the real Attention platform behind PyTorch's CUDA device API.""" + + if tensor.device.type != "cuda": + raise RuntimeError("strict Attention R/R requires CUDA or ROCm GPU tensors") + return "rocm" if torch.version.hip is not None else "cuda" + + +def _strict_attention_platform_contract(platform: str) -> tuple[str, str, str]: + if platform == "rocm": + return ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, + "rccl_ag_rs", + ) + if platform == "cuda": + return ( + STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_FA4_SCHEDULE_ID, + "cuda_ag_rs", + ) + raise RuntimeError(f"unsupported strict Attention platform {platform!r}") + + class SemanticOperatorHandle: """Resolve one exact semantic backend once for one framework process.""" @@ -478,7 +504,10 @@ def __call__( if query.ndim != expected_ndim or key.ndim != expected_ndim or value.ndim != expected_ndim: layout = "[T, H, D]" if packed_seq_params is not None else "[S, B, H, D]" raise RuntimeError(f"Megatron Attention Q/K/V must use {layout}") - _require_nvidia_cuda(query, "Attention") + runtime_platform = _require_attention_accelerator(query) + strict_core_id, strict_schedule, communication_backend = ( + _strict_attention_platform_contract(runtime_platform) + ) parallel_state = _megatron_parallel_state() cp_world = int(parallel_state.get_context_parallel_world_size()) @@ -494,7 +523,7 @@ def __call__( "context_parallel_size": cp_world, }, ) - operator.bind_cuda_runtime(process_group=cp_group) + operator.bind_accelerator_runtime(query, process_group=cp_group) scale = float(getattr(module, "softmax_scale", query.size(-1) ** -0.5)) def execute_sequence( @@ -533,11 +562,11 @@ def execute_sequence( local_block_offsets=block_offsets, ), config=AttentionAblationConfig( - strict_core_id=STRICT_ATTENTION_PRODUCTION_CORE_ID, - strict_schedule=STRICT_ATTENTION_FA4_SCHEDULE_ID, + strict_core_id=strict_core_id, + strict_schedule=strict_schedule, ), return_lse=True, - communication_backend="cuda_ag_rs" if cp_world > 1 else "none", + communication_backend=communication_backend if cp_world > 1 else "none", query_position_ids=position_ids, key_position_ids=position_ids, scale=scale, @@ -632,10 +661,10 @@ def execute_sequence( if packed_seq_params is not None else "megatron_sbh_zigzag_cp" ), - "materialization": "owner_local_zigzag_cuda_ag_rs", + "materialization": f"owner_local_zigzag_{communication_backend}", "cp_world_size": cp_world, "tp_world_size": tp_world, - "runtime_platform": "cuda", + "runtime_platform": runtime_platform, "triton_used": False, **execution_provenance, } @@ -748,16 +777,68 @@ def _vllm_kv_cache_views( kv_cache: torch.Tensor, *, head_size: int, + num_kv_heads: int | None = None, + platform: str | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Return vLLM's paged cache as [blocks, block, kv_heads, head].""" + """Return CUDA and ROCm vLLM caches as [blocks, block, kv_heads, head].""" + + def normalize(plane: torch.Tensor) -> torch.Tensor: + if plane.ndim != 4 or plane.size(-1) != head_size: + raise RuntimeError("vLLM K/V cache planes must be 4-D with a head-size tail") + if num_kv_heads is None or plane.size(2) == num_kv_heads: + # [blocks, block, heads, head] (LBHNC after selecting K/V). + return plane + if plane.size(1) == num_kv_heads: + # [blocks, heads, block, head]. + return plane.permute(0, 2, 1, 3) + if plane.size(0) == num_kv_heads: + # [heads, blocks, block, head] (LHBNC after selecting K/V). + return plane.permute(1, 2, 0, 3) + raise RuntimeError( + "vLLM K/V cache layout does not expose the declared number of KV heads" + ) + if ( + kv_cache.ndim == 4 + and platform == "rocm" + and kv_cache.size(1) == 2 + and num_kv_heads is not None + and kv_cache.size(-1) == num_kv_heads * head_size + ): + key_cache, value_cache = kv_cache.unbind(1) + blocks, block_size, _packed_heads = key_cache.shape + return ( + key_cache.view(blocks, block_size, num_kv_heads, head_size), + value_cache.view(blocks, block_size, num_kv_heads, head_size), + ) if kv_cache.ndim == 4 and kv_cache.size(-1) == 2 * head_size: return kv_cache.transpose(1, 2).split(head_size, dim=-1) + # The K/V axis is leading in CUDA FlashAttention caches and follows the + # block axis in the ROCm AITER layouts. Prefer the platform convention in + # the ambiguous two-block case where both dimensions happen to equal two. + if kv_cache.ndim == 5 and platform == "rocm" and kv_cache.size(1) == 2: + key_cache, value_cache = kv_cache.unbind(1) + return normalize(key_cache), normalize(value_cache) if kv_cache.ndim == 5 and kv_cache.size(0) == 2: key_cache, value_cache = kv_cache.unbind(0) - return key_cache, value_cache + return normalize(key_cache), normalize(value_cache) + if kv_cache.ndim == 5 and kv_cache.size(1) == 2: + key_cache, value_cache = kv_cache.unbind(1) + return normalize(key_cache), normalize(value_cache) + if ( + kv_cache.ndim == 4 + and kv_cache.size(1) == 2 + and num_kv_heads is not None + and kv_cache.size(-1) == num_kv_heads * head_size + ): + key_cache, value_cache = kv_cache.unbind(1) + blocks, block_size, _packed_heads = key_cache.shape + return ( + key_cache.view(blocks, block_size, num_kv_heads, head_size), + value_cache.view(blocks, block_size, num_kv_heads, head_size), + ) raise RuntimeError( - "vLLM FlashAttention KV cache must use " "[blocks, kv_heads, block, 2 * head_size]" + "vLLM Attention KV cache does not match a supported CUDA/ROCm paged layout" ) @@ -931,10 +1012,12 @@ def __call__( return output.zero_() if query.ndim != 3: raise RuntimeError("vLLM query must use [tokens, heads, head_dim]") - _require_nvidia_cuda(query, "Attention") + runtime_platform = _require_attention_accelerator(query) key_cache, value_cache = _vllm_kv_cache_views( kv_cache, head_size=int(impl.head_size), + num_kv_heads=int(impl.num_kv_heads), + platform=runtime_platform, ) if key_cache.dtype != query.dtype or value_cache.dtype != query.dtype: raise RuntimeError("strict vLLM Attention requires an unquantized KV cache") @@ -964,7 +1047,7 @@ def __call__( "context_parallel_size": 1, }, ) - runtime = operator.bind_cuda_runtime() + runtime = operator.bind_accelerator_runtime(query) groups, metadata_summary = self._materialization_groups( attn_metadata, query=query, @@ -1013,10 +1096,14 @@ def __call__( last_operator_provenance = _compact_attention_provenance(result.provenance) self._last_provenance = { "framework_layout": "vllm_paged_kv", - "materialization": "direct_paged_fa4", + "materialization": ( + "logical_paged_kv_to_aiter_ck_dense" + if runtime_platform == "rocm" + else "direct_paged_fa4" + ), "tp_world_size": tp_world, "tp_group_bound": tp_group is not None, - "runtime_platform": "cuda", + "runtime_platform": runtime_platform, "triton_used": False, "direct_output_buffer": direct_output_buffer, **metadata_summary, diff --git a/rl_engine/integrations/rocm_ablation.py b/rl_engine/integrations/rocm_ablation.py new file mode 100644 index 00000000..4a1fb040 --- /dev/null +++ b/rl_engine/integrations/rocm_ablation.py @@ -0,0 +1,588 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""End-to-end ROCm rollout/training Attention ablation runner. + +PR230 defines a four-cell production/RL-Kernel matrix. This module executes +that matrix against a real orchestration command (normally Vime), one fresh +process per cell, and validates the runtime readbacks emitted independently by +Megatron training and vLLM rollout workers. It deliberately does not execute +synthetic tensors or manufacture a checked-in result payload. + +The replay identity is frozen before the first case. Only case selection and +case-local artifact paths may differ between subprocesses. A successful exit +without both framework readbacks is a failure, as is an R-side readback that +does not prove ROCm execution through the strict AITER/CK runtime. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import subprocess +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import Any, Iterable, Mapping, Sequence + +from rl_engine.integrations.ablation import Implementation, IntegrationPlan +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, +) + +ROCM_ABLATION_SCHEMA_VERSION = "rlkernel.rocm.e2e_attention_ablation.v1" +ROCM_ATTENTION_CASE_IDS = ("P/P", "P/R", "R/P", "R/R") +STRICT_ROCM_ATTENTION_RUNTIME = "rlkernel.rocm.attention.aiter_ck_ag_rs.v1" +STRICT_ROCM_ATTENTION_CORE = STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID +STRICT_ROCM_ATTENTION_SCHEDULE = STRICT_ATTENTION_ROCM_SCHEDULE_ID + +# These values identify the immutable input and sampling stream. They must be +# present for an actual matrix execution; a dry run may omit them so launchers +# can be reviewed before expensive GPU allocation. +REQUIRED_REPLAY_ENV = ( + "MODEL_ROOT", + "TORCH_DIST_ROOT", + "PROMPT_DATA", +) +FROZEN_REPLAY_ENV = ( + *REQUIRED_REPLAY_ENV, + "VIME_CKPT", + "TOKENIZER_PATH", + "DATASET_SEED", + "PYTHONHASHSEED", + "ROLLOUT_SEED", + "TRAIN_SEED", + "NUM_ROLLOUT", + "ROLLOUT_BATCH_SIZE", + "N_SAMPLES_PER_PROMPT", + "MAX_PROMPT_LENGTH", + "MAX_RESPONSE_LENGTH", +) + +_FRAMEWORK_TARGETS = MappingProxyType( + { + "training": ("megatron", "training"), + "rollout": ("vllm", "rollout"), + } +) + + +@dataclass(frozen=True) +class RocmAttentionAblationCase: + """One PR230 Attention implementation pairing.""" + + case_id: str + plan: IntegrationPlan + + def __post_init__(self) -> None: + if self.case_id not in ROCM_ATTENTION_CASE_IDS: + raise ValueError(f"unknown ROCm Attention ablation case {self.case_id!r}") + if self.plan.cases["attention"].case_id != self.case_id: + raise ValueError("case_id must match the Attention case in the integration plan") + if self.plan.cases["ffn"].case_id != "P/P" or self.plan.cases["logp"].case_id != "P/P": + raise ValueError("the Attention matrix must hold FFN and Logp at P/P") + + @property + def slug(self) -> str: + return self.case_id.lower().replace("/", "-") + + def implementation_for(self, target: str) -> Implementation: + return self.plan.implementation_for("attention", target) + + +@dataclass(frozen=True) +class FrameworkRouteEvidence: + """Aggregated route evidence from one framework side of one case.""" + + framework: str + target: str + implementation: str + backend_ids: tuple[str, ...] + call_count: int + readback_count: int + runtime_platforms: tuple[str, ...] + actual_backends: tuple[str, ...] + + +@dataclass(frozen=True) +class RocmAblationCaseResult: + """Outcome of one real Vime subprocess.""" + + case_id: str + status: str + returncode: int | None + log_path: Path + readback_dir: Path + routes: Mapping[str, FrameworkRouteEvidence] = field(default_factory=dict) + errors: tuple[str, ...] = () + + def __post_init__(self) -> None: + if self.status not in {"not_run", "passed", "failed"}: + raise ValueError(f"unknown case status {self.status!r}") + object.__setattr__(self, "routes", MappingProxyType(dict(self.routes))) + + +def rocm_attention_ablation_matrix( + case_ids: Iterable[str] = ROCM_ATTENTION_CASE_IDS, +) -> tuple[RocmAttentionAblationCase, ...]: + """Build the PR230 matrix while freezing non-Attention modules.""" + + normalized = tuple(str(case_id).strip().upper() for case_id in case_ids) + if not normalized: + raise ValueError("at least one Attention ablation case is required") + if len(set(normalized)) != len(normalized): + raise ValueError("Attention ablation case IDs must be unique") + unknown = [case_id for case_id in normalized if case_id not in ROCM_ATTENTION_CASE_IDS] + if unknown: + raise ValueError(f"unknown Attention ablation cases: {', '.join(unknown)}") + return tuple( + RocmAttentionAblationCase( + case_id=case_id, + plan=IntegrationPlan.from_case_ids( + attention=case_id, + ffn="P/P", + logp="P/P", + ), + ) + for case_id in normalized + ) + + +def validate_rocm_host() -> dict[str, Any]: + """Verify that a real ROCm device is available before allocating a run.""" + + import torch + + hip = getattr(torch.version, "hip", None) + if hip is None: + raise RuntimeError("end-to-end ROCm ablation requires a ROCm PyTorch build") + if not torch.cuda.is_available() or torch.cuda.device_count() <= 0: + raise RuntimeError("end-to-end ROCm ablation requires at least one visible AMD GPU") + devices = tuple(torch.cuda.get_device_name(index) for index in range(torch.cuda.device_count())) + return { + "hip_runtime": str(hip), + "pytorch": str(torch.__version__), + "device_count": len(devices), + "devices": devices, + } + + +def validate_replay_environment(environment: Mapping[str, str]) -> None: + """Reject a run that cannot identify its immutable replay inputs.""" + + missing = [name for name in REQUIRED_REPLAY_ENV if not str(environment.get(name, "")).strip()] + if missing: + raise RuntimeError("missing frozen replay environment: " + ", ".join(missing)) + num_rollout = str(environment.get("NUM_ROLLOUT", "1")).strip() + try: + count = int(num_rollout) + except ValueError as exc: + raise RuntimeError("NUM_ROLLOUT must be an integer") from exc + if count != 1: + raise RuntimeError( + "the ablation matrix requires NUM_ROLLOUT=1 so every cell starts from " + "the same pre-update state" + ) + + +def replay_identity( + command: Sequence[str], + environment: Mapping[str, str], +) -> dict[str, Any]: + """Return one stable identity shared by every matrix cell.""" + + if not command or any(not isinstance(item, str) or not item for item in command): + raise ValueError("command must contain non-empty argument strings") + frozen = { + name: str(environment[name]) + for name in FROZEN_REPLAY_ENV + if str(environment.get(name, "")).strip() + } + payload = { + "command": list(command), + "environment": frozen, + } + canonical = json.dumps(payload, sort_keys=True, separators=(",", ":")) + payload["sha256"] = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + return payload + + +def build_case_environment( + base_environment: Mapping[str, str], + case: RocmAttentionAblationCase, + *, + case_dir: Path, +) -> dict[str, str]: + """Materialize one worker environment without mutating the parent process.""" + + environment = {str(key): str(value) for key, value in base_environment.items()} + readback_dir = case_dir / "readbacks" + environment.update( + { + "RL_KERNEL_ATTENTION_CASE": case.case_id, + "RL_KERNEL_FFN_CASE": "P/P", + "RL_KERNEL_LOGP_CASE": "P/P", + "RL_KERNEL_VLLM_INTEGRATION": "1", + "RL_KERNEL_READBACK_DIR": str(readback_dir), + "RL_KERNEL_ROUTE_REPORT": "1", + "RL_KERNEL_PLATFORM": "rocm", + "RL_KERNEL_ROCM_STRICT_ATTENTION": "1", + "VLLM_ATTENTION_BACKEND": "ROCM_AITER_FA", + "RL_KERNEL_ABLATION_CASE": case.case_id, + "RL_KERNEL_ABLATION_OUTPUT_DIR": str(case_dir), + "NUM_ROLLOUT": "1", + } + ) + return environment + + +def load_runtime_readbacks(directory: Path) -> list[dict[str, Any]]: + """Load framework-owned evidence and reject malformed artifacts.""" + + payloads: list[dict[str, Any]] = [] + if not directory.is_dir(): + return payloads + for path in sorted(directory.glob("*.json")): + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"invalid runtime readback {path}: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"runtime readback must contain an object: {path}") + value["_source_path"] = str(path) + payloads.append(value) + return payloads + + +def _nested_strings(value: Any, key: str) -> set[str]: + found: set[str] = set() + if isinstance(value, Mapping): + direct = value.get(key) + if isinstance(direct, str) and direct.strip(): + found.add(direct.strip()) + for item in value.values(): + found.update(_nested_strings(item, key)) + elif isinstance(value, (list, tuple)): + for item in value: + found.update(_nested_strings(item, key)) + return found + + +def _contains_truthy_fallback(value: Any) -> bool: + if isinstance(value, Mapping): + for key, item in value.items(): + if str(key).strip().lower() in { + "fallback", + "fallback_used", + "used_fallback", + "split_kv_fallback", + } and item not in (False, None, "", 0): + return True + if _contains_truthy_fallback(item): + return True + elif isinstance(value, (list, tuple)): + return any(_contains_truthy_fallback(item) for item in value) + return False + + +def _contains_triton(value: Any) -> bool: + if isinstance(value, str): + return "triton" in value.lower() + if isinstance(value, Mapping): + for key, item in value.items(): + if str(key).strip().lower() == "triton_used" and item is True: + return True + if _contains_triton(item): + return True + elif isinstance(value, (list, tuple)): + return any(_contains_triton(item) for item in value) + return False + + +def _attention_plan_case(payload: Mapping[str, Any]) -> str | None: + plan = payload.get("plan") + cases = plan.get("cases") if isinstance(plan, Mapping) else None + attention = cases.get("attention") if isinstance(cases, Mapping) else None + value = attention.get("case_id") if isinstance(attention, Mapping) else None + return str(value) if isinstance(value, str) else None + + +def _validate_route( + case: RocmAttentionAblationCase, + payloads: Sequence[Mapping[str, Any]], + *, + side: str, +) -> tuple[FrameworkRouteEvidence | None, list[str]]: + framework, target = _FRAMEWORK_TARGETS[side] + expected = case.implementation_for(target) + label = f"{framework}/{target}" + matching = [ + payload + for payload in payloads + if payload.get("framework") == framework and payload.get("target") == target + ] + if not matching: + return None, [f"missing {label} runtime readback"] + + errors: list[str] = [] + records: list[Mapping[str, Any]] = [] + for payload in matching: + if _attention_plan_case(payload) != case.case_id: + errors.append(f"{label} readback used a different Attention case") + if payload.get("fallbacks"): + errors.append(f"{label} recorded fallback: {payload['fallbacks']}") + operators = payload.get("operators") + record = operators.get("attention") if isinstance(operators, Mapping) else None + if isinstance(record, Mapping): + try: + call_count = int(record.get("call_count", 0)) + except (TypeError, ValueError): + errors.append(f"{label} Attention has an invalid call count") + continue + if call_count > 0: + records.append(record) + if not records: + errors.append(f"{label} Attention had zero executed calls") + return None, errors + + implementations = {str(record.get("implementation", "")) for record in records} + if implementations != {expected.value}: + errors.append( + f"{label} implementation mismatch: expected {expected.value}, " + f"observed {sorted(implementations)}" + ) + backend_ids = tuple(sorted({str(record.get("backend_id", "")) for record in records})) + provenance = [record.get("provenance", {}) for record in records] + runtime_platforms = tuple( + sorted( + set().union( + *(_nested_strings(item, "runtime_platform") for item in provenance) + ) + ) + ) + actual_backends = tuple( + sorted( + set().union(*(_nested_strings(item, "actual_backend") for item in provenance)) + ) + ) + strict_core_ids = set().union( + *(_nested_strings(item, "strict_core_id") for item in provenance) + ) + strict_schedules = set().union( + *(_nested_strings(item, "strict_schedule") for item in provenance) + ) + + if any(_contains_truthy_fallback(item) for item in provenance): + errors.append(f"{label} Attention provenance contains a fallback") + if expected is Implementation.PRODUCTION: + native_backend = f"{framework}.production.attention" + if any(backend != native_backend for backend in backend_ids): + errors.append( + f"{label} did not execute framework-native Attention: {backend_ids}" + ) + else: + if any(not backend.startswith("rlkernel.attention.") for backend in backend_ids): + errors.append(f"{label} did not execute the RL-Kernel Attention wrapper") + if runtime_platforms != ("rocm",): + errors.append(f"{label} did not prove ROCm execution: {runtime_platforms}") + if STRICT_ROCM_ATTENTION_RUNTIME not in actual_backends: + errors.append( + f"{label} did not prove strict AITER/CK runtime " + f"{STRICT_ROCM_ATTENTION_RUNTIME!r}" + ) + if STRICT_ROCM_ATTENTION_CORE not in strict_core_ids: + errors.append( + f"{label} did not prove strict AITER/CK core " + f"{STRICT_ROCM_ATTENTION_CORE!r}" + ) + if STRICT_ROCM_ATTENTION_SCHEDULE not in strict_schedules: + errors.append( + f"{label} did not prove fixed no-Split-KV schedule " + f"{STRICT_ROCM_ATTENTION_SCHEDULE!r}" + ) + if any(_contains_triton(item) for item in provenance): + errors.append(f"{label} strict ROCm Attention used Triton") + + evidence = FrameworkRouteEvidence( + framework=framework, + target=target, + implementation=expected.value, + backend_ids=backend_ids, + call_count=sum(int(record.get("call_count", 0)) for record in records), + readback_count=len(matching), + runtime_platforms=runtime_platforms, + actual_backends=actual_backends, + ) + return evidence, errors + + +def validate_case_readbacks( + case: RocmAttentionAblationCase, + payloads: Sequence[Mapping[str, Any]], +) -> tuple[dict[str, FrameworkRouteEvidence], tuple[str, ...]]: + """Require executed and correctly routed training and rollout evidence.""" + + routes: dict[str, FrameworkRouteEvidence] = {} + errors: list[str] = [] + for side in _FRAMEWORK_TARGETS: + route, route_errors = _validate_route(case, payloads, side=side) + if route is not None: + routes[side] = route + errors.extend(route_errors) + return routes, tuple(errors) + + +def run_rocm_attention_ablation( + command: Sequence[str], + *, + output_dir: Path, + base_environment: Mapping[str, str] | None = None, + case_ids: Iterable[str] = ROCM_ATTENTION_CASE_IDS, + execute: bool = False, +) -> tuple[RocmAblationCaseResult, ...]: + """Execute each case in a fresh process and validate its real readbacks.""" + + matrix = rocm_attention_ablation_matrix(case_ids) + environment = dict(os.environ if base_environment is None else base_environment) + identity = replay_identity(command, environment) + if execute: + validate_replay_environment(environment) + occupied = [output_dir / case.slug for case in matrix if (output_dir / case.slug).exists()] + if occupied: + joined = ", ".join(str(path) for path in occupied) + raise FileExistsError( + "refusing to mix ROCm ablation evidence with existing case directories: " + + joined + ) + validate_rocm_host() + output_dir.mkdir(parents=True, exist_ok=True) + + results: list[RocmAblationCaseResult] = [] + for case in matrix: + case_dir = output_dir / case.slug + readback_dir = case_dir / "readbacks" + log_path = case_dir / "run.log" + case_environment = build_case_environment(environment, case, case_dir=case_dir) + if not execute: + results.append( + RocmAblationCaseResult( + case_id=case.case_id, + status="not_run", + returncode=None, + log_path=log_path, + readback_dir=readback_dir, + ) + ) + continue + + case_dir.mkdir(parents=True, exist_ok=False) + readback_dir.mkdir(parents=True, exist_ok=False) + with log_path.open("w", encoding="utf-8") as log_handle: + process = subprocess.run( + list(command), + env=case_environment, + stdout=log_handle, + stderr=subprocess.STDOUT, + check=False, + ) + try: + payloads = load_runtime_readbacks(readback_dir) + routes, validation_errors = validate_case_readbacks(case, payloads) + except ValueError as exc: + routes = {} + validation_errors = (str(exc),) + errors = list(validation_errors) + if process.returncode != 0: + errors.insert(0, f"orchestration command exited with {process.returncode}") + results.append( + RocmAblationCaseResult( + case_id=case.case_id, + status="passed" if not errors else "failed", + returncode=process.returncode, + log_path=log_path, + readback_dir=readback_dir, + routes=routes, + errors=tuple(errors), + ) + ) + + write_markdown_summary( + output_dir / "summary.md", + results, + replay=identity, + executed=execute, + ) + return tuple(results) + + +def write_markdown_summary( + path: Path, + results: Sequence[RocmAblationCaseResult], + *, + replay: Mapping[str, Any], + executed: bool, +) -> None: + """Write a compact human-readable report; raw JSON remains worker evidence.""" + + lines = [ + "# ROCm rollout/training Attention ablation", + "", + f"- Schema: `{ROCM_ABLATION_SCHEMA_VERSION}`", + f"- Replay identity: `{replay['sha256']}`", + f"- Executed: `{'yes' if executed else 'no'}`", + "- FFN / Logp: fixed at `P/P`", + "", + "| Attention | Megatron training | vLLM rollout | Calls | Result |", + "|---|---|---|---:|---|", + ] + for result in results: + training = result.routes.get("training") + rollout = result.routes.get("rollout") + calls = sum(route.call_count for route in result.routes.values()) + lines.append( + "| {case} | {training} | {rollout} | {calls} | {status} |".format( + case=result.case_id, + training="—" if training is None else training.implementation, + rollout="—" if rollout is None else rollout.implementation, + calls=calls, + status=result.status, + ) + ) + failures = [result for result in results if result.errors] + if failures: + lines.extend(["", "## Failures", ""]) + for result in failures: + lines.append(f"- `{result.case_id}`: {'; '.join(result.errors)}") + lines.extend( + [ + "", + "The report is derived from Megatron and vLLM runtime readbacks. A zero-exit", + "command with missing route evidence is reported as failed.", + "", + ] + ) + path.write_text("\n".join(lines), encoding="utf-8") + + +__all__ = [ + "FROZEN_REPLAY_ENV", + "REQUIRED_REPLAY_ENV", + "ROCM_ABLATION_SCHEMA_VERSION", + "ROCM_ATTENTION_CASE_IDS", + "STRICT_ROCM_ATTENTION_CORE", + "STRICT_ROCM_ATTENTION_RUNTIME", + "STRICT_ROCM_ATTENTION_SCHEDULE", + "RocmAblationCaseResult", + "RocmAttentionAblationCase", + "build_case_environment", + "load_runtime_readbacks", + "replay_identity", + "rocm_attention_ablation_matrix", + "run_rocm_attention_ablation", + "validate_case_readbacks", + "validate_replay_environment", + "validate_rocm_host", + "write_markdown_summary", +] diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py index db651bce..ae65a2f5 100644 --- a/rl_engine/integrations/runtime.py +++ b/rl_engine/integrations/runtime.py @@ -251,8 +251,8 @@ def assert_strict_ready(self) -> None: wrong_backends.append(f"{module}={operator['backend_id']}") elif _contains_triton(operator): wrong_backends.append(f"{module}=triton") - elif _runtime_platform(operator.get("provenance")) != "cuda": - wrong_backends.append(f"{module}=non-cuda") + elif _runtime_platform(operator.get("provenance")) not in {"cuda", "rocm"}: + wrong_backends.append(f"{module}=non-cuda/rocm") failures: list[str] = [] if missing_hooks: failures.append("missing hooks: " + ", ".join(missing_hooks)) diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index 6351f0ab..37e2c521 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -105,7 +105,9 @@ def configure_vllm_environment(plan: IntegrationPlan, *, readback_dir: str | Non os.environ["RL_KERNEL_VLLM_INTEGRATION"] = "1" configure_integration_environment(plan, readback_dir=readback_dir) if plan.implementation_for("attention", "rollout") is Implementation.RL_KERNEL: - os.environ["VLLM_ATTENTION_BACKEND"] = "FLASH_ATTN" + os.environ["VLLM_ATTENTION_BACKEND"] = ( + "ROCM_AITER_FA" if torch.version.hip is not None else "FLASH_ATTN" + ) if plan.implementation_for("logp", "rollout") is Implementation.RL_KERNEL: real_vocab = os.getenv("RL_KERNEL_VLLM_REAL_VOCAB_SIZE", "").strip() padded_vocab = os.getenv("RL_KERNEL_VLLM_PADDED_VOCAB_SIZE", "").strip() @@ -602,19 +604,39 @@ def native(_sampler: Any, *call_args: Any, **call_kwargs: Any) -> Any: def _register_attention_backend(integration: VllmIntegration) -> None: global _RLK_ATTENTION_BACKEND, _RLK_ATTENTION_BUILDER, _RLK_ATTENTION_IMPL - from vllm.v1.attention.backends.flash_attn import ( - FlashAttentionBackend, - FlashAttentionImpl, - FlashAttentionMetadataBuilder, - ) from vllm.v1.attention.backends.registry import AttentionBackendEnum, register_backend + if torch.version.hip is not None: + from vllm.v1.attention.backends.rocm_aiter_fa import ( + AiterFlashAttentionBackend as PlatformAttentionBackend, + ) + from vllm.v1.attention.backends.rocm_aiter_fa import ( + AiterFlashAttentionImpl as PlatformAttentionImpl, + ) + from vllm.v1.attention.backends.rocm_aiter_fa import ( + AiterFlashAttentionMetadataBuilder as PlatformAttentionMetadataBuilder, + ) + + selected_backend = AttentionBackendEnum.ROCM_AITER_FA + else: + from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionBackend as PlatformAttentionBackend, + ) + from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionImpl as PlatformAttentionImpl, + ) + from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionMetadataBuilder as PlatformAttentionMetadataBuilder, + ) + + selected_backend = AttentionBackendEnum.FLASH_ATTN + operator: VllmAttentionOperator | None = None if integration.plan.implementation_for("attention", "rollout") is Implementation.RL_KERNEL: operator = VllmAttentionOperator() integration.install_operator("attention", operator) - class RlKernelFlashAttentionImpl(FlashAttentionImpl): + class RlKernelAttentionImpl(PlatformAttentionImpl): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) if operator is not None: @@ -626,42 +648,41 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: raise RuntimeError("vLLM Attention executed without its installed integration") def native(_impl: Any, *call_args: Any, **call_kwargs: Any) -> Any: - return FlashAttentionImpl.forward(self, *call_args, **call_kwargs) + return PlatformAttentionImpl.forward(self, *call_args, **call_kwargs) return integration.execute("attention", native, self, *args, **kwargs) - class RlKernelFlashAttentionMetadataBuilder(FlashAttentionMetadataBuilder): + class RlKernelAttentionMetadataBuilder(PlatformAttentionMetadataBuilder): pass - class RlKernelFlashAttentionBackend(FlashAttentionBackend): + class RlKernelAttentionBackend(PlatformAttentionBackend): @staticmethod def get_impl_cls() -> type[Any]: - return RlKernelFlashAttentionImpl + return RlKernelAttentionImpl @staticmethod def get_builder_cls() -> type[Any]: - return RlKernelFlashAttentionMetadataBuilder - - RlKernelFlashAttentionImpl.__module__ = __name__ - RlKernelFlashAttentionImpl.__qualname__ = "RlKernelFlashAttentionImpl" - RlKernelFlashAttentionMetadataBuilder.__module__ = __name__ - RlKernelFlashAttentionMetadataBuilder.__qualname__ = "RlKernelFlashAttentionMetadataBuilder" - RlKernelFlashAttentionBackend.__module__ = __name__ - RlKernelFlashAttentionBackend.__qualname__ = "RlKernelFlashAttentionBackend" - _RLK_ATTENTION_IMPL = RlKernelFlashAttentionImpl - _RLK_ATTENTION_BUILDER = RlKernelFlashAttentionMetadataBuilder - _RLK_ATTENTION_BACKEND = RlKernelFlashAttentionBackend - globals()["RlKernelFlashAttentionImpl"] = RlKernelFlashAttentionImpl - globals()["RlKernelFlashAttentionMetadataBuilder"] = RlKernelFlashAttentionMetadataBuilder - globals()["RlKernelFlashAttentionBackend"] = RlKernelFlashAttentionBackend - # vLLM 0.27 selects FLASH_ATTN for Qwen on CUDA before custom third-party - # names are considered. Override the selected enum in-place so the launcher - # does not need to edit vLLM source or pass version-specific CLI flags. + return RlKernelAttentionMetadataBuilder + + RlKernelAttentionImpl.__module__ = __name__ + RlKernelAttentionImpl.__qualname__ = "RlKernelAttentionImpl" + RlKernelAttentionMetadataBuilder.__module__ = __name__ + RlKernelAttentionMetadataBuilder.__qualname__ = "RlKernelAttentionMetadataBuilder" + RlKernelAttentionBackend.__module__ = __name__ + RlKernelAttentionBackend.__qualname__ = "RlKernelAttentionBackend" + _RLK_ATTENTION_IMPL = RlKernelAttentionImpl + _RLK_ATTENTION_BUILDER = RlKernelAttentionMetadataBuilder + _RLK_ATTENTION_BACKEND = RlKernelAttentionBackend + globals()["RlKernelAttentionImpl"] = RlKernelAttentionImpl + globals()["RlKernelAttentionMetadataBuilder"] = RlKernelAttentionMetadataBuilder + globals()["RlKernelAttentionBackend"] = RlKernelAttentionBackend + # Override the platform-selected enum so vLLM keeps its native metadata and + # cache update path while RL-Kernel owns the attention arithmetic call. register_backend( - AttentionBackendEnum.FLASH_ATTN, - f"{__name__}.RlKernelFlashAttentionBackend", + selected_backend, + f"{__name__}.RlKernelAttentionBackend", ) - integration.record_installed_hook("attention", f"{__name__}.RlKernelFlashAttentionBackend") + integration.record_installed_hook("attention", f"{__name__}.RlKernelAttentionBackend") def install_vllm_integration(plan: IntegrationPlan) -> VllmIntegration: diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index d594c164..72342686 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -163,27 +163,74 @@ def __init__( # fallback to the PyTorch reference when AG/RS is required. self.cp_backend = cp_backend self.communication_backend = communication_backend.strip() - self._cuda_runtime_group: Any = None - self._cuda_runtime_bound = False - - def bind_cuda_runtime(self, *, process_group: Any = None) -> Any: - """Bind the shared production CUDA core/transport once per process.""" - - if self._cuda_runtime_bound: - if process_group is not self._cuda_runtime_group: - raise AttentionContractError( - "Attention CUDA runtime is already bound to another process group" - ) - return self.core - from rl_engine.kernels.ops.cuda.attention.strict_runtime import StrictCUDAAttentionRuntime - - runtime = StrictCUDAAttentionRuntime(process_group=process_group) - self.core = runtime - self.cp_backend = runtime - self._cuda_runtime_group = process_group - self._cuda_runtime_bound = True - return runtime - + self._runtime_group: Any = None + self._runtime_platform: str | None = None + + def bind_accelerator_runtime( + self, + tensor: Tensor, + *, + process_group: Any = None, + ) -> Any: + """Bind the strict runtime matching a CUDA or ROCm tensor. + + PyTorch intentionally exposes AMD GPUs through the ``cuda`` device + type. The HIP build marker, rather than ``tensor.device.type``, is + therefore the platform discriminator used by every framework bridge. + """ + + if tensor.device.type != "cuda": + raise AttentionContractError( + "production Attention requires a CUDA or ROCm accelerator tensor" + ) + platform = "rocm" if torch.version.hip is not None else "cuda" + return self.bind_runtime(platform=platform, process_group=process_group) + + def bind_runtime(self, *, platform: str, process_group: Any = None) -> Any: + """Bind one platform runtime and reject process-local identity drift.""" + + normalized = platform.strip().lower() + if normalized not in {"cuda", "rocm"}: + raise AttentionContractError("Attention runtime platform must be 'cuda' or 'rocm'") + if self._runtime_platform is not None: + if normalized != self._runtime_platform: + raise AttentionContractError( + "Attention runtime is already bound to another accelerator platform" + ) + if process_group is not self._runtime_group: + raise AttentionContractError( + "Attention runtime is already bound to another process group" + ) + return self.core + + if normalized == "rocm": + from rl_engine.kernels.ops.rocm.attention.strict_runtime import ( + StrictRocmAttentionRuntime, + ) + + runtime = StrictRocmAttentionRuntime(process_group=process_group) + else: + from rl_engine.kernels.ops.cuda.attention.strict_runtime import ( + StrictCUDAAttentionRuntime, + ) + + runtime = StrictCUDAAttentionRuntime(process_group=process_group) + self.core = runtime + self.cp_backend = runtime + self._runtime_group = process_group + self._runtime_platform = normalized + return runtime + + def bind_cuda_runtime(self, *, process_group: Any = None) -> Any: + """Compatibility entry point for CUDA-only callers.""" + + return self.bind_runtime(platform="cuda", process_group=process_group) + + def bind_rocm_runtime(self, *, process_group: Any = None) -> Any: + """Bind the AITER/CK + RCCL production runtime once per process.""" + + return self.bind_runtime(platform="rocm", process_group=process_group) + def __call__( self, q: Tensor, @@ -395,13 +442,17 @@ def _select_backend( ) return self.cp_backend, _callable_backend_id(self.cp_backend) return self._reference_backend(), REFERENCE_BACKEND_ID - if q.device.type != "cuda" or torch.version.hip is not None: - return self._reference_backend(), REFERENCE_BACKEND_ID - if self.core is None: - from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( - DeterministicAttentionOp, - ) - + if q.device.type != "cuda": + return self._reference_backend(), REFERENCE_BACKEND_ID + if torch.version.hip is not None: + if self.core is None: + self.bind_rocm_runtime() + return self.core, _callable_backend_id(self.core) + if self.core is None: + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + self.core = DeterministicAttentionOp() return self.core, BACKEND_ID diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index e60edb27..23286d94 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -264,7 +264,14 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "deterministic": True, "split_kv": "contract_bound", "reduction_order": "global_block_index", - "strict_schedule": "single_batch_single_query_global_kv_blocks", + "strict_schedules": { + "cuda": "single_batch_fa4_kernel_sm100_no_splitkv", + "rocm": "single_batch_aiter_ck_dense_mha_no_splitkv", + }, + "platform_runtimes": { + "cuda": "rlkernel.cuda.attention.fa4_ag_rs.v1", + "rocm": "rlkernel.rocm.attention.aiter_ck_ag_rs.v1", + }, "strict_observable": True, }, lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, @@ -272,7 +279,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "rl_engine.kernels.ops.pytorch.attention.ablation.AttentionAblationOp" ), fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint="AttentionAblationOp-bitwise-v2", + version_or_build_fingerprint="AttentionAblationOp-platform-runtime-v3", ), OperatorBackendDescriptor( semantic_op="attention", diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index 3edb0a43..ce9ac462 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -139,6 +139,45 @@ def forward_with_lse(self, q, k, v, *, causal, scale): assert torch.equal(result.out, q) +def test_attention_wrapper_binds_rocm_runtime_once(monkeypatch): + import rl_engine.kernels.ops.rocm.attention.strict_runtime as rocm_runtime + + calls = [] + + class Runtime: + def __init__(self, *, process_group=None): + calls.append(process_group) + + monkeypatch.setattr(rocm_runtime, "StrictRocmAttentionRuntime", Runtime) + group = object() + operator = AttentionAblationOp() + + first = operator.bind_rocm_runtime(process_group=group) + second = operator.bind_runtime(platform="rocm", process_group=group) + + assert first is second + assert calls == [group] + assert operator.core is first and operator.cp_backend is first + + +def test_attention_wrapper_rejects_runtime_platform_or_group_drift(monkeypatch): + import rl_engine.kernels.ops.rocm.attention.strict_runtime as rocm_runtime + + class Runtime: + def __init__(self, *, process_group=None): + self.process_group = process_group + + monkeypatch.setattr(rocm_runtime, "StrictRocmAttentionRuntime", Runtime) + group = object() + operator = AttentionAblationOp() + operator.bind_rocm_runtime(process_group=group) + + with pytest.raises(AttentionContractError, match="platform"): + operator.bind_cuda_runtime(process_group=group) + with pytest.raises(AttentionContractError, match="process group"): + operator.bind_rocm_runtime(process_group=object()) + + def test_wrapper_owned_deterministic_core_does_not_require_external_provenance(): q, k, v = _qkv() diff --git a/tests/test_framework_runtime_adapters.py b/tests/test_framework_runtime_adapters.py index 27d54eed..5e140b2a 100644 --- a/tests/test_framework_runtime_adapters.py +++ b/tests/test_framework_runtime_adapters.py @@ -5,10 +5,11 @@ import ast import os +import sys from collections import namedtuple from dataclasses import dataclass from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import pytest import torch @@ -22,6 +23,7 @@ from rl_engine.integrations.framework_operators import ( MegatronAttentionOperator, SemanticOperatorHandle, + VllmAttentionOperator, VllmLogpOperator, _megatron_zigzag_layout, _packed_local_sequence_layout, @@ -35,6 +37,7 @@ from rl_engine.integrations.state import clear_active_integration from rl_engine.integrations.vllm_runtime import ( _patch_qwen3_strict_model, + _register_attention_backend, configure_vllm_environment, ) from rl_engine.kernels.attention_contract import ( @@ -129,6 +132,79 @@ def test_vllm_rlkernel_attention_overrides_selected_flash_attn_backend( assert os.environ["VLLM_ATTENTION_BACKEND"] == "FLASH_ATTN" +def test_vllm_rocm_attention_selects_aiter_metadata_backend(monkeypatch): + monkeypatch.delenv("VLLM_ATTENTION_BACKEND", raising=False) + monkeypatch.setattr(torch.version, "hip", "7.1") + plan = IntegrationPlan.from_case_ids(attention="P/R") + + configure_vllm_environment(plan) + + assert os.environ["VLLM_ATTENTION_BACKEND"] == "ROCM_AITER_FA" + + +def test_vllm_rocm_registration_wraps_aiter_backend(monkeypatch): + selected = [] + + class BackendEnum: + ROCM_AITER_FA = "ROCM_AITER_FA" + FLASH_ATTN = "FLASH_ATTN" + + class AiterImpl: + def __init__(self, *args, **kwargs): + del args, kwargs + + def forward(self, *args, **kwargs): + return args, kwargs + + class AiterBuilder: + pass + + class AiterBackend: + pass + + registry = ModuleType("vllm.v1.attention.backends.registry") + registry.AttentionBackendEnum = BackendEnum + registry.register_backend = lambda backend, path: selected.append((backend, path)) + aiter = ModuleType("vllm.v1.attention.backends.rocm_aiter_fa") + aiter.AiterFlashAttentionBackend = AiterBackend + aiter.AiterFlashAttentionImpl = AiterImpl + aiter.AiterFlashAttentionMetadataBuilder = AiterBuilder + monkeypatch.setitem(sys.modules, registry.__name__, registry) + monkeypatch.setitem(sys.modules, aiter.__name__, aiter) + monkeypatch.setattr(torch.version, "hip", "7.1") + + plan = IntegrationPlan.from_case_ids(attention="P/R") + + class Integration: + def __init__(self): + self.plan = plan + self.installed = {} + self.hooks = [] + + def install_operator(self, module, operator): + self.installed[module] = operator + + def record_installed_hook(self, module, hook): + self.hooks.append((module, hook)) + + integration = Integration() + _register_attention_backend(integration) + + assert selected == [ + ( + BackendEnum.ROCM_AITER_FA, + "rl_engine.integrations.vllm_runtime.RlKernelAttentionBackend", + ) + ] + assert "attention" in integration.installed + assert integration.hooks == [ + ( + "attention", + "rl_engine.integrations.vllm_runtime.RlKernelAttentionBackend", + ) + ] + + def test_megatron_install_is_idempotent_in_one_actor(): class Attention: def forward(self, value): @@ -187,7 +263,8 @@ def test_megatron_packed_attention_runs_each_sequence_in_thd_order(monkeypatch): calls: list[dict[str, object]] = [] class Operator: - def bind_cuda_runtime(self, *, process_group=None): + def bind_accelerator_runtime(self, tensor, *, process_group=None): + assert tensor is query assert process_group == "cp-group" def __call__(self, q, k, v, **kwargs): @@ -219,7 +296,11 @@ def get(self, tensor, *, topology): get_context_parallel_group=lambda: "cp-group", ) monkeypatch.setattr(framework_operators, "_megatron_parallel_state", lambda: parallel_state) - monkeypatch.setattr(framework_operators, "_require_nvidia_cuda", lambda tensor, module: None) + monkeypatch.setattr( + framework_operators, + "_require_attention_accelerator", + lambda tensor: "cuda", + ) packed = SimpleNamespace( qkv_format="thd", cu_seqlens_q=torch.tensor([0, 8, 16], dtype=torch.int32), @@ -246,7 +327,117 @@ def get(self, tensor, *, topology): assert calls[0]["query_position_ids"].tolist() == [ [0, 1, 6, 7], [0, 1, 6, 7], - ] + ] + + +def test_megatron_attention_binds_rocm_core_and_schedule(monkeypatch): + calls = [] + + class Operator: + def bind_accelerator_runtime(self, tensor, *, process_group=None): + calls.append((tensor, process_group)) + + def __call__(self, q, k, v, **kwargs): + del k, v + calls.append(kwargs) + return SimpleNamespace( + out=q.clone(), + provenance={ + "actual_backend": "rlkernel.rocm.attention.aiter_ck_ag_rs.v1", + "fallback": False, + }, + ) + + operator = Operator() + + class Handle: + provenance = {} + + def get(self, tensor, *, topology): + assert topology["context_parallel_size"] == 1 + return operator + + parallel_state = SimpleNamespace( + get_context_parallel_world_size=lambda: 1, + get_context_parallel_rank=lambda: 0, + get_tensor_model_parallel_world_size=lambda: 2, + get_tensor_model_parallel_rank=lambda: 0, + ) + monkeypatch.setattr(framework_operators, "_megatron_parallel_state", lambda: parallel_state) + monkeypatch.setattr( + framework_operators, + "_require_attention_accelerator", + lambda tensor: "rocm", + ) + query = torch.zeros(4, 1, 2, 8, dtype=torch.bfloat16) + key = torch.zeros(4, 1, 1, 8, dtype=torch.bfloat16) + adapter = MegatronAttentionOperator(handle=Handle()) + + output = adapter(SimpleNamespace(softmax_scale=0.25), query, key, key, None) + + assert output.shape == (4, 1, 16) + assert calls[0] == (query, None) + config = calls[1]["config"] + assert config.strict_core_id == "rlkernel.attention.rocm.aiter_ck_dense_mha.v1" + assert config.strict_schedule == "single_batch_aiter_ck_dense_mha_no_splitkv" + assert calls[1]["communication_backend"] == "none" + assert adapter.provenance["execution"]["runtime_platform"] == "rocm" + + +def test_vllm_attention_routes_paged_cache_to_rocm_runtime(monkeypatch): + runtime_calls = [] + + class Runtime: + def forward_paged_with_lse(self, q, k, v, **kwargs): + runtime_calls.append((q, k, v, kwargs)) + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:-1], dtype=torch.float32), + provenance={ + "actual_backend": "rlkernel.rocm.attention.aiter_ck_ag_rs.v1", + "fallback": False, + }, + ) + + class Operator: + def bind_accelerator_runtime(self, tensor, *, process_group=None): + assert process_group is None + return Runtime() + + class Handle: + provenance = {} + + def get(self, tensor, *, topology): + assert topology["context_parallel_size"] == 1 + return Operator() + + monkeypatch.setattr( + framework_operators, + "_require_attention_accelerator", + lambda tensor: "rocm", + ) + query = torch.zeros(1, 2, 8, dtype=torch.bfloat16) + kv_cache = torch.zeros(2, 1, 4, 16, dtype=torch.bfloat16) + metadata = SimpleNamespace( + block_table=torch.tensor([[0]], dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([1], dtype=torch.int32), + num_actual_tokens=1, + max_seq_len=1, + ) + impl = SimpleNamespace(head_size=8, num_heads=2, num_kv_heads=1, scale=8**-0.5) + adapter = VllmAttentionOperator(handle=Handle()) + + output = adapter(impl, object(), query, query, query, kv_cache, metadata) + + assert output.shape == (1, 16) + assert len(runtime_calls) == 1 + assert runtime_calls[0][0].shape == (1, 2, 1, 8) + assert adapter.provenance["execution"]["runtime_platform"] == "rocm" + assert ( + adapter.provenance["execution"]["materialization"] + == "logical_paged_kv_to_aiter_ck_dense" + ) def test_vllm_current_flash_attention_kv_cache_layout_is_materialized(): @@ -260,6 +451,49 @@ def test_vllm_current_flash_attention_kv_cache_layout_is_materialized(): assert torch.equal(value, cache.transpose(1, 2)[..., 5:]) +def test_vllm_rocm_kv_cache_pair_axis_is_materialized(): + cache = torch.arange(3 * 2 * 4 * 1 * 5).reshape(3, 2, 4, 1, 5) + + key, value = _vllm_kv_cache_views( + cache, + head_size=5, + num_kv_heads=1, + platform="rocm", + ) + + assert key.shape == (3, 4, 1, 5) + assert value.shape == (3, 4, 1, 5) + assert torch.equal(key, cache[:, 0]) + assert torch.equal(value, cache[:, 1]) + + +def test_vllm_rocm_lhbnc_kv_cache_is_normalized_to_block_major(): + cache = torch.arange(2 * 1 * 3 * 4 * 5).reshape(2, 1, 3, 4, 5) + + key, value = _vllm_kv_cache_views(cache, head_size=5, num_kv_heads=1) + + assert key.shape == (3, 4, 1, 5) + assert value.shape == (3, 4, 1, 5) + assert torch.equal(key, cache[0].permute(1, 2, 0, 3)) + assert torch.equal(value, cache[1].permute(1, 2, 0, 3)) + + +def test_vllm_rocm_flattened_kv_cache_is_unpacked(): + cache = torch.arange(3 * 2 * 4 * 10).reshape(3, 2, 4, 10) + + key, value = _vllm_kv_cache_views( + cache, + head_size=5, + num_kv_heads=2, + platform="rocm", + ) + + assert key.shape == (3, 4, 2, 5) + assert value.shape == (3, 4, 2, 5) + assert torch.equal(key.flatten(2), cache[:, 0]) + assert torch.equal(value.flatten(2), cache[:, 1]) + + def test_megatron_strict_attention_projections_install_without_debug_environment( monkeypatch, ): @@ -519,3 +753,25 @@ def test_strict_readback_accepts_cuda_without_triton(): integration.execute("attention", lambda value: value, "x") integration.assert_strict_ready() + + +def test_strict_readback_accepts_rocm_without_triton(): + plan = IntegrationPlan.from_case_ids(attention="R/R") + integration = FrameworkOperatorIntegration( + framework="megatron", + target="training", + plan=plan, + rl_kernel_operators={ + "attention": _ReadbackOperator( + { + "runtime_platform": "rocm", + "actual_backend": "rlkernel.rocm.attention.aiter_ck_ag_rs.v1", + "triton_used": False, + } + ) + }, + ) + integration.record_installed_hook("attention", "test.attention") + integration.execute("attention", lambda value: value, "x") + + integration.assert_strict_ready() diff --git a/tests/test_rocm_e2e_ablation.py b/tests/test_rocm_e2e_ablation.py new file mode 100644 index 00000000..295936af --- /dev/null +++ b/tests/test_rocm_e2e_ablation.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import rl_engine.integrations.rocm_ablation as rocm_ablation +from rl_engine.integrations.ablation import Implementation +from rl_engine.integrations.rocm_ablation import ( + STRICT_ROCM_ATTENTION_CORE, + STRICT_ROCM_ATTENTION_RUNTIME, + STRICT_ROCM_ATTENTION_SCHEDULE, + build_case_environment, + replay_identity, + rocm_attention_ablation_matrix, + run_rocm_attention_ablation, + validate_case_readbacks, + validate_replay_environment, +) + + +def _environment() -> dict[str, str]: + return { + "MODEL_ROOT": "/models/Qwen3-8B", + "TORCH_DIST_ROOT": "/models/Qwen3-8B_torch_dist", + "VIME_CKPT": "/checkpoints/pre-update", + "PROMPT_DATA": "/data/prompts.jsonl", + "NUM_ROLLOUT": "1", + "TRAIN_SEED": "1234", + "ROLLOUT_SEED": "42", + } + + +def _readback(case, *, side: str, calls: int = 1, strict: bool | None = None): + if side == "training": + framework, target = "megatron", "training" + else: + framework, target = "vllm", "rollout" + implementation = case.implementation_for(target) + is_strict = implementation is Implementation.RL_KERNEL if strict is None else strict + backend = ( + "rlkernel.attention.deterministic.v1" + if is_strict + else f"{framework}.production.attention" + ) + provenance = ( + { + "execution": { + "runtime_platform": "rocm", + "operator": { + "actual_backend": STRICT_ROCM_ATTENTION_RUNTIME, + "strict_core_id": STRICT_ROCM_ATTENTION_CORE, + "strict_schedule": STRICT_ROCM_ATTENTION_SCHEDULE, + "fallback": False, + "triton_used": False, + }, + } + } + if is_strict + else {} + ) + return { + "framework": framework, + "target": target, + "plan": case.plan.to_dict(), + "installed_hooks": {"attention": f"{framework}.attention"}, + "fallbacks": [], + "operators": { + "attention": { + "implementation": implementation.value, + "backend_id": backend, + "call_count": calls, + "provenance": provenance, + } + }, + } + + +def test_matrix_is_pr230_four_cell_attention_matrix(): + matrix = rocm_attention_ablation_matrix() + + assert tuple(case.case_id for case in matrix) == ("P/P", "P/R", "R/P", "R/R") + assert [ + ( + case.implementation_for("training"), + case.implementation_for("rollout"), + ) + for case in matrix + ] == [ + (Implementation.PRODUCTION, Implementation.PRODUCTION), + (Implementation.PRODUCTION, Implementation.RL_KERNEL), + (Implementation.RL_KERNEL, Implementation.PRODUCTION), + (Implementation.RL_KERNEL, Implementation.RL_KERNEL), + ] + assert all(case.plan.cases["ffn"].case_id == "P/P" for case in matrix) + assert all(case.plan.cases["logp"].case_id == "P/P" for case in matrix) + + +def test_matrix_rejects_unknown_and_duplicate_cells(): + with pytest.raises(ValueError, match="unknown"): + rocm_attention_ablation_matrix(["R/X"]) + with pytest.raises(ValueError, match="unique"): + rocm_attention_ablation_matrix(["R/R", "R/R"]) + + +def test_replay_identity_changes_only_when_frozen_inputs_change(): + environment = _environment() + first = replay_identity(["bash", "run.sh"], environment) + environment["RL_KERNEL_ATTENTION_CASE"] = "R/R" + second = replay_identity(["bash", "run.sh"], environment) + environment["PROMPT_DATA"] = "/data/other.jsonl" + changed = replay_identity(["bash", "run.sh"], environment) + + assert first == second + assert changed["sha256"] != first["sha256"] + + +def test_executable_replay_requires_inputs_and_single_rollout(): + environment = _environment() + validate_replay_environment(environment) + environment.pop("PROMPT_DATA") + with pytest.raises(RuntimeError, match="PROMPT_DATA"): + validate_replay_environment(environment) + environment = _environment() + environment["NUM_ROLLOUT"] = "2" + with pytest.raises(RuntimeError, match="NUM_ROLLOUT=1"): + validate_replay_environment(environment) + + +def test_case_environment_propagates_plan_to_megatron_vllm_and_ray(tmp_path): + case = rocm_attention_ablation_matrix(["P/R"])[0] + environment = build_case_environment(_environment(), case, case_dir=tmp_path / "p-r") + + assert environment["RL_KERNEL_ATTENTION_CASE"] == "P/R" + assert environment["RL_KERNEL_FFN_CASE"] == "P/P" + assert environment["RL_KERNEL_LOGP_CASE"] == "P/P" + assert environment["RL_KERNEL_VLLM_INTEGRATION"] == "1" + assert environment["RL_KERNEL_PLATFORM"] == "rocm" + assert environment["VLLM_ATTENTION_BACKEND"] == "ROCM_AITER_FA" + assert environment["RL_KERNEL_READBACK_DIR"].endswith("p-r/readbacks") + assert environment["NUM_ROLLOUT"] == "1" + + +@pytest.mark.parametrize("case_id", ["P/P", "P/R", "R/P", "R/R"]) +def test_case_readbacks_validate_both_real_framework_sides(case_id): + case = rocm_attention_ablation_matrix([case_id])[0] + routes, errors = validate_case_readbacks( + case, + [ + _readback(case, side="training", calls=11), + _readback(case, side="rollout", calls=17), + ], + ) + + assert errors == () + assert set(routes) == {"training", "rollout"} + assert routes["training"].call_count == 11 + assert routes["rollout"].call_count == 17 + + +def test_strict_side_requires_rocm_runtime_identity(): + case = rocm_attention_ablation_matrix(["R/R"])[0] + bad = _readback(case, side="rollout") + bad["operators"]["attention"]["provenance"]["execution"]["operator"][ + "actual_backend" + ] = "rlkernel.cuda.attention.fa4_ag_rs.v1" + + _routes, errors = validate_case_readbacks( + case, + [_readback(case, side="training"), bad], + ) + + assert any("strict AITER/CK runtime" in error for error in errors) + + +def test_strict_side_rejects_triton_and_missing_fixed_schedule(): + case = rocm_attention_ablation_matrix(["P/R"])[0] + bad = _readback(case, side="rollout") + operator = bad["operators"]["attention"]["provenance"]["execution"]["operator"] + operator["triton_used"] = True + operator["core_actual_backends"] = ["triton.attention"] + operator.pop("strict_schedule") + + _routes, errors = validate_case_readbacks( + case, + [_readback(case, side="training"), bad], + ) + + assert any("fixed no-Split-KV schedule" in error for error in errors) + assert any("used Triton" in error for error in errors) + + +def test_production_side_requires_framework_native_backend_identity(): + case = rocm_attention_ablation_matrix(["P/P"])[0] + bad = _readback(case, side="training") + bad["operators"]["attention"]["backend_id"] = ( + "rlkernel.attention.deterministic.v1" + ) + + _routes, errors = validate_case_readbacks( + case, + [bad, _readback(case, side="rollout")], + ) + + assert any("framework-native Attention" in error for error in errors) + + +def test_zero_exit_without_framework_evidence_is_a_failure(): + case = rocm_attention_ablation_matrix(["R/R"])[0] + _routes, errors = validate_case_readbacks(case, []) + + assert errors == ( + "missing megatron/training runtime readback", + "missing vllm/rollout runtime readback", + ) + + +def test_malformed_call_count_is_a_validation_error(): + case = rocm_attention_ablation_matrix(["P/P"])[0] + bad = _readback(case, side="training") + bad["operators"]["attention"]["call_count"] = "not-an-integer" + + _routes, errors = validate_case_readbacks( + case, + [bad, _readback(case, side="rollout")], + ) + + assert any("invalid call count" in error for error in errors) + + +def test_dry_run_writes_human_summary_without_fabricating_results(tmp_path): + results = run_rocm_attention_ablation( + ["bash", "run.sh"], + output_dir=tmp_path, + base_environment={}, + case_ids=["P/P", "R/R"], + execute=False, + ) + + assert [result.status for result in results] == ["not_run", "not_run"] + summary = (tmp_path / "summary.md").read_text(encoding="utf-8") + assert "Executed: `no`" in summary + assert "| P/P |" in summary and "| R/R |" in summary + assert not list(tmp_path.glob("*.json")) + + +def test_execute_rejects_stale_case_evidence_before_gpu_probe(monkeypatch, tmp_path): + (tmp_path / "r-r").mkdir() + probed = False + + def probe(): + nonlocal probed + probed = True + + monkeypatch.setattr(rocm_ablation, "validate_rocm_host", probe) + + with pytest.raises(FileExistsError, match="existing case directories"): + run_rocm_attention_ablation( + ["bash", "run.sh"], + output_dir=tmp_path, + base_environment=_environment(), + case_ids=["R/R"], + execute=True, + ) + + assert probed is False + + +def test_execute_runs_each_case_fresh_and_checks_emitted_readbacks(monkeypatch, tmp_path): + invocations: list[str] = [] + monkeypatch.setattr(rocm_ablation, "validate_rocm_host", lambda: {"hip_runtime": "7.1"}) + + def fake_run(command, *, env, stdout, stderr, check): + del command, stdout, stderr, check + case = rocm_attention_ablation_matrix([env["RL_KERNEL_ATTENTION_CASE"]])[0] + readback_dir = Path(env["RL_KERNEL_READBACK_DIR"]) + invocations.append(case.case_id) + for side in ("training", "rollout"): + payload = _readback(case, side=side) + (readback_dir / f"{side}.json").write_text( + json.dumps(payload), + encoding="utf-8", + ) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(rocm_ablation.subprocess, "run", fake_run) + results = run_rocm_attention_ablation( + ["bash", "run.sh"], + output_dir=tmp_path, + base_environment=_environment(), + case_ids=["P/R", "R/P"], + execute=True, + ) + + assert invocations == ["P/R", "R/P"] + assert [result.status for result in results] == ["passed", "passed"] + assert all(result.log_path.is_file() for result in results) + assert "Executed: `yes`" in (tmp_path / "summary.md").read_text(encoding="utf-8")