diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py
index ade13cab..6351f0ab 100644
--- a/rl_engine/integrations/vllm_runtime.py
+++ b/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:
diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn.py b/rl_engine/kernels/ops/cuda/attention/flash_attn.py
index 9ad510b3..e57cdeb5 100644
--- a/rl_engine/kernels/ops/cuda/attention/flash_attn.py
+++ b/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")
 
diff --git a/rl_engine/kernels/ops/triton/activation/swiglu.py b/rl_engine/kernels/ops/triton/activation/swiglu.py
index 6fb66313..25e5db77 100644
--- a/rl_engine/kernels/ops/triton/activation/swiglu.py
+++ b/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}'."
diff --git a/rl_engine/kernels/ops/triton/linear/embedding.py b/rl_engine/kernels/ops/triton/linear/embedding.py
index 99fa27a5..f18673b0 100644
--- a/rl_engine/kernels/ops/triton/linear/embedding.py
+++ b/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
diff --git a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py
index 66b99757..b6d5b0bc 100644
--- a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py
+++ b/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(
diff --git a/rl_engine/kernels/ops/triton/loss/grpo_loss.py b/rl_engine/kernels/ops/triton/loss/grpo_loss.py
index d6b402d8..26049c45 100644
--- a/rl_engine/kernels/ops/triton/loss/grpo_loss.py
+++ b/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].")
 
diff --git a/rl_engine/kernels/ops/triton/loss/linear_logp.py b/rl_engine/kernels/ops/triton/loss/linear_logp.py
index 561d81e4..a0528399 100644
--- a/rl_engine/kernels/ops/triton/loss/linear_logp.py
+++ b/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:
diff --git a/rl_engine/kernels/ops/triton/loss/ratio_kl.py b/rl_engine/kernels/ops/triton/loss/ratio_kl.py
index 6e256c67..e6c6429b 100644
--- a/rl_engine/kernels/ops/triton/loss/ratio_kl.py
+++ b/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),
diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py
index 50025db2..2c1ef9ac 100644
--- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py
+++ b/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
diff --git a/rl_engine/kernels/ops/triton/rmsnorm_triton.py b/rl_engine/kernels/ops/triton/rmsnorm_triton.py
index 63ba53c6..4fc2f8a5 100644
--- a/rl_engine/kernels/ops/triton/rmsnorm_triton.py
+++ b/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
diff --git a/rl_engine/kernels/ops/triton/rotary_embedding/rope.py b/rl_engine/kernels/ops/triton/rotary_embedding/rope.py
index a584e08c..c549039e 100644
--- a/rl_engine/kernels/ops/triton/rotary_embedding/rope.py
+++ b/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)
diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py
index 12ea9b21..b2dca65a 100644
--- a/rl_engine/kernels/registry.py
+++ b/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,7 +567,7 @@ 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],
@@ -572,20 +575,40 @@ def __init__(self):
                 "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],
diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py
index 388c4aec..e98cddab 100644
--- a/rl_engine/tests/test_dispatch.py
+++ b/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):
diff --git a/tests/test_batch_invariant_logp.py b/tests/test_batch_invariant_logp.py
index 0a81308f..d01e5e84 100644
--- a/tests/test_batch_invariant_logp.py
+++ b/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)
diff --git a/tests/test_grpo_loss.py b/tests/test_grpo_loss.py
index cb7ced30..ad2133fe 100644
--- a/tests/test_grpo_loss.py
+++ b/tests/test_grpo_loss.py
@@ -6,6 +6,7 @@
 
 from rl_engine.kernels.ops.pytorch.loss.grpo_loss import NativeGRPOLossOp
 from rl_engine.kernels.ops.triton.loss.grpo_loss import TritonGRPOLossOp
