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..d2454e674 100644 --- a/rl_engine/kernels/ops/pytorch/attention/__init__.py +++ b/rl_engine/kernels/ops/pytorch/attention/__init__.py @@ -0,0 +1,49 @@ +# 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: + # 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, + 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..a9781cfbc --- /dev/null +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +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.""" + return "triton" + + +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() + 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, OSError, RuntimeError) 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") + 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() + + 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..7c86834f6 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,26 +79,51 @@ 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.ROCM_FLASH_ATTN, + OpBackend.PYTORCH_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_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() + 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 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.", + rocm_attn_backend, + ) def _adjust_priority_for_hardware(self): """Prioritize the fused TMA LogP kernel only when it is compiled into the diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index d15b1ae4c..7f2839e18 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_flash_attention_by_default(monkeypatch): + monkeypatch.delenv("RL_KERNEL_ROCM_ATTN_BACKEND", raising=False) + + registry = KernelRegistry() + + assert registry._priority_map["rocm"]["attn"][0] == OpBackend.ROCM_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.PYTORCH_ATTN + assert registry._priority_map["rocm"]["attn"][1] == OpBackend.ROCM_FLASH_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..ae46496bf --- /dev/null +++ b/scripts/check_rocm_env.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import importlib.util +import os + + +def _fail(message: str) -> None: + raise SystemExit(f"ERROR: {message}") + + +def main() -> None: + try: + import torch + except ImportError as 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__}") + + if not torch.cuda.is_available(): + _fail("ROCm GPU is not available to PyTorch") + + device_name = torch.cuda.get_device_name(0) + triton_available = importlib.util.find_spec("triton") is not None + flash_attn_func_available = False + # 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 not flash_attn_func_available: + _fail("flash_attn AMD Triton backend 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..d68ad3cb7 --- /dev/null +++ b/tests/test_attention_correctness.py @@ -0,0 +1,442 @@ +# 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. + + 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. + """ + 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, 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_k, headdim, device=device, dtype=dtype) + v = torch.randn(batch, seqlen, nheads_k, 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( + "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", + (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, + ) + + +@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)