feat: add ROCm FlashAttention backend

Author: FED4

.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

docs/getting_started/installation.md

@@ -28,6 +28,38 @@ pip install -e ".[vllm]"
 Install the vLLM extra only on rollout or benchmark environments that need the
 vLLM runtime. Core CI and mocked integration tests do not require it.
 
+### 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 --require-flash-attn
+FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE python - <<'PY'
+from flash_attn import flash_attn_func
+
+print(flash_attn_func)
+PY
+```
+
+RL-Kernel uses PyTorch SDPA as the default ROCm attention path. To opt in to
+external FlashAttention for ROCm attention dispatch, set:
+
+```bash
+export RL_KERNEL_ROCM_ATTN_BACKEND=flash_attn
+```
+
 ## Development Dependencies
 
 ```bash

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"]

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",
+]

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)

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,13 +33,15 @@ 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.triton_grpo_loss.TritonGRPOLossOp"
     PYTORCH_GRPO_LOSS = "rl_engine.kernels.ops.pytorch.loss.grpo_loss.NativeGRPOLossOp"
 
     # 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"
 
 
@@ -72,24 +75,43 @@ 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],
                 # 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],
             },
             "cpu": {
                 "logp": [OpBackend.PYTORCH_NATIVE],
-                "attn": [OpBackend.PYTORCH_NATIVE],
+                "attn": [OpBackend.PYTORCH_ATTN],
                 "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS],
             },
         }
         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)."""

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)