From 9c47c85375f84b9edf11345ac9ffe5d8d96404f7 Mon Sep 17 00:00:00 2001 From: FED4 Date: Sat, 13 Jun 2026 11:59:31 +0800 Subject: [PATCH 1/2] feat: add ROCm FlashAttention backend --- .github/workflows/ci.yml | 4 + docs/getting_started/installation.md | 27 ++ .../kernels/ops/pytorch/attention/__init__.py | 34 ++ .../kernels/ops/rocm/attention/__init__.py | 8 + .../kernels/ops/rocm/attention/flash_attn.py | 75 ++++ rl_engine/kernels/registry.py | 28 +- rl_engine/tests/test_dispatch.py | 19 +- scripts/check_rocm_env.py | 138 +++++++ tests/test_attention_correctness.py | 351 ++++++++++++++++++ 9 files changed, 680 insertions(+), 4 deletions(-) create mode 100644 rl_engine/kernels/ops/rocm/attention/__init__.py create mode 100644 rl_engine/kernels/ops/rocm/attention/flash_attn.py create mode 100644 scripts/check_rocm_env.py create mode 100644 tests/test_attention_correctness.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 93c41a132..755f426d4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,6 +9,9 @@ on: pull_request: branches: [ main ] +permissions: + contents: read + jobs: linting: runs-on: ubuntu-latest @@ -63,6 +66,7 @@ jobs: - name: Run Mocked Hardware Discovery Tests run: | python -m pytest rl_engine/tests/test_dispatch.py -v + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rs docs: runs-on: ubuntu-latest diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index 5ea232c7e..d45635173 100644 --- a/docs/getting_started/installation.md +++ b/docs/getting_started/installation.md @@ -31,6 +31,33 @@ vLLM runtime. Core CI and mocked integration tests do not require it. For common CUDA, ROCm, vLLM, fallback, and CI questions, see the [FAQ](faq.md). +### ROCm Backend + +Use a ROCm PyTorch build that matches the installed ROCm toolchain. Then install +FlashAttention with an AMD backend: + +```bash +python -m pip install ninja packaging wheel psutil einops +git clone --recurse-submodules https://github.com/Dao-AILab/flash-attention.git +cd flash-attention +FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE \ + python -m pip install --no-build-isolation --no-deps . +cd .. +``` + +Verify the environment from the RL-Kernel checkout: + +```bash +python scripts/check_rocm_env.py +``` + +RL-Kernel uses external FlashAttention as the default ROCm attention path. To +fall back to PyTorch SDPA for ROCm attention dispatch, set: + +```bash +export RL_KERNEL_ROCM_ATTN_BACKEND=sdpa +``` + ## Development Dependencies ```bash diff --git a/rl_engine/kernels/ops/pytorch/attention/__init__.py b/rl_engine/kernels/ops/pytorch/attention/__init__.py index e69de29bb..d315e5be4 100644 --- a/rl_engine/kernels/ops/pytorch/attention/__init__.py +++ b/rl_engine/kernels/ops/pytorch/attention/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import torch +import torch.nn.functional as F + + +class NativeAttentionOp: + """PyTorch SDPA fallback for FlashAttention-layout tensors.""" + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + causal: bool = False, + ) -> torch.Tensor: + q_ref = q.transpose(1, 2) + k_ref = k.transpose(1, 2) + v_ref = v.transpose(1, 2) + out = F.scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + dropout_p=dropout_p, + is_causal=causal, + scale=softmax_scale, + ) + return out.transpose(1, 2) + + +__all__ = ["NativeAttentionOp"] diff --git a/rl_engine/kernels/ops/rocm/attention/__init__.py b/rl_engine/kernels/ops/rocm/attention/__init__.py new file mode 100644 index 000000000..150c937f5 --- /dev/null +++ b/rl_engine/kernels/ops/rocm/attention/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .flash_attn import RocmFlashAttentionOp + +__all__ = [ + "RocmFlashAttentionOp", +] diff --git a/rl_engine/kernels/ops/rocm/attention/flash_attn.py b/rl_engine/kernels/ops/rocm/attention/flash_attn.py new file mode 100644 index 000000000..56b2e8c6c --- /dev/null +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os + +import torch + +from rl_engine.utils.logger import logger + + +def _select_flash_attn_backend() -> str: + """Select the installed FlashAttention ROCm backend before importing flash_attn.""" + if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE", "").upper() == "TRUE": + return "triton" + if importlib.util.find_spec("flash_attn_2_cuda") is None: + os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE" + return "triton" + return "ck" + + +class RocmFlashAttentionOp: + """ + Standard FlashAttention wrapper for ROCm. + Demonstrates the reference structure for adding new operator families. + """ + + def __init__(self): + if torch.version.hip is None: + raise RuntimeError("RocmFlashAttentionOp requires a ROCm PyTorch build.") + + backend = _select_flash_attn_backend() + try: + from flash_attn import flash_attn_func + + self.op = flash_attn_func + logger.info("Successfully linked to external flash_attn library (%s backend).", backend) + except ImportError as exc: + raise RuntimeError( + "ROCm FlashAttention requires a ROCm-compatible flash-attn installation. " + "See docs/getting_started/installation.md#rocm-backend." + ) from exc + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + causal: bool = False, + ) -> torch.Tensor: + """ + Standard attention forward pass. + Args: + q: (batch, seqlen, nheads, headdim) + k: (batch, seqlen, nheads_k, headdim) + v: (batch, seqlen, nheads_k, headdim) + """ + valid_dtypes = (torch.float16, torch.bfloat16) + if ( + q.dtype not in valid_dtypes + or k.dtype not in valid_dtypes + or v.dtype not in valid_dtypes + ): + raise TypeError("FlashAttention requires FP16 or BF16 for q/k/v") + # PyTorch uses the CUDA device API for both CUDA and ROCm tensors. + if not (q.is_cuda and k.is_cuda and v.is_cuda): + raise ValueError("Inputs must be on a CUDA/ROCm GPU device") + if not (q.device == k.device == v.device): + raise ValueError("q, k, and v must be on the same device") + + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + + return self.op(q, k, v, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=causal) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 7aae08fbc..7831c4ef3 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -2,6 +2,7 @@ # Copyright (c) 2026 RL-Kernel Contributors import importlib +import os from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type @@ -32,6 +33,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # AMD ROCm optimized stack ROCM_AITER = "rl_engine.kernels.ops.rocm.aiter.AiterOp" ROCM_CK = "rl_engine.kernels.ops.rocm.composable_kernel.CKOp" + ROCM_FLASH_ATTN = "rl_engine.kernels.ops.rocm.attention.flash_attn.RocmFlashAttentionOp" # GRPO loss (group reward normalization + clipped surrogate + KL) TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" @@ -43,6 +45,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # Generic fallback TRITON_GENERIC = "rl_engine.kernels.ops.triton.generic.TritonOp" + PYTORCH_ATTN = "rl_engine.kernels.ops.pytorch.attention.NativeAttentionOp" PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" @@ -76,27 +79,46 @@ def __init__(self): OpBackend.CUDA_FUSED_LOGP_GENERIC, OpBackend.PYTORCH_NATIVE, ], - "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], + "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], # Default dispatch logic for new operators }, "rocm": { "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], - "attn": [OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], + "attn": [ + OpBackend.PYTORCH_ATTN, + OpBackend.ROCM_FLASH_ATTN, + OpBackend.TRITON_GENERIC, + ], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], - "attn": [OpBackend.PYTORCH_NATIVE], + "attn": [OpBackend.PYTORCH_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], }, } logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") + self._adjust_priority_from_env() self._adjust_priority_for_hardware() + def _adjust_priority_from_env(self): + rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() + if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: + self._priority_map["rocm"]["attn"] = [ + OpBackend.ROCM_FLASH_ATTN, + OpBackend.PYTORCH_ATTN, + OpBackend.TRITON_GENERIC, + ] + elif rocm_attn_backend and rocm_attn_backend not in {"native", "pytorch", "sdpa"}: + logger.warning( + "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.", + rocm_attn_backend, + ) + def _adjust_priority_for_hardware(self): """Prioritize the fused TMA LogP kernel only when it is compiled into the extension and the device is TMA-capable (SM90/100/120).""" diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index d15b1ae4c..3fa361a6a 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -4,7 +4,7 @@ import torch from rl_engine.executors.rollout import RolloutExecutor -from rl_engine.kernels.registry import kernel_registry +from rl_engine.kernels.registry import KernelRegistry, OpBackend, kernel_registry from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -25,6 +25,23 @@ def test_device_and_registry(): logger.info(f"Retrieved Attention Operator: {attn_op}") +def test_rocm_attention_uses_native_sdpa_by_default(monkeypatch): + monkeypatch.delenv("RL_KERNEL_ROCM_ATTN_BACKEND", raising=False) + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"][0] == OpBackend.PYTORCH_ATTN + + +def test_rocm_attention_flash_attn_opt_in(monkeypatch): + monkeypatch.setenv("RL_KERNEL_ROCM_ATTN_BACKEND", " flash_attn ") + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"][0] == OpBackend.ROCM_FLASH_ATTN + assert registry._priority_map["rocm"]["attn"][1] == OpBackend.PYTORCH_ATTN + + def test_executor_flow(): executor = RolloutExecutor() mock_input_ids = torch.ones((1, 16), dtype=torch.long) diff --git a/scripts/check_rocm_env.py b/scripts/check_rocm_env.py new file mode 100644 index 000000000..0819d2fcd --- /dev/null +++ b/scripts/check_rocm_env.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import argparse +import importlib.util +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + + +def _fail(message: str) -> None: + raise SystemExit(f"ERROR: {message}") + + +def _rocm_major_minor(version: str | None) -> tuple[str, str] | None: + if not version: + return None + match = re.search(r"([0-9]+)\.([0-9]+)", version) + if match is None: + return None + return match.group(1), match.group(2) + + +def _find_hipcc() -> str | None: + for env_name in ("ROCM_HOME", "HIP_PATH"): + env_path = os.environ.get(env_name) + if env_path: + hipcc = Path(env_path) / "bin" / "hipcc" + if hipcc.exists(): + return str(hipcc) + + hipcc = shutil.which("hipcc") + if hipcc: + return hipcc + + fallback = Path("/opt/rocm/bin/hipcc") + if fallback.exists(): + return str(fallback) + + return None + + +def _flash_attn_backend() -> str | None: + if importlib.util.find_spec("flash_attn") is None: + return None + if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE", "").upper() == "TRUE": + return "triton" + if importlib.util.find_spec("flash_attn_2_cuda") is not None: + return "ck" + return "triton-available-if-enabled" + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Check the local ROCm PyTorch environment.") + parser.add_argument( + "--require-flash-attn", + action="store_true", + help="fail if flash_attn_func cannot be imported", + ) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + + try: + import torch + except ImportError as exc: + _fail(f"PyTorch is not installed in {sys.executable}: {exc}") + + if torch.version.hip is None: + _fail(f"PyTorch is not a ROCm build: torch={torch.__version__}") + + if not torch.cuda.is_available(): + _fail("ROCm GPU is not available to PyTorch") + + hipcc = _find_hipcc() + if hipcc is None: + _fail("Could not find hipcc from ROCM_HOME, HIP_PATH, PATH, or /opt/rocm/bin/hipcc") + + try: + hipcc_output = subprocess.check_output([hipcc, "--version"], text=True) + except (subprocess.CalledProcessError, OSError) as exc: + _fail(f"Could not run {hipcc} --version: {exc}") + + torch_rocm = _rocm_major_minor(torch.version.hip) + hipcc_rocm = _rocm_major_minor(hipcc_output) + if torch_rocm is None: + _fail(f"Could not parse torch.version.hip={torch.version.hip!r}") + if hipcc_rocm is None: + _fail(f"Could not parse {hipcc} --version output") + if torch_rocm != hipcc_rocm: + _fail( + "ROCm version mismatch: " + f"torch.version.hip={torch.version.hip}, hipcc major/minor={'.'.join(hipcc_rocm)}" + ) + + device_name = torch.cuda.get_device_name(0) + capability = torch.cuda.get_device_capability(0) + triton_available = importlib.util.find_spec("triton") is not None + flash_attn_backend = _flash_attn_backend() + flash_attn_func_available = False + flash_attn_error = None + if flash_attn_backend is not None: + if flash_attn_backend == "triton-available-if-enabled": + os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE" + try: + from flash_attn import flash_attn_func + except (ImportError, OSError, RuntimeError) as exc: + flash_attn_error = str(exc) + else: + flash_attn_func_available = flash_attn_func is not None + + print(f"Python: {sys.executable}") + print(f"torch: {torch.__version__}") + print(f"torch.version.hip: {torch.version.hip}") + print(f"hipcc: {hipcc}") + print(f"hipcc ROCm: {'.'.join(hipcc_rocm)}") + print(f"GPU: {device_name}") + print(f"compute capability: {capability}") + print(f"triton: {'available' if triton_available else 'missing'}") + print(f"flash_attn backend: {flash_attn_backend or 'not installed'}") + print(f"flash_attn_func: {'available' if flash_attn_func_available else 'not available'}") + + if flash_attn_error: + print(f"flash_attn import error: {flash_attn_error}") + + if args.require_flash_attn and not flash_attn_func_available: + _fail("flash_attn_func is required but could not be imported") + + +if __name__ == "__main__": + main() diff --git a/tests/test_attention_correctness.py b/tests/test_attention_correctness.py new file mode 100644 index 000000000..767b140e7 --- /dev/null +++ b/tests/test_attention_correctness.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import os +from contextlib import contextmanager + +import pytest +import torch +import torch.nn.functional as F + +try: + from torch.nn.attention import SDPBackend, sdpa_kernel +except ImportError: + SDPBackend = None + sdpa_kernel = None + + +DTYPE_CASES = [ + pytest.param(torch.float16, 1e-3, 1e-3, id="fp16"), + pytest.param(torch.bfloat16, 2e-2, 2e-2, id="bf16"), +] + +AVAILABILITY_ERRORS = (ImportError, ModuleNotFoundError, OSError, RuntimeError) + +ATTENTION_SHAPES = [ + pytest.param(1, 128, 4, 64, id="b1-s128-h4-d64"), + pytest.param(2, 256, 8, 128, id="b2-s256-h8-d128"), + pytest.param(1, 512, 8, 256, id="b1-s512-h8-d256"), + pytest.param(1, 1024, 16, 64, id="b1-s1024-h16-d64"), + pytest.param(1, 2048, 16, 64, id="b1-s2048-h16-d64"), +] + + +@contextmanager +def sdpa_math_backend(): + if sdpa_kernel is not None and SDPBackend is not None: + with sdpa_kernel(SDPBackend.MATH): + yield + return + + if hasattr(torch.backends.cuda, "sdp_kernel"): + with torch.backends.cuda.sdp_kernel( + enable_flash=False, + enable_math=True, + enable_mem_efficient=False, + ): + yield + return + + raise RuntimeError("PyTorch SDPA backend selector is unavailable") + + +def should_print_attention_diff(): + return os.getenv("PRINT_ATTENTION_DIFF", "").lower() in {"1", "true", "yes"} + + +def pytorch_sdpa_reference(q, k, v, *, causal, softmax_scale): + """ + Compute a PyTorch SDPA math-backend reference for FlashAttention-layout inputs. + + Inputs use FlashAttention layout: (batch, seqlen, nheads, headdim). + PyTorch SDPA uses: (batch, nheads, seqlen, headdim), so this helper + transposes before and after the reference call. + """ + q_ref = q.transpose(1, 2).contiguous() + k_ref = k.transpose(1, 2).contiguous() + v_ref = v.transpose(1, 2).contiguous() + + with sdpa_math_backend(): + expected = F.scaled_dot_product_attention( + q_ref.float(), + k_ref.float(), + v_ref.float(), + dropout_p=0.0, + is_causal=causal, + scale=softmax_scale, + ) + + return expected.transpose(1, 2).contiguous() + + +def make_qkv(batch, seqlen, nheads, headdim, device, dtype): + torch.manual_seed(0) + q = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype) + k = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype) + v = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype) + return q, k, v + + +def describe_attention_diff(actual, expected, *, dtype, atol, rtol, causal, softmax_scale): + actual_f = actual.float() + expected_f = expected.float() + abs_diff = (actual_f - expected_f).abs() + rel_diff = abs_diff / expected_f.abs().clamp_min(1e-12) + + return ( + "FlashAttention vs PyTorch SDPA math backend diff: " + f"dtype={dtype}, shape={tuple(actual.shape)}, causal={causal}, " + f"softmax_scale={softmax_scale}, atol={atol}, rtol={rtol}, " + f"max_abs_diff={abs_diff.max().item():.6g}, " + f"mean_abs_diff={abs_diff.mean().item():.6g}, " + f"max_rel_diff={rel_diff.max().item():.6g}" + ) + + +def print_attention_diff(actual, expected, *, dtype, atol, rtol, causal, softmax_scale): + if should_print_attention_diff(): + print( + describe_attention_diff( + actual, + expected, + dtype=dtype, + atol=atol, + rtol=rtol, + causal=causal, + softmax_scale=softmax_scale, + ) + ) + + +def is_cuda_platform(): + return torch.cuda.is_available() and torch.version.hip is None + + +def is_rocm_platform(): + return torch.cuda.is_available() and torch.version.hip is not None + + +def cuda_flash_attention_availability(): + if not torch.cuda.is_available(): + return ( + False, + "CUDA is not available, check CUDA device, driver/runtime compatibility, " + "and torch CUDA build", + ) + if not is_cuda_platform(): + return False, "current torch build is not CUDA platform" + try: + from rl_engine.kernels.ops.cuda.attention.flash_attn import FlashAttentionOp + + FlashAttentionOp() + except AVAILABILITY_ERRORS as exc: + return False, f"CUDA FlashAttentionOp is unavailable: {exc}" + return True, "" + + +def rocm_flash_attention_availability(): + if not torch.cuda.is_available(): + return ( + False, + "ROCm is not available, check AMD GPU device, driver/runtime compatibility, " + "and torch ROCm build", + ) + if not is_rocm_platform(): + return False, "current torch build is not ROCm platform" + try: + from rl_engine.kernels.ops.rocm.attention.flash_attn import RocmFlashAttentionOp + + RocmFlashAttentionOp() + except AVAILABILITY_ERRORS as exc: + return False, f"ROCm FlashAttentionOp is unavailable: {exc}" + return True, "" + + +def native_attention_availability(): + if not torch.cuda.is_available(): + return False, "CUDA/ROCm GPU is not available" + try: + from rl_engine.kernels.ops.pytorch.attention import NativeAttentionOp + + NativeAttentionOp() + except AVAILABILITY_ERRORS as exc: + return False, f"NativeAttentionOp is unavailable: {exc}" + return True, "" + + +def assert_flash_attention_matches_sdpa( + op, + dtype, + atol, + rtol, + causal, + use_explicit_scale, + batch, + seqlen, + nheads, + headdim, +): + # PyTorch exposes both NVIDIA CUDA and AMD ROCm GPUs through the "cuda" device API. + device = torch.device("cuda") + q, k, v = make_qkv(batch, seqlen, nheads, headdim, device, dtype) + softmax_scale = (1.0 / headdim**0.5) if use_explicit_scale else None + + actual = op( + q, + k, + v, + dropout_p=0.0, + softmax_scale=softmax_scale, + causal=causal, + ) + + expected = pytorch_sdpa_reference(q, k, v, causal=causal, softmax_scale=softmax_scale) + print_attention_diff( + actual, + expected, + dtype=dtype, + atol=atol, + rtol=rtol, + causal=causal, + softmax_scale=softmax_scale, + ) + + torch.testing.assert_close( + actual.float(), + expected.float(), + atol=atol, + rtol=rtol, + msg=describe_attention_diff( + actual, + expected, + dtype=dtype, + atol=atol, + rtol=rtol, + causal=causal, + softmax_scale=softmax_scale, + ), + ) + + +@pytest.mark.parametrize(("dtype", "atol", "rtol"), DTYPE_CASES) +@pytest.mark.parametrize( + "causal", + (pytest.param(False, id="noncausal"), pytest.param(True, id="causal")), +) +@pytest.mark.parametrize( + "use_explicit_scale", + (pytest.param(False, id="default-scale"), pytest.param(True, id="explicit-scale")), +) +@pytest.mark.parametrize(("batch", "seqlen", "nheads", "headdim"), ATTENTION_SHAPES) +def test_cuda_flash_attention_matches_sdpa( + dtype, + atol, + rtol, + causal, + use_explicit_scale, + batch, + seqlen, + nheads, + headdim, +): + available, reason = cuda_flash_attention_availability() + if not available: + pytest.skip(reason) + + from rl_engine.kernels.ops.cuda.attention.flash_attn import FlashAttentionOp + + assert_flash_attention_matches_sdpa( + FlashAttentionOp(), + dtype, + atol, + rtol, + causal, + use_explicit_scale, + batch, + seqlen, + nheads, + headdim, + ) + + +@pytest.mark.parametrize(("dtype", "atol", "rtol"), DTYPE_CASES) +@pytest.mark.parametrize( + "causal", + (pytest.param(False, id="noncausal"), pytest.param(True, id="causal")), +) +@pytest.mark.parametrize( + "use_explicit_scale", + (pytest.param(False, id="default-scale"), pytest.param(True, id="explicit-scale")), +) +@pytest.mark.parametrize(("batch", "seqlen", "nheads", "headdim"), ATTENTION_SHAPES) +def test_rocm_flash_attention_matches_sdpa( + dtype, + atol, + rtol, + causal, + use_explicit_scale, + batch, + seqlen, + nheads, + headdim, +): + available, reason = rocm_flash_attention_availability() + if not available: + pytest.skip(reason) + + from rl_engine.kernels.ops.rocm.attention.flash_attn import RocmFlashAttentionOp + + assert_flash_attention_matches_sdpa( + RocmFlashAttentionOp(), + dtype, + atol, + rtol, + causal, + use_explicit_scale, + batch, + seqlen, + nheads, + headdim, + ) + + +@pytest.mark.parametrize(("dtype", "atol", "rtol"), DTYPE_CASES) +@pytest.mark.parametrize( + "causal", + (pytest.param(False, id="noncausal"), pytest.param(True, id="causal")), +) +@pytest.mark.parametrize( + "use_explicit_scale", + (pytest.param(False, id="default-scale"), pytest.param(True, id="explicit-scale")), +) +@pytest.mark.parametrize(("batch", "seqlen", "nheads", "headdim"), ATTENTION_SHAPES) +def test_native_attention_matches_sdpa( + dtype, + atol, + rtol, + causal, + use_explicit_scale, + batch, + seqlen, + nheads, + headdim, +): + available, reason = native_attention_availability() + if not available: + pytest.skip(reason) + + from rl_engine.kernels.ops.pytorch.attention import NativeAttentionOp + + assert_flash_attention_matches_sdpa( + NativeAttentionOp(), + dtype, + atol, + rtol, + causal, + use_explicit_scale, + batch, + seqlen, + nheads, + headdim, + ) From aa8468f02ffee53309f0742445951edcd2495724 Mon Sep 17 00:00:00 2001 From: FED4 Date: Thu, 18 Jun 2026 17:37:32 +0800 Subject: [PATCH 2/2] fix: address ROCm attention review feedback --- .../kernels/ops/pytorch/attention/__init__.py | 15 ++ .../kernels/ops/rocm/attention/flash_attn.py | 35 +++-- rl_engine/kernels/registry.py | 10 +- rl_engine/tests/test_dispatch.py | 12 +- scripts/check_rocm_env.py | 128 +++--------------- tests/test_attention_correctness.py | 97 ++++++++++++- tests/test_kernel_registry.py | 84 ++++++++++++ 7 files changed, 255 insertions(+), 126 deletions(-) create mode 100644 tests/test_kernel_registry.py diff --git a/rl_engine/kernels/ops/pytorch/attention/__init__.py b/rl_engine/kernels/ops/pytorch/attention/__init__.py index d315e5be4..d2454e674 100644 --- a/rl_engine/kernels/ops/pytorch/attention/__init__.py +++ b/rl_engine/kernels/ops/pytorch/attention/__init__.py @@ -17,9 +17,24 @@ def __call__( softmax_scale: float | None = None, causal: bool = False, ) -> torch.Tensor: + # Convert FlashAttention layout to PyTorch SDPA layout: + # (batch, seqlen, nheads, headdim) -> (batch, nheads, seqlen, headdim) q_ref = q.transpose(1, 2) k_ref = k.transpose(1, 2) v_ref = v.transpose(1, 2) + + q_head_num = q_ref.shape[1] + k_head_num = k_ref.shape[1] + if k_head_num != v_ref.shape[1]: + raise ValueError("k and v must have the same number of heads") + + if q_head_num != k_head_num: + if q_head_num % k_head_num != 0: + raise ValueError("q heads must be divisible by k/v heads for GQA/MQA") + repeat = q_head_num // k_head_num + k_ref = k_ref.repeat_interleave(repeat, dim=1) + v_ref = v_ref.repeat_interleave(repeat, dim=1) + out = F.scaled_dot_product_attention( q_ref, k_ref, diff --git a/rl_engine/kernels/ops/rocm/attention/flash_attn.py b/rl_engine/kernels/ops/rocm/attention/flash_attn.py index 56b2e8c6c..a9781cfbc 100644 --- a/rl_engine/kernels/ops/rocm/attention/flash_attn.py +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -1,22 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -import importlib.util import os import torch from rl_engine.utils.logger import logger +_MAX_TESTED_ROCM_TRITON_HEAD_DIM = 512 + def _select_flash_attn_backend() -> str: - """Select the installed FlashAttention ROCm backend before importing flash_attn.""" - if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE", "").upper() == "TRUE": - return "triton" - if importlib.util.find_spec("flash_attn_2_cuda") is None: - os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE" - return "triton" - return "ck" + """Select the installed FlashAttention ROCm backend.""" + return "triton" class RocmFlashAttentionOp: @@ -30,12 +26,15 @@ def __init__(self): raise RuntimeError("RocmFlashAttentionOp requires a ROCm PyTorch build.") backend = _select_flash_attn_backend() + if backend == "triton": + # flash-attn selects the ROCm CK/Triton backend at import time. + os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE" try: from flash_attn import flash_attn_func self.op = flash_attn_func logger.info("Successfully linked to external flash_attn library (%s backend).", backend) - except ImportError as exc: + except (ImportError, OSError, RuntimeError) as exc: raise RuntimeError( "ROCm FlashAttention requires a ROCm-compatible flash-attn installation. " "See docs/getting_started/installation.md#rocm-backend." @@ -69,6 +68,24 @@ def __call__( raise ValueError("Inputs must be on a CUDA/ROCm GPU device") if not (q.device == k.device == v.device): raise ValueError("q, k, and v must be on the same device") + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError( + "q, k, and v must be rank-4 tensors: (batch, seqlen, nheads, head_dim)" + ) + + head_dim = q.shape[-1] + if head_dim == 0: + raise ValueError("head_dim must be positive") + if k.shape[-1] != head_dim or v.shape[-1] != head_dim: + raise ValueError("q, k, and v must have the same head_dim") + if head_dim > _MAX_TESTED_ROCM_TRITON_HEAD_DIM: + raise NotImplementedError( + "RL-Kernel's ROCm FlashAttention wrapper currently supports " + f"head_dim <= {_MAX_TESTED_ROCM_TRITON_HEAD_DIM}; got {head_dim}" + ) + + if softmax_scale is None: + softmax_scale = q.shape[-1] ** -0.5 q, k, v = q.contiguous(), k.contiguous(), v.contiguous() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 7831c4ef3..7c86834f6 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -87,8 +87,8 @@ def __init__(self): "rocm": { "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], "attn": [ - OpBackend.PYTORCH_ATTN, OpBackend.ROCM_FLASH_ATTN, + OpBackend.PYTORCH_ATTN, OpBackend.TRITON_GENERIC, ], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], @@ -102,8 +102,8 @@ def __init__(self): }, } logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") - self._adjust_priority_from_env() self._adjust_priority_for_hardware() + self._adjust_priority_from_env() def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() @@ -113,6 +113,12 @@ def _adjust_priority_from_env(self): OpBackend.PYTORCH_ATTN, OpBackend.TRITON_GENERIC, ] + elif rocm_attn_backend in {"native", "pytorch", "sdpa"}: + self._priority_map["rocm"]["attn"] = [ + OpBackend.PYTORCH_ATTN, + OpBackend.ROCM_FLASH_ATTN, + OpBackend.TRITON_GENERIC, + ] elif rocm_attn_backend and rocm_attn_backend not in {"native", "pytorch", "sdpa"}: logger.warning( "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.", diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 3fa361a6a..7f2839e18 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -25,21 +25,21 @@ def test_device_and_registry(): logger.info(f"Retrieved Attention Operator: {attn_op}") -def test_rocm_attention_uses_native_sdpa_by_default(monkeypatch): +def test_rocm_attention_uses_flash_attention_by_default(monkeypatch): monkeypatch.delenv("RL_KERNEL_ROCM_ATTN_BACKEND", raising=False) registry = KernelRegistry() - assert registry._priority_map["rocm"]["attn"][0] == OpBackend.PYTORCH_ATTN + assert registry._priority_map["rocm"]["attn"][0] == OpBackend.ROCM_FLASH_ATTN -def test_rocm_attention_flash_attn_opt_in(monkeypatch): - monkeypatch.setenv("RL_KERNEL_ROCM_ATTN_BACKEND", " flash_attn ") +def test_rocm_attention_native_sdpa_opt_out(monkeypatch): + monkeypatch.setenv("RL_KERNEL_ROCM_ATTN_BACKEND", " sdpa ") registry = KernelRegistry() - assert registry._priority_map["rocm"]["attn"][0] == OpBackend.ROCM_FLASH_ATTN - assert registry._priority_map["rocm"]["attn"][1] == OpBackend.PYTORCH_ATTN + assert registry._priority_map["rocm"]["attn"][0] == OpBackend.PYTORCH_ATTN + assert registry._priority_map["rocm"]["attn"][1] == OpBackend.ROCM_FLASH_ATTN def test_executor_flow(): diff --git a/scripts/check_rocm_env.py b/scripts/check_rocm_env.py index 0819d2fcd..ae46496bf 100644 --- a/scripts/check_rocm_env.py +++ b/scripts/check_rocm_env.py @@ -3,75 +3,19 @@ from __future__ import annotations -import argparse import importlib.util import os -import re -import shutil -import subprocess -import sys -from pathlib import Path def _fail(message: str) -> None: raise SystemExit(f"ERROR: {message}") -def _rocm_major_minor(version: str | None) -> tuple[str, str] | None: - if not version: - return None - match = re.search(r"([0-9]+)\.([0-9]+)", version) - if match is None: - return None - return match.group(1), match.group(2) - - -def _find_hipcc() -> str | None: - for env_name in ("ROCM_HOME", "HIP_PATH"): - env_path = os.environ.get(env_name) - if env_path: - hipcc = Path(env_path) / "bin" / "hipcc" - if hipcc.exists(): - return str(hipcc) - - hipcc = shutil.which("hipcc") - if hipcc: - return hipcc - - fallback = Path("/opt/rocm/bin/hipcc") - if fallback.exists(): - return str(fallback) - - return None - - -def _flash_attn_backend() -> str | None: - if importlib.util.find_spec("flash_attn") is None: - return None - if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE", "").upper() == "TRUE": - return "triton" - if importlib.util.find_spec("flash_attn_2_cuda") is not None: - return "ck" - return "triton-available-if-enabled" - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Check the local ROCm PyTorch environment.") - parser.add_argument( - "--require-flash-attn", - action="store_true", - help="fail if flash_attn_func cannot be imported", - ) - return parser.parse_args() - - def main() -> None: - args = _parse_args() - try: import torch except ImportError as exc: - _fail(f"PyTorch is not installed in {sys.executable}: {exc}") + _fail(f"PyTorch is not installed: {exc}") if torch.version.hip is None: _fail(f"PyTorch is not a ROCm build: torch={torch.__version__}") @@ -79,59 +23,31 @@ def main() -> None: if not torch.cuda.is_available(): _fail("ROCm GPU is not available to PyTorch") - hipcc = _find_hipcc() - if hipcc is None: - _fail("Could not find hipcc from ROCM_HOME, HIP_PATH, PATH, or /opt/rocm/bin/hipcc") - - try: - hipcc_output = subprocess.check_output([hipcc, "--version"], text=True) - except (subprocess.CalledProcessError, OSError) as exc: - _fail(f"Could not run {hipcc} --version: {exc}") - - torch_rocm = _rocm_major_minor(torch.version.hip) - hipcc_rocm = _rocm_major_minor(hipcc_output) - if torch_rocm is None: - _fail(f"Could not parse torch.version.hip={torch.version.hip!r}") - if hipcc_rocm is None: - _fail(f"Could not parse {hipcc} --version output") - if torch_rocm != hipcc_rocm: - _fail( - "ROCm version mismatch: " - f"torch.version.hip={torch.version.hip}, hipcc major/minor={'.'.join(hipcc_rocm)}" - ) - device_name = torch.cuda.get_device_name(0) - capability = torch.cuda.get_device_capability(0) triton_available = importlib.util.find_spec("triton") is not None - flash_attn_backend = _flash_attn_backend() flash_attn_func_available = False - flash_attn_error = None - if flash_attn_backend is not None: - if flash_attn_backend == "triton-available-if-enabled": - os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE" - try: - from flash_attn import flash_attn_func - except (ImportError, OSError, RuntimeError) as exc: - flash_attn_error = str(exc) - else: - flash_attn_func_available = flash_attn_func is not None - - print(f"Python: {sys.executable}") - print(f"torch: {torch.__version__}") - print(f"torch.version.hip: {torch.version.hip}") - print(f"hipcc: {hipcc}") - print(f"hipcc ROCm: {'.'.join(hipcc_rocm)}") - print(f"GPU: {device_name}") - print(f"compute capability: {capability}") - print(f"triton: {'available' if triton_available else 'missing'}") - print(f"flash_attn backend: {flash_attn_backend or 'not installed'}") - print(f"flash_attn_func: {'available' if flash_attn_func_available else 'not available'}") - - if flash_attn_error: - print(f"flash_attn import error: {flash_attn_error}") + # flash-attn selects the ROCm CK/Triton backend at import time. + os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE" + try: + from flash_attn import flash_attn_func + except (ImportError, OSError, RuntimeError) as exc: + flash_attn_status = f"not available ({exc})" + else: + flash_attn_func_available = flash_attn_func is not None + flash_attn_status = "available" if flash_attn_func_available else "not available" + + print("backend availability:") + print( + " ROCm PyTorch runtime: " + f"available (torch={torch.__version__}, hip={torch.version.hip}, GPU={device_name})" + ) + print(" PyTorch SDPA fallback: available") + print(f" Triton package: {'available' if triton_available else 'not available'}") + print(f" flash-attn AMD Triton: {flash_attn_status}") + print(" ROCm CK: not selected by this checker") - if args.require_flash_attn and not flash_attn_func_available: - _fail("flash_attn_func is required but could not be imported") + if not flash_attn_func_available: + _fail("flash_attn AMD Triton backend is required but could not be imported") if __name__ == "__main__": diff --git a/tests/test_attention_correctness.py b/tests/test_attention_correctness.py index 767b140e7..d68ad3cb7 100644 --- a/tests/test_attention_correctness.py +++ b/tests/test_attention_correctness.py @@ -58,6 +58,10 @@ def pytorch_sdpa_reference(q, k, v, *, causal, softmax_scale): """ Compute a PyTorch SDPA math-backend reference for FlashAttention-layout inputs. + The reference runs in fp32 and validates each low-precision backend against + that golden value. It does not guarantee CUDA and ROCm backends are + numerically identical to each other. + Inputs use FlashAttention layout: (batch, seqlen, nheads, headdim). PyTorch SDPA uses: (batch, nheads, seqlen, headdim), so this helper transposes before and after the reference call. @@ -79,11 +83,13 @@ def pytorch_sdpa_reference(q, k, v, *, causal, softmax_scale): return expected.transpose(1, 2).contiguous() -def make_qkv(batch, seqlen, nheads, headdim, device, dtype): +def make_qkv(batch, seqlen, nheads, headdim, device, dtype, nheads_k=None): torch.manual_seed(0) + if nheads_k is None: + nheads_k = nheads q = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype) - k = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype) - v = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype) + k = torch.randn(batch, seqlen, nheads_k, headdim, device=device, dtype=dtype) + v = torch.randn(batch, seqlen, nheads_k, headdim, device=device, dtype=dtype) return q, k, v @@ -310,6 +316,30 @@ def test_rocm_flash_attention_matches_sdpa( ) +@pytest.mark.parametrize( + "causal", + (pytest.param(False, id="noncausal"), pytest.param(True, id="causal")), +) +def test_rocm_flash_attention_rejects_unsupported_head_dim(causal): + available, reason = rocm_flash_attention_availability() + if not available: + pytest.skip(reason) + + from rl_engine.kernels.ops.rocm.attention.flash_attn import RocmFlashAttentionOp + + q, k, v = make_qkv( + batch=1, + seqlen=64, + nheads=2, + headdim=513, + device=torch.device("cuda"), + dtype=torch.float16, + ) + + with pytest.raises(NotImplementedError, match="head_dim <= 512"): + RocmFlashAttentionOp()(q, k, v, causal=causal) + + @pytest.mark.parametrize(("dtype", "atol", "rtol"), DTYPE_CASES) @pytest.mark.parametrize( "causal", @@ -349,3 +379,64 @@ def test_native_attention_matches_sdpa( nheads, headdim, ) + + +@pytest.mark.parametrize( + ("nheads", "nheads_k"), + (pytest.param(8, 4, id="gqa"), pytest.param(8, 1, id="mqa")), +) +@pytest.mark.parametrize( + "causal", + (pytest.param(False, id="noncausal"), pytest.param(True, id="causal")), +) +def test_native_attention_supports_gqa_mqa(nheads, nheads_k, causal): + available, reason = native_attention_availability() + if not available: + pytest.skip(reason) + + from rl_engine.kernels.ops.pytorch.attention import NativeAttentionOp + + device = torch.device("cuda") + dtype = torch.float16 + q, k, v = make_qkv( + batch=1, + seqlen=128, + nheads=nheads, + nheads_k=nheads_k, + headdim=64, + device=device, + dtype=dtype, + ) + + actual = NativeAttentionOp()(q, k, v, dropout_p=0.0, causal=causal) + repeat = nheads // nheads_k + expected = pytorch_sdpa_reference( + q, + k.repeat_interleave(repeat, dim=2), + v.repeat_interleave(repeat, dim=2), + causal=causal, + softmax_scale=None, + ) + + torch.testing.assert_close(actual.float(), expected.float(), atol=1e-3, rtol=1e-3) + + +def test_native_attention_rejects_invalid_gqa_head_ratio(): + available, reason = native_attention_availability() + if not available: + pytest.skip(reason) + + from rl_engine.kernels.ops.pytorch.attention import NativeAttentionOp + + q, k, v = make_qkv( + batch=1, + seqlen=128, + nheads=10, + nheads_k=3, + headdim=64, + device=torch.device("cuda"), + dtype=torch.float16, + ) + + with pytest.raises(ValueError, match="q heads must be divisible"): + NativeAttentionOp()(q, k, v) diff --git a/tests/test_kernel_registry.py b/tests/test_kernel_registry.py new file mode 100644 index 000000000..991688a8c --- /dev/null +++ b/tests/test_kernel_registry.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest + +from rl_engine.kernels import registry as registry_module +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +def test_rocm_attention_defaults_to_flash_attention(monkeypatch): + monkeypatch.delenv("RL_KERNEL_ROCM_ATTN_BACKEND", raising=False) + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"] == [ + OpBackend.ROCM_FLASH_ATTN, + OpBackend.PYTORCH_ATTN, + OpBackend.TRITON_GENERIC, + ] + + +@pytest.mark.parametrize("value", ["FLASH_ATTN", "flash-attn", "Flash_Attention", " flash_attn "]) +def test_rocm_attention_flash_opt_in_aliases(monkeypatch, value): + monkeypatch.setenv("RL_KERNEL_ROCM_ATTN_BACKEND", value) + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"] == [ + OpBackend.ROCM_FLASH_ATTN, + OpBackend.PYTORCH_ATTN, + OpBackend.TRITON_GENERIC, + ] + + +@pytest.mark.parametrize("value", ["native", "PYTORCH", " sdpa "]) +def test_rocm_attention_can_opt_out_to_sdpa(monkeypatch, value): + monkeypatch.setenv("RL_KERNEL_ROCM_ATTN_BACKEND", value) + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"] == [ + OpBackend.PYTORCH_ATTN, + OpBackend.ROCM_FLASH_ATTN, + OpBackend.TRITON_GENERIC, + ] + + +def test_rocm_attention_env_override_wins_after_hardware_adjustment(monkeypatch): + def fake_hardware_adjustment(registry): + registry._priority_map["rocm"]["attn"] = [ + OpBackend.PYTORCH_ATTN, + OpBackend.ROCM_FLASH_ATTN, + OpBackend.TRITON_GENERIC, + ] + + monkeypatch.setenv("RL_KERNEL_ROCM_ATTN_BACKEND", "flash_attn") + monkeypatch.setattr(KernelRegistry, "_adjust_priority_for_hardware", fake_hardware_adjustment) + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"] == [ + OpBackend.ROCM_FLASH_ATTN, + OpBackend.PYTORCH_ATTN, + OpBackend.TRITON_GENERIC, + ] + + +def test_rocm_attention_unknown_env_value_uses_default_and_warns(monkeypatch): + warnings = [] + + def fake_warning(message, *args): + warnings.append(message % args) + + monkeypatch.setenv("RL_KERNEL_ROCM_ATTN_BACKEND", "unknown") + monkeypatch.setattr(registry_module.logger, "warning", fake_warning) + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"] == [ + OpBackend.ROCM_FLASH_ATTN, + OpBackend.PYTORCH_ATTN, + OpBackend.TRITON_GENERIC, + ] + assert any("Unknown RL_KERNEL_ROCM_ATTN_BACKEND=unknown" in warning for warning in warnings)