feat: enable Triton kernels on MUSA

rl_engine/integrations/vllm_runtime.py

@@ -475,11 +475,7 @@ def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any:
                 )[instance.tp_rank].contiguous()
 
             assert instance.quant_method is not None
-            bias_ = (
-                None
-                if (instance.tp_rank > 0 or instance.skip_bias_add)
-                else instance.bias
-            )
+            bias_ = None if (instance.tp_rank > 0 or instance.skip_bias_add) else instance.bias
             output_parallel = instance.quant_method.apply(instance, input_parallel, bias_)
 
             if instance.reduce_results and instance.tp_size > 1:

rl_engine/kernels/ops/cuda/attention/flash_attn.py

@@ -130,9 +130,7 @@ def precompile_training(
         tensors, RNG state, or distributed collectives.
         """
         if torch.version.hip is not None:
-            raise StrictFlashAttentionUnavailable(
-                "FA4 CUDA precompile is unavailable on ROCm"
-            )
+            raise StrictFlashAttentionUnavailable("FA4 CUDA precompile is unavailable on ROCm")
         if not torch.cuda.is_available():
             raise StrictFlashAttentionUnavailable(
                 "FA4 CUDA precompile requires an available CUDA device"
@@ -144,11 +142,7 @@ def precompile_training(
         if head_dim <= 0 or sequence_length <= 0:
             raise ValueError("head_dim and sequence_length must be positive")
 
-        target = (
-            torch.device("cuda", torch.cuda.current_device())
-            if device is None
-            else device
-        )
+        target = torch.device("cuda", torch.cuda.current_device()) if device is None else device
         if target.type != "cuda":
             raise ValueError("strict FA4 training precompile requires a CUDA device")
 

rl_engine/kernels/ops/triton/activation/swiglu.py

@@ -180,13 +180,13 @@ def __call__(self, x: Tensor) -> Tensor:
         return self.forward(x)
 
     def forward(self, x: Tensor) -> Tensor:
-        if x.device.type not in ("cuda", "hip", "xpu"):
+        if x.device.type not in ("cuda", "hip", "xpu", "musa"):
             raise RuntimeError(f"TritonSiLUOp requires a GPU tensor, got device '{x.device}'.")
         _validate_dtype(x, "x")
         return _SiLUTritonFunction.apply(x)
 
     def forward_fp32(self, x: Tensor) -> Tensor:
-        if x.device.type not in ("cuda", "hip", "xpu"):
+        if x.device.type not in ("cuda", "hip", "xpu", "musa"):
             raise RuntimeError(f"TritonSiLUOp requires a GPU tensor, got device '{x.device}'.")
         _validate_dtype(x, "x")
         return _SiLUTritonFunction.apply(x.float())
@@ -201,10 +201,11 @@ def __call__(self, gate: Tensor, up: Tensor) -> Tensor:
         return self.forward(gate, up)
 
     def forward(self, gate: Tensor, up: Tensor) -> Tensor:
-        if gate.device.type not in ("cuda", "hip", "xpu") or up.device.type not in (
+        if gate.device.type not in ("cuda", "hip", "xpu", "musa") or up.device.type not in (
             "cuda",
             "hip",
             "xpu",
+            "musa",
         ):
             raise RuntimeError(
                 f"TritonSwiGLUOp requires GPU tensors, got gate='{gate.device}', up='{up.device}'."
@@ -226,10 +227,11 @@ def forward(self, gate: Tensor, up: Tensor) -> Tensor:
         return _SwiGLUTritonFunction.apply(gate, up)
 
     def forward_fp32(self, gate: Tensor, up: Tensor) -> Tensor:
-        if gate.device.type not in ("cuda", "hip", "xpu") or up.device.type not in (
+        if gate.device.type not in ("cuda", "hip", "xpu", "musa") or up.device.type not in (
             "cuda",
             "hip",
             "xpu",
+            "musa",
         ):
             raise RuntimeError(
                 f"TritonSwiGLUOp requires GPU tensors, got gate='{gate.device}', up='{up.device}'."

rl_engine/kernels/ops/triton/linear/embedding.py

@@ -238,8 +238,14 @@ def __call__(self, token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tenso
         return self.forward(token_ids, weight)
 
     def forward(self, token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
-        if not token_ids.is_cuda or not weight.is_cuda:
-            raise RuntimeError("TritonEmbeddingOp requires CUDA tensors")
+        supported_devices = ("cuda", "hip", "xpu", "musa")
+        if (
+            token_ids.device.type not in supported_devices
+            or weight.device.type not in supported_devices
+        ):
+            raise RuntimeError(
+                "TritonEmbeddingOp requires accelerator tensors " "(CUDA / ROCm / XPU / MUSA)"
+            )
         return _TritonEmbeddingFunction.apply(token_ids, weight)
 
     @staticmethod

rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py

@@ -185,7 +185,7 @@ class TritonBatchInvariantLogpOp:
     Computes ``logits[t, target_ids[t]] - logsumexp(logits[t, :])`` using a
     one-pass online softmax Triton kernel with locked reduction order.
 
-    Requires a GPU tensor (CUDA / ROCm).
+    Requires a GPU tensor (CUDA / ROCm / XPU / MUSA).
     """
 
     def __init__(self) -> None:
@@ -209,10 +209,10 @@ def apply(
         *,
         validate: bool = False,
     ) -> torch.Tensor:
-        if logits.device.type not in ("cuda", "xpu", "hip"):
+        if logits.device.type not in ("cuda", "xpu", "hip", "musa"):
             raise RuntimeError(
                 "TritonBatchInvariantLogpOp requires a GPU tensor "
-                f"(CUDA / ROCm / XPU), got device '{logits.device}'."
+                f"(CUDA / ROCm / XPU / MUSA), got device '{logits.device}'."
             )
         if logits.dim() < 2:
             raise ValueError(

rl_engine/kernels/ops/triton/loss/grpo_loss.py

@@ -145,8 +145,10 @@ def group_advantages(
         eps: float = 1e-6,
     ) -> torch.Tensor:
         """Per-sequence reward normalization, computed by the Triton group kernel."""
-        if not rewards.is_cuda:
-            raise RuntimeError("TritonGRPOLossOp requires CUDA tensors.")
+        if rewards.device.type not in ("cuda", "hip", "xpu", "musa"):
+            raise RuntimeError(
+                "TritonGRPOLossOp requires accelerator tensors " "(CUDA / ROCm / XPU / MUSA)."
+            )
         flat = rewards.reshape(-1).to(torch.float32)
         n = flat.numel()
         bounds, max_group = self._build_bounds(n, flat.device, samples_per_prompt, group_boundaries)
@@ -195,8 +197,10 @@ def apply(
         beta: float = 0.0,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Evaluate the loss from logits + per-sequence advantages."""
-        if not policy_logits.is_cuda:
-            raise RuntimeError("TritonGRPOLossOp requires CUDA tensors.")
+        if policy_logits.device.type not in ("cuda", "hip", "xpu", "musa"):
+            raise RuntimeError(
+                "TritonGRPOLossOp requires accelerator tensors " "(CUDA / ROCm / XPU / MUSA)."
+            )
         if completion_mask.ndim != 2:
             raise ValueError("completion_mask must be 2D [num_sequences, completion_len].")
 

rl_engine/kernels/ops/triton/loss/linear_logp.py

@@ -198,9 +198,9 @@ def apply(
         vocab_start_index: int = 0,
         global_vocab_size: Optional[int] = None,
     ) -> torch.Tensor:
-        if hidden.device.type not in ("cuda", "xpu", "hip"):
+        if hidden.device.type not in ("cuda", "xpu", "hip", "musa"):
             raise RuntimeError(
-                "TritonLinearLogpOp requires a GPU tensor (CUDA / ROCm / XPU), got "
+                "TritonLinearLogpOp requires a GPU tensor (CUDA / ROCm / XPU / MUSA), got "
                 f"device '{hidden.device}'."
             )
         if hidden.shape[:-1] != target_ids.shape:

rl_engine/kernels/ops/triton/loss/ratio_kl.py

@@ -130,8 +130,10 @@ def _ratio_kl_bwd_kernel(
 class _RatioKLFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, policy_logits, ref_logits, action_ids, attention_mask, old_logps):
-        if not policy_logits.is_cuda:
-            raise RuntimeError("TritonRatioKLOp requires CUDA/ROCm tensors.")
+        if policy_logits.device.type not in ("cuda", "hip", "xpu", "musa"):
+            raise RuntimeError(
+                "TritonRatioKLOp requires accelerator tensors " "(CUDA / ROCm / XPU / MUSA)."
+            )
         expected = tuple(policy_logits.shape[:-1])
         for name, t in (
             ("action_ids", action_ids),

rl_engine/kernels/ops/triton/matmul/det_gemm.py

@@ -146,7 +146,9 @@ def __init__(self):
 
     def __call__(self, a, b):
         assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only"
-        assert a.is_cuda and b.is_cuda, "CUDA only"
+        assert a.device.type in ("cuda", "musa")
+        assert b.device.type in ("cuda", "musa")
+        assert a.device == b.device, "GEMM operands must be on the same device"
         return _TritonDetGemmFn.apply(a, b, False)
 
     def forward_fp32(self, a, b):
@@ -155,7 +157,9 @@ def forward_fp32(self, a, b):
             torch.float32,
         ):
             raise TypeError("FP32-output Triton GEMM requires BF16 or FP32 inputs")
-        assert a.is_cuda and b.is_cuda, "CUDA only"
+        assert a.device.type in ("cuda", "musa")
+        assert b.device.type in ("cuda", "musa")
+        assert a.device == b.device, "GEMM operands must be on the same device"
         return _TritonDetGemmFn.apply(a, b, True)
 
     forward_accum_fp32 = forward_fp32

rl_engine/kernels/ops/triton/rmsnorm_triton.py

@@ -76,7 +76,8 @@ def _require_triton():
 
 def rmsnorm_triton_forward_with_rstd(x, weight, eps: float = 1e-6):
     _require_triton()
-    assert x.is_cuda and weight.is_cuda
+    assert x.device.type in ("cuda", "hip", "xpu", "musa")
+    assert weight.device.type in ("cuda", "hip", "xpu", "musa")
     assert x.dim() == 2 and weight.dim() == 1
     rows, hidden = x.shape
     assert weight.numel() == hidden
@@ -112,7 +113,8 @@ def rmsnorm_triton_backward_rows(grad_out, x, weight, rstd):
 class RMSNormTriton(torch.autograd.Function):
     @staticmethod
     def forward(ctx, x, weight, eps: float = 1e-6):
-        assert x.is_cuda and weight.is_cuda
+        assert x.device.type in ("cuda", "hip", "xpu", "musa")
+        assert weight.device.type in ("cuda", "hip", "xpu", "musa")
         assert x.dim() == 2 and weight.dim() == 1
         T, H = x.shape
         assert weight.numel() == H

rl_engine/kernels/ops/triton/rotary_embedding/rope.py

@@ -192,6 +192,6 @@ def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0)
         return self.forward(x, positions, theta=theta)
 
     def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor:
-        if x.device.type not in ("cuda", "hip", "xpu"):
+        if x.device.type not in ("cuda", "hip", "xpu", "musa"):
             raise RuntimeError(f"TritonRoPEOp requires a GPU tensor, got device '{x.device}'.")
         return _RoPEFunction.apply(x, positions, theta)

rl_engine/kernels/registry.py

@@ -113,11 +113,13 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     )
 
     # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth)
+    TRITON_RMS_NORM = "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp"
     PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp"
 
     # Generic fallback
     TRITON_GENERIC = "rl_engine.kernels.ops.triton.generic.TritonOp"
     PYTORCH_ATTN = "rl_engine.kernels.ops.pytorch.attention.NativeAttentionOp"
+    TRITON_LOGP = "rl_engine.kernels.ops.triton.loss.logp.TritonLogpOp"
     PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp"
     PYTORCH_NATIVE_MATMUL = "rl_engine.kernels.ops.pytorch.linear.matmul.NativeMatmulOp"
     PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp"
@@ -149,6 +151,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     # WS1 pure-PyTorch ground-truth linear ops
     PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp"
     # WS1 pure-PyTorch ground-truth embedding ops
+    TRITON_EMBEDDING = "rl_engine.kernels.ops.triton.linear.embedding.TritonEmbeddingOp"
     PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp"
     CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp"
     CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp"
@@ -564,28 +567,48 @@ def __init__(self):
                 "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU],
             },
             "musa": {
-                "logp": [OpBackend.PYTORCH_NATIVE],
+                "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE],
                 "logp_indexed": [OpBackend.PYTORCH_NATIVE],
                 "logp_online": [OpBackend.PYTORCH_NATIVE],
                 "logp_online_indexed": [OpBackend.PYTORCH_NATIVE],
                 "logp_deterministic": [OpBackend.PYTORCH_NATIVE],
                 "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE],
                 "attn": [OpBackend.PYTORCH_ATTN],
                 "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION],
+                "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION],
+                "ws2_attention": [
+                    OpBackend.PYTORCH_CP_ATTENTION,
+                    OpBackend.PYTORCH_NATIVE_ATTENTION,
+                ],
                 "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN],
-                "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS],
-                "rope": [OpBackend.PYTORCH_NATIVE_ROPE],
-                "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP],
-                "ratio_kl": [OpBackend.PYTORCH_RATIO_KL],
+                "grpo_loss": [
+                    OpBackend.TRITON_GRPO_LOSS,
+                    OpBackend.PYTORCH_GRPO_LOSS,
+                ],
+                "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE],
+                "linear_logp": [
+                    OpBackend.TRITON_LINEAR_LOGP,
+                    OpBackend.PYTORCH_LINEAR_LOGP,
+                ],
+                "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
                 "pack": [OpBackend.PYTORCH_PACK],
-                "det_gemm": [],
-                "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP],
+                "det_gemm": [OpBackend.TRITON_DET_GEMM],
+                "batch_invariant_logp": [
+                    OpBackend.TRITON_BATCH_INVARIANT_LOGP,
+                    OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
+                ],
                 "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL],
-                "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM],
+                "rms_norm": [
+                    OpBackend.TRITON_RMS_NORM,
+                    OpBackend.PYTORCH_NATIVE_RMS_NORM,
+                ],
                 "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
-                "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING],
-                "silu": [OpBackend.PYTORCH_NATIVE_SILU],
-                "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU],
+                "embedding": [
+                    OpBackend.TRITON_EMBEDDING,
+                    OpBackend.PYTORCH_NATIVE_EMBEDDING,
+                ],
+                "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU],
+                "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU],
             },
             "cpu": {
                 "logp": [OpBackend.PYTORCH_NATIVE],

rl_engine/tests/test_dispatch.py

@@ -4,7 +4,6 @@
 import sys
 from types import ModuleType
 
-import pytest
 import torch
 
 import rl_engine.platforms.device as device_module
@@ -106,20 +105,54 @@ def is_available():
         assert context.is_musa is True
         assert context.device.type == "musa"
 
-    def test_musa_dispatch_uses_only_pytorch_fallbacks(self, monkeypatch):
+    def test_musa_dispatch_prefers_validated_triton_backends(self, monkeypatch):
         self._mock_musa_device(monkeypatch)
         registry = KernelRegistry()
 
         assert registry._platform_for_device("musa") == "musa"
-        for candidates in registry._priority_map["musa"].values():
+        expected = {
+            "logp": [
+                OpBackend.TRITON_LOGP,
+                OpBackend.PYTORCH_NATIVE,
+            ],
+            "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE],
+            "linear_logp": [
+                OpBackend.TRITON_LINEAR_LOGP,
+                OpBackend.PYTORCH_LINEAR_LOGP,
+            ],
+            "grpo_loss": [
+                OpBackend.TRITON_GRPO_LOSS,
+                OpBackend.PYTORCH_GRPO_LOSS,
+            ],
+            "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
+            "det_gemm": [OpBackend.TRITON_DET_GEMM],
+            "batch_invariant_logp": [
+                OpBackend.TRITON_BATCH_INVARIANT_LOGP,
+                OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
+            ],
+            "rms_norm": [
+                OpBackend.TRITON_RMS_NORM,
+                OpBackend.PYTORCH_NATIVE_RMS_NORM,
+            ],
+            "embedding": [
+                OpBackend.TRITON_EMBEDDING,
+                OpBackend.PYTORCH_NATIVE_EMBEDDING,
+            ],
+            "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU],
+            "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU],
+        }
+        for op_name, candidates in expected.items():
+            assert registry._priority_map["musa"][op_name] == candidates
+        for op_name, candidates in registry._priority_map["musa"].items():
+            if op_name in expected:
+                continue
             assert all(candidate.name.startswith("PYTORCH_") for candidate in candidates)
 
-    def test_musa_det_gemm_fails_closed(self, monkeypatch):
+    def test_musa_det_gemm_uses_only_validated_triton_backend(self, monkeypatch):
         self._mock_musa_device(monkeypatch)
         registry = KernelRegistry()
 
-        with pytest.raises(RuntimeError, match="No functional backend"):
-            registry.get_op("det_gemm", device="musa")
+        assert registry._priority_map["musa"]["det_gemm"] == [OpBackend.TRITON_DET_GEMM]
 
 
 def test_registry_explicit_device_selects_device_platform(monkeypatch):

tests/test_batch_invariant_logp.py

@@ -14,7 +14,7 @@
 
 from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp
 from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp
-from rl_engine.platforms.device import _npu_available
+from rl_engine.platforms.device import _npu_available, device_ctx
 
 _V = 300
 
@@ -1309,7 +1309,7 @@ def test_registry_dispatches_correctly():
         or type(op).__name__ == "BatchInvariantLogpSM90Op"
         or type(op).__name__ == "BatchInvariantLogpAscendOp"
     )
-    logits = torch.randn(4, _V, device="cuda" if torch.cuda.is_available() else "cpu")
+    logits = torch.randn(4, _V, device=device_ctx.device)
     target = torch.randint(0, _V, (4,), device=logits.device)
     out = op(logits, target)
     ref = _reference_logp(logits, target)