+from rl_engine.platforms.device import device_ctx
 from rl_engine.testing import (
     compute_policy_ratio,
     compute_reference_kl,
@@ -349,7 +350,7 @@ def test_registry_dispatches_grpo_loss():
 
     op = kernel_registry.get_op("grpo_loss")
     assert hasattr(op, "forward") and hasattr(op, "group_advantages")
-    if _HAS_TRITON and torch.cuda.is_available():
+    if _HAS_TRITON and (torch.cuda.is_available() or device_ctx.is_musa):
         assert isinstance(op, TritonGRPOLossOp)
     else:
         assert isinstance(op, NativeGRPOLossOp)
diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py
index 1135bad8..90e7e21b 100644
--- a/tests/test_linear_logp.py
+++ b/tests/test_linear_logp.py
@@ -29,6 +29,11 @@
     reason="Triton linear log-prob requires a CUDA device and Triton.",
 )
 
+requires_triton_musa = pytest.mark.skipif(
+    not (_HAS_TRITON and hasattr(torch, "musa") and torch.musa.is_available()),
+    reason="Triton linear log-prob requires a MUSA device and Triton.",
+)
+
 
 def _sm90_available():
     """SM90 forward needs a Hopper GPU and the kernel compiled into the extension."""
@@ -502,6 +507,32 @@ def test_triton_forward_matches_native_bf16():
     assert torch.allclose(out, ref, atol=2e-2)
 
 
+@requires_triton_musa
+def test_triton_musa_forward_backward_matches_fp32_reference():
+    from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp
+
+    hidden, weight, target, bias = _inputs(2, device="musa", dtype=torch.bfloat16)
+    gen = torch.Generator(device="musa").manual_seed(3)
+    grad_out = torch.randn(_N, device="musa", generator=gen)
+
+    trit_hidden = hidden.detach().clone().requires_grad_(True)
+    trit_weight = weight.detach().clone().requires_grad_(True)
+    trit_bias = bias.detach().clone().requires_grad_(True)
+    out = TritonLinearLogpOp()(trit_hidden, trit_weight, target, trit_bias)
+    out.backward(grad_out)
+
+    ref_hidden = hidden.float().requires_grad_(True)
+    ref_weight = weight.float().requires_grad_(True)
+    ref_bias = bias.float().requires_grad_(True)
+    ref = NativeLinearLogpOp()(ref_hidden, ref_weight, target, ref_bias)
+    ref.backward(grad_out)
+
+    assert torch.allclose(out.float(), ref, atol=2e-2)
+    assert torch.allclose(trit_hidden.grad.float(), ref_hidden.grad, atol=2e-2)
+    assert torch.allclose(trit_weight.grad.float(), ref_weight.grad, atol=2e-2)
+    assert torch.allclose(trit_bias.grad.float(), ref_bias.grad, atol=2e-2)
+
+
 @requires_triton_cuda
 @pytest.mark.parametrize("use_bias", [True, False])
 def test_triton_backward_matches_native(use_bias):
@@ -1261,8 +1292,11 @@ def test_registry_dispatch_matches_native():
     from rl_engine.platforms.device import device_ctx
 
     op = kernel_registry.get_op("linear_logp")
-    device = device_ctx.device if device_ctx.device_type == "cuda" else "cpu"
+    device = (
+        device_ctx.device if device_ctx.device_type in {"cuda", "hip", "xpu", "musa"} else "cpu"
+    )
     hidden, weight, target, bias = _inputs(6, device=device)
     out = op(hidden, weight, target, bias)
     ref = NativeLinearLogpOp()(hidden, weight, target, bias)
-    assert torch.allclose(out.cpu(), ref.cpu(), atol=1e-3)
+    atol = 5e-2 if device_ctx.device_type == "musa" else 1e-3
+    assert torch.allclose(out.cpu(), ref.cpu(), atol=atol)
diff --git a/tests/test_logp.py b/tests/test_logp.py
index bcb0e3f9..22ea0bad 100644
--- a/tests/test_logp.py
+++ b/tests/test_logp.py
@@ -9,6 +9,7 @@
 import torch
 
 from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp
