From e59c88ab985e34f365e36febd4a46a9d5612a384 Mon Sep 17 00:00:00 2001 From: Arlo-mt Date: Mon, 31 Aug 2026 17:49:26 +0800 Subject: [PATCH 1/6] feat: add MUSA native kernel build path --- csrc/musa/activation.mu | 366 +++++++++++ csrc/musa/deterministic_attention.mu | 687 ++++++++++++++++++++ csrc/musa/deterministic_logp_kernel.mu | 360 ++++++++++ csrc/musa/embedding_lm_head.mu | 303 +++++++++ csrc/musa/fused_logp_kernel.mu | 129 ++++ csrc/musa/ops.cpp | 112 ++++ csrc/musa/rmsnorm.mu | 379 +++++++++++ csrc/musa/rope.mu | 102 +++ rl_engine/kernels/ops/musa/.gitkeep | 1 + rl_engine/kernels/ops/musa/__init__.py | 6 + rl_engine/kernels/ops/musa/loss/__init__.py | 6 + rl_engine/kernels/ops/musa/loss/logp.py | 44 ++ rl_engine/kernels/registry.py | 44 +- rl_engine/platforms/constants.py | 2 + rl_engine/platforms/device.py | 37 +- setup.py | 61 +- tests/test_musa_extension.py | 33 + 17 files changed, 2656 insertions(+), 16 deletions(-) create mode 100644 csrc/musa/activation.mu create mode 100644 csrc/musa/deterministic_attention.mu create mode 100644 csrc/musa/deterministic_logp_kernel.mu create mode 100644 csrc/musa/embedding_lm_head.mu create mode 100644 csrc/musa/fused_logp_kernel.mu create mode 100644 csrc/musa/ops.cpp create mode 100644 csrc/musa/rmsnorm.mu create mode 100644 csrc/musa/rope.mu create mode 100644 rl_engine/kernels/ops/musa/.gitkeep create mode 100644 rl_engine/kernels/ops/musa/__init__.py create mode 100644 rl_engine/kernels/ops/musa/loss/__init__.py create mode 100644 rl_engine/kernels/ops/musa/loss/logp.py create mode 100644 tests/test_musa_extension.py diff --git a/csrc/musa/activation.mu b/csrc/musa/activation.mu new file mode 100644 index 00000000..e7ef88ec --- /dev/null +++ b/csrc/musa/activation.mu @@ -0,0 +1,366 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant SiLU / SwiGLU CUDA kernels (WS1 elementwise activations). +// +// Semantics match NativeSiLUOp / NativeSwiGLUOp: +// silu(x) = x * sigmoid(x) (math in fp32) +// swiglu(g, u) = silu(g) * u (math in fp32) +// +// Pure elementwise / token-local: no cross-row reduction, so batch size and +// padding cannot change a row's result (Axis-A bitwise invariance). + +#include +#include +#include +#include + +namespace { + +__device__ __forceinline__ float silu_f32(float x) { + // sigmoid(x) = 1 / (1 + exp(-x)); use expf for device fp32. + const float s = 1.0f / (1.0f + expf(-x)); + return x * s; +} + +__device__ __forceinline__ float silu_grad_f32(float x) { + // d/dx [x * s] = s + x * s * (1 - s) = s * (1 + x * (1 - s)), s = sigmoid(x) + const float s = 1.0f / (1.0f + expf(-x)); + return s * (1.0f + x * (1.0f - s)); +} + +template +__global__ void silu_forward_kernel( + const scalar_t* __restrict__ x, + scalar_t* __restrict__ y, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float xv = static_cast(x[idx]); + y[idx] = static_cast(silu_f32(xv)); +} + +template +__global__ void silu_backward_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ x, + scalar_t* __restrict__ dx, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float dyv = static_cast(dy[idx]); + const float xv = static_cast(x[idx]); + dx[idx] = static_cast(dyv * silu_grad_f32(xv)); +} + +template +__global__ void swiglu_forward_kernel( + const scalar_t* __restrict__ gate, + const scalar_t* __restrict__ up, + scalar_t* __restrict__ y, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float gv = static_cast(gate[idx]); + const float uv = static_cast(up[idx]); + y[idx] = static_cast(silu_f32(gv) * uv); +} + +template +__global__ void swiglu_backward_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ gate, + const scalar_t* __restrict__ up, + scalar_t* __restrict__ d_gate, + scalar_t* __restrict__ d_up, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float dyv = static_cast(dy[idx]); + const float gv = static_cast(gate[idx]); + const float uv = static_cast(up[idx]); + const float s = silu_f32(gv); + // d_up = dy * silu(gate); d_gate = dy * up * silu'(gate) + d_up[idx] = static_cast(dyv * s); + d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); +} + +template +__global__ void swiglu_packed_forward_kernel( + const scalar_t* __restrict__ gate_up, + scalar_t* __restrict__ y, + const int64_t n, + const int64_t width) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const int64_t row = idx / width; + const int64_t column = idx - row * width; + const int64_t gate_index = row * (2 * width) + column; + const float gv = static_cast(gate_up[gate_index]); + const float uv = static_cast(gate_up[gate_index + width]); + y[idx] = static_cast(silu_f32(gv) * uv); +} + +template +__global__ void swiglu_packed_backward_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ gate_up, + scalar_t* __restrict__ d_gate, + scalar_t* __restrict__ d_up, + const int64_t n, + const int64_t width) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const int64_t row = idx / width; + const int64_t column = idx - row * width; + const int64_t gate_index = row * (2 * width) + column; + const float dyv = static_cast(dy[idx]); + const float gv = static_cast(gate_up[gate_index]); + const float uv = static_cast(gate_up[gate_index + width]); + const float s = silu_f32(gv); + d_up[idx] = static_cast(dyv * s); + d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); +} + +static void launch_1d(int64_t n, int& threads, int64_t& blocks) { + threads = 256; + blocks = (n + threads - 1) / threads; + if (blocks == 0) { + blocks = 1; + } +} + +static void check_cuda_contig(const torch::Tensor& t, const char* name) { + TORCH_CHECK(t.device().type() == c10::kPrivateUse1, name, " must be a CUDA tensor"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + // Supported activation dtypes only: fp16 / bf16 / fp32 (reject float64). + TORCH_CHECK( + t.scalar_type() == at::kHalf || t.scalar_type() == at::kBFloat16 || + t.scalar_type() == at::kFloat, + name, + " must be fp16, bf16, or fp32, got ", + t.scalar_type()); +} + +static void check_same_device( + const torch::Tensor& lhs, + const torch::Tensor& rhs, + const char* lhs_name, + const char* rhs_name) { + TORCH_CHECK( + lhs.device() == rhs.device(), + lhs_name, + " and ", + rhs_name, + " must be on the same CUDA device, got ", + lhs.device(), + " and ", + rhs.device()); +} + +} // namespace + +torch::Tensor silu_forward_cuda(torch::Tensor x) { + check_cuda_contig(x, "x"); + const at::musa::OptionalMUSAGuard device_guard(device_of(x)); + auto y = torch::empty_like(x); + const int64_t n = x.numel(); + if (n == 0) { + return y; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "silu_forward_cuda", [&] { + silu_forward_kernel<<>>( + x.data_ptr(), y.data_ptr(), n); + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return y; +} + +torch::Tensor silu_backward_cuda(torch::Tensor dy, torch::Tensor x) { + check_cuda_contig(dy, "dy"); + check_cuda_contig(x, "x"); + check_same_device(dy, x, "dy", "x"); + TORCH_CHECK(dy.sizes() == x.sizes(), "dy and x must share shape"); + TORCH_CHECK(dy.scalar_type() == x.scalar_type(), "dy and x must share dtype"); + const at::musa::OptionalMUSAGuard device_guard(device_of(x)); + auto dx = torch::empty_like(x); + const int64_t n = x.numel(); + if (n == 0) { + return dx; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "silu_backward_cuda", [&] { + silu_backward_kernel<<>>( + dy.data_ptr(), x.data_ptr(), dx.data_ptr(), n); + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return dx; +} + +torch::Tensor swiglu_forward_cuda(torch::Tensor gate, torch::Tensor up) { + check_cuda_contig(gate, "gate"); + check_cuda_contig(up, "up"); + check_same_device(gate, up, "gate", "up"); + TORCH_CHECK(gate.sizes() == up.sizes(), "gate and up must share shape"); + TORCH_CHECK(gate.scalar_type() == up.scalar_type(), "gate and up must share dtype"); + const at::musa::OptionalMUSAGuard device_guard(device_of(gate)); + auto y = torch::empty_like(gate); + const int64_t n = gate.numel(); + if (n == 0) { + return y; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate.scalar_type(), + "swiglu_forward_cuda", + [&] { + swiglu_forward_kernel<<>>( + gate.data_ptr(), up.data_ptr(), y.data_ptr(), n); + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return y; +} + +std::vector swiglu_backward_cuda( + torch::Tensor dy, + torch::Tensor gate, + torch::Tensor up) { + check_cuda_contig(dy, "dy"); + check_cuda_contig(gate, "gate"); + check_cuda_contig(up, "up"); + check_same_device(dy, gate, "dy", "gate"); + check_same_device(gate, up, "gate", "up"); + TORCH_CHECK(gate.sizes() == up.sizes(), "gate and up must share shape"); + TORCH_CHECK(dy.sizes() == gate.sizes(), "dy and gate must share shape"); + TORCH_CHECK(dy.scalar_type() == gate.scalar_type(), "dy and gate must share dtype"); + TORCH_CHECK(up.scalar_type() == gate.scalar_type(), "up and gate must share dtype"); + const at::musa::OptionalMUSAGuard device_guard(device_of(gate)); + auto d_gate = torch::empty_like(gate); + auto d_up = torch::empty_like(up); + const int64_t n = gate.numel(); + if (n == 0) { + return {d_gate, d_up}; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate.scalar_type(), + "swiglu_backward_cuda", + [&] { + swiglu_backward_kernel<<>>( + dy.data_ptr(), + gate.data_ptr(), + up.data_ptr(), + d_gate.data_ptr(), + d_up.data_ptr(), + n); + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return {d_gate, d_up}; +} + +torch::Tensor swiglu_packed_forward_cuda(torch::Tensor gate_up) { + check_cuda_contig(gate_up, "gate_up"); + TORCH_CHECK(gate_up.dim() == 2, "gate_up must be [rows, 2 * intermediate]"); + TORCH_CHECK(gate_up.size(1) % 2 == 0, "gate_up width must be even"); + const at::musa::OptionalMUSAGuard device_guard(device_of(gate_up)); + const int64_t width = gate_up.size(1) / 2; + auto y = torch::empty({gate_up.size(0), width}, gate_up.options()); + const int64_t n = y.numel(); + if (n == 0) { + return y; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate_up.scalar_type(), + "swiglu_packed_forward_cuda", + [&] { + swiglu_packed_forward_kernel<<>>( + gate_up.data_ptr(), y.data_ptr(), n, width); + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return y; +} + +std::vector swiglu_packed_backward_cuda( + torch::Tensor dy, + torch::Tensor gate_up) { + check_cuda_contig(dy, "dy"); + check_cuda_contig(gate_up, "gate_up"); + check_same_device(dy, gate_up, "dy", "gate_up"); + TORCH_CHECK(gate_up.dim() == 2, "gate_up must be [rows, 2 * intermediate]"); + TORCH_CHECK(gate_up.size(1) % 2 == 0, "gate_up width must be even"); + TORCH_CHECK(dy.dim() == 2, "dy must be [rows, intermediate]"); + TORCH_CHECK( + dy.size(0) == gate_up.size(0) && dy.size(1) * 2 == gate_up.size(1), + "dy shape must match the packed gate/up halves"); + TORCH_CHECK(dy.scalar_type() == gate_up.scalar_type(), "dy and gate_up must share dtype"); + const at::musa::OptionalMUSAGuard device_guard(device_of(gate_up)); + auto d_gate = torch::empty_like(dy); + auto d_up = torch::empty_like(dy); + const int64_t n = dy.numel(); + if (n == 0) { + return {d_gate, d_up}; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate_up.scalar_type(), + "swiglu_packed_backward_cuda", + [&] { + swiglu_packed_backward_kernel<<>>( + dy.data_ptr(), + gate_up.data_ptr(), + d_gate.data_ptr(), + d_up.data_ptr(), + n, + dy.size(1)); + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return {d_gate, d_up}; +} diff --git a/csrc/musa/deterministic_attention.mu b/csrc/musa/deterministic_attention.mu new file mode 100644 index 00000000..98ec4f35 --- /dev/null +++ b/csrc/musa/deterministic_attention.mu @@ -0,0 +1,687 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic standard-softmax attention (issue #147). +// Forward: QK kernel → masked softmax+LSE kernel → PV kernel. +// All reductions use fixed order; no split-KV or dynamic dispatch. + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int64_t kDeterministicAttentionHeadDim = 128; +constexpr int kSoftmaxThreads = 256; + +// --------------------------------------------------------------------------- +// QK Kernel: scores[b, hq, q, k] = scale * sum_{d=0}^{D-1} Q[b,hq,q,d]*K[b,kv_head,k,d] +// Grid: (Skv_blocks, Sq_blocks, B * Hq) +// Block: (TILE_K, TILE_Q) threads, each thread computes one score element. +// --------------------------------------------------------------------------- +constexpr int kQKTileQ = 16; +constexpr int kQKTileK = 16; + +template +__global__ void qk_kernel( + const scalar_t* __restrict__ Q, // [B, Hq, Sq, D] + const scalar_t* __restrict__ K, // [B, Hkv, Skv, D] + float* __restrict__ scores, // [B, Hq, Sq, Skv] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D, + float scale) { + + const int k_idx = blockIdx.x * kQKTileK + threadIdx.x; + const int q_idx = blockIdx.y * kQKTileQ + threadIdx.y; + const int bh = blockIdx.z; // flattened (b * Hq + hq) + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || k_idx >= Skv) return; + + const int kv_head = hq / (Hq / Hkv); + + const scalar_t* q_ptr = Q + ((int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D); + const scalar_t* k_ptr = K + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D + (int64_t)k_idx * D); + + float acc = 0.0f; + #pragma unroll 8 + for (int64_t d = 0; d < D; ++d) { + acc += (float)q_ptr[d] * (float)k_ptr[d]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv + k_idx; + scores[out_idx] = scale * acc; +} + +// --------------------------------------------------------------------------- +// Masked Softmax + LSE Kernel +// One CTA per (b, hq, q) row. Fixed 256 threads. +// Applies causal + padding mask, computes max, sum-exp, writes P and LSE. +// --------------------------------------------------------------------------- +__global__ void masked_softmax_lse_kernel( + float* __restrict__ scores, // [B, Hq, Sq, Skv] in-place -> P + float* __restrict__ lse, // [B, Hq, Sq] + const bool* __restrict__ pad_mask, // [B, Skv] or nullptr + int64_t B, int64_t Hq, int64_t Sq, int64_t Skv, + bool causal) { + + const int row_idx = blockIdx.x; // flattened (b * Hq * Sq + hq * Sq + q) + const int b = row_idx / (Hq * Sq); + const int hq = (row_idx / Sq) % Hq; + const int q = row_idx % Sq; + + float* row = scores + (int64_t)row_idx * Skv; + + // Causal boundary: key_index <= Skv - Sq + q + const int64_t causal_limit = causal ? (Skv - Sq + q + 1) : Skv; + + // Phase 1: Apply masks and find max + __shared__ float smax[kSoftmaxThreads]; + float thread_max = -INFINITY; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + bool valid = (k < causal_limit); + if (valid && pad_mask != nullptr) { + valid = pad_mask[(int64_t)b * Skv + k]; + } + if (!valid) { + row[k] = -INFINITY; + } + if (valid) { + thread_max = fmaxf(thread_max, row[k]); + } + } + + // Warp reduction for max + smax[threadIdx.x] = thread_max; + __syncthreads(); + // Tree reduction + for (int stride = kSoftmaxThreads / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + smax[threadIdx.x] = fmaxf(smax[threadIdx.x], smax[threadIdx.x + stride]); + } + __syncthreads(); + } + float row_max = smax[0]; + + // Phase 2: Compute sum of exp(s - max) + __shared__ float ssum[kSoftmaxThreads]; + float thread_sum = 0.0f; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + bool valid = (k < causal_limit); + if (valid && pad_mask != nullptr) { + valid = pad_mask[(int64_t)b * Skv + k]; + } + if (valid) { + float val = expf(row[k] - row_max); + row[k] = val; + thread_sum += val; + } else { + row[k] = 0.0f; + } + } + + ssum[threadIdx.x] = thread_sum; + __syncthreads(); + for (int stride = kSoftmaxThreads / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + ssum[threadIdx.x] += ssum[threadIdx.x + stride]; + } + __syncthreads(); + } + float row_sum = ssum[0]; + + // Phase 3: Normalize to get P, compute LSE + float lse_val; + if (row_sum == 0.0f) { + // Fully masked row + lse_val = -INFINITY; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + row[k] = 0.0f; + } + } else { + lse_val = row_max + logf(row_sum); + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + row[k] /= row_sum; + } + } + + if (threadIdx.x == 0) { + lse[row_idx] = lse_val; + } +} + +// --------------------------------------------------------------------------- +// PV Kernel: out[b, hq, q, d] = sum_{k=0}^{Skv-1} P[b,hq,q,k] * V[b,kv_head,k,d] +// Grid: (D_blocks, Sq_blocks, B * Hq) +// Each thread computes one output element with sequential k accumulation. +// --------------------------------------------------------------------------- +constexpr int kPVTileQ = 16; +constexpr int kPVTileD = 16; + +template +__global__ void pv_kernel( + const float* __restrict__ P, // [B, Hq, Sq, Skv] + const input_t* __restrict__ V, // [B, Hkv, Skv, D] + output_t* __restrict__ out, // [B, Hq, Sq, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int q_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int bh = blockIdx.z; + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || d_idx >= D) return; + + const int kv_head = hq / (Hq / Hkv); + + const float* p_row = P + ((int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv); + const input_t* v_base = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); + + float acc = 0.0f; + for (int64_t k = 0; k < Skv; ++k) { + acc += p_row[k] * (float)v_base[k * D + d_idx]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D + d_idx; + out[out_idx] = (output_t)acc; +} + +void check_deterministic_attention_inputs( + const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::optional& key_padding_mask) { + TORCH_CHECK(q.device().type() == c10::kPrivateUse1 && k.device().type() == c10::kPrivateUse1 && v.device().type() == c10::kPrivateUse1, + "deterministic_attention: q, k, v must be CUDA tensors"); + TORCH_CHECK(q.device() == k.device() && q.device() == v.device(), + "deterministic_attention: q, k, v must be on the same device"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, + "deterministic_attention: q/k/v must be 4-D [B, H, S, D]"); + TORCH_CHECK( + q.scalar_type() == at::kHalf || q.scalar_type() == at::kBFloat16, + "deterministic_attention: only FP16 and BF16 are supported, got ", + q.scalar_type()); + TORCH_CHECK(k.scalar_type() == q.scalar_type() && v.scalar_type() == q.scalar_type(), + "deterministic_attention: q, k, v must share the same dtype"); + + const int64_t B = q.size(0); + const int64_t Hq = q.size(1); + const int64_t Sq = q.size(2); + const int64_t D = q.size(3); + const int64_t Hkv = k.size(1); + const int64_t Skv = k.size(2); + + TORCH_CHECK(D == kDeterministicAttentionHeadDim, + "deterministic_attention: head dim D must be ", + kDeterministicAttentionHeadDim, + ", got ", + D); + TORCH_CHECK(k.size(0) == B && v.size(0) == B, + "deterministic_attention: batch size mismatch between q/k/v"); + TORCH_CHECK(v.size(1) == Hkv && v.size(2) == Skv && k.size(3) == D && v.size(3) == D, + "deterministic_attention: k/v shape mismatch"); + TORCH_CHECK(Hq % Hkv == 0, + "deterministic_attention: Hq (", + Hq, + ") must be divisible by Hkv (", + Hkv, + ") for GQA"); + TORCH_CHECK(Sq >= 1 && Skv >= 1, + "deterministic_attention: Sq and Skv must be positive"); + + if (key_padding_mask.has_value() && key_padding_mask->defined()) { + const auto& mask = *key_padding_mask; + TORCH_CHECK(mask.device().type() == c10::kPrivateUse1, "deterministic_attention: key_padding_mask must be CUDA"); + TORCH_CHECK(mask.device() == q.device(), + "deterministic_attention: key_padding_mask must match q device"); + TORCH_CHECK(mask.scalar_type() == at::kBool, + "deterministic_attention: key_padding_mask must be bool"); + TORCH_CHECK(mask.dim() == 2 && mask.size(0) == B && mask.size(1) == Skv, + "deterministic_attention: key_padding_mask must be [B, Skv]"); + } +} + +} // namespace + +// Returns {out, lse, P}: +// out: [B, Hq, Sq, D] same dtype as q +// lse: [B, Hq, Sq] FP32 +// P: [B, Hq, Sq, Skv] FP32 (softmax probabilities, saved for backward) +std::vector deterministic_attention_forward_impl( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + torch::optional key_padding_mask, + bool output_fp32) { + check_deterministic_attention_inputs(q, k, v, key_padding_mask); + + const at::musa::OptionalMUSAGuard device_guard(at::device_of(q)); + + auto q_contig = q.contiguous(); + auto k_contig = k.contiguous(); + auto v_contig = v.contiguous(); + torch::optional mask_contig; + if (key_padding_mask.has_value() && key_padding_mask->defined()) { + mask_contig = key_padding_mask->contiguous(); + } + + const int64_t B = q_contig.size(0); + const int64_t Hq = q_contig.size(1); + const int64_t Sq = q_contig.size(2); + const int64_t D = q_contig.size(3); + const int64_t Hkv = k_contig.size(1); + const int64_t Skv = k_contig.size(2); + + auto stream = at::musa::getCurrentMUSAStream(); + + // Allocate scores [B, Hq, Sq, Skv] FP32 + auto scores = torch::empty({B, Hq, Sq, Skv}, q_contig.options().dtype(at::kFloat)); + + // --- Launch QK kernel --- + { + dim3 block(kQKTileK, kQKTileQ); + dim3 grid( + (Skv + kQKTileK - 1) / kQKTileK, + (Sq + kQKTileQ - 1) / kQKTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "qk_kernel", [&] { + qk_kernel<<>>( + q_contig.data_ptr(), + k_contig.data_ptr(), + scores.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + }); + } + + // --- Launch Masked Softmax + LSE kernel --- + auto lse = torch::empty({B, Hq, Sq}, q_contig.options().dtype(at::kFloat)); + { + const int64_t num_rows = B * Hq * Sq; + dim3 block(kSoftmaxThreads); + dim3 grid(num_rows); + const bool* pad_mask_ptr = nullptr; + if (mask_contig.has_value()) { + pad_mask_ptr = mask_contig->data_ptr(); + } + masked_softmax_lse_kernel<<>>( + scores.data_ptr(), + lse.data_ptr(), + pad_mask_ptr, + B, Hq, Sq, Skv, causal); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + } + + // --- Launch PV kernel --- + auto out = output_fp32 + ? torch::empty(q_contig.sizes(), q_contig.options().dtype(at::kFloat)) + : torch::empty_like(q_contig); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Sq + kPVTileQ - 1) / kPVTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "pv_kernel", [&] { + if (output_fp32) { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } else { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } + C10_MUSA_KERNEL_LAUNCH_CHECK(); + }); + } + + return {out, lse, scores}; +} + +std::vector deterministic_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask) { + return deterministic_attention_forward_impl( + q, k, v, causal, scale, key_padding_mask, false); +} + +std::vector deterministic_attention_forward_fp32( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask) { + return deterministic_attention_forward_impl( + q, k, v, causal, scale, key_padding_mask, true); +} + +// =========================================================================== +// BACKWARD +// =========================================================================== +namespace { + +// --------------------------------------------------------------------------- +// dP kernel: dP[b,hq,q,k] = sum_{d=0}^{D-1} dO[b,hq,q,d] * V[b,kv_head,k,d] +// Grid: (Skv_blocks, Sq_blocks, B*Hq) +// --------------------------------------------------------------------------- +template +__global__ void dp_kernel( + const scalar_t* __restrict__ dO, // [B, Hq, Sq, D] + const scalar_t* __restrict__ V, // [B, Hkv, Skv, D] + float* __restrict__ dP, // [B, Hq, Sq, Skv] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D) { + + const int k_idx = blockIdx.x * kQKTileK + threadIdx.x; + const int q_idx = blockIdx.y * kQKTileQ + threadIdx.y; + const int bh = blockIdx.z; + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || k_idx >= Skv) return; + + const int kv_head = hq / (Hq / Hkv); + + const scalar_t* do_ptr = dO + ((int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D); + const scalar_t* v_ptr = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D + (int64_t)k_idx * D); + + float acc = 0.0f; + #pragma unroll 8 + for (int64_t d = 0; d < D; ++d) { + acc += (float)do_ptr[d] * (float)v_ptr[d]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv + k_idx; + dP[out_idx] = acc; +} + +// --------------------------------------------------------------------------- +// Softmax backward kernel: one CTA per (b, hq, q) row +// delta[row] = sum_k(dP[row,k] * P[row,k]) +// dS[row,k] = P[row,k] * (dP[row,k] - delta) +// Writes dS in-place over the dP buffer. +// --------------------------------------------------------------------------- +__global__ void softmax_backward_kernel( + float* __restrict__ dP_dS, // [B, Hq, Sq, Skv] - input dP, output dS + const float* __restrict__ P, // [B, Hq, Sq, Skv] + int64_t Skv) { + + const int row_idx = blockIdx.x; + float* ds_row = dP_dS + (int64_t)row_idx * Skv; + const float* p_row = P + (int64_t)row_idx * Skv; + + // Compute delta = sum_k(dP * P) + __shared__ float sdelta[kSoftmaxThreads]; + float thread_delta = 0.0f; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + thread_delta += ds_row[k] * p_row[k]; + } + sdelta[threadIdx.x] = thread_delta; + __syncthreads(); + for (int stride = kSoftmaxThreads / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + sdelta[threadIdx.x] += sdelta[threadIdx.x + stride]; + } + __syncthreads(); + } + float delta = sdelta[0]; + + // dS = P * (dP - delta) + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + ds_row[k] = p_row[k] * (ds_row[k] - delta); + } +} + +// --------------------------------------------------------------------------- +// dQ kernel: dQ[b,hq,q,d] = scale * sum_{k=0}^{Skv-1} dS[b,hq,q,k] * K[b,kv_head,k,d] +// Grid: (D_blocks, Sq_blocks, B*Hq) +// --------------------------------------------------------------------------- +template +__global__ void dq_kernel( + const float* __restrict__ dS, // [B, Hq, Sq, Skv] + const scalar_t* __restrict__ K, // [B, Hkv, Skv, D] + scalar_t* __restrict__ dQ, // [B, Hq, Sq, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D, + float scale) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int q_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int bh = blockIdx.z; + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || d_idx >= D) return; + + const int kv_head = hq / (Hq / Hkv); + + const float* ds_row = dS + ((int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv); + const scalar_t* k_base = K + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); + + float acc = 0.0f; + for (int64_t k = 0; k < Skv; ++k) { + acc += ds_row[k] * (float)k_base[k * D + d_idx]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D + d_idx; + dQ[out_idx] = (scalar_t)(scale * acc); +} + +// --------------------------------------------------------------------------- +// dK kernel: dK[b,hkv,k,d] = scale * sum_{local=0..g-1} sum_{q=0..Sq-1} dS[b,hq,q,k]*Q[b,hq,q,d] +// Grid: (D_blocks, Skv_blocks, B*Hkv) +// Each thread: single writer for one dK element (§4.1 fixed order). +// --------------------------------------------------------------------------- +template +__global__ void dk_kernel( + const float* __restrict__ dS, // [B, Hq, Sq, Skv] + const scalar_t* __restrict__ Q, // [B, Hq, Sq, D] + scalar_t* __restrict__ dK, // [B, Hkv, Skv, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D, + float scale) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int k_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int b_hkv = blockIdx.z; + const int b = b_hkv / Hkv; + const int hkv = b_hkv % Hkv; + + if (k_idx >= Skv || d_idx >= D) return; + + const int64_t g = Hq / Hkv; + + float acc = 0.0f; + for (int64_t local = 0; local < g; ++local) { + int64_t hq = hkv * g + local; + for (int64_t qi = 0; qi < Sq; ++qi) { + float ds_val = dS[(int64_t)b * Hq * Sq * Skv + hq * Sq * Skv + qi * Skv + k_idx]; + float q_val = (float)Q[(int64_t)b * Hq * Sq * D + hq * Sq * D + qi * D + d_idx]; + acc += ds_val * q_val; + } + } + + const int64_t out_idx = (int64_t)b * Hkv * Skv * D + (int64_t)hkv * Skv * D + (int64_t)k_idx * D + d_idx; + dK[out_idx] = (scalar_t)(scale * acc); +} + +// --------------------------------------------------------------------------- +// dV kernel: dV[b,hkv,k,d] = sum_{local=0..g-1} sum_{q=0..Sq-1} P[b,hq,q,k]*dO[b,hq,q,d] +// Grid: (D_blocks, Skv_blocks, B*Hkv) +// Each thread: single writer for one dV element (§4.1 fixed order). +// --------------------------------------------------------------------------- +template +__global__ void dv_kernel( + const float* __restrict__ P, // [B, Hq, Sq, Skv] + const scalar_t* __restrict__ dO, // [B, Hq, Sq, D] + scalar_t* __restrict__ dV, // [B, Hkv, Skv, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int k_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int b_hkv = blockIdx.z; + const int b = b_hkv / Hkv; + const int hkv = b_hkv % Hkv; + + if (k_idx >= Skv || d_idx >= D) return; + + const int64_t g = Hq / Hkv; + + float acc = 0.0f; + for (int64_t local = 0; local < g; ++local) { + int64_t hq = hkv * g + local; + for (int64_t qi = 0; qi < Sq; ++qi) { + float p_val = P[(int64_t)b * Hq * Sq * Skv + hq * Sq * Skv + qi * Skv + k_idx]; + float do_val = (float)dO[(int64_t)b * Hq * Sq * D + hq * Sq * D + qi * D + d_idx]; + acc += p_val * do_val; + } + } + + const int64_t out_idx = (int64_t)b * Hkv * Skv * D + (int64_t)hkv * Skv * D + (int64_t)k_idx * D + d_idx; + dV[out_idx] = (scalar_t)acc; +} + +} // namespace (backward kernels) + +// Returns {dQ, dK, dV} +std::vector deterministic_attention_backward( + torch::Tensor grad_output, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor P, // saved from forward [B, Hq, Sq, Skv] FP32 + bool causal, + double scale, + torch::optional key_padding_mask) { + + const at::musa::OptionalMUSAGuard device_guard(at::device_of(q)); + + auto dO = grad_output.contiguous(); + auto q_c = q.contiguous(); + auto k_c = k.contiguous(); + auto v_c = v.contiguous(); + auto P_c = P.contiguous(); + + const int64_t B = q_c.size(0); + const int64_t Hq = q_c.size(1); + const int64_t Sq = q_c.size(2); + const int64_t D = q_c.size(3); + const int64_t Hkv = k_c.size(1); + const int64_t Skv = k_c.size(2); + + auto stream = at::musa::getCurrentMUSAStream(); + + // dP = dO @ V^T [B, Hq, Sq, Skv] + auto dP = torch::empty({B, Hq, Sq, Skv}, q_c.options().dtype(at::kFloat)); + { + dim3 block(kQKTileK, kQKTileQ); + dim3 grid( + (Skv + kQKTileK - 1) / kQKTileK, + (Sq + kQKTileQ - 1) / kQKTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dp_kernel", [&] { + dp_kernel<<>>( + dO.data_ptr(), + v_c.data_ptr(), + dP.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + }); + } + + // Softmax backward: dS = P * (dP - delta), writes in-place over dP + { + const int64_t num_rows = B * Hq * Sq; + softmax_backward_kernel<<>>( + dP.data_ptr(), + P_c.data_ptr(), + Skv); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + } + // dP buffer now contains dS + + // dQ = scale * dS @ K + auto dQ = torch::empty_like(q_c); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Sq + kPVTileQ - 1) / kPVTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dq_kernel", [&] { + dq_kernel<<>>( + dP.data_ptr(), + k_c.data_ptr(), + dQ.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + }); + } + + // dK = scale * dS^T @ Q (per kv_head, accumulate over query heads in group) + auto dK = torch::empty_like(k_c); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Skv + kPVTileQ - 1) / kPVTileQ, + B * Hkv); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dk_kernel", [&] { + dk_kernel<<>>( + dP.data_ptr(), + q_c.data_ptr(), + dK.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + }); + } + + // dV = P^T @ dO (per kv_head, accumulate over query heads in group) + auto dV = torch::empty_like(v_c); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Skv + kPVTileQ - 1) / kPVTileQ, + B * Hkv); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dv_kernel", [&] { + dv_kernel<<>>( + P_c.data_ptr(), + dO.data_ptr(), + dV.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + }); + } + + return {dQ, dK, dV}; +} diff --git a/csrc/musa/deterministic_logp_kernel.mu b/csrc/musa/deterministic_logp_kernel.mu new file mode 100644 index 00000000..bcb7a967 --- /dev/null +++ b/csrc/musa/deterministic_logp_kernel.mu @@ -0,0 +1,360 @@ +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kDeterministicLogpSmallBlockSize = 128; +constexpr int kDeterministicLogpMediumBlockSize = 256; +constexpr int kDeterministicLogpLargeBlockSize = 512; +constexpr int kDeterministicLogpSmallVocabLimit = 128; +constexpr int kDeterministicLogpMediumVocabLimit = 4096; +constexpr int kDeterministicLogpWarpSize = 32; +constexpr float kDeterministicLogpNegInf = -3.4028234663852886e38F; + +template +__device__ __forceinline__ T deterministic_logp_shfl_down_32(T value, unsigned int delta) { +#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__) + return __shfl_down(value, delta, kDeterministicLogpWarpSize); +#else + return __shfl_down_sync(0xffffffffu, value, delta, kDeterministicLogpWarpSize); +#endif +} + +template +struct DeterministicLogpBlockTraits { + static_assert( + BlockSize == kDeterministicLogpSmallBlockSize || + BlockSize == kDeterministicLogpMediumBlockSize || + BlockSize == kDeterministicLogpLargeBlockSize, + "deterministic logp reduction topology requires a supported fixed block size"); + static_assert(BlockSize % kDeterministicLogpWarpSize == 0, "block size must be warp-aligned"); + static constexpr int WarpCount = BlockSize / kDeterministicLogpWarpSize; +}; + +template +__device__ __forceinline__ float deterministicBlockReduceMax(float val) { + constexpr int WarpCount = DeterministicLogpBlockTraits::WarpCount; + __shared__ float shared[WarpCount]; + + int lane = threadIdx.x & (kDeterministicLogpWarpSize - 1); + int wid = threadIdx.x / kDeterministicLogpWarpSize; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset)); + } + + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + + const bool has_warp_value = threadIdx.x < WarpCount; + const int shared_idx = has_warp_value ? threadIdx.x : 0; + val = has_warp_value ? shared[shared_idx] : kDeterministicLogpNegInf; + if (wid == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset)); + } + } + return val; +} + +template +__device__ __forceinline__ float deterministicBlockReduceSum(float val) { + constexpr int WarpCount = DeterministicLogpBlockTraits::WarpCount; + __shared__ float shared[WarpCount]; + + int lane = threadIdx.x & (kDeterministicLogpWarpSize - 1); + int wid = threadIdx.x / kDeterministicLogpWarpSize; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + val += deterministic_logp_shfl_down_32(val, offset); + } + + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + + const bool has_warp_value = threadIdx.x < WarpCount; + const int shared_idx = has_warp_value ? threadIdx.x : 0; + val = has_warp_value ? shared[shared_idx] : 0.0f; + if (wid == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + val += deterministic_logp_shfl_down_32(val, offset); + } + } + return val; +} + +template +__global__ void __launch_bounds__(BlockSize) deterministic_logp_forward_kernel( + const input_t* __restrict__ logits, + const int64_t* __restrict__ token_ids, + output_t* __restrict__ output, + const int64_t* __restrict__ row_indices, + int64_t total_rows, + int vocab_size) { + int64_t row = row_indices == nullptr ? blockIdx.x : row_indices[blockIdx.x]; + if (row < 0 || row >= total_rows) { + return; + } + + const input_t* row_logits = logits + row * vocab_size; + + float local_max = kDeterministicLogpNegInf; + for (int col = threadIdx.x; col < vocab_size; col += BlockSize) { + local_max = fmaxf(local_max, static_cast(row_logits[col])); + } + + float max_val = deterministicBlockReduceMax(local_max); + + __shared__ float row_max; + if (threadIdx.x == 0) { + row_max = max_val; + } + __syncthreads(); + + float local_sum = 0.0f; + for (int col = threadIdx.x; col < vocab_size; col += BlockSize) { + local_sum += expf(static_cast(row_logits[col]) - row_max); + } + + float sum_val = deterministicBlockReduceSum(local_sum); + + __shared__ float row_sum; + if (threadIdx.x == 0) { + row_sum = sum_val; + } + __syncthreads(); + + // Indexed mode may launch duplicate row ids. The writes are idempotent: + // every duplicate writer computes and stores the same deterministic value. + if (threadIdx.x == 0) { + int64_t target_id = token_ids[row]; + if (target_id >= 0 && target_id < vocab_size) { + float target_logit = static_cast(row_logits[target_id]); + output[row] = static_cast(target_logit - row_max - logf(row_sum)); + } else { + output[row] = static_cast(0.0f); + } + } +} + +void check_deterministic_logp_inputs( + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& output) { + TORCH_CHECK(logits.device().type() == c10::kPrivateUse1, "logits must be a CUDA tensor"); + TORCH_CHECK(token_ids.device().type() == c10::kPrivateUse1, "token_ids must be a CUDA tensor"); + TORCH_CHECK(output.device().type() == c10::kPrivateUse1, "output must be a CUDA tensor"); + TORCH_CHECK( + logits.device() == token_ids.device(), + "logits and token_ids must be on the same CUDA device"); + TORCH_CHECK( + logits.device() == output.device(), + "logits and output must be on the same CUDA device"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + TORCH_CHECK(token_ids.dim() == 1, "token_ids must be a 1D tensor"); + TORCH_CHECK(output.dim() == 1, "output must be a 1D tensor"); + TORCH_CHECK(token_ids.scalar_type() == at::ScalarType::Long, "token_ids must be int64"); + TORCH_CHECK( + token_ids.numel() == logits.size(0), + "token_ids length must match logits rows"); + TORCH_CHECK(output.numel() == logits.size(0), "output length must match logits rows"); + TORCH_CHECK(output.is_contiguous(), "output must be contiguous"); + TORCH_CHECK(logits.size(1) > 0, "logits vocab dimension must be non-empty"); + TORCH_CHECK( + logits.size(0) <= std::numeric_limits::max(), + "logits row count exceeds CUDA grid-x limit"); + TORCH_CHECK( + logits.size(1) <= std::numeric_limits::max(), + "logits vocab dimension exceeds int32 kernel limit"); + TORCH_CHECK( + output.scalar_type() == at::ScalarType::Float || + output.scalar_type() == at::ScalarType::Double || + output.scalar_type() == at::ScalarType::Half || + output.scalar_type() == at::ScalarType::BFloat16, + "output dtype must be float64, float32, float16, or bfloat16"); +} + +void check_deterministic_logp_indices( + const torch::Tensor& logits, + const torch::Tensor& row_indices) { + TORCH_CHECK( + reinterpret_cast(logits.data_ptr()) % 16 == 0, + "logits must be 16-byte aligned"); + TORCH_CHECK(row_indices.device().type() == c10::kPrivateUse1, "row_indices must be a CUDA tensor"); + TORCH_CHECK( + logits.device() == row_indices.device(), + "logits and row_indices must be on the same CUDA device"); + TORCH_CHECK(row_indices.dim() == 1, "row_indices must be a 1D tensor"); + TORCH_CHECK(row_indices.scalar_type() == at::ScalarType::Long, "row_indices must be int64"); + TORCH_CHECK( + row_indices.numel() <= std::numeric_limits::max(), + "row_indices length exceeds CUDA grid-x limit"); +} + +void launch_deterministic_logp_kernel( + const torch::Tensor& logits, + const torch::Tensor& token_ids, + const torch::Tensor& output, + const int64_t* row_indices_ptr, + int64_t launch_rows, + int64_t total_rows, + int64_t vocab_size) { + if (launch_rows == 0) { + return; + } + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + logits.scalar_type(), + "deterministic_logp_kernel", + ([&] { + using input_t = scalar_t; + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + output.scalar_type(), + "deterministic_logp_output_kernel", + ([&] { + using output_t = scalar_t; + const int vocab_size_i32 = static_cast(vocab_size); + const int launch_rows_i32 = static_cast(launch_rows); + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + if (vocab_size <= kDeterministicLogpSmallVocabLimit) { + deterministic_logp_forward_kernel< + input_t, + output_t, + kDeterministicLogpSmallBlockSize><<< + launch_rows_i32, + kDeterministicLogpSmallBlockSize, + 0, + stream>>>( + logits.data_ptr(), + token_ids.data_ptr(), + output.data_ptr(), + row_indices_ptr, + total_rows, + vocab_size_i32); + } else if (vocab_size <= kDeterministicLogpMediumVocabLimit) { + deterministic_logp_forward_kernel< + input_t, + output_t, + kDeterministicLogpMediumBlockSize><<< + launch_rows_i32, + kDeterministicLogpMediumBlockSize, + 0, + stream>>>( + logits.data_ptr(), + token_ids.data_ptr(), + output.data_ptr(), + row_indices_ptr, + total_rows, + vocab_size_i32); + } else { + deterministic_logp_forward_kernel< + input_t, + output_t, + kDeterministicLogpLargeBlockSize><<< + launch_rows_i32, + kDeterministicLogpLargeBlockSize, + 0, + stream>>>( + logits.data_ptr(), + token_ids.data_ptr(), + output.data_ptr(), + row_indices_ptr, + total_rows, + vocab_size_i32); + } + })); + })); + + C10_MUSA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + +torch::Tensor deterministic_logp_forward_out( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor output) { + check_deterministic_logp_inputs(logits, token_ids, output); + + auto logits_contig = logits.contiguous(); + auto token_ids_contig = token_ids.contiguous(); + + int64_t total_rows = logits_contig.size(0); + int64_t vocab_size = logits_contig.size(1); + launch_deterministic_logp_kernel( + logits_contig, + token_ids_contig, + output, + nullptr, + total_rows, + total_rows, + vocab_size); + + return output; +} + +torch::Tensor deterministic_logp_forward_indexed_out( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices, + torch::Tensor output) { + check_deterministic_logp_inputs(logits, token_ids, output); + check_deterministic_logp_indices(logits, row_indices); + + auto logits_contig = logits.contiguous(); + auto token_ids_contig = token_ids.contiguous(); + auto row_indices_contig = row_indices.contiguous(); + + int64_t total_rows = logits_contig.size(0); + int64_t vocab_size = logits_contig.size(1); + int64_t valid_rows = row_indices_contig.numel(); + + launch_deterministic_logp_kernel( + logits_contig, + token_ids_contig, + output, + row_indices_contig.data_ptr(), + valid_rows, + total_rows, + vocab_size); + + return output; +} + +torch::Tensor deterministic_logp_forward(torch::Tensor logits, torch::Tensor token_ids) { + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + auto output = torch::empty({logits.size(0)}, logits.options()); + return deterministic_logp_forward_out(logits, token_ids, output); +} + +torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids) { + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + auto output = torch::empty({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + return deterministic_logp_forward_out(logits, token_ids, output); +} + +torch::Tensor deterministic_logp_forward_indexed_fp32( + torch::Tensor logits, + torch::Tensor token_ids, + torch::Tensor row_indices) { + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + auto output = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); + return deterministic_logp_forward_indexed_out(logits, token_ids, row_indices, output); +} diff --git a/csrc/musa/embedding_lm_head.mu b/csrc/musa/embedding_lm_head.mu new file mode 100644 index 00000000..c0aa1fc4 --- /dev/null +++ b/csrc/musa/embedding_lm_head.mu @@ -0,0 +1,303 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Single-card SM90 batch-invariant embedding and LM-head reference kernels. +// +// These kernels intentionally avoid Split-K. For LM-head, each output element +// owns the full hidden-dimension reduction inside one CTA, so the K traversal +// and reduction tree depend only on hidden_size, not on batch/sequence layout. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxCachedDevices = 64; + +// Reference path: one CTA per logit is intentionally throughput-heavy for large +// vocab prefill. It exists to preserve a fixed K reduction order for WS1. + +template +__device__ __forceinline__ float to_float(T value) { + return static_cast(value); +} + +template +__device__ __forceinline__ T from_float(float value) { + return static_cast(value); +} + +template +__device__ __forceinline__ float block_sum(float value) { + static_assert(BlockSize % 32 == 0, "BlockSize must be a warp multiple"); + __shared__ float shared[BlockSize / 32]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + + if (lane == 0) { + shared[warp] = value; + } + __syncthreads(); + + value = threadIdx.x < (BlockSize / 32) ? shared[lane] : 0.0f; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + } + return value; +} + +template +__global__ void embedding_sm90_forward_kernel(const int64_t *__restrict__ token_ids, + const input_t *__restrict__ weight, + output_t *__restrict__ output, + int64_t num_tokens, int64_t hidden_size, + int64_t /*vocab_size*/) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = num_tokens * hidden_size; + if (idx >= total) { + return; + } + + const int64_t token_row = idx / hidden_size; + const int64_t hidden_col = idx - token_row * hidden_size; + const int64_t token_id = token_ids[token_row]; + output[idx] = static_cast(weight[token_id * hidden_size + hidden_col]); +} + +template +__global__ void lm_head_sm90_forward_kernel(const scalar_t *__restrict__ hidden, + const scalar_t *__restrict__ weight, + const float *__restrict__ bias, + output_t *__restrict__ output, + int64_t num_tokens, int64_t hidden_size, + int64_t vocab_size) { + const int64_t out_idx = static_cast(blockIdx.x); + const int64_t total = num_tokens * vocab_size; + if (out_idx >= total) { + return; + } + + const int64_t token_row = out_idx / vocab_size; + const int64_t vocab_col = out_idx - token_row * vocab_size; + const scalar_t *hidden_row = hidden + token_row * hidden_size; + const scalar_t *weight_row = weight + vocab_col * hidden_size; + + float acc = 0.0f; + for (int64_t k = threadIdx.x; k < hidden_size; k += blockDim.x) { + acc += to_float(hidden_row[k]) * to_float(weight_row[k]); + } + acc = block_sum(acc); + + if (threadIdx.x == 0) { + if (bias != nullptr) { + acc += bias[vocab_col]; + } + output[out_idx] = from_float(acc); + } +} + +std::vector embedding_output_sizes(torch::Tensor token_ids, int64_t hidden_size) { + std::vector sizes; + sizes.reserve(static_cast(token_ids.dim()) + 1); + for (int64_t i = 0; i < token_ids.dim(); ++i) { + sizes.push_back(token_ids.size(i)); + } + sizes.push_back(hidden_size); + return sizes; +} + +std::vector lm_head_output_sizes(torch::Tensor hidden, int64_t vocab_size) { + std::vector sizes; + sizes.reserve(static_cast(hidden.dim())); + for (int64_t i = 0; i < hidden.dim() - 1; ++i) { + sizes.push_back(hidden.size(i)); + } + sizes.push_back(vocab_size); + return sizes; +} + +bool is_supported_float_dtype(at::ScalarType dtype) { + return dtype == at::kFloat || dtype == at::kHalf || dtype == at::kBFloat16; +} + +void check_sm90_device() { + // The port uses generic warp/reduction instructions, not NVIDIA SM90 PTX. +} + +void check_cuda_same_device(torch::Tensor a, torch::Tensor b, const char *a_name, + const char *b_name) { + TORCH_CHECK(a.device().type() == c10::kPrivateUse1 && b.device().type() == c10::kPrivateUse1, a_name, " and ", b_name, + " must be CUDA tensors"); + TORCH_CHECK(a.device() == b.device(), a_name, " and ", b_name, + " must be on the same CUDA device"); +} + +torch::Tensor embedding_sm90_forward_impl(torch::Tensor token_ids, torch::Tensor weight, + bool output_fp32) { + TORCH_CHECK(token_ids.device().type() == c10::kPrivateUse1 && weight.device().type() == c10::kPrivateUse1, + "token_ids and weight must be CUDA tensors"); + TORCH_CHECK(token_ids.device() == weight.device(), + "token_ids and weight must be on the same CUDA device"); + TORCH_CHECK(weight.dim() == 2, "embedding weight must be [vocab, hidden]"); + TORCH_CHECK(weight.is_contiguous(), "embedding weight must be contiguous"); + TORCH_CHECK(is_supported_float_dtype(weight.scalar_type()), + "embedding_sm90 supports fp32, fp16, and bf16 weights"); + + c10::musa::MUSAGuard device_guard(weight.device()); + check_sm90_device(); + const int64_t vocab_size = weight.size(0); + const int64_t hidden_size = weight.size(1); + const int64_t num_tokens = token_ids.numel(); + auto ids = token_ids.reshape({num_tokens}).to(at::kLong).contiguous(); + if (num_tokens > 0) { + const int64_t min_id = ids.min().item(); + const int64_t max_id = ids.max().item(); + TORCH_CHECK(min_id >= 0 && max_id < vocab_size, + "embedding_sm90 token ids must be in [0, ", vocab_size - 1, + "], got [", min_id, ", ", max_id, "]"); + } + auto out_options = weight.options().dtype(output_fp32 ? at::kFloat : weight.scalar_type()); + auto output = torch::empty(embedding_output_sizes(token_ids, hidden_size), out_options); + if (num_tokens == 0 || hidden_size == 0) { + return output; + } + + const int64_t total = num_tokens * hidden_size; + TORCH_CHECK(total <= std::numeric_limits::max(), + "embedding_sm90 launch size exceeds CUDA grid limit"); + const int blocks = static_cast((total + kThreads - 1) / kThreads); + auto output_2d = output.reshape({num_tokens, hidden_size}); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, weight.scalar_type(), "embedding_sm90_forward", [&] { + if (output_fp32) { + embedding_sm90_forward_kernel + <<>>( + ids.data_ptr(), weight.data_ptr(), + output_2d.data_ptr(), num_tokens, hidden_size, vocab_size); + } else { + embedding_sm90_forward_kernel + <<>>( + ids.data_ptr(), weight.data_ptr(), + output_2d.data_ptr(), num_tokens, hidden_size, vocab_size); + } + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return output; +} + +torch::Tensor lm_head_sm90_forward_impl(torch::Tensor hidden, torch::Tensor weight, + torch::optional bias, + bool output_fp32) { + check_cuda_same_device(hidden, weight, "hidden", "weight"); + TORCH_CHECK(hidden.dim() >= 2, "hidden must have shape [..., hidden]"); + TORCH_CHECK(weight.dim() == 2, "lm_head weight must be [vocab, hidden]"); + TORCH_CHECK(hidden.size(-1) == weight.size(1), "hidden/weight hidden-dim mismatch"); + TORCH_CHECK(is_supported_float_dtype(hidden.scalar_type()), + "lm_head_sm90 supports fp32, fp16, and bf16 hidden states"); + TORCH_CHECK(is_supported_float_dtype(weight.scalar_type()), + "lm_head_sm90 supports fp32, fp16, and bf16 weights"); + + c10::musa::MUSAGuard device_guard(hidden.device()); + check_sm90_device(); + const int64_t hidden_size = hidden.size(-1); + TORCH_CHECK(hidden_size > 0, "lm_head hidden dimension must be non-zero"); + const int64_t vocab_size = weight.size(0); + const int64_t num_tokens = hidden.numel() / hidden_size; + const at::ScalarType compute_dtype = output_fp32 ? at::kFloat : hidden.scalar_type(); + + auto hidden_2d = hidden.reshape({num_tokens, hidden_size}).to(compute_dtype).contiguous(); + auto weight_2d = weight.to(compute_dtype).contiguous(); + torch::Tensor bias_f; + const float *bias_ptr = nullptr; + if (bias.has_value()) { + TORCH_CHECK(bias->device().type() == c10::kPrivateUse1, "lm_head bias must be a CUDA tensor"); + TORCH_CHECK(bias->device() == hidden.device(), + "lm_head bias must be on the same CUDA device as hidden"); + TORCH_CHECK(bias->dim() == 1, "lm_head bias must be 1-D [vocab]"); + TORCH_CHECK(bias->numel() == vocab_size, "lm_head bias must have vocab elements"); + TORCH_CHECK(is_supported_float_dtype(bias->scalar_type()), + "lm_head_sm90 supports fp32, fp16, and bf16 bias"); + bias_f = bias->reshape({vocab_size}).to(at::kFloat).contiguous(); + bias_ptr = bias_f.data_ptr(); + } + + auto out_options = hidden.options().dtype(output_fp32 ? at::kFloat : hidden.scalar_type()); + auto output = torch::empty(lm_head_output_sizes(hidden, vocab_size), out_options); + if (num_tokens == 0 || vocab_size == 0) { + return output; + } + + const int64_t total = num_tokens * vocab_size; + TORCH_CHECK(total <= std::numeric_limits::max(), + "lm_head_sm90 launch size exceeds CUDA grid limit"); + auto output_2d = output.reshape({num_tokens, vocab_size}); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, compute_dtype, "lm_head_sm90_forward", [&] { + if (output_fp32) { + lm_head_sm90_forward_kernel + <<(total), kThreads, 0, at::musa::getCurrentMUSAStream()>>>( + hidden_2d.data_ptr(), weight_2d.data_ptr(), + bias_ptr, output_2d.data_ptr(), num_tokens, hidden_size, + vocab_size); + } else { + lm_head_sm90_forward_kernel + <<(total), kThreads, 0, at::musa::getCurrentMUSAStream()>>>( + hidden_2d.data_ptr(), weight_2d.data_ptr(), + bias_ptr, output_2d.data_ptr(), num_tokens, hidden_size, + vocab_size); + } + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return output; +} + +} // namespace + +torch::Tensor embedding_sm90_forward(torch::Tensor token_ids, torch::Tensor weight) { + return embedding_sm90_forward_impl(token_ids, weight, false); +} + +torch::Tensor embedding_sm90_forward_fp32(torch::Tensor token_ids, torch::Tensor weight) { + return embedding_sm90_forward_impl(token_ids, weight, true); +} + +torch::Tensor lm_head_sm90_forward(torch::Tensor hidden, torch::Tensor weight, + torch::optional bias) { + return lm_head_sm90_forward_impl(hidden, weight, bias, false); +} + +torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::Tensor weight, + torch::optional bias) { + return lm_head_sm90_forward_impl(hidden, weight, bias, true); +} + +torch::Tensor det_gemm_rowwise_fwd_fp32(torch::Tensor a, torch::Tensor b) { + TORCH_CHECK(a.dim() == 2 && b.dim() == 2, + "det_gemm_rowwise_fwd_fp32 expects [M,K] @ [K,N]"); + TORCH_CHECK(a.size(1) == b.size(0), + "det_gemm_rowwise_fwd_fp32: K mismatch"); + // lm_head_sm90_forward_impl computes one output element per CTA using a + // fixed 256-thread block reduction. Passing B^T as [N,K] exposes that + // deterministic rowwise reduction as a general GEMM configuration. + return lm_head_sm90_forward_impl( + a, b.transpose(0, 1).contiguous(), torch::optional{}, true); +} diff --git a/csrc/musa/fused_logp_kernel.mu b/csrc/musa/fused_logp_kernel.mu new file mode 100644 index 00000000..cf554975 --- /dev/null +++ b/csrc/musa/fused_logp_kernel.mu @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#include +#include +#include +#include + +#include + +namespace { + +constexpr int kBlockSize = 256; + +__device__ __forceinline__ float block_reduce_max(float value) { + __shared__ float partial[32]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = fmaxf(value, __shfl_down_sync(0xffffffffu, value, offset, 32)); + } + if (lane == 0) { + partial[warp] = value; + } + __syncthreads(); + + value = threadIdx.x < (kBlockSize / 32) ? partial[lane] : -FLT_MAX; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = fmaxf(value, __shfl_down_sync(0xffffffffu, value, offset, 32)); + } + } + if (threadIdx.x == 0) { + partial[0] = value; + } + __syncthreads(); + return partial[0]; +} + +__device__ __forceinline__ float block_reduce_sum(float value) { + __shared__ float partial[32]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset, 32); + } + if (lane == 0) { + partial[warp] = value; + } + __syncthreads(); + + value = threadIdx.x < (kBlockSize / 32) ? partial[lane] : 0.0f; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset, 32); + } + } + if (threadIdx.x == 0) { + partial[0] = value; + } + __syncthreads(); + return partial[0]; +} + +template +__global__ void fused_logp_kernel( + const scalar_t* __restrict__ logits, + const int64_t* __restrict__ token_ids, + scalar_t* __restrict__ output, + int rows, + int vocab) { + const int row = blockIdx.x; + if (row >= rows) { + return; + } + + const scalar_t* row_logits = logits + static_cast(row) * vocab; + float row_max = -FLT_MAX; + for (int col = threadIdx.x; col < vocab; col += blockDim.x) { + row_max = fmaxf(row_max, static_cast(row_logits[col])); + } + row_max = block_reduce_max(row_max); + + float row_sum = 0.0f; + for (int col = threadIdx.x; col < vocab; col += blockDim.x) { + row_sum += expf(static_cast(row_logits[col]) - row_max); + } + row_sum = block_reduce_sum(row_sum); + + if (threadIdx.x == 0) { + const int64_t target = token_ids[row]; + const float target_logit = static_cast(row_logits[target]); + output[row] = static_cast(target_logit - row_max - logf(row_sum)); + } +} + +} // namespace + +torch::Tensor fused_logp_forward_musa(torch::Tensor logits, torch::Tensor token_ids) { + auto output = torch::empty({logits.size(0)}, logits.options()); + const int rows = static_cast(logits.size(0)); + const int vocab = static_cast(logits.size(1)); + if (rows == 0) { + return output; + } + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + logits.scalar_type(), + "musa_fused_logp", + [&] { + fused_logp_kernel<<>>( + logits.data_ptr(), + token_ids.data_ptr(), + output.data_ptr(), + rows, + vocab); + }); + C10_MUSA_KERNEL_LAUNCH_CHECK(); + return output; +} diff --git a/csrc/musa/ops.cpp b/csrc/musa/ops.cpp new file mode 100644 index 00000000..51b05970 --- /dev/null +++ b/csrc/musa/ops.cpp @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#include + +#include + +torch::Tensor fused_logp_forward_musa(torch::Tensor logits, torch::Tensor token_ids); +torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); +torch::Tensor silu_forward_cuda(torch::Tensor x); +torch::Tensor silu_backward_cuda(torch::Tensor dy, torch::Tensor x); +torch::Tensor swiglu_forward_cuda(torch::Tensor gate, torch::Tensor up); +std::vector swiglu_backward_cuda( + torch::Tensor dy, torch::Tensor gate, torch::Tensor up); +torch::Tensor swiglu_packed_forward_cuda(torch::Tensor gate_up); +std::vector swiglu_packed_backward_cuda( + torch::Tensor dy, torch::Tensor gate_up); +void rmsnorm_forward_cuda( + torch::Tensor x, torch::Tensor weight, torch::Tensor y, torch::Tensor rstd, double eps); +void rmsnorm_backward_dx_cuda( + torch::Tensor dy, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, torch::Tensor dx); +void rmsnorm_backward_partial_dw_cuda( + torch::Tensor dy, torch::Tensor x, torch::Tensor rstd, torch::Tensor mask, + torch::Tensor partial_dw); +void rmsnorm_backward_reduce_dw_cuda(torch::Tensor partial_dw, torch::Tensor dw); +void reduce_rows_fp32_left_fold_cuda(torch::Tensor rows, torch::Tensor output); +torch::Tensor embedding_sm90_forward_fp32(torch::Tensor token_ids, torch::Tensor weight); +torch::Tensor lm_head_sm90_forward_fp32( + torch::Tensor hidden, + torch::Tensor weight, + torch::optional bias); +torch::Tensor rope_apply_sm90( + torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin, + double sin_sign); +std::vector deterministic_attention_forward_fp32( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + torch::optional key_padding_mask); +std::vector deterministic_attention_backward( + torch::Tensor grad_output, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor probs, + bool causal, + double scale, + torch::optional key_padding_mask); + +torch::Tensor fused_logp_forward(torch::Tensor logits, torch::Tensor token_ids) { + TORCH_CHECK( + logits.device().type() == c10::kPrivateUse1, + "logits must be a MUSA tensor, got ", + logits.device()); + TORCH_CHECK( + token_ids.device().type() == c10::kPrivateUse1, + "token_ids must be a MUSA tensor, got ", + token_ids.device()); + TORCH_CHECK(logits.device() == token_ids.device(), "logits and token_ids must share a device"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + TORCH_CHECK(token_ids.dim() == 1, "token_ids must be a 1D tensor"); + TORCH_CHECK(token_ids.scalar_type() == at::ScalarType::Long, "token_ids must be int64"); + TORCH_CHECK(token_ids.numel() == logits.size(0), "token_ids length must match logits rows"); + TORCH_CHECK(logits.size(0) <= std::numeric_limits::max(), "too many logits rows"); + TORCH_CHECK(logits.size(1) > 0, "logits vocabulary dimension must be non-empty"); + if (token_ids.numel() > 0) { + TORCH_CHECK( + token_ids.min().item() >= 0 && + token_ids.max().item() < logits.size(1), + "token_ids must be within the logits vocabulary dimension"); + } + TORCH_CHECK( + logits.scalar_type() == at::ScalarType::Float || + logits.scalar_type() == at::ScalarType::Half || + logits.scalar_type() == at::ScalarType::BFloat16, + "MUSA fused_logp supports float32, float16, and bfloat16 logits"); + + auto logits_contiguous = logits.contiguous(); + auto token_ids_contiguous = token_ids.contiguous(); + return fused_logp_forward_musa(logits_contiguous, token_ids_contiguous); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fused_logp", &fused_logp_forward, "MUSA fused selected-token log-probability"); + m.def("deterministic_logp_fp32", &deterministic_logp_forward_fp32); + m.def("silu_forward", &silu_forward_cuda); + m.def("silu_backward", &silu_backward_cuda); + m.def("swiglu_forward", &swiglu_forward_cuda); + m.def("swiglu_backward", &swiglu_backward_cuda); + m.def("swiglu_packed_forward", &swiglu_packed_forward_cuda); + m.def("swiglu_packed_backward", &swiglu_packed_backward_cuda); + m.def("rmsnorm_forward", &rmsnorm_forward_cuda); + m.def("rmsnorm_backward_dx", &rmsnorm_backward_dx_cuda); + m.def("rmsnorm_backward_partial_dw", &rmsnorm_backward_partial_dw_cuda); + m.def("rmsnorm_backward_reduce_dw", &rmsnorm_backward_reduce_dw_cuda); + m.def("reduce_rows_fp32_left_fold", &reduce_rows_fp32_left_fold_cuda); + m.def("embedding_fp32", &embedding_sm90_forward_fp32); + m.def("lm_head_fp32", &lm_head_sm90_forward_fp32, py::arg("hidden"), + py::arg("weight"), py::arg("bias") = py::none()); + m.def("rope_apply", &rope_apply_sm90); + m.def("deterministic_attention_fp32", &deterministic_attention_forward_fp32, + py::arg("q"), py::arg("k"), py::arg("v"), py::arg("causal"), + py::arg("scale"), py::arg("key_padding_mask") = py::none()); + m.def("deterministic_attention_backward", &deterministic_attention_backward, + py::arg("grad_output"), py::arg("q"), py::arg("k"), py::arg("v"), + py::arg("probs"), py::arg("causal"), py::arg("scale"), + py::arg("key_padding_mask") = py::none()); +} diff --git a/csrc/musa/rmsnorm.mu b/csrc/musa/rmsnorm.mu new file mode 100644 index 00000000..88ef9127 --- /dev/null +++ b/csrc/musa/rmsnorm.mu @@ -0,0 +1,379 @@ +#include +#include +#if !defined(USE_ROCM) +#include +#include +#endif +#include +#include +#include +#include +#include + +template +__device__ __forceinline__ float load_as_float(const scalar_t* ptr) { + return static_cast(*ptr); +} + +template <> +__device__ __forceinline__ float load_as_float(const at::Half* ptr) { + const __half* p = reinterpret_cast(ptr); + return __half2float(*p); +} + +template <> +__device__ __forceinline__ float load_as_float(const at::BFloat16* ptr) { + const __mt_bfloat16* p = reinterpret_cast(ptr); + return __bfloat162float(*p); +} + + +template +__device__ __forceinline__ void store_from_float(scalar_t* ptr, float v) { + *ptr = static_cast(v); +} + +template <> +__device__ __forceinline__ void store_from_float(at::Half* ptr, float v) { + __half* p = reinterpret_cast<__half*>(ptr); + *p = __float2half(v); +} + +template <> +__device__ __forceinline__ void store_from_float(at::BFloat16* ptr, float v) { + __mt_bfloat16* p = reinterpret_cast<__mt_bfloat16*>(ptr); + *p = __float2bfloat16(v); +} + + +__device__ __forceinline__ float block_reduce_sum(float v) { + extern __shared__ float smem[]; + int tid = threadIdx.x; + assert((blockDim.x & (blockDim.x - 1)) == 0); + + smem[tid] = v; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + smem[tid] += smem[tid + stride]; + } + __syncthreads(); + } + + return smem[0]; +} + + +static int choose_threads(int H) { + if (H <= 64) return 64; + if (H <= 128) return 128; + if (H <= 256) return 256; + return 512; +} + +constexpr int RMSNORM_DW_ROWS_PER_CHUNK = 256; +constexpr int RMSNORM_DW_H_TILE = 128; + +#if !defined(USE_ROCM) +constexpr int FP32_LEFT_FOLD_THREADS = 256; + +__global__ void reduce_rows_fp32_left_fold_kernel( + const float* __restrict__ rows, + float* __restrict__ output, + int64_t row_count, + int64_t columns +) { + int64_t column = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (column >= columns) { + return; + } + + float acc = 0.0f; +#pragma unroll 1 + for (int64_t row = 0; row < row_count; ++row) { + acc = __fadd_rn(acc, rows[row * columns + column]); + } + output[column] = acc; +} +#endif + +int64_t rmsnorm_backward_dw_chunks_cuda(int64_t rows) { + return (rows + RMSNORM_DW_ROWS_PER_CHUNK - 1) / RMSNORM_DW_ROWS_PER_CHUNK; +} + +template +__global__ void rmsnorm_fwd_kernel( + const scalar_t* __restrict__ x, + const weight_t* __restrict__ weight, + scalar_t* __restrict__ y, + float* __restrict__ rstd, + int T, + int H, + float eps +) { + int row = blockIdx.x; + int tid = threadIdx.x; + + const scalar_t* x_row = x + row * H; + scalar_t* y_row = y + row * H; + + float local_sum = 0.0f; + + // 计算 sum(x^2)，每个 thread 负责若干列。 + for (int col = tid; col < H; col += blockDim.x) { + float xv = load_as_float(x_row + col); + local_sum += xv * xv; + } + + // 固定 block reduction。 + float sum = block_reduce_sum(local_sum); + + float row_rstd = rsqrtf(sum / static_cast(H) + eps); + + if (tid == 0) { + rstd[row] = row_rstd; + } + + __syncthreads(); + + // 写出 y = x * rstd * weight。 + for (int col = tid; col < H; col += blockDim.x) { + float xv = load_as_float(x_row + col); + float wv = load_as_float(weight + col); + float out = xv * row_rstd * wv; + store_from_float(y_row + col, out); + } +} + + +template +__global__ void rmsnorm_bwd_dx_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ x, + const weight_t* __restrict__ weight, + const float* __restrict__ rstd, + scalar_t* __restrict__ dx, + int T, + int H +) { + int row = blockIdx.x; + int tid = threadIdx.x; + + const scalar_t* dy_row = dy + row * H; + const scalar_t* x_row = x + row * H; + scalar_t* dx_row = dx + row * H; + + float local_dot = 0.0f; + + for (int col = tid; col < H; col += blockDim.x) { + float dyv = load_as_float(dy_row + col); + float xv = load_as_float(x_row + col); + float wv = load_as_float(weight + col); + local_dot += dyv * wv * xv; + } + + float dot = block_reduce_sum(local_dot); + + float r = rstd[row]; + float coeff = dot * r * r * r / static_cast(H); + + for (int col = tid; col < H; col += blockDim.x) { + float dyv = load_as_float(dy_row + col); + float xv = load_as_float(x_row + col); + float wv = load_as_float(weight + col); + + float out = r * dyv * wv - xv * coeff; + store_from_float(dx_row + col, out); + } +} + + +template +__global__ void rmsnorm_partial_dw_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ x, + const float* __restrict__ rstd, + const bool* __restrict__ mask, + float* __restrict__ partial_dw, + int T, + int H +) { + int chunk = blockIdx.x; + int h = blockIdx.y * RMSNORM_DW_H_TILE + threadIdx.x; + if (h >= H) { + return; + } + + int t0 = chunk * RMSNORM_DW_ROWS_PER_CHUNK; + float acc = 0.0f; + +#pragma unroll + for (int i = 0; i < RMSNORM_DW_ROWS_PER_CHUNK; ++i) { + int row = t0 + i; + float contrib = 0.0f; + if (row < T && mask[row]) { + int idx = row * H + h; + float dyv = load_as_float(dy + idx); + float xv = load_as_float(x + idx); + contrib = dyv * xv * rstd[row]; + } + acc += contrib; + } + + partial_dw[chunk * H + h] = acc; +} + + +__global__ void rmsnorm_reduce_dw_kernel( + const float* __restrict__ partial_dw, + float* __restrict__ dw, + int chunks, + int H +) { + int h = blockIdx.x * RMSNORM_DW_H_TILE + threadIdx.x; + if (h >= H) { + return; + } + + float acc = 0.0f; + for (int chunk = 0; chunk < chunks; ++chunk) { + acc += partial_dw[chunk * H + h]; + } + dw[h] = acc; +} + + +void rmsnorm_forward_cuda( + torch::Tensor x, + torch::Tensor weight, + torch::Tensor y, + torch::Tensor rstd, + double eps +) { + int T = x.size(0); + int H = x.size(1); + int threads = choose_threads(H); + size_t smem = threads * sizeof(float); + + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, x.scalar_type(), "rmsnorm_forward_cuda", [&] { + using x_t = scalar_t; + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, weight.scalar_type(), "rmsnorm_forward_weight_cuda", [&] { + using w_t = scalar_t; + rmsnorm_fwd_kernel<<>>( + x.data_ptr(), + weight.data_ptr(), + y.data_ptr(), + rstd.data_ptr(), + T, + H, + static_cast(eps) + ); + }); + }); +} + + +void rmsnorm_backward_dx_cuda( + torch::Tensor dy, + torch::Tensor x, + torch::Tensor weight, + torch::Tensor rstd, + torch::Tensor dx +) { + int T = x.size(0); + int H = x.size(1); + int threads = choose_threads(H); + size_t smem = threads * sizeof(float); + + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, x.scalar_type(), "rmsnorm_backward_dx_cuda", [&] { + using x_t = scalar_t; + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, weight.scalar_type(), "rmsnorm_backward_dx_weight_cuda", [&] { + using w_t = scalar_t; + rmsnorm_bwd_dx_kernel<<>>( + dy.data_ptr(), + x.data_ptr(), + weight.data_ptr(), + rstd.data_ptr(), + dx.data_ptr(), + T, + H + ); + }); + }); +} + + +void rmsnorm_backward_partial_dw_cuda( + torch::Tensor dy, + torch::Tensor x, + torch::Tensor rstd, + torch::Tensor mask, + torch::Tensor partial_dw +) { + int T = x.size(0); + int H = x.size(1); + + int chunks = (T + RMSNORM_DW_ROWS_PER_CHUNK - 1) / RMSNORM_DW_ROWS_PER_CHUNK; + dim3 blocks(chunks, (H + RMSNORM_DW_H_TILE - 1) / RMSNORM_DW_H_TILE); + dim3 threads(RMSNORM_DW_H_TILE); + + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, x.scalar_type(), "rmsnorm_partial_dw_cuda", [&] { + rmsnorm_partial_dw_kernel<<>>( + dy.data_ptr(), + x.data_ptr(), + rstd.data_ptr(), + mask.data_ptr(), + partial_dw.data_ptr(), + T, + H + ); + }); +} + + +void rmsnorm_backward_reduce_dw_cuda( + torch::Tensor partial_dw, + torch::Tensor dw +) { + int chunks = partial_dw.size(0); + int H = partial_dw.size(1); + + dim3 blocks((H + RMSNORM_DW_H_TILE - 1) / RMSNORM_DW_H_TILE); + dim3 threads(RMSNORM_DW_H_TILE); + + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + rmsnorm_reduce_dw_kernel<<>>( + partial_dw.data_ptr(), + dw.data_ptr(), + chunks, + H + ); +} + +#if !defined(USE_ROCM) +void reduce_rows_fp32_left_fold_cuda( + torch::Tensor rows, + torch::Tensor output +) { + const c10::musa::MUSAGuard device_guard(rows.device()); + int64_t columns = rows.size(1); + int64_t blocks = (columns + FP32_LEFT_FOLD_THREADS - 1) / FP32_LEFT_FOLD_THREADS; + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + reduce_rows_fp32_left_fold_kernel<<>>( + rows.data_ptr(), + output.data_ptr(), + rows.size(0), + columns + ); + C10_MUSA_KERNEL_LAUNCH_CHECK(); +} +#endif diff --git a/csrc/musa/rope.mu b/csrc/musa/rope.mu new file mode 100644 index 00000000..3d70d9fc --- /dev/null +++ b/csrc/musa/rope.mu @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// CUDA RoPE kernel for SM90 (GPT-NeoX / HF rotate-half), matching NativeRoPEOp. +// +// For a row (one [B, H, S] token vector of width D, half = D / 2) at sequence +// index s = row % S and pair index i in [0, half): +// +// c = cos[s, i] (fp32, precomputed to match the reference math) +// sn = sin[s, i] * sin_sign (sin_sign = +1 forward, -1 backward) +// out[i] = x[i] * c - x[i + half] * sn +// out[i + half] = x[i + half] * c + x[i] * sn +// +// The elementwise rotation is done in fp32 and rounded back to the input dtype. +// Backward reuses the same kernel with sin_sign = -1 (RoPE is an orthogonal +// per-position rotation, so grad_x = grad_out * cos - rotate_half(grad_out) * sin). + +#include +#include +#include + +namespace { + +template +__global__ void rope_apply_sm90_kernel( + const scalar_t* __restrict__ x, // [n_rows, D] + const float* __restrict__ cos, // [S, half] + const float* __restrict__ sin, // [S, half] + scalar_t* __restrict__ out, // [n_rows, D] + const int64_t n_rows, + const int S, + const int half, + const float sin_sign) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + const int64_t total = n_rows * static_cast(half); + if (idx >= total) { + return; + } + + const int64_t row = idx / half; + const int i = static_cast(idx % half); + const int seq = static_cast(row % S); + + const float c = cos[seq * half + i]; + const float sn = sin[seq * half + i] * sin_sign; + + const int64_t base = row * (2LL * half); + const float x1 = static_cast(x[base + i]); + const float x2 = static_cast(x[base + i + half]); + + out[base + i] = static_cast(x1 * c - x2 * sn); + out[base + i + half] = static_cast(x2 * c + x1 * sn); +} + +} // namespace + +// x: [n_rows, D] contiguous (any float dtype); cos/sin: [S, half] fp32 contiguous. +torch::Tensor rope_apply_sm90( + torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin, + double sin_sign) { + TORCH_CHECK(x.device().type() == c10::kPrivateUse1, "rope: x must be a CUDA tensor"); + TORCH_CHECK(x.dim() == 2, "rope: x must be 2-D [n_rows, D]"); + TORCH_CHECK(x.is_contiguous(), "rope: x must be contiguous"); + TORCH_CHECK(cos.device().type() == c10::kPrivateUse1 && sin.device().type() == c10::kPrivateUse1, "rope: cos/sin must be CUDA tensors"); + TORCH_CHECK(cos.scalar_type() == torch::kFloat32 && sin.scalar_type() == torch::kFloat32, + "rope: cos/sin must be fp32"); + TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(), "rope: cos/sin must be contiguous"); + + const int64_t n_rows = x.size(0); + const int64_t D = x.size(1); + TORCH_CHECK(D % 2 == 0, "rope: head_dim must be even"); + const int half = static_cast(D / 2); + const int S = static_cast(cos.size(0)); + TORCH_CHECK(cos.size(1) == half && sin.size(1) == half, + "rope: cos/sin last dim must equal head_dim/2"); + TORCH_CHECK(S > 0 && n_rows % S == 0, + "rope: n_rows must be divisible by seq length S"); + + const at::musa::OptionalMUSAGuard device_guard(device_of(x)); + auto out = torch::empty_like(x); + + const int64_t total = n_rows * static_cast(half); + const int threads = 256; + const int64_t blocks = (total + threads - 1) / threads; + auto stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "rope_apply_sm90", [&] { + rope_apply_sm90_kernel<<>>( + x.data_ptr(), + cos.data_ptr(), + sin.data_ptr(), + out.data_ptr(), + n_rows, + S, + half, + static_cast(sin_sign)); + }); + return out; +} diff --git a/rl_engine/kernels/ops/musa/.gitkeep b/rl_engine/kernels/ops/musa/.gitkeep new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/rl_engine/kernels/ops/musa/.gitkeep @@ -0,0 +1 @@ + diff --git a/rl_engine/kernels/ops/musa/__init__.py b/rl_engine/kernels/ops/musa/__init__.py new file mode 100644 index 00000000..2bbc36bb --- /dev/null +++ b/rl_engine/kernels/ops/musa/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .loss import MusaFusedLogpOp + +__all__ = ["MusaFusedLogpOp"] diff --git a/rl_engine/kernels/ops/musa/loss/__init__.py b/rl_engine/kernels/ops/musa/loss/__init__.py new file mode 100644 index 00000000..5af48e21 --- /dev/null +++ b/rl_engine/kernels/ops/musa/loss/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .logp import MusaFusedLogpOp + +__all__ = ["MusaFusedLogpOp"] diff --git a/rl_engine/kernels/ops/musa/loss/logp.py b/rl_engine/kernels/ops/musa/loss/logp.py new file mode 100644 index 00000000..3e96e486 --- /dev/null +++ b/rl_engine/kernels/ops/musa/loss/logp.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MUSA native fused selected-token log-probability backend.""" + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + +class _MusaFusedLogpFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + token_ids_1d = token_ids.reshape(-1).to(device=logits.device, dtype=torch.long).contiguous() + output = _C.fused_logp(logits_2d, token_ids_1d) + ctx.save_for_backward(logits_2d, token_ids_1d) + ctx.input_shape = tuple(logits.shape) + ctx.input_dtype = logits.dtype + return output.reshape(logits.shape[:-1]) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + logits, token_ids = ctx.saved_tensors + probs = torch.softmax(logits.float(), dim=-1) + rows = torch.arange(logits.size(0), device=logits.device) + probs[rows, token_ids] -= 1.0 + grad = -grad_output.reshape(-1, 1).float() * probs + return grad.to(ctx.input_dtype).reshape(ctx.input_shape), None + + +class MusaFusedLogpOp: + """Generic MUSA fused LogP; backward uses a portable PyTorch formula.""" + + is_fused_logp = True + + def __init__(self): + if not _EXT_AVAILABLE or not hasattr(_C, "fused_logp"): + raise RuntimeError("MUSA fused_logp extension is unavailable") + + def __call__(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + if logits.device.type != "musa": + raise RuntimeError(f"MusaFusedLogpOp requires a MUSA tensor, got {logits.device}") + return _MusaFusedLogpFunction.apply(logits, token_ids) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index efde5c25..c3810eab 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -42,6 +42,9 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ROCM_CK = "rl_engine.kernels.ops.rocm.composable_kernel.CKOp" ROCM_FLASH_ATTN = "rl_engine.kernels.ops.rocm.attention.flash_attn.RocmFlashAttentionOp" + # Moore Threads MUSA native generic LogP kernel. + MUSA_FUSED_LOGP = "rl_engine.kernels.ops.musa.loss.logp.MusaFusedLogpOp" + # GRPO loss (group reward normalization + clipped surrogate + KL) TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" PYTORCH_GRPO_LOSS = "rl_engine.kernels.ops.pytorch.loss.grpo_loss.NativeGRPOLossOp" @@ -161,7 +164,7 @@ def resolve_logp_op_type( class KernelRegistry: """ Central dispatcher for high-performance kernels. - Handles dynamic routing between ROCm and CUDA backends at runtime. + Handles dynamic routing between CUDA, ROCm, and MUSA backends at runtime. """ def __init__(self): @@ -196,7 +199,11 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], - "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], + "attn": [ + OpBackend.FLASH_ATTN, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_ATTN, + ], "attention": [ OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, @@ -252,7 +259,10 @@ def __init__(self): "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "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], + "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": [OpBackend.TRITON_DET_GEMM], @@ -267,6 +277,30 @@ def __init__(self): "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], }, + "musa": { + "logp": [OpBackend.MUSA_FUSED_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], + "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], + "pack": [OpBackend.PYTORCH_PACK], + "det_gemm": [OpBackend.PYTORCH_GEMM], + "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP], + "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], + "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], + }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], "logp_deterministic": [OpBackend.PYTORCH_NATIVE], @@ -405,6 +439,8 @@ def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: if device_ctx.is_rocm: return "rocm" + if device_ctx.is_musa: + return "musa" if device_ctx.device_type == "cuda": return "cuda" return "cpu" @@ -412,6 +448,8 @@ def _platform_for_device(self, device: torch.device | str | None) -> str: resolved = torch.device(device) if resolved.type == "cuda": return "rocm" if torch.version.hip is not None else "cuda" + if resolved.type == "musa": + return "musa" if resolved.type in self._priority_map: return resolved.type return "cpu" diff --git a/rl_engine/platforms/constants.py b/rl_engine/platforms/constants.py index fbf9b692..6b81d4cd 100644 --- a/rl_engine/platforms/constants.py +++ b/rl_engine/platforms/constants.py @@ -7,6 +7,7 @@ class DeviceType(Enum): CUDA = "cuda" # NVIDIA GPU ROCM = "rocm" # AMD GPU + MUSA = "musa" # Moore Threads GPU NPU = "npu" # Huawei Ascend TPU = "tpu" # Google TPU XPU = "xpu" # Intel GPU @@ -17,6 +18,7 @@ class BackendLib(Enum): FLASHINFER = "flashinfer" AITER = "aiter" TRITON = "triton" + MUSA = "musa" # Moore Threads MUSA stack CANN = "cann" # Huawei NPU NATIVE = "native" # PyTorch Default diff --git a/rl_engine/platforms/device.py b/rl_engine/platforms/device.py index 5a494d71..a55d8b80 100644 --- a/rl_engine/platforms/device.py +++ b/rl_engine/platforms/device.py @@ -7,23 +7,44 @@ from rl_engine.utils.logger import logger +def is_musa_available() -> bool: + """Return whether torch_musa is installed and a MUSA device is usable.""" + try: + import torch_musa # noqa: F401 + except ImportError: + return False + return bool(hasattr(torch, "musa") and torch.musa.is_available()) + + class DeviceContext: """ Hardware-aware context manager for high-performance RL tasks. - Provides transparent support for both AMD (ROCm/HIP) and NVIDIA (CUDA) - architectures to ensure backend-agnostic scaling for RL operators. + Provides transparent support for AMD (ROCm/HIP), NVIDIA (CUDA), and + Moore Threads (MUSA) architectures. """ def __init__(self): - self.device = torch.device( - DeviceType.CUDA.value if torch.cuda.is_available() else DeviceType.CPU.value - ) + if is_musa_available(): + self.device = torch.device(DeviceType.MUSA.value) + elif torch.cuda.is_available(): + self.device = torch.device(DeviceType.CUDA.value) + else: + self.device = torch.device(DeviceType.CPU.value) self.is_rocm = False + self.is_musa = False self.backend_version = "N/A" self.device_type = DeviceType.CPU.value - if self.device.type == DeviceType.CUDA.value: + if self.device.type == DeviceType.MUSA.value: + self.is_musa = True + self.device_type = DeviceType.MUSA.value + self.backend_version = str(getattr(torch.version, "musa", "N/A")) + logger.info_once( + "RL-Engine initialized with Moore Threads MUSA backend" + f" (Version: {self.backend_version})" + ) + elif self.device.type == DeviceType.CUDA.value: # Distinct detection for AMD HIP and NVIDIA CUDA if hasattr(torch.version, "hip") and torch.version.hip is not None: self.is_rocm = True @@ -47,9 +68,9 @@ def __init__(self): def get_preferred_dtype(self): """ Returns the optimal data type for the current hardware. - AMD ROCm typically yields better performance with bfloat16 in RL workloads. + AMD ROCm and Moore Threads MUSA typically use bfloat16 for RL workloads. """ - return torch.bfloat16 if self.is_rocm else torch.float16 + return torch.bfloat16 if self.is_rocm or self.is_musa else torch.float16 device_ctx = DeviceContext() diff --git a/setup.py b/setup.py index 6a35bcb3..14409d87 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,25 @@ def _load_envs_module(): envs = _load_envs_module() +def _musa_build_available(torch) -> bool: + """Return whether the installed PyTorch stack can build MUSA extensions.""" + try: + import torch_musa # noqa: F401 + except ImportError: + return False + + try: + return bool( + hasattr(torch, "musa") + and ( + torch.musa.is_available() + or bool(os.environ.get("TORCH_MUSA_ARCH_LIST", "").strip()) + ) + ) + except Exception: + return False + + def _load_torch_extension_tools(): try: import torch @@ -30,6 +49,11 @@ def _load_torch_extension_tools(): raise return None, None, None + if _musa_build_available(torch): + from torch_musa.utils.musa_extension import BuildExtension, MUSAExtension + + return torch, BuildExtension, MUSAExtension + from torch.utils.cpp_extension import BuildExtension, CUDAExtension # CUDAExtension is also the supported extension entry point for ROCm @@ -44,7 +68,9 @@ def _native_extension_required() -> bool: envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or bool(os.environ.get("TORCH_MUSA_ARCH_LIST", "").strip()) or envs.env_flag("FORCE_CUDA") + or envs.env_flag("FORCE_MUSA") ) @@ -93,7 +119,7 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() + torch, _, Extension = _load_torch_extension_tools() if torch is None: message = ( "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " @@ -111,12 +137,37 @@ def get_extensions(): return [] extensions = [] + is_musa = _musa_build_available(torch) torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") is_rocm = getattr(torch.version, "hip", None) is not None + if is_musa: + extensions.append( + Extension( + name="rl_engine._C", + sources=[ + "csrc/musa/ops.cpp", + "csrc/musa/fused_logp_kernel.mu", + "csrc/musa/deterministic_logp_kernel.mu", + "csrc/musa/activation.mu", + "csrc/musa/rmsnorm.mu", + "csrc/musa/deterministic_attention.mu", + "csrc/musa/rope.mu", + "csrc/musa/embedding_lm_head.mu", + ], + include_dirs=[], + extra_compile_args={ + "cxx": ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_MUSA"], + "mcc": ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_MUSA"], + }, + extra_link_args=list(torch_rpath), + ) + ) + return extensions + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add @@ -249,7 +300,7 @@ def get_extensions(): nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) extensions.append( - CUDAExtension( + Extension( name="rl_engine._C", sources=cuda_sources, include_dirs=[], @@ -263,9 +314,9 @@ def get_extensions(): if _native_extension_required() and not extensions: raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." + "rl_engine._C was requested but no CUDA/ROCm/MUSA build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm or MUSA build, " + "set the corresponding architecture variable." ) return extensions diff --git a/tests/test_musa_extension.py b/tests/test_musa_extension.py new file mode 100644 index 00000000..b7caa839 --- /dev/null +++ b/tests/test_musa_extension.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest +import torch + + +try: + import torch_musa # noqa: F401 +except ImportError: + torch_musa = None + + +requires_musa_extension = pytest.mark.skipif( + torch_musa is None or not hasattr(torch, "musa") or not torch.musa.is_available(), + reason="requires a MUSA runtime", +) + + +@requires_musa_extension +def test_musa_fused_logp_extension_matches_reference(): + from rl_engine import _C + + if not hasattr(_C, "fused_logp"): + pytest.fail("MUSA extension is present but does not export fused_logp") + + torch.manual_seed(7) + logits = torch.randn(5, 37, device="musa", dtype=torch.bfloat16) + token_ids = torch.tensor([0, 3, 11, 29, 36], device="musa", dtype=torch.long) + + actual = _C.fused_logp(logits, token_ids) + reference = torch.log_softmax(logits.float(), dim=-1).gather(1, token_ids[:, None]).squeeze(1) + torch.testing.assert_close(actual.float(), reference, rtol=2e-2, atol=2e-2) From 12376c9c6c0a9ec652bf4dc604fa2c1fae43e19a Mon Sep 17 00:00:00 2001 From: mt Date: Tue, 1 Sep 2026 10:32:39 +0800 Subject: [PATCH 2/6] test: adapt MUSA dispatch checks --- tests/test_batch_invariant_logp.py | 3 ++- tests/test_grpo_loss.py | 3 ++- tests/test_logp.py | 12 +++++++++++- tests/test_op_accuracy.py | 3 +-- tests/test_ratio_kl.py | 3 ++- 5 files changed, 18 insertions(+), 6 deletions(-) diff --git a/tests/test_batch_invariant_logp.py b/tests/test_batch_invariant_logp.py index 8ae06b90..c3f351e8 100644 --- a/tests/test_batch_invariant_logp.py +++ b/tests/test_batch_invariant_logp.py @@ -14,6 +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 device_ctx _V = 300 @@ -1027,7 +1028,7 @@ def test_registry_dispatches_correctly(): or type(op).__name__ == "TritonBatchInvariantLogpOp" or type(op).__name__ == "BatchInvariantLogpSM90Op" ) - 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..9cfae257 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() and not device_ctx.is_musa: assert isinstance(op, TritonGRPOLossOp) else: assert isinstance(op, NativeGRPOLossOp) diff --git a/tests/test_logp.py b/tests/test_logp.py index bcb0e3f9..9792d45b 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,13 @@ 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.base import _EXT_AVAILABLE + from rl_engine.kernels.ops.musa.loss.logp import MusaFusedLogpOp + + if _EXT_AVAILABLE: + assert isinstance(op, MusaFusedLogpOp) + else: + assert isinstance(op, NativeLogpOp) + 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..b529957c 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() and not device_ctx.is_musa: assert isinstance(op, TritonRatioKLOp) else: assert isinstance(op, NativeRatioKLOp) From 48beca4efad6268c9e8bf383ae639036890255d5 Mon Sep 17 00:00:00 2001 From: mt Date: Tue, 1 Sep 2026 11:11:27 +0800 Subject: [PATCH 3/6] feat: dispatch MUSA native kernels --- csrc/musa/ops.cpp | 2 + rl_engine/kernels/ops/musa/__init__.py | 24 +- rl_engine/kernels/ops/musa/native.py | 365 +++++++++++++++++++++++++ rl_engine/kernels/registry.py | 39 ++- tests/test_attention.py | 6 + tests/test_batch_invariant_logp.py | 1 + tests/test_rms_norm.py | 6 +- 7 files changed, 430 insertions(+), 13 deletions(-) create mode 100644 rl_engine/kernels/ops/musa/native.py diff --git a/csrc/musa/ops.cpp b/csrc/musa/ops.cpp index 51b05970..95c3eeb4 100644 --- a/csrc/musa/ops.cpp +++ b/csrc/musa/ops.cpp @@ -29,6 +29,7 @@ torch::Tensor lm_head_sm90_forward_fp32( torch::Tensor hidden, torch::Tensor weight, torch::optional bias); +torch::Tensor det_gemm_rowwise_fwd_fp32(torch::Tensor a, torch::Tensor b); torch::Tensor rope_apply_sm90( torch::Tensor x, torch::Tensor cos, @@ -101,6 +102,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("embedding_fp32", &embedding_sm90_forward_fp32); m.def("lm_head_fp32", &lm_head_sm90_forward_fp32, py::arg("hidden"), py::arg("weight"), py::arg("bias") = py::none()); + m.def("det_gemm_rowwise_fwd_fp32", &det_gemm_rowwise_fwd_fp32); m.def("rope_apply", &rope_apply_sm90); m.def("deterministic_attention_fp32", &deterministic_attention_forward_fp32, py::arg("q"), py::arg("k"), py::arg("v"), py::arg("causal"), diff --git a/rl_engine/kernels/ops/musa/__init__.py b/rl_engine/kernels/ops/musa/__init__.py index 2bbc36bb..06da8712 100644 --- a/rl_engine/kernels/ops/musa/__init__.py +++ b/rl_engine/kernels/ops/musa/__init__.py @@ -2,5 +2,27 @@ # Copyright (c) 2026 RL-Kernel Contributors from .loss import MusaFusedLogpOp +from .native import ( + MusaDetGemmOp, + MusaDeterministicAttentionOp, + MusaDeterministicLogpOp, + MusaEmbeddingOp, + MusaLMHeadOp, + MusaRMSNormOp, + MusaRoPEOp, + MusaSiLUOp, + MusaSwiGLUOp, +) -__all__ = ["MusaFusedLogpOp"] +__all__ = [ + "MusaFusedLogpOp", + "MusaDetGemmOp", + "MusaDeterministicAttentionOp", + "MusaDeterministicLogpOp", + "MusaEmbeddingOp", + "MusaLMHeadOp", + "MusaRMSNormOp", + "MusaRoPEOp", + "MusaSiLUOp", + "MusaSwiGLUOp", +] diff --git a/rl_engine/kernels/ops/musa/native.py b/rl_engine/kernels/ops/musa/native.py new file mode 100644 index 00000000..9b686455 --- /dev/null +++ b/rl_engine/kernels/ops/musa/native.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Wrappers for the native MUSA kernels exposed by ``rl_engine._C``.""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + +def _require(symbol: str) -> None: + if not _EXT_AVAILABLE or _C is None or not hasattr(_C, symbol): + raise RuntimeError(f"MUSA native symbol {symbol!r} is unavailable") + + +def _check_musa(*tensors: torch.Tensor) -> None: + if any(t.device.type != "musa" for t in tensors): + raise RuntimeError("MUSA native kernels require MUSA tensors") + + +class _MusaLogpFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, logits: torch.Tensor, token_ids: torch.Tensor): + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + ids = token_ids.reshape(-1).to(device=logits.device, dtype=torch.long).contiguous() + output = _C.deterministic_logp_fp32(logits_2d, ids) + ctx.save_for_backward(logits_2d, ids) + ctx.input_shape = tuple(logits.shape) + ctx.input_dtype = logits.dtype + return output.reshape(logits.shape[:-1]) + + @staticmethod + def backward(ctx, grad_output): + logits, ids = ctx.saved_tensors + probs = torch.softmax(logits.float(), dim=-1) + rows = torch.arange(logits.size(0), device=logits.device) + probs[rows, ids] -= 1.0 + grad = -grad_output.reshape(-1, 1).float() * probs + return grad.to(ctx.input_dtype).reshape(ctx.input_shape), None + + +class MusaDeterministicLogpOp: + is_batch_invariant = True + + def __init__(self): + _require("deterministic_logp_fp32") + + def __call__(self, logits, token_ids): + return self.apply_fp32(logits, token_ids) + + def apply(self, logits, token_ids): + return self.apply_fp32(logits, token_ids) + + def apply_fp32(self, logits, token_ids): + _check_musa(logits, token_ids) + return _MusaLogpFunction.apply(logits, token_ids) + + forward = apply_fp32 + forward_fp32 = apply_fp32 + online_fp32 = apply_fp32 + + def indexed_fp32(self, logits, token_ids, row_indices): + output = torch.zeros(logits.shape[:-1], device=logits.device, dtype=torch.float32) + return self.indexed_out(logits, token_ids, row_indices, output) + + def indexed_out(self, logits, token_ids, row_indices, output): + full = self.apply_fp32(logits, token_ids) + output.reshape(-1).index_copy_( + 0, + row_indices.reshape(-1).to(device=output.device, dtype=torch.long), + full.reshape(-1).index_select( + 0, row_indices.reshape(-1).to(device=full.device, dtype=torch.long) + ).to(output.dtype), + ) + return output + + out = indexed_out + online_indexed_fp32 = indexed_fp32 + online_indexed_out = indexed_out + + +class _MusaSiLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x) + return _C.silu_forward(x.contiguous()) + + @staticmethod + def backward(ctx, grad_output): + (x,) = ctx.saved_tensors + return _C.silu_backward(grad_output.contiguous(), x) + + +class MusaSiLUOp: + def __init__(self): + _require("silu_forward") + + def __call__(self, x): + return self.forward(x) + + def forward(self, x): + _check_musa(x) + return _MusaSiLUFunction.apply(x) + + def forward_fp32(self, x): + _check_musa(x) + return _MusaSiLUFunction.apply(x.float()) + + +class _MusaSwiGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, gate, up): + ctx.save_for_backward(gate, up) + return _C.swiglu_forward(gate.contiguous(), up.contiguous()) + + @staticmethod + def backward(ctx, grad_output): + gate, up = ctx.saved_tensors + d_gate, d_up = _C.swiglu_backward(grad_output.contiguous(), gate, up) + return d_gate, d_up + + +class MusaSwiGLUOp: + def __init__(self): + _require("swiglu_forward") + + def __call__(self, gate, up): + return self.forward(gate, up) + + def forward(self, gate, up): + _check_musa(gate, up) + return _MusaSwiGLUFunction.apply(gate, up) + + def forward_fp32(self, gate, up): + _check_musa(gate, up) + return _MusaSwiGLUFunction.apply(gate.float(), up.float()) + + +class _MusaRMSNormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, eps): + x_2d = x.reshape(-1, x.size(-1)).contiguous() + weight = weight.contiguous() + y = torch.empty_like(x_2d) + rstd = torch.empty(x_2d.size(0), device=x.device, dtype=torch.float32) + _C.rmsnorm_forward(x_2d, weight, y, rstd, float(eps)) + ctx.save_for_backward(x_2d, weight, rstd) + ctx.input_shape = tuple(x.shape) + ctx.eps = eps + return y.reshape_as(x) + + @staticmethod + def backward(ctx, grad_output): + x, weight, rstd = ctx.saved_tensors + dy = grad_output.reshape_as(x).reshape(-1, x.size(-1)).contiguous() + dx = torch.empty_like(x) + _C.rmsnorm_backward_dx(dy, x, weight, rstd, dx) + dw = (dy.float() * x.float() * rstd[:, None]).sum(dim=0).to(weight.dtype) + return dx.reshape(ctx.input_shape), dw, None + + +class MusaRMSNormOp: + def __init__(self): + _require("rmsnorm_forward") + _require("rmsnorm_backward_dx") + + def __call__(self, x, weight, *, eps=1e-6): + return self.forward(x, weight, eps=eps) + + def forward(self, x, weight, *, eps=1e-6): + _check_musa(x, weight) + return _MusaRMSNormFunction.apply(x, weight, eps) + + def forward_fp32(self, x, weight, *, eps=1e-6): + _check_musa(x, weight) + return _MusaRMSNormFunction.apply(x.float(), weight.float(), eps) + + +class _MusaEmbeddingFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, token_ids, weight): + ids = token_ids.reshape(-1).to(dtype=torch.long).contiguous() + output = _C.embedding_fp32(ids, weight.contiguous()).to(weight.dtype) + ctx.save_for_backward(ids) + ctx.weight_shape = tuple(weight.shape) + ctx.weight_dtype = weight.dtype + ctx.output_shape = tuple(token_ids.shape) + (weight.size(1),) + return output.reshape(ctx.output_shape) + + @staticmethod + def backward(ctx, grad_output): + (ids,) = ctx.saved_tensors + grad_weight = torch.zeros( + ctx.weight_shape, device=grad_output.device, dtype=ctx.weight_dtype + ) + grad_weight.index_add_(0, ids, grad_output.reshape(-1, ctx.weight_shape[1])) + return None, grad_weight + + +class MusaEmbeddingOp: + def __init__(self): + _require("embedding_fp32") + + def __call__(self, token_ids, weight): + return self.forward(token_ids, weight) + + def forward(self, token_ids, weight): + _check_musa(token_ids, weight) + return _MusaEmbeddingFunction.apply(token_ids, weight) + + def forward_fp32(self, token_ids, weight): + _check_musa(token_ids, weight) + return _C.embedding_fp32(token_ids.reshape(-1).long(), weight).reshape( + *token_ids.shape, weight.size(1) + ) + + +class _MusaLMHeadFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, hidden, weight, bias, output_dtype): + output = _C.lm_head_fp32(hidden, weight, bias) + ctx.save_for_backward(hidden, weight, bias if bias is not None else torch.empty(0)) + ctx.has_bias = bias is not None + ctx.output_dtype = output_dtype + return output.to(output_dtype) + + @staticmethod + def backward(ctx, grad_output): + hidden, weight, bias = ctx.saved_tensors + grad = grad_output.float().reshape(-1, grad_output.size(-1)) + hidden_2d = hidden.float().reshape(-1, hidden.size(-1)) + grad_hidden = (grad @ weight.float()).reshape_as(hidden) + grad_weight = grad.t() @ hidden_2d + grad_bias = grad.sum(dim=0) if ctx.has_bias else None + return grad_hidden.to(hidden.dtype), grad_weight.to(weight.dtype), grad_bias, None + + +class MusaLMHeadOp: + def __init__(self): + _require("lm_head_fp32") + + def __call__(self, hidden, weight, *, bias=None): + return self.forward(hidden, weight, bias=bias) + + def forward(self, hidden, weight, *, bias=None): + _check_musa(hidden, weight) + if bias is not None: + _check_musa(bias) + return _MusaLMHeadFunction.apply(hidden, weight, bias, hidden.dtype) + + def forward_fp32(self, hidden, weight, *, bias=None): + _check_musa(hidden, weight) + return _C.lm_head_fp32(hidden, weight, bias) + + +def _cos_sin(positions, half, theta, device): + inv = 1.0 / (theta ** (torch.arange(half, device=device, dtype=torch.float32) / half)) + freqs = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) * inv + return freqs.cos().contiguous(), freqs.sin().contiguous() + + +class _MusaRoPEFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, positions, theta): + if positions.dim() != 1: + raise NotImplementedError("MUSA native RoPE currently supports positions [S]") + half = x.size(-1) // 2 + cos, sin = _cos_sin(positions, half, theta, x.device) + x_2d = x.contiguous().reshape(-1, x.size(-1)) + output = _C.rope_apply(x_2d, cos, sin, 1.0).reshape_as(x) + ctx.save_for_backward(cos, sin) + ctx.seq_len = positions.numel() + return output + + @staticmethod + def backward(ctx, grad_output): + cos, sin = ctx.saved_tensors + grad = grad_output.contiguous().reshape(-1, grad_output.size(-1)) + return _C.rope_apply(grad, cos, sin, -1.0).reshape_as(grad_output), None, None + + +class MusaRoPEOp: + def __init__(self): + _require("rope_apply") + + def __call__(self, x, positions, *, theta=1_000_000.0): + return self.forward(x, positions, theta=theta) + + def forward(self, x, positions, *, theta=1_000_000.0): + _check_musa(x) + return _MusaRoPEFunction.apply(x, positions, float(theta)) + + def forward_fp32(self, x, positions, *, theta=1_000_000.0): + _check_musa(x) + return _MusaRoPEFunction.apply(x.float(), positions, float(theta)) + + +class _MusaAttentionFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, causal, scale, mask): + mask_arg = mask if mask.numel() else None + out, lse, probs = _C.deterministic_attention_fp32( + q.contiguous(), k.contiguous(), v.contiguous(), causal, float(scale), mask_arg + ) + ctx.save_for_backward(q, k, v, probs, mask) + ctx.has_mask = mask.numel() != 0 + ctx.causal = causal + ctx.scale = scale + return out.to(q.dtype), lse + + @staticmethod + def backward(ctx, grad_out, grad_lse): + q, k, v, probs, mask = ctx.saved_tensors + mask_arg = mask if ctx.has_mask else None + dq, dk, dv = _C.deterministic_attention_backward( + grad_out.contiguous(), q, k, v, probs, ctx.causal, float(ctx.scale), mask_arg + ) + return dq, dk, dv, None, None, None + + +class MusaDeterministicAttentionOp: + def __init__(self): + _require("deterministic_attention_fp32") + _require("deterministic_attention_backward") + + def __call__(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): + return self.forward(q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask) + + def forward(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): + _check_musa(q, k, v) + resolved_scale = scale if scale is not None else 1.0 / math.sqrt(q.size(-1)) + mask = ( + key_padding_mask.contiguous() + if key_padding_mask is not None + else torch.empty(0, device=q.device, dtype=torch.bool) + ) + out, _lse = _MusaAttentionFunction.apply(q, k, v, causal, resolved_scale, mask) + return out + + def forward_with_lse(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): + _check_musa(q, k, v) + resolved_scale = scale if scale is not None else 1.0 / math.sqrt(q.size(-1)) + mask = ( + key_padding_mask.contiguous() + if key_padding_mask is not None + else torch.empty(0, device=q.device, dtype=torch.bool) + ) + return _MusaAttentionFunction.apply(q, k, v, causal, resolved_scale, mask) + + +class MusaDetGemmOp: + def __init__(self): + _require("det_gemm_rowwise_fwd_fp32") + + def __call__(self, a, b): + _check_musa(a, b) + return _C.det_gemm_rowwise_fwd_fp32(a, b) + + forward = __call__ + forward_fp32 = __call__ diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index c3810eab..d7a275d5 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -44,6 +44,17 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # Moore Threads MUSA native generic LogP kernel. MUSA_FUSED_LOGP = "rl_engine.kernels.ops.musa.loss.logp.MusaFusedLogpOp" + MUSA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.musa.native.MusaDeterministicLogpOp" + MUSA_DETERMINISTIC_ATTENTION = ( + "rl_engine.kernels.ops.musa.native.MusaDeterministicAttentionOp" + ) + MUSA_DET_GEMM = "rl_engine.kernels.ops.musa.native.MusaDetGemmOp" + MUSA_RMS_NORM = "rl_engine.kernels.ops.musa.native.MusaRMSNormOp" + MUSA_EMBEDDING = "rl_engine.kernels.ops.musa.native.MusaEmbeddingOp" + MUSA_LM_HEAD = "rl_engine.kernels.ops.musa.native.MusaLMHeadOp" + MUSA_ROPE = "rl_engine.kernels.ops.musa.native.MusaRoPEOp" + MUSA_SILU = "rl_engine.kernels.ops.musa.native.MusaSiLUOp" + MUSA_SWIGLU = "rl_engine.kernels.ops.musa.native.MusaSwiGLUOp" # GRPO loss (group reward normalization + clipped surrogate + KL) TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" @@ -282,24 +293,30 @@ def __init__(self): "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], + "logp_deterministic": [OpBackend.MUSA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE], + "logp_deterministic_indexed": [ + OpBackend.MUSA_DETERMINISTIC_LOGP, + OpBackend.PYTORCH_NATIVE, + ], "attn": [OpBackend.PYTORCH_ATTN], - "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "attention": [OpBackend.MUSA_DETERMINISTIC_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], + "rope": [OpBackend.MUSA_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], - "det_gemm": [OpBackend.PYTORCH_GEMM], - "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP], + "det_gemm": [OpBackend.MUSA_DET_GEMM, OpBackend.PYTORCH_GEMM], + "batch_invariant_logp": [ + OpBackend.MUSA_DETERMINISTIC_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], - "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], + "rms_norm": [OpBackend.MUSA_RMS_NORM, OpBackend.PYTORCH_NATIVE_RMS_NORM], + "lm_head": [OpBackend.MUSA_LM_HEAD, OpBackend.PYTORCH_NATIVE_LM_HEAD], + "embedding": [OpBackend.MUSA_EMBEDDING, OpBackend.PYTORCH_NATIVE_EMBEDDING], + "silu": [OpBackend.MUSA_SILU, OpBackend.PYTORCH_NATIVE_SILU], + "swiglu": [OpBackend.MUSA_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], diff --git a/tests/test_attention.py b/tests/test_attention.py index 469c6d30..2eb4735e 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -29,6 +29,7 @@ from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp from rl_engine.kernels.registry import kernel_registry +from rl_engine.platforms.device import device_ctx # Qwen3-8B attention dims (synthetic tensors, no checkpoint). Unlike embedding / # lm_head (whose multi-GB weight forces shrinking), attention's cost is the @@ -436,6 +437,11 @@ def test_gradient_matches_reference(): def test_registry_dispatches_native_attention_op(): """Resolve attention to the deterministic CUDA op or native fallback.""" op = kernel_registry.get_op("attention") + if device_ctx.is_musa: + from rl_engine.kernels.ops.musa.native import MusaDeterministicAttentionOp + + assert isinstance(op, MusaDeterministicAttentionOp) + return # On CUDA with the extension built, the registry prefers DeterministicAttentionOp. # On CPU or without the CUDA extension, it falls back to NativeAttentionOp. from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp diff --git a/tests/test_batch_invariant_logp.py b/tests/test_batch_invariant_logp.py index c3f351e8..bb5c05ab 100644 --- a/tests/test_batch_invariant_logp.py +++ b/tests/test_batch_invariant_logp.py @@ -1027,6 +1027,7 @@ def test_registry_dispatches_correctly(): isinstance(op, NativeBatchInvariantLogpOp) or type(op).__name__ == "TritonBatchInvariantLogpOp" or type(op).__name__ == "BatchInvariantLogpSM90Op" + or type(op).__name__ == "MusaDeterministicLogpOp" ) logits = torch.randn(4, _V, device=device_ctx.device) target = torch.randint(0, _V, (4,), device=logits.device) diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 43a9cf85..d6401472 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -6,6 +6,7 @@ import torch.nn.functional as F from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda +from rl_engine.kernels.ops.musa.native import MusaRMSNormOp from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton @@ -236,7 +237,10 @@ def test_registry_dispatches_rms_norm(): from rl_engine.kernels.registry import kernel_registry op = kernel_registry.get_op("rms_norm") - assert isinstance(op, NativeRMSNormOp) + if torch.musa.is_available(): + assert isinstance(op, MusaRMSNormOp) + else: + assert isinstance(op, NativeRMSNormOp) assert hasattr(op, "forward") and hasattr(op, "forward_fp32") From 67f126c0da8de6ebe3d95a79fea32bb8d963cca3 Mon Sep 17 00:00:00 2001 From: mt Date: Tue, 1 Sep 2026 11:29:15 +0800 Subject: [PATCH 4/6] refactor: align MUSA ops layout with CUDA --- rl_engine/kernels/ops/musa/__init__.py | 26 +- rl_engine/kernels/ops/musa/_common.py | 18 + .../kernels/ops/musa/activation/__init__.py | 6 + .../kernels/ops/musa/activation/swiglu.py | 66 ++++ .../kernels/ops/musa/attention/__init__.py | 6 + .../ops/musa/attention/deterministic_attn.py | 80 ++++ rl_engine/kernels/ops/musa/linear/__init__.py | 7 + .../kernels/ops/musa/linear/embedding.py | 47 +++ rl_engine/kernels/ops/musa/linear/lm_head.py | 49 +++ rl_engine/kernels/ops/musa/loss/__init__.py | 3 +- .../ops/musa/loss/batch_invariant_logp.py | 71 ++++ rl_engine/kernels/ops/musa/matmul/__init__.py | 6 + rl_engine/kernels/ops/musa/matmul/det_gemm.py | 19 + rl_engine/kernels/ops/musa/native.py | 365 ------------------ rl_engine/kernels/ops/musa/norm/__init__.py | 6 + rl_engine/kernels/ops/musa/norm/rmsnorm.py | 48 +++ .../ops/musa/rotary_embedding/__init__.py | 6 + .../kernels/ops/musa/rotary_embedding/rope.py | 50 +++ rl_engine/kernels/registry.py | 20 +- tests/test_attention.py | 4 +- tests/test_rms_norm.py | 2 +- 21 files changed, 504 insertions(+), 401 deletions(-) create mode 100644 rl_engine/kernels/ops/musa/_common.py create mode 100644 rl_engine/kernels/ops/musa/activation/__init__.py create mode 100644 rl_engine/kernels/ops/musa/activation/swiglu.py create mode 100644 rl_engine/kernels/ops/musa/attention/__init__.py create mode 100644 rl_engine/kernels/ops/musa/attention/deterministic_attn.py create mode 100644 rl_engine/kernels/ops/musa/linear/__init__.py create mode 100644 rl_engine/kernels/ops/musa/linear/embedding.py create mode 100644 rl_engine/kernels/ops/musa/linear/lm_head.py create mode 100644 rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py create mode 100644 rl_engine/kernels/ops/musa/matmul/__init__.py create mode 100644 rl_engine/kernels/ops/musa/matmul/det_gemm.py delete mode 100644 rl_engine/kernels/ops/musa/native.py create mode 100644 rl_engine/kernels/ops/musa/norm/__init__.py create mode 100644 rl_engine/kernels/ops/musa/norm/rmsnorm.py create mode 100644 rl_engine/kernels/ops/musa/rotary_embedding/__init__.py create mode 100644 rl_engine/kernels/ops/musa/rotary_embedding/rope.py diff --git a/rl_engine/kernels/ops/musa/__init__.py b/rl_engine/kernels/ops/musa/__init__.py index 06da8712..dced5845 100644 --- a/rl_engine/kernels/ops/musa/__init__.py +++ b/rl_engine/kernels/ops/musa/__init__.py @@ -1,28 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .loss import MusaFusedLogpOp -from .native import ( - MusaDetGemmOp, - MusaDeterministicAttentionOp, - MusaDeterministicLogpOp, - MusaEmbeddingOp, - MusaLMHeadOp, - MusaRMSNormOp, - MusaRoPEOp, - MusaSiLUOp, - MusaSwiGLUOp, -) +from . import activation, attention, linear, loss, matmul, norm, rotary_embedding -__all__ = [ - "MusaFusedLogpOp", - "MusaDetGemmOp", - "MusaDeterministicAttentionOp", - "MusaDeterministicLogpOp", - "MusaEmbeddingOp", - "MusaLMHeadOp", - "MusaRMSNormOp", - "MusaRoPEOp", - "MusaSiLUOp", - "MusaSwiGLUOp", -] +__all__ = ["activation", "attention", "linear", "loss", "matmul", "norm", "rotary_embedding"] diff --git a/rl_engine/kernels/ops/musa/_common.py b/rl_engine/kernels/ops/musa/_common.py new file mode 100644 index 00000000..6741e95f --- /dev/null +++ b/rl_engine/kernels/ops/musa/_common.py @@ -0,0 +1,18 @@ +# 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 + + +def require_musa_symbol(symbol: str) -> None: + if not _EXT_AVAILABLE or _C is None or not hasattr(_C, symbol): + raise RuntimeError(f"MUSA native symbol {symbol!r} is unavailable") + + +def check_musa(*tensors: torch.Tensor) -> None: + if any(t.device.type != "musa" for t in tensors): + raise RuntimeError("MUSA native kernels require MUSA tensors") diff --git a/rl_engine/kernels/ops/musa/activation/__init__.py b/rl_engine/kernels/ops/musa/activation/__init__.py new file mode 100644 index 00000000..2ae13f49 --- /dev/null +++ b/rl_engine/kernels/ops/musa/activation/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .swiglu import MusaSiLUOp, MusaSwiGLUOp + +__all__ = ["MusaSiLUOp", "MusaSwiGLUOp"] diff --git a/rl_engine/kernels/ops/musa/activation/swiglu.py b/rl_engine/kernels/ops/musa/activation/swiglu.py new file mode 100644 index 00000000..5982c77e --- /dev/null +++ b/rl_engine/kernels/ops/musa/activation/swiglu.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol +from rl_engine.kernels.ops.base import _C + + +class _MusaSiLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x) + return _C.silu_forward(x.contiguous()) + + @staticmethod + def backward(ctx, grad_output): + (x,) = ctx.saved_tensors + return _C.silu_backward(grad_output.contiguous(), x) + + +class MusaSiLUOp: + def __init__(self): + require_musa_symbol("silu_forward") + + def __call__(self, x): + return self.forward(x) + + def forward(self, x): + check_musa(x) + return _MusaSiLUFunction.apply(x) + + def forward_fp32(self, x): + check_musa(x) + return _MusaSiLUFunction.apply(x.float()) + + +class _MusaSwiGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, gate, up): + ctx.save_for_backward(gate, up) + return _C.swiglu_forward(gate.contiguous(), up.contiguous()) + + @staticmethod + def backward(ctx, grad_output): + gate, up = ctx.saved_tensors + d_gate, d_up = _C.swiglu_backward(grad_output.contiguous(), gate, up) + return d_gate, d_up + + +class MusaSwiGLUOp: + def __init__(self): + require_musa_symbol("swiglu_forward") + + def __call__(self, gate, up): + return self.forward(gate, up) + + def forward(self, gate, up): + check_musa(gate, up) + return _MusaSwiGLUFunction.apply(gate, up) + + def forward_fp32(self, gate, up): + check_musa(gate, up) + return _MusaSwiGLUFunction.apply(gate.float(), up.float()) diff --git a/rl_engine/kernels/ops/musa/attention/__init__.py b/rl_engine/kernels/ops/musa/attention/__init__.py new file mode 100644 index 00000000..f9761323 --- /dev/null +++ b/rl_engine/kernels/ops/musa/attention/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .deterministic_attn import MusaDeterministicAttentionOp + +__all__ = ["MusaDeterministicAttentionOp"] diff --git a/rl_engine/kernels/ops/musa/attention/deterministic_attn.py b/rl_engine/kernels/ops/musa/attention/deterministic_attn.py new file mode 100644 index 00000000..6e4a898d --- /dev/null +++ b/rl_engine/kernels/ops/musa/attention/deterministic_attn.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import math + +import torch + +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol + + +class _MusaDeterministicAttentionFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, causal, scale, mask): + mask_arg = mask if mask.numel() else None + out, lse, probs = _C.deterministic_attention_fp32( + q.contiguous(), k.contiguous(), v.contiguous(), causal, float(scale), mask_arg + ) + ctx.save_for_backward(q, k, v, probs, mask) + ctx.has_mask = mask.numel() != 0 + ctx.causal = causal + ctx.scale = scale + return out.to(q.dtype), lse + + @staticmethod + def backward(ctx, grad_out, grad_lse): + q, k, v, probs, mask = ctx.saved_tensors + mask_arg = mask if ctx.has_mask else None + dq, dk, dv = _C.deterministic_attention_backward( + grad_out.contiguous(), q, k, v, probs, ctx.causal, float(ctx.scale), mask_arg + ) + return dq, dk, dv, None, None, None + + +class MusaDeterministicAttentionOp: + def __init__(self): + require_musa_symbol("deterministic_attention_fp32") + require_musa_symbol("deterministic_attention_backward") + + def __call__(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): + return self.forward(q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask) + + def _apply(self, q, k, v, causal, scale, key_padding_mask): + check_musa(q, k, v) + if key_padding_mask is not None: + check_musa(key_padding_mask) + resolved_scale = scale if scale is not None else 1.0 / math.sqrt(q.size(-1)) + mask = ( + key_padding_mask.contiguous() + if key_padding_mask is not None + else torch.empty(0, device=q.device, dtype=torch.bool) + ) + return _MusaDeterministicAttentionFunction.apply( + q, k, v, causal, resolved_scale, mask + ) + + def forward(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): + out, _lse = self._apply(q, k, v, causal, scale, key_padding_mask) + return out + + def forward_with_lse(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): + return self._apply(q, k, v, causal, scale, key_padding_mask) + + def forward_fp32(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): + check_musa(q, k, v) + if key_padding_mask is not None: + check_musa(key_padding_mask) + resolved_scale = scale if scale is not None else 1.0 / math.sqrt(q.size(-1)) + mask = ( + key_padding_mask.contiguous() + if key_padding_mask is not None + else torch.empty(0, device=q.device, dtype=torch.bool) + ) + out, lse, _probs = _C.deterministic_attention_fp32( + q.contiguous(), k.contiguous(), v.contiguous(), + causal, float(resolved_scale), mask if mask.numel() else None + ) + return out, lse diff --git a/rl_engine/kernels/ops/musa/linear/__init__.py b/rl_engine/kernels/ops/musa/linear/__init__.py new file mode 100644 index 00000000..e0f673ef --- /dev/null +++ b/rl_engine/kernels/ops/musa/linear/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .embedding import MusaEmbeddingOp +from .lm_head import MusaLMHeadOp + +__all__ = ["MusaEmbeddingOp", "MusaLMHeadOp"] diff --git a/rl_engine/kernels/ops/musa/linear/embedding.py b/rl_engine/kernels/ops/musa/linear/embedding.py new file mode 100644 index 00000000..853d66b2 --- /dev/null +++ b/rl_engine/kernels/ops/musa/linear/embedding.py @@ -0,0 +1,47 @@ +# 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 +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol + + +class _MusaEmbeddingFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, token_ids, weight): + ids = token_ids.reshape(-1).to(dtype=torch.long).contiguous() + output = _C.embedding_fp32(ids, weight.contiguous()).to(weight.dtype) + ctx.save_for_backward(ids) + ctx.weight_shape = tuple(weight.shape) + ctx.weight_dtype = weight.dtype + ctx.output_shape = tuple(token_ids.shape) + (weight.size(1),) + return output.reshape(ctx.output_shape) + + @staticmethod + def backward(ctx, grad_output): + (ids,) = ctx.saved_tensors + grad_weight = torch.zeros( + ctx.weight_shape, device=grad_output.device, dtype=ctx.weight_dtype + ) + grad_weight.index_add_(0, ids, grad_output.reshape(-1, ctx.weight_shape[1])) + return None, grad_weight + + +class MusaEmbeddingOp: + def __init__(self): + require_musa_symbol("embedding_fp32") + + def __call__(self, token_ids, weight): + return self.forward(token_ids, weight) + + def forward(self, token_ids, weight): + check_musa(token_ids, weight) + return _MusaEmbeddingFunction.apply(token_ids, weight) + + def forward_fp32(self, token_ids, weight): + check_musa(token_ids, weight) + ids = token_ids.reshape(-1).to(device=weight.device, dtype=torch.long) + return _C.embedding_fp32(ids, weight).reshape(*token_ids.shape, weight.size(1)) diff --git a/rl_engine/kernels/ops/musa/linear/lm_head.py b/rl_engine/kernels/ops/musa/linear/lm_head.py new file mode 100644 index 00000000..958b818e --- /dev/null +++ b/rl_engine/kernels/ops/musa/linear/lm_head.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch +from typing import Optional + +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol + + +class _MusaLMHeadFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, hidden, weight, bias, output_dtype): + output = _C.lm_head_fp32(hidden, weight, bias) + ctx.save_for_backward(hidden, weight, bias if bias is not None else torch.empty(0)) + ctx.has_bias = bias is not None + return output.to(output_dtype) + + @staticmethod + def backward(ctx, grad_output): + hidden, weight, bias = ctx.saved_tensors + grad = grad_output.float().reshape(-1, grad_output.size(-1)) + hidden_2d = hidden.float().reshape(-1, hidden.size(-1)) + grad_hidden = (grad @ weight.float()).reshape_as(hidden) + grad_weight = grad.t() @ hidden_2d + grad_bias = grad.sum(dim=0) if ctx.has_bias else None + return grad_hidden.to(hidden.dtype), grad_weight.to(weight.dtype), grad_bias, None + + +class MusaLMHeadOp: + def __init__(self): + require_musa_symbol("lm_head_fp32") + + def __call__(self, hidden, weight, *, bias: Optional[torch.Tensor] = None): + return self.forward(hidden, weight, bias=bias) + + def forward(self, hidden, weight, *, bias: Optional[torch.Tensor] = None): + check_musa(hidden, weight) + if bias is not None: + check_musa(bias) + return _MusaLMHeadFunction.apply(hidden, weight, bias, hidden.dtype) + + def forward_fp32(self, hidden, weight, *, bias: Optional[torch.Tensor] = None): + check_musa(hidden, weight) + if bias is not None: + check_musa(bias) + return _C.lm_head_fp32(hidden, weight, bias) diff --git a/rl_engine/kernels/ops/musa/loss/__init__.py b/rl_engine/kernels/ops/musa/loss/__init__.py index 5af48e21..6d5dcbab 100644 --- a/rl_engine/kernels/ops/musa/loss/__init__.py +++ b/rl_engine/kernels/ops/musa/loss/__init__.py @@ -2,5 +2,6 @@ # Copyright (c) 2026 RL-Kernel Contributors from .logp import MusaFusedLogpOp +from .batch_invariant_logp import MusaDeterministicLogpOp -__all__ = ["MusaFusedLogpOp"] +__all__ = ["MusaFusedLogpOp", "MusaDeterministicLogpOp"] diff --git a/rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py new file mode 100644 index 00000000..f29558b8 --- /dev/null +++ b/rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py @@ -0,0 +1,71 @@ +# 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 +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol + + +class _MusaDeterministicLogpFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, logits: torch.Tensor, token_ids: torch.Tensor): + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + ids = token_ids.reshape(-1).to(device=logits.device, dtype=torch.long).contiguous() + output = _C.deterministic_logp_fp32(logits_2d, ids) + ctx.save_for_backward(logits_2d, ids) + ctx.input_shape = tuple(logits.shape) + ctx.input_dtype = logits.dtype + return output.reshape(logits.shape[:-1]) + + @staticmethod + def backward(ctx, grad_output): + logits, ids = ctx.saved_tensors + probs = torch.softmax(logits.float(), dim=-1) + rows = torch.arange(logits.size(0), device=logits.device) + probs[rows, ids] -= 1.0 + grad = -grad_output.reshape(-1, 1).float() * probs + return grad.to(ctx.input_dtype).reshape(ctx.input_shape), None + + +class MusaDeterministicLogpOp: + """Batch-invariant selected log-probability using the native MUSA kernel.""" + + is_batch_invariant = True + + def __init__(self): + require_musa_symbol("deterministic_logp_fp32") + + def __call__(self, logits, token_ids): + return self.apply_fp32(logits, token_ids) + + def apply(self, logits, token_ids): + return self.apply_fp32(logits, token_ids) + + def apply_fp32(self, logits, token_ids): + check_musa(logits, token_ids) + return _MusaDeterministicLogpFunction.apply(logits, token_ids) + + forward = apply_fp32 + forward_fp32 = apply_fp32 + online_fp32 = apply_fp32 + + def out(self, logits, token_ids, output): + output.copy_(self.apply_fp32(logits, token_ids).to(output.dtype)) + return output + + def indexed_fp32(self, logits, token_ids, row_indices): + output = torch.zeros(logits.shape[:-1], device=logits.device, dtype=torch.float32) + return self.indexed_out(logits, token_ids, row_indices, output) + + def indexed_out(self, logits, token_ids, row_indices, output): + full = self.apply_fp32(logits, token_ids).reshape(-1) + indices = row_indices.reshape(-1).to(device=full.device, dtype=torch.long) + output.reshape(-1).index_copy_(0, indices, full.index_select(0, indices).to(output.dtype)) + return output + + online_out = out + online_indexed_fp32 = indexed_fp32 + online_indexed_out = indexed_out diff --git a/rl_engine/kernels/ops/musa/matmul/__init__.py b/rl_engine/kernels/ops/musa/matmul/__init__.py new file mode 100644 index 00000000..d5ec4cc2 --- /dev/null +++ b/rl_engine/kernels/ops/musa/matmul/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .det_gemm import MusaDetGemmOp + +__all__ = ["MusaDetGemmOp"] diff --git a/rl_engine/kernels/ops/musa/matmul/det_gemm.py b/rl_engine/kernels/ops/musa/matmul/det_gemm.py new file mode 100644 index 00000000..c437214b --- /dev/null +++ b/rl_engine/kernels/ops/musa/matmul/det_gemm.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import torch + +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol + + +class MusaDetGemmOp: + def __init__(self): + require_musa_symbol("det_gemm_rowwise_fwd_fp32") + + def __call__(self, a, b): + check_musa(a, b) + return _C.det_gemm_rowwise_fwd_fp32(a, b) + + forward = __call__ + forward_fp32 = __call__ diff --git a/rl_engine/kernels/ops/musa/native.py b/rl_engine/kernels/ops/musa/native.py deleted file mode 100644 index 9b686455..00000000 --- a/rl_engine/kernels/ops/musa/native.py +++ /dev/null @@ -1,365 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Wrappers for the native MUSA kernels exposed by ``rl_engine._C``.""" - -from __future__ import annotations - -import math -from typing import Optional - -import torch - -from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - - -def _require(symbol: str) -> None: - if not _EXT_AVAILABLE or _C is None or not hasattr(_C, symbol): - raise RuntimeError(f"MUSA native symbol {symbol!r} is unavailable") - - -def _check_musa(*tensors: torch.Tensor) -> None: - if any(t.device.type != "musa" for t in tensors): - raise RuntimeError("MUSA native kernels require MUSA tensors") - - -class _MusaLogpFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, logits: torch.Tensor, token_ids: torch.Tensor): - logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() - ids = token_ids.reshape(-1).to(device=logits.device, dtype=torch.long).contiguous() - output = _C.deterministic_logp_fp32(logits_2d, ids) - ctx.save_for_backward(logits_2d, ids) - ctx.input_shape = tuple(logits.shape) - ctx.input_dtype = logits.dtype - return output.reshape(logits.shape[:-1]) - - @staticmethod - def backward(ctx, grad_output): - logits, ids = ctx.saved_tensors - probs = torch.softmax(logits.float(), dim=-1) - rows = torch.arange(logits.size(0), device=logits.device) - probs[rows, ids] -= 1.0 - grad = -grad_output.reshape(-1, 1).float() * probs - return grad.to(ctx.input_dtype).reshape(ctx.input_shape), None - - -class MusaDeterministicLogpOp: - is_batch_invariant = True - - def __init__(self): - _require("deterministic_logp_fp32") - - def __call__(self, logits, token_ids): - return self.apply_fp32(logits, token_ids) - - def apply(self, logits, token_ids): - return self.apply_fp32(logits, token_ids) - - def apply_fp32(self, logits, token_ids): - _check_musa(logits, token_ids) - return _MusaLogpFunction.apply(logits, token_ids) - - forward = apply_fp32 - forward_fp32 = apply_fp32 - online_fp32 = apply_fp32 - - def indexed_fp32(self, logits, token_ids, row_indices): - output = torch.zeros(logits.shape[:-1], device=logits.device, dtype=torch.float32) - return self.indexed_out(logits, token_ids, row_indices, output) - - def indexed_out(self, logits, token_ids, row_indices, output): - full = self.apply_fp32(logits, token_ids) - output.reshape(-1).index_copy_( - 0, - row_indices.reshape(-1).to(device=output.device, dtype=torch.long), - full.reshape(-1).index_select( - 0, row_indices.reshape(-1).to(device=full.device, dtype=torch.long) - ).to(output.dtype), - ) - return output - - out = indexed_out - online_indexed_fp32 = indexed_fp32 - online_indexed_out = indexed_out - - -class _MusaSiLUFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, x): - ctx.save_for_backward(x) - return _C.silu_forward(x.contiguous()) - - @staticmethod - def backward(ctx, grad_output): - (x,) = ctx.saved_tensors - return _C.silu_backward(grad_output.contiguous(), x) - - -class MusaSiLUOp: - def __init__(self): - _require("silu_forward") - - def __call__(self, x): - return self.forward(x) - - def forward(self, x): - _check_musa(x) - return _MusaSiLUFunction.apply(x) - - def forward_fp32(self, x): - _check_musa(x) - return _MusaSiLUFunction.apply(x.float()) - - -class _MusaSwiGLUFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, gate, up): - ctx.save_for_backward(gate, up) - return _C.swiglu_forward(gate.contiguous(), up.contiguous()) - - @staticmethod - def backward(ctx, grad_output): - gate, up = ctx.saved_tensors - d_gate, d_up = _C.swiglu_backward(grad_output.contiguous(), gate, up) - return d_gate, d_up - - -class MusaSwiGLUOp: - def __init__(self): - _require("swiglu_forward") - - def __call__(self, gate, up): - return self.forward(gate, up) - - def forward(self, gate, up): - _check_musa(gate, up) - return _MusaSwiGLUFunction.apply(gate, up) - - def forward_fp32(self, gate, up): - _check_musa(gate, up) - return _MusaSwiGLUFunction.apply(gate.float(), up.float()) - - -class _MusaRMSNormFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, x, weight, eps): - x_2d = x.reshape(-1, x.size(-1)).contiguous() - weight = weight.contiguous() - y = torch.empty_like(x_2d) - rstd = torch.empty(x_2d.size(0), device=x.device, dtype=torch.float32) - _C.rmsnorm_forward(x_2d, weight, y, rstd, float(eps)) - ctx.save_for_backward(x_2d, weight, rstd) - ctx.input_shape = tuple(x.shape) - ctx.eps = eps - return y.reshape_as(x) - - @staticmethod - def backward(ctx, grad_output): - x, weight, rstd = ctx.saved_tensors - dy = grad_output.reshape_as(x).reshape(-1, x.size(-1)).contiguous() - dx = torch.empty_like(x) - _C.rmsnorm_backward_dx(dy, x, weight, rstd, dx) - dw = (dy.float() * x.float() * rstd[:, None]).sum(dim=0).to(weight.dtype) - return dx.reshape(ctx.input_shape), dw, None - - -class MusaRMSNormOp: - def __init__(self): - _require("rmsnorm_forward") - _require("rmsnorm_backward_dx") - - def __call__(self, x, weight, *, eps=1e-6): - return self.forward(x, weight, eps=eps) - - def forward(self, x, weight, *, eps=1e-6): - _check_musa(x, weight) - return _MusaRMSNormFunction.apply(x, weight, eps) - - def forward_fp32(self, x, weight, *, eps=1e-6): - _check_musa(x, weight) - return _MusaRMSNormFunction.apply(x.float(), weight.float(), eps) - - -class _MusaEmbeddingFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, token_ids, weight): - ids = token_ids.reshape(-1).to(dtype=torch.long).contiguous() - output = _C.embedding_fp32(ids, weight.contiguous()).to(weight.dtype) - ctx.save_for_backward(ids) - ctx.weight_shape = tuple(weight.shape) - ctx.weight_dtype = weight.dtype - ctx.output_shape = tuple(token_ids.shape) + (weight.size(1),) - return output.reshape(ctx.output_shape) - - @staticmethod - def backward(ctx, grad_output): - (ids,) = ctx.saved_tensors - grad_weight = torch.zeros( - ctx.weight_shape, device=grad_output.device, dtype=ctx.weight_dtype - ) - grad_weight.index_add_(0, ids, grad_output.reshape(-1, ctx.weight_shape[1])) - return None, grad_weight - - -class MusaEmbeddingOp: - def __init__(self): - _require("embedding_fp32") - - def __call__(self, token_ids, weight): - return self.forward(token_ids, weight) - - def forward(self, token_ids, weight): - _check_musa(token_ids, weight) - return _MusaEmbeddingFunction.apply(token_ids, weight) - - def forward_fp32(self, token_ids, weight): - _check_musa(token_ids, weight) - return _C.embedding_fp32(token_ids.reshape(-1).long(), weight).reshape( - *token_ids.shape, weight.size(1) - ) - - -class _MusaLMHeadFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, hidden, weight, bias, output_dtype): - output = _C.lm_head_fp32(hidden, weight, bias) - ctx.save_for_backward(hidden, weight, bias if bias is not None else torch.empty(0)) - ctx.has_bias = bias is not None - ctx.output_dtype = output_dtype - return output.to(output_dtype) - - @staticmethod - def backward(ctx, grad_output): - hidden, weight, bias = ctx.saved_tensors - grad = grad_output.float().reshape(-1, grad_output.size(-1)) - hidden_2d = hidden.float().reshape(-1, hidden.size(-1)) - grad_hidden = (grad @ weight.float()).reshape_as(hidden) - grad_weight = grad.t() @ hidden_2d - grad_bias = grad.sum(dim=0) if ctx.has_bias else None - return grad_hidden.to(hidden.dtype), grad_weight.to(weight.dtype), grad_bias, None - - -class MusaLMHeadOp: - def __init__(self): - _require("lm_head_fp32") - - def __call__(self, hidden, weight, *, bias=None): - return self.forward(hidden, weight, bias=bias) - - def forward(self, hidden, weight, *, bias=None): - _check_musa(hidden, weight) - if bias is not None: - _check_musa(bias) - return _MusaLMHeadFunction.apply(hidden, weight, bias, hidden.dtype) - - def forward_fp32(self, hidden, weight, *, bias=None): - _check_musa(hidden, weight) - return _C.lm_head_fp32(hidden, weight, bias) - - -def _cos_sin(positions, half, theta, device): - inv = 1.0 / (theta ** (torch.arange(half, device=device, dtype=torch.float32) / half)) - freqs = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) * inv - return freqs.cos().contiguous(), freqs.sin().contiguous() - - -class _MusaRoPEFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, x, positions, theta): - if positions.dim() != 1: - raise NotImplementedError("MUSA native RoPE currently supports positions [S]") - half = x.size(-1) // 2 - cos, sin = _cos_sin(positions, half, theta, x.device) - x_2d = x.contiguous().reshape(-1, x.size(-1)) - output = _C.rope_apply(x_2d, cos, sin, 1.0).reshape_as(x) - ctx.save_for_backward(cos, sin) - ctx.seq_len = positions.numel() - return output - - @staticmethod - def backward(ctx, grad_output): - cos, sin = ctx.saved_tensors - grad = grad_output.contiguous().reshape(-1, grad_output.size(-1)) - return _C.rope_apply(grad, cos, sin, -1.0).reshape_as(grad_output), None, None - - -class MusaRoPEOp: - def __init__(self): - _require("rope_apply") - - def __call__(self, x, positions, *, theta=1_000_000.0): - return self.forward(x, positions, theta=theta) - - def forward(self, x, positions, *, theta=1_000_000.0): - _check_musa(x) - return _MusaRoPEFunction.apply(x, positions, float(theta)) - - def forward_fp32(self, x, positions, *, theta=1_000_000.0): - _check_musa(x) - return _MusaRoPEFunction.apply(x.float(), positions, float(theta)) - - -class _MusaAttentionFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, q, k, v, causal, scale, mask): - mask_arg = mask if mask.numel() else None - out, lse, probs = _C.deterministic_attention_fp32( - q.contiguous(), k.contiguous(), v.contiguous(), causal, float(scale), mask_arg - ) - ctx.save_for_backward(q, k, v, probs, mask) - ctx.has_mask = mask.numel() != 0 - ctx.causal = causal - ctx.scale = scale - return out.to(q.dtype), lse - - @staticmethod - def backward(ctx, grad_out, grad_lse): - q, k, v, probs, mask = ctx.saved_tensors - mask_arg = mask if ctx.has_mask else None - dq, dk, dv = _C.deterministic_attention_backward( - grad_out.contiguous(), q, k, v, probs, ctx.causal, float(ctx.scale), mask_arg - ) - return dq, dk, dv, None, None, None - - -class MusaDeterministicAttentionOp: - def __init__(self): - _require("deterministic_attention_fp32") - _require("deterministic_attention_backward") - - def __call__(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): - return self.forward(q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask) - - def forward(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): - _check_musa(q, k, v) - resolved_scale = scale if scale is not None else 1.0 / math.sqrt(q.size(-1)) - mask = ( - key_padding_mask.contiguous() - if key_padding_mask is not None - else torch.empty(0, device=q.device, dtype=torch.bool) - ) - out, _lse = _MusaAttentionFunction.apply(q, k, v, causal, resolved_scale, mask) - return out - - def forward_with_lse(self, q, k, v, *, causal=True, scale=None, key_padding_mask=None): - _check_musa(q, k, v) - resolved_scale = scale if scale is not None else 1.0 / math.sqrt(q.size(-1)) - mask = ( - key_padding_mask.contiguous() - if key_padding_mask is not None - else torch.empty(0, device=q.device, dtype=torch.bool) - ) - return _MusaAttentionFunction.apply(q, k, v, causal, resolved_scale, mask) - - -class MusaDetGemmOp: - def __init__(self): - _require("det_gemm_rowwise_fwd_fp32") - - def __call__(self, a, b): - _check_musa(a, b) - return _C.det_gemm_rowwise_fwd_fp32(a, b) - - forward = __call__ - forward_fp32 = __call__ diff --git a/rl_engine/kernels/ops/musa/norm/__init__.py b/rl_engine/kernels/ops/musa/norm/__init__.py new file mode 100644 index 00000000..cbd971bd --- /dev/null +++ b/rl_engine/kernels/ops/musa/norm/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .rmsnorm import MusaRMSNormOp + +__all__ = ["MusaRMSNormOp"] diff --git a/rl_engine/kernels/ops/musa/norm/rmsnorm.py b/rl_engine/kernels/ops/musa/norm/rmsnorm.py new file mode 100644 index 00000000..268c9f04 --- /dev/null +++ b/rl_engine/kernels/ops/musa/norm/rmsnorm.py @@ -0,0 +1,48 @@ +# 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 +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol + + +class _MusaRMSNormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, eps): + x_2d = x.reshape(-1, x.size(-1)).contiguous() + weight = weight.contiguous() + y = torch.empty_like(x_2d) + rstd = torch.empty(x_2d.size(0), device=x.device, dtype=torch.float32) + _C.rmsnorm_forward(x_2d, weight, y, rstd, float(eps)) + ctx.save_for_backward(x_2d, weight, rstd) + ctx.input_shape = tuple(x.shape) + return y.reshape_as(x) + + @staticmethod + def backward(ctx, grad_output): + x, weight, rstd = ctx.saved_tensors + dy = grad_output.reshape(-1, x.size(-1)).contiguous() + dx = torch.empty_like(x) + _C.rmsnorm_backward_dx(dy, x, weight, rstd, dx.reshape_as(x)) + dw = (dy.float() * x.float() * rstd[:, None]).sum(dim=0).to(weight.dtype) + return dx.reshape(ctx.input_shape), dw, None + + +class MusaRMSNormOp: + def __init__(self): + require_musa_symbol("rmsnorm_forward") + require_musa_symbol("rmsnorm_backward_dx") + + def __call__(self, x, weight, *, eps=1e-6): + return self.forward(x, weight, eps=eps) + + def forward(self, x, weight, *, eps=1e-6): + check_musa(x, weight) + return _MusaRMSNormFunction.apply(x, weight, eps) + + def forward_fp32(self, x, weight, *, eps=1e-6): + check_musa(x, weight) + return _MusaRMSNormFunction.apply(x.float(), weight.float(), eps) diff --git a/rl_engine/kernels/ops/musa/rotary_embedding/__init__.py b/rl_engine/kernels/ops/musa/rotary_embedding/__init__.py new file mode 100644 index 00000000..421dfccb --- /dev/null +++ b/rl_engine/kernels/ops/musa/rotary_embedding/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .rope import MusaRoPEOp + +__all__ = ["MusaRoPEOp"] diff --git a/rl_engine/kernels/ops/musa/rotary_embedding/rope.py b/rl_engine/kernels/ops/musa/rotary_embedding/rope.py new file mode 100644 index 00000000..c3e3ca67 --- /dev/null +++ b/rl_engine/kernels/ops/musa/rotary_embedding/rope.py @@ -0,0 +1,50 @@ +# 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 +from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol + + +def _cos_sin(positions, half, theta, device): + inv = 1.0 / (theta ** (torch.arange(half, device=device, dtype=torch.float32) / half)) + freqs = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) * inv + return freqs.cos().contiguous(), freqs.sin().contiguous() + + +class _MusaRoPEFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, positions, theta): + if positions.dim() != 1: + raise NotImplementedError("MUSA native RoPE currently supports positions [S]") + half = x.size(-1) // 2 + cos, sin = _cos_sin(positions, half, theta, x.device) + x_2d = x.contiguous().reshape(-1, x.size(-1)) + output = _C.rope_apply(x_2d, cos, sin, 1.0).reshape_as(x) + ctx.save_for_backward(cos, sin) + return output + + @staticmethod + def backward(ctx, grad_output): + cos, sin = ctx.saved_tensors + grad = grad_output.contiguous().reshape(-1, grad_output.size(-1)) + return _C.rope_apply(grad, cos, sin, -1.0).reshape_as(grad_output), None, None + + +class MusaRoPEOp: + def __init__(self): + require_musa_symbol("rope_apply") + + def __call__(self, x, positions, *, theta=1_000_000.0): + return self.forward(x, positions, theta=theta) + + def forward(self, x, positions, *, theta=1_000_000.0): + check_musa(x) + return _MusaRoPEFunction.apply(x, positions, float(theta)) + + def forward_fp32(self, x, positions, *, theta=1_000_000.0): + check_musa(x) + return _MusaRoPEFunction.apply(x.float(), positions, float(theta)) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index d7a275d5..9f239b1b 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -44,17 +44,19 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # Moore Threads MUSA native generic LogP kernel. MUSA_FUSED_LOGP = "rl_engine.kernels.ops.musa.loss.logp.MusaFusedLogpOp" - MUSA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.musa.native.MusaDeterministicLogpOp" + MUSA_DETERMINISTIC_LOGP = ( + "rl_engine.kernels.ops.musa.loss.batch_invariant_logp.MusaDeterministicLogpOp" + ) MUSA_DETERMINISTIC_ATTENTION = ( - "rl_engine.kernels.ops.musa.native.MusaDeterministicAttentionOp" + "rl_engine.kernels.ops.musa.attention.deterministic_attn.MusaDeterministicAttentionOp" ) - MUSA_DET_GEMM = "rl_engine.kernels.ops.musa.native.MusaDetGemmOp" - MUSA_RMS_NORM = "rl_engine.kernels.ops.musa.native.MusaRMSNormOp" - MUSA_EMBEDDING = "rl_engine.kernels.ops.musa.native.MusaEmbeddingOp" - MUSA_LM_HEAD = "rl_engine.kernels.ops.musa.native.MusaLMHeadOp" - MUSA_ROPE = "rl_engine.kernels.ops.musa.native.MusaRoPEOp" - MUSA_SILU = "rl_engine.kernels.ops.musa.native.MusaSiLUOp" - MUSA_SWIGLU = "rl_engine.kernels.ops.musa.native.MusaSwiGLUOp" + MUSA_DET_GEMM = "rl_engine.kernels.ops.musa.matmul.det_gemm.MusaDetGemmOp" + MUSA_RMS_NORM = "rl_engine.kernels.ops.musa.norm.rmsnorm.MusaRMSNormOp" + MUSA_EMBEDDING = "rl_engine.kernels.ops.musa.linear.embedding.MusaEmbeddingOp" + MUSA_LM_HEAD = "rl_engine.kernels.ops.musa.linear.lm_head.MusaLMHeadOp" + MUSA_ROPE = "rl_engine.kernels.ops.musa.rotary_embedding.rope.MusaRoPEOp" + MUSA_SILU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSiLUOp" + MUSA_SWIGLU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSwiGLUOp" # GRPO loss (group reward normalization + clipped surrogate + KL) TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" diff --git a/tests/test_attention.py b/tests/test_attention.py index 2eb4735e..4e23cf2d 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -438,7 +438,9 @@ def test_registry_dispatches_native_attention_op(): """Resolve attention to the deterministic CUDA op or native fallback.""" op = kernel_registry.get_op("attention") if device_ctx.is_musa: - from rl_engine.kernels.ops.musa.native import MusaDeterministicAttentionOp + from rl_engine.kernels.ops.musa.attention.deterministic_attn import ( + MusaDeterministicAttentionOp, + ) assert isinstance(op, MusaDeterministicAttentionOp) return diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index d6401472..14c5d4a0 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -6,7 +6,7 @@ import torch.nn.functional as F from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda -from rl_engine.kernels.ops.musa.native import MusaRMSNormOp +from rl_engine.kernels.ops.musa.norm.rmsnorm import MusaRMSNormOp from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton From a73f3dd44464b7d90179048b025aa6d41f9eccfe Mon Sep 17 00:00:00 2001 From: mt Date: Wed, 2 Sep 2026 17:02:06 +0800 Subject: [PATCH 5/6] refactor: simplify MUSA backend names Signed-off-by: mt --- rl_engine/kernels/registry.py | 44 +++++++++++++++++------------------ 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 9f239b1b..2bc38b66 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -43,20 +43,20 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ROCM_FLASH_ATTN = "rl_engine.kernels.ops.rocm.attention.flash_attn.RocmFlashAttentionOp" # Moore Threads MUSA native generic LogP kernel. - MUSA_FUSED_LOGP = "rl_engine.kernels.ops.musa.loss.logp.MusaFusedLogpOp" - MUSA_DETERMINISTIC_LOGP = ( + FUSED_LOGP = "rl_engine.kernels.ops.musa.loss.logp.MusaFusedLogpOp" + DETERMINISTIC_LOGP = ( "rl_engine.kernels.ops.musa.loss.batch_invariant_logp.MusaDeterministicLogpOp" ) - MUSA_DETERMINISTIC_ATTENTION = ( + DETERMINISTIC_ATTENTION = ( "rl_engine.kernels.ops.musa.attention.deterministic_attn.MusaDeterministicAttentionOp" ) - MUSA_DET_GEMM = "rl_engine.kernels.ops.musa.matmul.det_gemm.MusaDetGemmOp" - MUSA_RMS_NORM = "rl_engine.kernels.ops.musa.norm.rmsnorm.MusaRMSNormOp" - MUSA_EMBEDDING = "rl_engine.kernels.ops.musa.linear.embedding.MusaEmbeddingOp" - MUSA_LM_HEAD = "rl_engine.kernels.ops.musa.linear.lm_head.MusaLMHeadOp" - MUSA_ROPE = "rl_engine.kernels.ops.musa.rotary_embedding.rope.MusaRoPEOp" - MUSA_SILU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSiLUOp" - MUSA_SWIGLU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSwiGLUOp" + DET_GEMM = "rl_engine.kernels.ops.musa.matmul.det_gemm.MusaDetGemmOp" + RMS_NORM = "rl_engine.kernels.ops.musa.norm.rmsnorm.MusaRMSNormOp" + EMBEDDING = "rl_engine.kernels.ops.musa.linear.embedding.MusaEmbeddingOp" + LM_HEAD = "rl_engine.kernels.ops.musa.linear.lm_head.MusaLMHeadOp" + ROPE = "rl_engine.kernels.ops.musa.rotary_embedding.rope.MusaRoPEOp" + SILU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSiLUOp" + SWIGLU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSwiGLUOp" # GRPO loss (group reward normalization + clipped surrogate + KL) TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" @@ -291,34 +291,34 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], }, "musa": { - "logp": [OpBackend.MUSA_FUSED_LOGP, OpBackend.PYTORCH_NATIVE], + "logp": [OpBackend.FUSED_LOGP, OpBackend.PYTORCH_NATIVE], "logp_indexed": [OpBackend.PYTORCH_NATIVE], "logp_online": [OpBackend.PYTORCH_NATIVE], "logp_online_indexed": [OpBackend.PYTORCH_NATIVE], - "logp_deterministic": [OpBackend.MUSA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE], + "logp_deterministic": [OpBackend.DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE], "logp_deterministic_indexed": [ - OpBackend.MUSA_DETERMINISTIC_LOGP, + OpBackend.DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], "attn": [OpBackend.PYTORCH_ATTN], - "attention": [OpBackend.MUSA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION], + "attention": [OpBackend.DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], - "rope": [OpBackend.MUSA_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], + "rope": [OpBackend.ROPE, OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], - "det_gemm": [OpBackend.MUSA_DET_GEMM, OpBackend.PYTORCH_GEMM], + "det_gemm": [OpBackend.DET_GEMM, OpBackend.PYTORCH_GEMM], "batch_invariant_logp": [ - OpBackend.MUSA_DETERMINISTIC_LOGP, + OpBackend.DETERMINISTIC_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], - "rms_norm": [OpBackend.MUSA_RMS_NORM, OpBackend.PYTORCH_NATIVE_RMS_NORM], - "lm_head": [OpBackend.MUSA_LM_HEAD, OpBackend.PYTORCH_NATIVE_LM_HEAD], - "embedding": [OpBackend.MUSA_EMBEDDING, OpBackend.PYTORCH_NATIVE_EMBEDDING], - "silu": [OpBackend.MUSA_SILU, OpBackend.PYTORCH_NATIVE_SILU], - "swiglu": [OpBackend.MUSA_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "rms_norm": [OpBackend.RMS_NORM, OpBackend.PYTORCH_NATIVE_RMS_NORM], + "lm_head": [OpBackend.LM_HEAD, OpBackend.PYTORCH_NATIVE_LM_HEAD], + "embedding": [OpBackend.EMBEDDING, OpBackend.PYTORCH_NATIVE_EMBEDDING], + "silu": [OpBackend.SILU, OpBackend.PYTORCH_NATIVE_SILU], + "swiglu": [OpBackend.SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], From 8823e3255454d64f0694e4c23b3c960958cb646d Mon Sep 17 00:00:00 2001 From: mt Date: Wed, 2 Sep 2026 17:16:42 +0800 Subject: [PATCH 6/6] refactor: remove MUSA prefix from op classes Signed-off-by: mt --- .../kernels/ops/musa/activation/__init__.py | 4 +- .../kernels/ops/musa/activation/swiglu.py | 4 +- .../kernels/ops/musa/attention/__init__.py | 4 +- .../ops/musa/attention/deterministic_attn.py | 2 +- rl_engine/kernels/ops/musa/linear/__init__.py | 6 +-- .../kernels/ops/musa/linear/embedding.py | 2 +- rl_engine/kernels/ops/musa/linear/lm_head.py | 2 +- rl_engine/kernels/ops/musa/loss/__init__.py | 6 +-- .../ops/musa/loss/batch_invariant_logp.py | 2 +- rl_engine/kernels/ops/musa/loss/logp.py | 4 +- rl_engine/kernels/ops/musa/matmul/__init__.py | 4 +- rl_engine/kernels/ops/musa/matmul/det_gemm.py | 2 +- rl_engine/kernels/ops/musa/norm/__init__.py | 4 +- rl_engine/kernels/ops/musa/norm/rmsnorm.py | 2 +- .../ops/musa/rotary_embedding/__init__.py | 4 +- .../kernels/ops/musa/rotary_embedding/rope.py | 2 +- rl_engine/kernels/registry.py | 54 ++++++++++--------- tests/test_attention.py | 4 +- tests/test_batch_invariant_logp.py | 2 +- tests/test_logp.py | 4 +- tests/test_rms_norm.py | 4 +- 21 files changed, 64 insertions(+), 58 deletions(-) diff --git a/rl_engine/kernels/ops/musa/activation/__init__.py b/rl_engine/kernels/ops/musa/activation/__init__.py index 2ae13f49..ab20d986 100644 --- a/rl_engine/kernels/ops/musa/activation/__init__.py +++ b/rl_engine/kernels/ops/musa/activation/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .swiglu import MusaSiLUOp, MusaSwiGLUOp +from .swiglu import SiLUOp, SwiGLUOp -__all__ = ["MusaSiLUOp", "MusaSwiGLUOp"] +__all__ = ["SiLUOp", "SwiGLUOp"] diff --git a/rl_engine/kernels/ops/musa/activation/swiglu.py b/rl_engine/kernels/ops/musa/activation/swiglu.py index 5982c77e..90a7ddee 100644 --- a/rl_engine/kernels/ops/musa/activation/swiglu.py +++ b/rl_engine/kernels/ops/musa/activation/swiglu.py @@ -21,7 +21,7 @@ def backward(ctx, grad_output): return _C.silu_backward(grad_output.contiguous(), x) -class MusaSiLUOp: +class SiLUOp: def __init__(self): require_musa_symbol("silu_forward") @@ -50,7 +50,7 @@ def backward(ctx, grad_output): return d_gate, d_up -class MusaSwiGLUOp: +class SwiGLUOp: def __init__(self): require_musa_symbol("swiglu_forward") diff --git a/rl_engine/kernels/ops/musa/attention/__init__.py b/rl_engine/kernels/ops/musa/attention/__init__.py index f9761323..fe78dd6d 100644 --- a/rl_engine/kernels/ops/musa/attention/__init__.py +++ b/rl_engine/kernels/ops/musa/attention/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .deterministic_attn import MusaDeterministicAttentionOp +from .deterministic_attn import DeterministicAttentionOp -__all__ = ["MusaDeterministicAttentionOp"] +__all__ = ["DeterministicAttentionOp"] diff --git a/rl_engine/kernels/ops/musa/attention/deterministic_attn.py b/rl_engine/kernels/ops/musa/attention/deterministic_attn.py index 6e4a898d..2216eb82 100644 --- a/rl_engine/kernels/ops/musa/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/musa/attention/deterministic_attn.py @@ -34,7 +34,7 @@ def backward(ctx, grad_out, grad_lse): return dq, dk, dv, None, None, None -class MusaDeterministicAttentionOp: +class DeterministicAttentionOp: def __init__(self): require_musa_symbol("deterministic_attention_fp32") require_musa_symbol("deterministic_attention_backward") diff --git a/rl_engine/kernels/ops/musa/linear/__init__.py b/rl_engine/kernels/ops/musa/linear/__init__.py index e0f673ef..abc64787 100644 --- a/rl_engine/kernels/ops/musa/linear/__init__.py +++ b/rl_engine/kernels/ops/musa/linear/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .embedding import MusaEmbeddingOp -from .lm_head import MusaLMHeadOp +from .embedding import EmbeddingOp +from .lm_head import LMHeadOp -__all__ = ["MusaEmbeddingOp", "MusaLMHeadOp"] +__all__ = ["EmbeddingOp", "LMHeadOp"] diff --git a/rl_engine/kernels/ops/musa/linear/embedding.py b/rl_engine/kernels/ops/musa/linear/embedding.py index 853d66b2..1c96be0f 100644 --- a/rl_engine/kernels/ops/musa/linear/embedding.py +++ b/rl_engine/kernels/ops/musa/linear/embedding.py @@ -30,7 +30,7 @@ def backward(ctx, grad_output): return None, grad_weight -class MusaEmbeddingOp: +class EmbeddingOp: def __init__(self): require_musa_symbol("embedding_fp32") diff --git a/rl_engine/kernels/ops/musa/linear/lm_head.py b/rl_engine/kernels/ops/musa/linear/lm_head.py index 958b818e..96ca3fae 100644 --- a/rl_engine/kernels/ops/musa/linear/lm_head.py +++ b/rl_engine/kernels/ops/musa/linear/lm_head.py @@ -29,7 +29,7 @@ def backward(ctx, grad_output): return grad_hidden.to(hidden.dtype), grad_weight.to(weight.dtype), grad_bias, None -class MusaLMHeadOp: +class LMHeadOp: def __init__(self): require_musa_symbol("lm_head_fp32") diff --git a/rl_engine/kernels/ops/musa/loss/__init__.py b/rl_engine/kernels/ops/musa/loss/__init__.py index 6d5dcbab..25919573 100644 --- a/rl_engine/kernels/ops/musa/loss/__init__.py +++ b/rl_engine/kernels/ops/musa/loss/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .logp import MusaFusedLogpOp -from .batch_invariant_logp import MusaDeterministicLogpOp +from .logp import FusedLogpOp +from .batch_invariant_logp import DeterministicLogpOp -__all__ = ["MusaFusedLogpOp", "MusaDeterministicLogpOp"] +__all__ = ["FusedLogpOp", "DeterministicLogpOp"] diff --git a/rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py index f29558b8..3d2f2693 100644 --- a/rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/musa/loss/batch_invariant_logp.py @@ -30,7 +30,7 @@ def backward(ctx, grad_output): return grad.to(ctx.input_dtype).reshape(ctx.input_shape), None -class MusaDeterministicLogpOp: +class DeterministicLogpOp: """Batch-invariant selected log-probability using the native MUSA kernel.""" is_batch_invariant = True diff --git a/rl_engine/kernels/ops/musa/loss/logp.py b/rl_engine/kernels/ops/musa/loss/logp.py index 3e96e486..d0ac02d6 100644 --- a/rl_engine/kernels/ops/musa/loss/logp.py +++ b/rl_engine/kernels/ops/musa/loss/logp.py @@ -29,7 +29,7 @@ def backward(ctx, grad_output: torch.Tensor): return grad.to(ctx.input_dtype).reshape(ctx.input_shape), None -class MusaFusedLogpOp: +class FusedLogpOp: """Generic MUSA fused LogP; backward uses a portable PyTorch formula.""" is_fused_logp = True @@ -40,5 +40,5 @@ def __init__(self): def __call__(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: if logits.device.type != "musa": - raise RuntimeError(f"MusaFusedLogpOp requires a MUSA tensor, got {logits.device}") + raise RuntimeError(f"FusedLogpOp requires a MUSA tensor, got {logits.device}") return _MusaFusedLogpFunction.apply(logits, token_ids) diff --git a/rl_engine/kernels/ops/musa/matmul/__init__.py b/rl_engine/kernels/ops/musa/matmul/__init__.py index d5ec4cc2..786499c2 100644 --- a/rl_engine/kernels/ops/musa/matmul/__init__.py +++ b/rl_engine/kernels/ops/musa/matmul/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .det_gemm import MusaDetGemmOp +from .det_gemm import DetGemmOp -__all__ = ["MusaDetGemmOp"] +__all__ = ["DetGemmOp"] diff --git a/rl_engine/kernels/ops/musa/matmul/det_gemm.py b/rl_engine/kernels/ops/musa/matmul/det_gemm.py index c437214b..b0f1b3d5 100644 --- a/rl_engine/kernels/ops/musa/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/musa/matmul/det_gemm.py @@ -7,7 +7,7 @@ from rl_engine.kernels.ops.musa._common import check_musa, require_musa_symbol -class MusaDetGemmOp: +class DetGemmOp: def __init__(self): require_musa_symbol("det_gemm_rowwise_fwd_fp32") diff --git a/rl_engine/kernels/ops/musa/norm/__init__.py b/rl_engine/kernels/ops/musa/norm/__init__.py index cbd971bd..12eb7720 100644 --- a/rl_engine/kernels/ops/musa/norm/__init__.py +++ b/rl_engine/kernels/ops/musa/norm/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .rmsnorm import MusaRMSNormOp +from .rmsnorm import RMSNormOp -__all__ = ["MusaRMSNormOp"] +__all__ = ["RMSNormOp"] diff --git a/rl_engine/kernels/ops/musa/norm/rmsnorm.py b/rl_engine/kernels/ops/musa/norm/rmsnorm.py index 268c9f04..0384d6ca 100644 --- a/rl_engine/kernels/ops/musa/norm/rmsnorm.py +++ b/rl_engine/kernels/ops/musa/norm/rmsnorm.py @@ -31,7 +31,7 @@ def backward(ctx, grad_output): return dx.reshape(ctx.input_shape), dw, None -class MusaRMSNormOp: +class RMSNormOp: def __init__(self): require_musa_symbol("rmsnorm_forward") require_musa_symbol("rmsnorm_backward_dx") diff --git a/rl_engine/kernels/ops/musa/rotary_embedding/__init__.py b/rl_engine/kernels/ops/musa/rotary_embedding/__init__.py index 421dfccb..06157f39 100644 --- a/rl_engine/kernels/ops/musa/rotary_embedding/__init__.py +++ b/rl_engine/kernels/ops/musa/rotary_embedding/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .rope import MusaRoPEOp +from .rope import RoPEOp -__all__ = ["MusaRoPEOp"] +__all__ = ["RoPEOp"] diff --git a/rl_engine/kernels/ops/musa/rotary_embedding/rope.py b/rl_engine/kernels/ops/musa/rotary_embedding/rope.py index c3e3ca67..96d42caa 100644 --- a/rl_engine/kernels/ops/musa/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/musa/rotary_embedding/rope.py @@ -34,7 +34,7 @@ def backward(ctx, grad_output): return _C.rope_apply(grad, cos, sin, -1.0).reshape_as(grad_output), None, None -class MusaRoPEOp: +class RoPEOp: def __init__(self): require_musa_symbol("rope_apply") diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 2bc38b66..ad22fb6a 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -43,20 +43,20 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ROCM_FLASH_ATTN = "rl_engine.kernels.ops.rocm.attention.flash_attn.RocmFlashAttentionOp" # Moore Threads MUSA native generic LogP kernel. - FUSED_LOGP = "rl_engine.kernels.ops.musa.loss.logp.MusaFusedLogpOp" - DETERMINISTIC_LOGP = ( - "rl_engine.kernels.ops.musa.loss.batch_invariant_logp.MusaDeterministicLogpOp" + MUSA_FUSED_LOGP = "rl_engine.kernels.ops.musa.loss.logp.FusedLogpOp" + MUSA_DETERMINISTIC_LOGP = ( + "rl_engine.kernels.ops.musa.loss.batch_invariant_logp.DeterministicLogpOp" ) - DETERMINISTIC_ATTENTION = ( - "rl_engine.kernels.ops.musa.attention.deterministic_attn.MusaDeterministicAttentionOp" + MUSA_DETERMINISTIC_ATTENTION = ( + "rl_engine.kernels.ops.musa.attention.deterministic_attn.DeterministicAttentionOp" ) - DET_GEMM = "rl_engine.kernels.ops.musa.matmul.det_gemm.MusaDetGemmOp" - RMS_NORM = "rl_engine.kernels.ops.musa.norm.rmsnorm.MusaRMSNormOp" - EMBEDDING = "rl_engine.kernels.ops.musa.linear.embedding.MusaEmbeddingOp" - LM_HEAD = "rl_engine.kernels.ops.musa.linear.lm_head.MusaLMHeadOp" - ROPE = "rl_engine.kernels.ops.musa.rotary_embedding.rope.MusaRoPEOp" - SILU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSiLUOp" - SWIGLU = "rl_engine.kernels.ops.musa.activation.swiglu.MusaSwiGLUOp" + MUSA_DET_GEMM = "rl_engine.kernels.ops.musa.matmul.det_gemm.DetGemmOp" + MUSA_RMS_NORM = "rl_engine.kernels.ops.musa.norm.rmsnorm.RMSNormOp" + MUSA_EMBEDDING = "rl_engine.kernels.ops.musa.linear.embedding.EmbeddingOp" + MUSA_LM_HEAD = "rl_engine.kernels.ops.musa.linear.lm_head.LMHeadOp" + MUSA_ROPE = "rl_engine.kernels.ops.musa.rotary_embedding.rope.RoPEOp" + MUSA_SILU = "rl_engine.kernels.ops.musa.activation.swiglu.SiLUOp" + MUSA_SWIGLU = "rl_engine.kernels.ops.musa.activation.swiglu.SwiGLUOp" # GRPO loss (group reward normalization + clipped surrogate + KL) TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" @@ -291,34 +291,40 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], }, "musa": { - "logp": [OpBackend.FUSED_LOGP, OpBackend.PYTORCH_NATIVE], + "logp": [OpBackend.MUSA_FUSED_LOGP, OpBackend.PYTORCH_NATIVE], "logp_indexed": [OpBackend.PYTORCH_NATIVE], "logp_online": [OpBackend.PYTORCH_NATIVE], "logp_online_indexed": [OpBackend.PYTORCH_NATIVE], - "logp_deterministic": [OpBackend.DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE], + "logp_deterministic": [ + OpBackend.MUSA_DETERMINISTIC_LOGP, + OpBackend.PYTORCH_NATIVE, + ], "logp_deterministic_indexed": [ - OpBackend.DETERMINISTIC_LOGP, + OpBackend.MUSA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], "attn": [OpBackend.PYTORCH_ATTN], - "attention": [OpBackend.DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION], + "attention": [ + OpBackend.MUSA_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], - "rope": [OpBackend.ROPE, OpBackend.PYTORCH_NATIVE_ROPE], + "rope": [OpBackend.MUSA_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], - "det_gemm": [OpBackend.DET_GEMM, OpBackend.PYTORCH_GEMM], + "det_gemm": [OpBackend.MUSA_DET_GEMM, OpBackend.PYTORCH_GEMM], "batch_invariant_logp": [ - OpBackend.DETERMINISTIC_LOGP, + OpBackend.MUSA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], - "rms_norm": [OpBackend.RMS_NORM, OpBackend.PYTORCH_NATIVE_RMS_NORM], - "lm_head": [OpBackend.LM_HEAD, OpBackend.PYTORCH_NATIVE_LM_HEAD], - "embedding": [OpBackend.EMBEDDING, OpBackend.PYTORCH_NATIVE_EMBEDDING], - "silu": [OpBackend.SILU, OpBackend.PYTORCH_NATIVE_SILU], - "swiglu": [OpBackend.SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "rms_norm": [OpBackend.MUSA_RMS_NORM, OpBackend.PYTORCH_NATIVE_RMS_NORM], + "lm_head": [OpBackend.MUSA_LM_HEAD, OpBackend.PYTORCH_NATIVE_LM_HEAD], + "embedding": [OpBackend.MUSA_EMBEDDING, OpBackend.PYTORCH_NATIVE_EMBEDDING], + "silu": [OpBackend.MUSA_SILU, OpBackend.PYTORCH_NATIVE_SILU], + "swiglu": [OpBackend.MUSA_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], diff --git a/tests/test_attention.py b/tests/test_attention.py index 4e23cf2d..6ecf13be 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -439,10 +439,10 @@ def test_registry_dispatches_native_attention_op(): op = kernel_registry.get_op("attention") if device_ctx.is_musa: from rl_engine.kernels.ops.musa.attention.deterministic_attn import ( - MusaDeterministicAttentionOp, + DeterministicAttentionOp, ) - assert isinstance(op, MusaDeterministicAttentionOp) + assert isinstance(op, DeterministicAttentionOp) return # On CUDA with the extension built, the registry prefers DeterministicAttentionOp. # On CPU or without the CUDA extension, it falls back to NativeAttentionOp. diff --git a/tests/test_batch_invariant_logp.py b/tests/test_batch_invariant_logp.py index bb5c05ab..3d21a573 100644 --- a/tests/test_batch_invariant_logp.py +++ b/tests/test_batch_invariant_logp.py @@ -1027,7 +1027,7 @@ def test_registry_dispatches_correctly(): isinstance(op, NativeBatchInvariantLogpOp) or type(op).__name__ == "TritonBatchInvariantLogpOp" or type(op).__name__ == "BatchInvariantLogpSM90Op" - or type(op).__name__ == "MusaDeterministicLogpOp" + or type(op).__name__ == "DeterministicLogpOp" ) logits = torch.randn(4, _V, device=device_ctx.device) target = torch.randint(0, _V, (4,), device=logits.device) diff --git a/tests/test_logp.py b/tests/test_logp.py index 9792d45b..a523c9ba 100644 --- a/tests/test_logp.py +++ b/tests/test_logp.py @@ -158,10 +158,10 @@ def test_registry_returns_logp_op(self): op = kernel_registry.get_op("logp") if device_ctx.is_musa: from rl_engine.kernels.ops.base import _EXT_AVAILABLE - from rl_engine.kernels.ops.musa.loss.logp import MusaFusedLogpOp + from rl_engine.kernels.ops.musa.loss.logp import FusedLogpOp if _EXT_AVAILABLE: - assert isinstance(op, MusaFusedLogpOp) + assert isinstance(op, FusedLogpOp) else: assert isinstance(op, NativeLogpOp) else: diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 14c5d4a0..d7f15375 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -6,7 +6,7 @@ import torch.nn.functional as F from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda -from rl_engine.kernels.ops.musa.norm.rmsnorm import MusaRMSNormOp +from rl_engine.kernels.ops.musa.norm.rmsnorm import RMSNormOp from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton @@ -238,7 +238,7 @@ def test_registry_dispatches_rms_norm(): op = kernel_registry.get_op("rms_norm") if torch.musa.is_available(): - assert isinstance(op, MusaRMSNormOp) + assert isinstance(op, RMSNormOp) else: assert isinstance(op, NativeRMSNormOp) assert hasattr(op, "forward") and hasattr(op, "forward_fp32")