feat(musa): add native deterministic gemm kernel

csrc/musa/det_gemm.mu

@@ -0,0 +1,179 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+
+#include <musa_runtime.h>
+#include <c10/core/DeviceGuard.h>
+#include <torch/extension.h>
+#include <torch_musa/csrc/aten/musa/Exceptions.h>
+#include <torch_musa/csrc/aten/musa/MUSAContext.h>
+
+namespace {
+
+constexpr int kBlockSize = 256;
+
+template <
+    typename input_t,
+    typename output_t,
+    bool TransposeA,
+    bool TransposeB,
+    bool TransposeOutput>
+__global__ void det_gemm_kernel(
+    const input_t* a,
+    const input_t* b,
+    output_t* c,
+    int m,
+    int n,
+    int k) {
+    const int index = blockIdx.x * blockDim.x + threadIdx.x;
+    if (index >= m * n) {
+        return;
+    }
+    const int row = index / n;
+    const int column = index % n;
+    float accumulator = 0.0f;
+    for (int inner = 0; inner < k; ++inner) {
+        const int a_index = TransposeA ? inner * m + row : row * k + inner;
+        const int b_index = TransposeB ? column * k + inner : inner * n + column;
+        accumulator +=
+            static_cast<float>(a[a_index]) * static_cast<float>(b[b_index]);
+    }
+    const int output_index = TransposeOutput ? column * m + row : index;
+    c[output_index] = static_cast<output_t>(accumulator);
+}
+
+template <typename input_t, typename output_t>
+void launch(
+    torch::Tensor a,
+    torch::Tensor b,
+    torch::Tensor output,
+    int m,
+    int n,
+    int k,
+    bool transpose_a,
+    bool transpose_b,
+    bool transpose_output) {
+    const int blocks = (m * n + kBlockSize - 1) / kBlockSize;
+    auto stream = at::musa::getCurrentMUSAStream();
+    if (transpose_a) {
+        if (transpose_output) {
+            det_gemm_kernel<input_t, output_t, true, false, true>
+                <<<blocks, kBlockSize, 0, stream>>>(
+                    a.data_ptr<input_t>(),
+                    b.data_ptr<input_t>(),
+                    output.data_ptr<output_t>(),
+                    m,
+                    n,
+                    k);
+        } else {
+            det_gemm_kernel<input_t, output_t, true, false, false>
+                <<<blocks, kBlockSize, 0, stream>>>(
+                    a.data_ptr<input_t>(),
+                    b.data_ptr<input_t>(),
+                    output.data_ptr<output_t>(),
+                    m,
+                    n,
+                    k);
+        }
+    } else if (transpose_b) {
+        det_gemm_kernel<input_t, output_t, false, true, false>
+            <<<blocks, kBlockSize, 0, stream>>>(
+                a.data_ptr<input_t>(),
+                b.data_ptr<input_t>(),
+                output.data_ptr<output_t>(),
+                m,
+                n,
+                k);
+    } else {
+        det_gemm_kernel<input_t, output_t, false, false, false>
+            <<<blocks, kBlockSize, 0, stream>>>(
+                a.data_ptr<input_t>(),
+                b.data_ptr<input_t>(),
+                output.data_ptr<output_t>(),
+                m,
+                n,
+                k);
+    }
+}
+
+void check_inputs(torch::Tensor a, torch::Tensor b) {
+    TORCH_CHECK(
+        a.device().type() == c10::kPrivateUse1 &&
+            b.device().type() == c10::kPrivateUse1,
+        "det_gemm requires MUSA tensors");
+    TORCH_CHECK(a.device() == b.device(), "det_gemm tensors must share a device");
+    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "det_gemm expects 2-D tensors");
+    TORCH_CHECK(a.scalar_type() == b.scalar_type(), "det_gemm dtypes must match");
+    TORCH_CHECK(
+        a.scalar_type() == at::ScalarType::BFloat16,
+        "MUSA det_gemm currently supports bfloat16 inputs");
+}
+
+torch::Tensor dispatch(
+    torch::Tensor a,
+    torch::Tensor b,
+    int m,
+    int n,
+    int k,
+    bool transpose_a,
+    bool transpose_b,
+    bool transpose_output,
+    bool output_fp32) {
+    const c10::DeviceGuard device_guard(a.device());
+    a = a.contiguous();
+    b = b.contiguous();
+    auto options = output_fp32 ? a.options().dtype(torch::kFloat) : a.options();
+    auto output = torch::empty(
+        transpose_output ? std::vector<int64_t>{n, m}
+                         : std::vector<int64_t>{m, n},
+        options);
+    if (m == 0 || n == 0) {
+        return output;
+    }
+    if (output_fp32) {
+        launch<c10::BFloat16, float>(
+            a, b, output, m, n, k, transpose_a, transpose_b, transpose_output);
+    } else {
+        launch<c10::BFloat16, c10::BFloat16>(
+            a, b, output, m, n, k, transpose_a, transpose_b, transpose_output);
+    }
+    C10_MUSA_KERNEL_LAUNCH_CHECK();
+    return output;
+}
+
+}  // namespace
+
+torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b) {
+    check_inputs(a, b);
+    TORCH_CHECK(b.size(0) == a.size(1), "det_gemm_fwd: K mismatch");
+    return dispatch(a, b, a.size(0), b.size(1), a.size(1), false, false, false, false);
+}
+
+torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b) {
+    check_inputs(a, b);
+    TORCH_CHECK(b.size(0) == a.size(1), "det_gemm_fwd_fp32: K mismatch");
+    return dispatch(a, b, a.size(0), b.size(1), a.size(1), false, false, false, true);
+}
+
+torch::Tensor det_gemm_fwd_rhs_transposed(torch::Tensor a, torch::Tensor bt) {
+    check_inputs(a, bt);
+    TORCH_CHECK(bt.size(1) == a.size(1), "det_gemm_fwd_rhs_transposed: K mismatch");
+    return dispatch(a, bt, a.size(0), bt.size(0), a.size(1), false, true, false, false);
+}
+
+torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b) {
+    check_inputs(dc, b);
+    TORCH_CHECK(b.size(1) == dc.size(1), "det_gemm_da: N mismatch");
+    return dispatch(dc, b, dc.size(0), b.size(0), dc.size(1), false, true, false, false);
+}
+
+torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc) {
+    check_inputs(a, dc);
+    TORCH_CHECK(a.size(0) == dc.size(0), "det_gemm_db: M mismatch");
+    return dispatch(a, dc, a.size(1), dc.size(1), a.size(0), true, false, false, false);
+}
+
+torch::Tensor det_gemm_db_transposed(torch::Tensor a, torch::Tensor dc) {
+    check_inputs(a, dc);
+    TORCH_CHECK(a.size(0) == dc.size(0), "det_gemm_db_transposed: M mismatch");
+    return dispatch(a, dc, a.size(1), dc.size(1), a.size(0), true, false, true, false);
+}

csrc/musa/ops.cpp

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+
+#include <torch/extension.h>
+
+torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b);
+torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b);
+torch::Tensor det_gemm_fwd_rhs_transposed(torch::Tensor a, torch::Tensor bt);
+torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b);
+torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc);
+torch::Tensor det_gemm_db_transposed(torch::Tensor a, torch::Tensor dc);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("det_gemm_fwd", &det_gemm_fwd);
+    m.def("det_gemm_fwd_fp32", &det_gemm_fwd_fp32);
+    m.def("det_gemm_fwd_rhs_transposed", &det_gemm_fwd_rhs_transposed);
+    m.def("det_gemm_da", &det_gemm_da);
+    m.def("det_gemm_db", &det_gemm_db);
+    m.def("det_gemm_db_transposed", &det_gemm_db_transposed);
+}

rl_engine/kernels/ops/musa/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

rl_engine/kernels/ops/musa/matmul/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .det_gemm import MusaDetGemmOp
+
+__all__ = ["MusaDetGemmOp"]

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

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+from __future__ import annotations
+
+import torch
+
+from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
+
+
+class _MusaDetGemmFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, a: torch.Tensor, b: torch.Tensor, output_fp32: bool):
+        ctx.save_for_backward(a, b)
+        ctx.output_fp32 = bool(output_fp32)
+        if ctx.output_fp32:
+            return _C.det_gemm_fwd_fp32(a, b)
+        return _C.det_gemm_fwd(a, b)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        a, b = ctx.saved_tensors
+        grad_a = _C.det_gemm_da(grad_output, b) if ctx.needs_input_grad[0] else None
+        grad_b = _C.det_gemm_db(a, grad_output) if ctx.needs_input_grad[1] else None
+        return grad_a, grad_b, None
+
+
+class _MusaDetLinearFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, a: torch.Tensor, weight: torch.Tensor):
+        ctx.save_for_backward(a, weight)
+        return _C.det_gemm_fwd_rhs_transposed(a, weight)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        a, weight = ctx.saved_tensors
+        grad_a = _C.det_gemm_fwd(grad_output, weight) if ctx.needs_input_grad[0] else None
+        grad_weight = _C.det_gemm_db_transposed(a, grad_output) if ctx.needs_input_grad[1] else None
+        return grad_a, grad_weight
+
+
+class MusaDetGemmOp:
+    """Deterministic fixed-order GEMM for MUSA BF16 tensors."""
+
+    def __init__(self) -> None:
+        if not _EXT_AVAILABLE or _C is None:
+            raise RuntimeError("MUSA det_gemm requires the compiled extension")
+        required = (
+            "det_gemm_fwd",
+            "det_gemm_fwd_fp32",
+            "det_gemm_fwd_rhs_transposed",
+            "det_gemm_da",
+            "det_gemm_db",
+            "det_gemm_db_transposed",
+        )
+        missing = [name for name in required if not hasattr(_C, name)]
+        if missing:
+            raise RuntimeError(f"MUSA det_gemm extension is missing: {', '.join(missing)}")
+
+    @staticmethod
+    def _check(a: torch.Tensor, b: torch.Tensor) -> None:
+        if a.device.type != "musa" or b.device.type != "musa":
+            raise ValueError("MUSA det_gemm requires MUSA tensors")
+        if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16:
+            raise TypeError("MUSA det_gemm currently supports BF16 tensors")
+
+    def __call__(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        self._check(a, b)
+        if a.ndim != 2 or b.ndim != 2 or a.size(1) != b.size(0):
+            raise ValueError("det_gemm expects A[M,K] and B[K,N]")
+        return _MusaDetGemmFunction.apply(a.contiguous(), b.contiguous(), False)
+
+    def forward_fp32(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        self._check(a, b)
+        if a.ndim != 2 or b.ndim != 2 or a.size(1) != b.size(0):
+            raise ValueError("det_gemm expects A[M,K] and B[K,N]")
+        return _MusaDetGemmFunction.apply(a.contiguous(), b.contiguous(), True)
+
+    forward_accum_fp32 = forward_fp32
+
+    def linear(self, a: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+        self._check(a, weight)
+        if a.ndim != 2 or weight.ndim != 2 or a.size(1) != weight.size(1):
+            raise ValueError("linear expects A[M,K] and weight[N,K]")
+        return _MusaDetLinearFunction.apply(a.contiguous(), weight.contiguous())
+
+
+def deterministic_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    return MusaDetGemmOp()(a, b)

rl_engine/kernels/registry.py

@@ -108,6 +108,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     PYTORCH_PACK = "rl_engine.kernels.ops.pytorch.packing.pack.NativePackOp"
     # Batch-invariant deterministic GEMM (WS1 #146)
     CUDA_DET_GEMM = "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp"
+    MUSA_DET_GEMM = "rl_engine.kernels.ops.musa.matmul.det_gemm.MusaDetGemmOp"
     TRITON_DET_GEMM = "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp"
     ASCEND_DET_GEMM = "rl_engine.kernels.ops.ascend.matmul.det_gemm.DetGemmAscendOp"
     # NON-deterministic reference (torch.matmul); reference/benchmark ONLY,
@@ -676,7 +677,7 @@ def __init__(self):
                 ],
                 "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
                 "pack": [OpBackend.PYTORCH_PACK],
-                "det_gemm": [OpBackend.TRITON_DET_GEMM],
+                "det_gemm": [OpBackend.MUSA_DET_GEMM, OpBackend.TRITON_DET_GEMM],
                 "batch_invariant_logp": [
                     OpBackend.TRITON_BATCH_INVARIANT_LOGP,
                     OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,

rl_engine/tests/test_dispatch.py

@@ -125,7 +125,10 @@ def test_musa_dispatch_prefers_validated_triton_backends(self, monkeypatch):
                 OpBackend.PYTORCH_GRPO_LOSS,
             ],
             "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
-            "det_gemm": [OpBackend.TRITON_DET_GEMM],
+            "det_gemm": [
+                OpBackend.MUSA_DET_GEMM,
+                OpBackend.TRITON_DET_GEMM,
+            ],
             "batch_invariant_logp": [
                 OpBackend.TRITON_BATCH_INVARIANT_LOGP,
                 OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
@@ -152,7 +155,10 @@ def test_musa_det_gemm_uses_only_validated_triton_backend(self, monkeypatch):
         self._mock_musa_device(monkeypatch)
         registry = KernelRegistry()
 
-        assert registry._priority_map["musa"]["det_gemm"] == [OpBackend.TRITON_DET_GEMM]
+        assert registry._priority_map["musa"]["det_gemm"] == [
+            OpBackend.MUSA_DET_GEMM,
+            OpBackend.TRITON_DET_GEMM,
+        ]
 
 
 def test_registry_explicit_device_selects_device_platform(monkeypatch):