+from rl_engine.platforms.device import device_ctx
 
 
 def _make_inputs(
@@ -155,4 +156,9 @@ def test_registry_returns_logp_op(self):
         from rl_engine.kernels.registry import kernel_registry
 
         op = kernel_registry.get_op("logp")
-        assert isinstance(op, NativeLogpOp)
+        if device_ctx.is_musa:
+            from rl_engine.kernels.ops.triton.loss.logp import TritonLogpOp
+
+            assert isinstance(op, TritonLogpOp)
+        else:
+            assert isinstance(op, NativeLogpOp)
diff --git a/tests/test_op_accuracy.py b/tests/test_op_accuracy.py
index 56f7ab69..b8605b96 100644
--- a/tests/test_op_accuracy.py
+++ b/tests/test_op_accuracy.py
@@ -157,8 +157,6 @@ def test_accuracy():
     with torch.no_grad():
         ref_logp = torch.log_softmax(logits.float(), dim=-1)
         ref_logp = torch.gather(ref_logp, dim=-1, index=token_ids.unsqueeze(-1).long()).squeeze(-1)
-        ref_logp = ref_logp.to(dtype)
-
     if device.type == "cuda":
         torch.cuda.synchronize()
 
@@ -169,6 +167,7 @@ def test_accuracy():
         logger.error(f"Failed to execute FusedLogp: {e}")
         raise
 
+    ref_logp = ref_logp.to(custom_logp.dtype)
     diff = torch.abs(ref_logp.float() - custom_logp.float()).max().item()
     threshold = 1e-2 if dtype in (torch.bfloat16, torch.float16) else 1e-5
 
diff --git a/tests/test_ratio_kl.py b/tests/test_ratio_kl.py
index 2de2d7de..1705d1bb 100644
--- a/tests/test_ratio_kl.py
+++ b/tests/test_ratio_kl.py
@@ -11,6 +11,7 @@
     _ratio_kl_bwd_kernel,
     _ratio_kl_fwd_kernel,
 )
+from rl_engine.platforms.device import device_ctx
 from rl_engine.testing import make_synthetic_rl_kernel_batch, selected_logprobs_reference
 
 try:
@@ -515,7 +516,7 @@ def test_registry_dispatches_ratio_kl():
     from rl_engine.kernels.registry import kernel_registry
 
     op = kernel_registry.get_op("ratio_kl")
-    if _HAS_TRITON and torch.cuda.is_available():
+    if _HAS_TRITON and (torch.cuda.is_available() or device_ctx.is_musa):
         assert isinstance(op, TritonRatioKLOp)
     else:
         assert isinstance(op, NativeRatioKLOp)
diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py
index 6572603e..fb62993e 100644
--- a/tests/test_rms_norm.py
+++ b/tests/test_rms_norm.py
@@ -7,7 +7,12 @@
 
 from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp, rmsnorm_cuda
 from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp
-from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton
+from rl_engine.kernels.ops.triton.rmsnorm_triton import (
+    _TRITON_AVAILABLE,
+    RMSNormTritonOp,
+    rmsnorm_triton,
+)
+from rl_engine.platforms.device import device_ctx
 
 try:
     from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
@@ -239,7 +244,14 @@ def test_registry_dispatches_rms_norm():
     from rl_engine.kernels.registry import kernel_registry
 
     op = kernel_registry.get_op("rms_norm")
-    if torch.cuda.is_available() and _HAS_CUDA_RMSNORM:
+    if device_ctx.is_musa:
+        if _TRITON_AVAILABLE:
+            assert isinstance(op, RMSNormTritonOp)
+            assert hasattr(op, "forward")
+        else:
+            assert isinstance(op, NativeRMSNormOp)
+            assert hasattr(op, "forward") and hasattr(op, "forward_fp32")
+    elif torch.cuda.is_available() and _HAS_CUDA_RMSNORM:
         assert isinstance(op, RMSNormCudaOp)
         assert hasattr(op, "forward")
     else: