diff --git a/README.md b/README.md index e2788c6a..1a727ca0 100644 --- a/README.md +++ b/README.md @@ -123,6 +123,10 @@ python -m pip install -e . # Native CUDA or ROCm extension (install a matching PyTorch build first) RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e . python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)" + +# Ascend C extension (Linux host with matching CANN + torch_npu) +KERNEL_ALIGN_FORCE_ASCEND=1 python -m pip install --no-build-isolation -e . +python -c "import rl_engine._C_npu as C; assert hasattr(C, 'rope_apply_ascend'); print(C.__file__)" ``` ### Contributions diff --git a/csrc/ascend/batch_invariant_logp_ascend.asc b/csrc/ascend/batch_invariant_logp_ascend.asc index dead4cbe..49139240 100644 --- a/csrc/ascend/batch_invariant_logp_ascend.asc +++ b/csrc/ascend/batch_invariant_logp_ascend.asc @@ -307,10 +307,3 @@ std::vector batch_invariant_logp_ascend_forward(torch::Tensor log } return {logp, lse}; } - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("batch_invariant_logp_ascend", - &batch_invariant_logp_ascend_forward, - "Batch-invariant selected-token log-probability (Ascend C forward)"); -} diff --git a/csrc/ascend/npu_module.cpp b/csrc/ascend/npu_module.cpp new file mode 100644 index 00000000..4256baf8 --- /dev/null +++ b/csrc/ascend/npu_module.cpp @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#include + +#include + +std::vector batch_invariant_logp_ascend_forward(torch::Tensor logits, + torch::Tensor target, + int64_t ignore_index); + +torch::Tensor rope_apply_ascend_forward(torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin, + double sin_sign); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("batch_invariant_logp_ascend", + &batch_invariant_logp_ascend_forward, + "Batch-invariant selected-token log-probability (Ascend C forward)"); + m.def("rope_apply_ascend", + &rope_apply_ascend_forward, + "GPT-NeoX/HF rotate-half RoPE apply (Ascend C forward/backward primitive)"); +} diff --git a/csrc/ascend/rope_ascend.asc b/csrc/ascend/rope_ascend.asc new file mode 100644 index 00000000..c3606909 --- /dev/null +++ b/csrc/ascend/rope_ascend.asc @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Ascend C RoPE apply primitive (GPT-NeoX / Hugging Face rotate-half). +// +// The Python wrapper builds fp32 cos/sin tables from positions and theta, then +// flattens x to [n_rows, D]. For pair i in [0, D/2): +// +// out[i] = x[i] * cos[i] - x[i+D/2] * sin[i] * sin_sign +// out[i+D/2] = x[i+D/2] * cos[i] + x[i] * sin[i] * sin_sign +// +// sin_sign=+1 is the forward rotation and sin_sign=-1 is its transpose, used +// for grad_x. Each row is processed by one block with a fixed tile order, so +// adding or moving other batch rows cannot alter the instruction sequence. + +#include +#include + +#include "kernel_operator.h" + +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +constexpr uint32_t ROPE_TILE_HALF = 4096; +constexpr int64_t ROPE_MAX_BLOCKS = 128; + +template +class KernelRopeApply { +public: + __aicore__ inline KernelRopeApply(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR x, + GM_ADDR cos, + GM_ADDR sin, + GM_ADDR out, + int64_t numRows, + int64_t tableRows, + int64_t headDim, + float sinSign) + { + numRows_ = numRows; + tableRows_ = tableRows; + headDim_ = headDim; + halfDim_ = headDim / 2; + sinSign_ = sinSign; + xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x)); + cosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(cos)); + sinGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(sin)); + outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out)); + + pipe_->InitBuffer(x1InBuf_, ROPE_TILE_HALF * sizeof(T)); + pipe_->InitBuffer(x2InBuf_, ROPE_TILE_HALF * sizeof(T)); + pipe_->InitBuffer(x1FpBuf_, ROPE_TILE_HALF * sizeof(float)); + pipe_->InitBuffer(x2FpBuf_, ROPE_TILE_HALF * sizeof(float)); + pipe_->InitBuffer(cosBuf_, ROPE_TILE_HALF * sizeof(float)); + pipe_->InitBuffer(sinBuf_, ROPE_TILE_HALF * sizeof(float)); + pipe_->InitBuffer(out1FpBuf_, ROPE_TILE_HALF * sizeof(float)); + pipe_->InitBuffer(out2FpBuf_, ROPE_TILE_HALF * sizeof(float)); + pipe_->InitBuffer(tmpFpBuf_, ROPE_TILE_HALF * sizeof(float)); + pipe_->InitBuffer(out1Buf_, ROPE_TILE_HALF * sizeof(T)); + pipe_->InitBuffer(out2Buf_, ROPE_TILE_HALF * sizeof(T)); + + // Mark the reusable input and output buffers as initially available. + AscendC::SetFlag(0); + AscendC::SetFlag(0); + } + + __aicore__ inline void Process() + { + for (int64_t row = AscendC::GetBlockIdx(); row < numRows_; row += AscendC::GetBlockNum()) { + ProcessRow(row); + } + // Drain the final tile before the block exits and its UB is reclaimed. + AscendC::WaitFlag(0); + AscendC::WaitFlag(0); + } + +private: + __aicore__ inline void ProcessRow(int64_t row) + { + const int64_t tableRow = row % tableRows_; + for (int64_t start = 0; start < halfDim_; start += ROPE_TILE_HALF) { + const int64_t remaining = halfDim_ - start; + const uint32_t count = static_cast( + remaining < ROPE_TILE_HALF ? remaining : ROPE_TILE_HALF); + + // Previous vector reads are complete before MTE2 reuses input/cache buffers. + AscendC::WaitFlag(0); + CopyIn(row, tableRow, start, count); + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + + // Previous MTE3 reads are complete before V reuses output buffers. + AscendC::WaitFlag(0); + Compute(count); + + // The next MTE2 tile may reuse its buffers after all vector reads finish. + AscendC::SetFlag(0); + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + CopyOut(row, start, count); + AscendC::SetFlag(0); + } + } + + __aicore__ inline void CopyIn(int64_t row, + int64_t tableRow, + int64_t start, + uint32_t count) + { + AscendC::DataCopyExtParams xParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyExtParams fpParams{ + 1, static_cast(count * sizeof(float)), 0, 0, 0}; + AscendC::DataCopyPadExtParams xPad{false, 0, 0, 0}; + AscendC::DataCopyPadExtParams fpPad{false, 0, 0, 0}; + + const int64_t xBase = row * headDim_ + start; + if constexpr (std::is_same_v) { + AscendC::DataCopyPad(x1FpBuf_.Get(), xGm_[xBase], xParams, xPad); + AscendC::DataCopyPad( + x2FpBuf_.Get(), xGm_[xBase + halfDim_], xParams, xPad); + } else { + AscendC::DataCopyPad(x1InBuf_.Get(), xGm_[xBase], xParams, xPad); + AscendC::DataCopyPad(x2InBuf_.Get(), xGm_[xBase + halfDim_], xParams, xPad); + } + + const int64_t cacheBase = tableRow * halfDim_ + start; + AscendC::DataCopyPad(cosBuf_.Get(), cosGm_[cacheBase], fpParams, fpPad); + AscendC::DataCopyPad(sinBuf_.Get(), sinGm_[cacheBase], fpParams, fpPad); + } + + __aicore__ inline void Compute(uint32_t count) + { + AscendC::LocalTensor x1 = x1FpBuf_.Get(); + AscendC::LocalTensor x2 = x2FpBuf_.Get(); + if constexpr (!std::is_same_v) { + AscendC::Cast(x1, x1InBuf_.Get(), AscendC::RoundMode::CAST_NONE, count); + AscendC::Cast(x2, x2InBuf_.Get(), AscendC::RoundMode::CAST_NONE, count); + } + + AscendC::LocalTensor cos = cosBuf_.Get(); + AscendC::LocalTensor sin = sinBuf_.Get(); + AscendC::LocalTensor out1 = out1FpBuf_.Get(); + AscendC::LocalTensor out2 = out2FpBuf_.Get(); + AscendC::LocalTensor tmp = tmpFpBuf_.Get(); + + AscendC::Muls(sin, sin, sinSign_, count); + AscendC::Mul(out1, x1, cos, count); + AscendC::Mul(tmp, x2, sin, count); + AscendC::Sub(out1, out1, tmp, count); + AscendC::Mul(out2, x2, cos, count); + AscendC::Mul(tmp, x1, sin, count); + AscendC::Add(out2, out2, tmp, count); + + if constexpr (!std::is_same_v) { + AscendC::Cast( + out1Buf_.Get(), out1, AscendC::RoundMode::CAST_RINT, count); + AscendC::Cast( + out2Buf_.Get(), out2, AscendC::RoundMode::CAST_RINT, count); + } + } + + __aicore__ inline void CopyOut(int64_t row, int64_t start, uint32_t count) + { + AscendC::DataCopyExtParams outParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + const int64_t outBase = row * headDim_ + start; + if constexpr (std::is_same_v) { + AscendC::DataCopyPad(outGm_[outBase], out1FpBuf_.Get(), outParams); + AscendC::DataCopyPad( + outGm_[outBase + halfDim_], out2FpBuf_.Get(), outParams); + } else { + AscendC::DataCopyPad(outGm_[outBase], out1Buf_.Get(), outParams); + AscendC::DataCopyPad(outGm_[outBase + halfDim_], out2Buf_.Get(), outParams); + } + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor xGm_; + AscendC::GlobalTensor cosGm_; + AscendC::GlobalTensor sinGm_; + AscendC::GlobalTensor outGm_; + AscendC::TBuf x1InBuf_; + AscendC::TBuf x2InBuf_; + AscendC::TBuf x1FpBuf_; + AscendC::TBuf x2FpBuf_; + AscendC::TBuf cosBuf_; + AscendC::TBuf sinBuf_; + AscendC::TBuf out1FpBuf_; + AscendC::TBuf out2FpBuf_; + AscendC::TBuf tmpFpBuf_; + AscendC::TBuf out1Buf_; + AscendC::TBuf out2Buf_; + int64_t numRows_; + int64_t tableRows_; + int64_t headDim_; + int64_t halfDim_; + float sinSign_; +}; + +} // namespace + +extern "C" __global__ __vector__ void rope_apply_ascend_kernel_fp32( + GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR out, + int64_t numRows, int64_t tableRows, int64_t headDim, float sinSign) +{ + AscendC::TPipe pipe; + KernelRopeApply op(&pipe); + op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign); + op.Process(); +} + +extern "C" __global__ __vector__ void rope_apply_ascend_kernel_fp16( + GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR out, + int64_t numRows, int64_t tableRows, int64_t headDim, float sinSign) +{ + AscendC::TPipe pipe; + KernelRopeApply op(&pipe); + op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign); + op.Process(); +} + +extern "C" __global__ __vector__ void rope_apply_ascend_kernel_bf16( + GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR out, + int64_t numRows, int64_t tableRows, int64_t headDim, float sinSign) +{ + AscendC::TPipe pipe; + KernelRopeApply op(&pipe); + op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign); + op.Process(); +} + +torch::Tensor rope_apply_ascend_forward(torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin, + double sin_sign) +{ + TORCH_CHECK(x.is_privateuseone(), "rope: x must be on an NPU device"); + 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( + x.scalar_type() == at::kHalf || x.scalar_type() == at::kBFloat16 || + x.scalar_type() == at::kFloat, + "rope: x must be fp16, bf16, or fp32"); + TORCH_CHECK(cos.is_privateuseone() && sin.is_privateuseone(), + "rope: cos/sin must be on an NPU device"); + TORCH_CHECK(cos.device() == x.device() && sin.device() == x.device(), + "rope: x, cos, and sin must be on the same NPU device"); + TORCH_CHECK(cos.scalar_type() == at::kFloat && sin.scalar_type() == at::kFloat, + "rope: cos/sin must be fp32"); + TORCH_CHECK(cos.dim() == 2 && sin.dim() == 2, + "rope: cos/sin must be 2-D [table_rows, D/2]"); + TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(), + "rope: cos/sin must be contiguous"); + TORCH_CHECK(cos.sizes() == sin.sizes(), "rope: cos/sin shapes must match"); + + const int64_t numRows = x.size(0); + const int64_t headDim = x.size(1); + TORCH_CHECK(headDim > 0 && headDim % 2 == 0, + "rope: head_dim must be a positive even number"); + TORCH_CHECK(cos.size(1) == headDim / 2, + "rope: cos/sin last dimension must equal head_dim/2"); + + torch::Tensor out = at::empty_like(x); + if (numRows == 0) { + return out; + } + + const int64_t tableRows = cos.size(0); + TORCH_CHECK(tableRows > 0, "rope: cos/sin table must contain at least one row"); + TORCH_CHECK(numRows % tableRows == 0, + "rope: n_rows must be divisible by the cos/sin table row count"); + auto aclStream = c10_npu::getCurrentNPUStream().stream(true); + const uint32_t blockNum = static_cast(std::min(numRows, ROPE_MAX_BLOCKS)); + const float sign = static_cast(sin_sign); + + auto* xPtr = reinterpret_cast(x.mutable_data_ptr()); + auto* cosPtr = reinterpret_cast(cos.mutable_data_ptr()); + auto* sinPtr = reinterpret_cast(sin.mutable_data_ptr()); + auto* outPtr = reinterpret_cast(out.mutable_data_ptr()); + if (x.scalar_type() == at::kFloat) { + rope_apply_ascend_kernel_fp32<<>>( + xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sign); + } else if (x.scalar_type() == at::kHalf) { + rope_apply_ascend_kernel_fp16<<>>( + xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sign); + } else { + rope_apply_ascend_kernel_bf16<<>>( + xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sign); + } + return out; +} diff --git a/docs/operators/rope.md b/docs/operators/rope.md index cfc171b0..d3e3e49f 100644 --- a/docs/operators/rope.md +++ b/docs/operators/rope.md @@ -1,19 +1,18 @@ # RoPE RoPE applies rotary position embeddings to per-head query or key tensors. The -current implementation is a pure PyTorch reference operator for Issue #108 -ground-truth validation; it is not a fused CUDA or Triton kernel. - -This page documents the PyTorch baseline version. +project provides the pure PyTorch ground truth plus CUDA, Triton, and Ascend C +candidate backends using the same GPT-NeoX/Hugging Face rotate-half convention. ## Entry Point ```python from rl_engine.kernels.registry import kernel_registry +from rl_engine.kernels.ops.pytorch.rotary_embedding import NativeRoPEOp rope = kernel_registry.get_op("rope") output = rope.forward(x, positions, theta=1_000_000.0) -reference = rope.forward_fp32(x, positions, theta=1_000_000.0) +reference = NativeRoPEOp().forward_fp32(x, positions, theta=1_000_000.0) ``` The operator can also be imported directly: @@ -29,9 +28,12 @@ rope = NativeRoPEOp() | Backend | Wrapper | Native symbol | Notes | | --- | --- | --- | --- | | PyTorch native | `NativeRoPEOp` | None | Reference baseline for Qwen3-style RoPE. | +| Ascend C | `RoPEAscendOp` | `_C_npu.rope_apply_ascend` | FP16/BF16/FP32, FP32 rotation math, autograd through the inverse rotation. | +| CUDA SM90 | `RoPESM90Op` | `_C.rope_apply_sm90` | Hopper build only. | +| Triton | `TritonRoPEOp` | JIT kernel | CUDA/ROCm candidate. | -`kernel_registry.get_op("rope")` dispatches to the PyTorch native backend on CPU, -CUDA, and ROCm. CUDA/Triton fused RoPE kernels should compare against this reference. +`kernel_registry.get_op("rope")` prefers `RoPEAscendOp` on NPU, the SM90/Triton +candidates on CUDA, and the PyTorch implementation as the portable fallback. ## Tensor Contract @@ -93,6 +95,8 @@ as `[S]` and `[B, S]`, batch invariance, and Qwen3 query/key head shapes. ## Implementation Files - `rl_engine/kernels/ops/pytorch/rotary_embedding/rope.py` +- `rl_engine/kernels/ops/ascend/rotary_embedding/rope.py` +- `csrc/ascend/rope_ascend.asc` - `rl_engine/kernels/ops/pytorch/rotary_embedding/__init__.py` - `rl_engine/kernels/registry.py` - `tests/test_rope.py` diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..e990b6f4 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,10 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... + +def rope_apply_ascend( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + sin_sign: float, +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index ca4a462e..f56dd602 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -163,6 +163,7 @@ def _load_object(path: str) -> Any: "pytorch": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", "triton": "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp", "cuda-sm90": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op", + "ascend": "rl_engine.kernels.ops.ascend.rotary_embedding.rope.RoPEAscendOp", }, grad_input_names=("x",), ), diff --git a/rl_engine/kernels/ops/ascend/__init__.py b/rl_engine/kernels/ops/ascend/__init__.py index ab85458d..5b1eb253 100644 --- a/rl_engine/kernels/ops/ascend/__init__.py +++ b/rl_engine/kernels/ops/ascend/__init__.py @@ -2,3 +2,4 @@ # Copyright (c) 2026 RL-Kernel Contributors from . import loss # noqa: F401 +from . import rotary_embedding # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/rotary_embedding/__init__.py b/rl_engine/kernels/ops/ascend/rotary_embedding/__init__.py new file mode 100644 index 00000000..1c0317a2 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/rotary_embedding/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .rope import RoPEAscendOp diff --git a/rl_engine/kernels/ops/ascend/rotary_embedding/rope.py b/rl_engine/kernels/ops/ascend/rotary_embedding/rope.py new file mode 100644 index 00000000..89b93111 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/rotary_embedding/rope.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Ascend C RoPE backend (GPT-NeoX/Hugging Face rotate-half convention).""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import Tensor + +from rl_engine.utils.logger import logger + +_C_npu: Any = None +try: + from rl_engine import _C_npu + + _NPU_EXT_AVAILABLE = True +except ImportError: # pragma: no cover - Ascend extension not built + _NPU_EXT_AVAILABLE = False + + +def _build_cos_sin( + positions: Tensor, + half: int, + theta: float, + device: torch.device, +) -> tuple[Tensor, Tensor]: + """Build fp32 [table_rows, half] caches with the reference RoPE formula.""" + inv_freq = 1.0 / ( + theta ** (torch.arange(0, half, dtype=torch.float32, device=device) / half) + ) + freqs = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) * inv_freq + return freqs.cos().contiguous(), freqs.sin().contiguous() + + +def _rope_table( + x: Tensor, positions: Tensor, theta: float +) -> tuple[Tensor, Tensor, Tensor]: + """Flatten x so row modulo table length selects the correct position cache.""" + if x.dim() < 2: + raise ValueError( + f"x must have at least 2 dimensions, got shape {tuple(x.shape)}" + ) + dim = x.shape[-1] + if dim <= 0 or dim % 2 != 0: + raise ValueError(f"RoPE head_dim must be a positive even number, got {dim}") + + if positions.dim() == 1: + table_len = int(positions.shape[0]) + x_2d = x.contiguous().reshape(-1, dim) + if table_len == 0: + if x_2d.shape[0] != 0: + raise ValueError("positions cannot be empty when x contains rows") + elif x_2d.shape[0] % table_len != 0: + raise ValueError( + f"row count {x_2d.shape[0]} not divisible by seq length {table_len}; " + "expected a [..., S, D] contiguous layout." + ) + cos, sin = _build_cos_sin(positions, dim // 2, float(theta), x.device) + return x_2d, cos, sin + + if positions.dim() != 2: + raise ValueError( + f"positions must be [S] or [B, S], got shape {tuple(positions.shape)}" + ) + batch, seq = positions.shape + if x.shape[0] != batch or x.shape[-2] != seq: + raise ValueError( + f"positions {tuple(positions.shape)} is incompatible with x {tuple(x.shape)}; " + "expected x [B, ..., S, D]" + ) + if x.dim() == 4: + x_2d = x.permute(1, 0, 2, 3).contiguous().reshape(-1, dim) + elif x.dim() == 3: + x_2d = x.contiguous().reshape(-1, dim) + else: + raise ValueError( + f"RoPE [B, S] positions require x [B, S, D] or [B, H, S, D], got {x.dim()}D" + ) + + table_len = batch * seq + if table_len == 0: + if x_2d.shape[0] != 0: + raise ValueError("positions cannot be empty when x contains rows") + elif x_2d.shape[0] % table_len != 0: + raise ValueError(f"row count {x_2d.shape[0]} not divisible by B*S={table_len}") + cos, sin = _build_cos_sin(positions, dim // 2, float(theta), x.device) + return x_2d, cos, sin + + +def _restore_rope(out_2d: Tensor, x: Tensor, positions: Tensor) -> Tensor: + if positions.dim() == 1 or x.dim() != 4: + return out_2d.reshape(x.shape) + heads, batch, seq, dim = x.shape[1], x.shape[0], x.shape[2], x.shape[3] + return out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() + + +class _RoPEAscendFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: + x_2d, cos, sin = _rope_table(x, positions, theta) + ctx.save_for_backward(cos, sin) + ctx.x_shape = tuple(x.shape) + ctx.pos_dim = positions.dim() + out_2d = _C_npu.rope_apply_ascend(x_2d, cos, sin, 1.0) + return _restore_rope(out_2d, x, positions) + + @staticmethod + def backward(ctx, grad_out: Tensor): + cos, sin = ctx.saved_tensors + grad_x = None + if ctx.needs_input_grad[0]: + if ctx.pos_dim == 2 and len(ctx.x_shape) == 4: + grad_2d = ( + grad_out.permute(1, 0, 2, 3) + .contiguous() + .reshape(-1, ctx.x_shape[-1]) + ) + out_2d = _C_npu.rope_apply_ascend(grad_2d, cos, sin, -1.0) + heads, batch, seq, dim = ( + ctx.x_shape[1], + ctx.x_shape[0], + ctx.x_shape[2], + ctx.x_shape[3], + ) + grad_x = ( + out_2d.reshape(heads, batch, seq, dim) + .permute(1, 0, 2, 3) + .contiguous() + ) + else: + grad_2d = grad_out.contiguous().reshape(-1, grad_out.shape[-1]) + grad_x = _C_npu.rope_apply_ascend(grad_2d, cos, sin, -1.0).reshape( + grad_out.shape + ) + return grad_x, None, None + + +class RoPEAscendOp: + """Differentiable Ascend C RoPE backend for fp16, bf16, and fp32 inputs.""" + + op_class = "elementwise" + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "rope_apply_ascend"): + raise RuntimeError( + "rope_apply_ascend is not compiled into _C_npu. Rebuild on an Ascend host with " + "'KERNEL_ALIGN_FORCE_ASCEND=1 pip install --no-build-isolation -e .'." + ) + logger.info( + "Successfully linked to precompiled _C_npu.rope_apply_ascend kernel." + ) + + def __call__( + self, + x: Tensor, + positions: Tensor, + *, + theta: float = 1_000_000.0, + ) -> Tensor: + return self.forward(x, positions, theta=theta) + + def forward( + self, + x: Tensor, + positions: Tensor, + *, + theta: float = 1_000_000.0, + ) -> Tensor: + if x.device.type != "npu": + raise RuntimeError( + f"RoPEAscendOp requires an NPU tensor, got device '{x.device}'." + ) + if x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + f"RoPEAscendOp supports fp16, bf16, and fp32, got {x.dtype}." + ) + return _RoPEAscendFunction.apply(x, positions, float(theta)) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 7070728a..85f6002c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -151,6 +151,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op" + ASCEND_ROPE = "rl_engine.kernels.ops.ascend.rotary_embedding.rope.RoPEAscendOp" PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" CUDA_SILU = "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp" @@ -690,6 +691,10 @@ def __init__(self): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + self._priority_map["npu"]["rope"] = [ + OpBackend.ASCEND_ROPE, + OpBackend.PYTORCH_NATIVE_ROPE, + ] logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() self._adjust_priority_from_env() diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 388c4aec..dfc43d28 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -166,6 +166,10 @@ def fake_load_backend(backend): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + assert registry._priority_map["npu"]["rope"] == [ + OpBackend.ASCEND_ROPE, + OpBackend.PYTORCH_NATIVE_ROPE, + ] def test_npu_available_handles_runtime_failure(monkeypatch): diff --git a/setup.py b/setup.py index ee560db0..74c1c184 100644 --- a/setup.py +++ b/setup.py @@ -331,7 +331,13 @@ def _ascend_extensions(): asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) if not asc_srcs: raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") - return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] + sources: list[str] = asc_srcs + # Some kernels ship a C++ pybind host alongside the .asc sources; include + # it when present (compiled per-source by _bisheng_compile_cmd). + host_cpp = Path("csrc/ascend/npu_module.cpp") + if host_cpp.is_file(): + sources = [str(host_cpp), *asc_srcs] + return [Extension(name="rl_engine._C_npu", sources=sources, language="asc")] def _bisheng_compile_cmd(ext, ext_fullpath): @@ -366,6 +372,49 @@ def _bisheng_compile_cmd(ext, ext_fullpath): os.path.join(ascend_home, "lib64"), ] + if any(not str(src).endswith(".asc") for src in ext.sources): + # Mixed extension (C++ pybind host + .asc kernels): the -x asc driver + # cannot compile C++, so compile each source to an object and link them + # in a second step. + import subprocess + import tempfile + + build_temp = tempfile.mkdtemp(prefix="rl_kernel_ascend_") + objects = [] + for src in ext.sources: + src = str(src) + obj = os.path.join(build_temp, Path(src).name + ".o") + src_cmd = [ + "bisheng", + "-std=c++17", + "-O2", + "-fPIC", + "-c", + f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", + f"-DTORCH_EXTENSION_NAME={module_name}", + ] + if src.endswith(".asc"): + src_cmd += ["-x", "asc", f"--npu-arch={soc}"] + src_cmd += [f"-I{d}" for d in include_dirs if d] + src_cmd += [src, "-o", obj] + subprocess.check_call(src_cmd) + objects.append(obj) + cmd = [ + "bisheng", + "-shared", + *objects, + "-lascendcl", + "-ltorch_npu", + "-ltorch", + "-ltorch_cpu", + "-ltorch_python", + "-lc10", + "-o", + ext_fullpath, + ] + cmd += [f"-L{d}" for d in lib_dirs if d] + return cmd + cmd = [ "bisheng", "-x", diff --git a/tests/test_rope.py b/tests/test_rope.py index 67f4a294..38534483 100644 --- a/tests/test_rope.py +++ b/tests/test_rope.py @@ -15,6 +15,7 @@ import torch from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.platforms.device import _npu_available # --------------------------------------------------------------------------- # Fixtures & helpers @@ -74,7 +75,9 @@ def test_forward_fp32_returns_fp32_even_with_bf16_input(self): def test_call_equals_forward(self): op = NativeRoPEOp() x, pos = _make_inputs(2, 32, 16, QWEN3_HEAD_DIM) - assert torch.equal(op(x, pos, theta=QWEN3_THETA), op.forward(x, pos, theta=QWEN3_THETA)) + assert torch.equal( + op(x, pos, theta=QWEN3_THETA), op.forward(x, pos, theta=QWEN3_THETA) + ) def test_pure_function_no_inplace(self): op = NativeRoPEOp() @@ -143,7 +146,9 @@ def test_batch1_vs_batchN_bitwise(self): full_out = op.forward_fp32(x, pos) for i in range(x.shape[0]): single_out = op.forward_fp32(x[i : i + 1], pos) - assert torch.equal(full_out[i], single_out[0]), f"Batch invariance broken at row {i}" + assert torch.equal( + full_out[i], single_out[0] + ), f"Batch invariance broken at row {i}" def test_batch_invariance_with_padding(self): """Padded batch (extra rows) must not affect valid rows.""" @@ -230,7 +235,8 @@ def test_forward_vs_fp32_within_tolerance(self, dtype, atol, rtol): out_fp32 = op.forward_fp32(x_typed, pos) diff = (out_typed - out_fp32).abs().max().item() assert torch.allclose(out_typed, out_fp32, atol=atol, rtol=rtol), ( - f"dtype={dtype}, max_abs_error={diff:.3e} exceeds " f"atol={atol}, rtol={rtol}" + f"dtype={dtype}, max_abs_error={diff:.3e} exceeds " + f"atol={atol}, rtol={rtol}" ) @@ -285,7 +291,9 @@ def test_packed_logical_positions_match_per_sample_rope(self): assert not torch.equal(packed_out, naive_out) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="candidate RoPE requires CUDA") +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="candidate RoPE requires CUDA" +) class TestCandidateRoPELayouts: def _candidates(self): from rl_engine.kernels.ops.triton.rotary_embedding.rope import TritonRoPEOp @@ -328,3 +336,118 @@ def test_packed_logical_positions_on_candidates(self): for name, op in self._candidates(): got = op.forward(packed, packed_pos, theta=QWEN3_THETA).float() assert torch.allclose(got, gold, atol=2e-2, rtol=1.6e-2), name + + +# --------------------------------------------------------------------------- +# Ascend C candidate +# --------------------------------------------------------------------------- + + +def _ascend_rope_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine.kernels.ops.ascend.rotary_embedding.rope import ( + _C_npu, + _NPU_EXT_AVAILABLE, + ) + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "rope_apply_ascend") + + +requires_ascend_rope = pytest.mark.skipif( + not _ascend_rope_available(), + reason="rope_apply_ascend is not compiled (requires an Ascend NPU build).", +) + + +def _make_ascend_inputs(batch: int, heads: int, seq: int, dim: int, dtype: torch.dtype): + x, positions = _make_inputs(batch, heads, seq, dim, seed=2026) + return x.to(dtype=dtype, device="npu"), positions.to(device="npu") + + +def test_ascend_per_batch_table_layout_round_trips_without_an_npu(): + from rl_engine.kernels.ops.ascend.rotary_embedding.rope import ( + _restore_rope, + _rope_table, + ) + + x = torch.arange(2 * 3 * 4 * 8, dtype=torch.float32).reshape(2, 3, 4, 8) + positions = torch.stack([torch.arange(4), torch.arange(4) + 17]) + x_2d, cos, sin = _rope_table(x, positions, QWEN3_THETA) + + assert torch.equal(x_2d, x.permute(1, 0, 2, 3).reshape(-1, 8)) + assert cos.shape == sin.shape == (8, 4) + assert torch.equal(_restore_rope(x_2d, x, positions), x) + + +def test_ascend_per_batch_table_rejects_incompatible_shapes_without_an_npu(): + from rl_engine.kernels.ops.ascend.rotary_embedding.rope import _rope_table + + with pytest.raises(ValueError, match="incompatible"): + _rope_table( + torch.randn(2, 3, 4, 8), + torch.zeros(3, 4, dtype=torch.long), + QWEN3_THETA, + ) + + +@requires_ascend_rope +class TestRoPEAscend: + def _op(self): + from rl_engine.kernels.ops.ascend.rotary_embedding.rope import RoPEAscendOp + + return RoPEAscendOp() + + @pytest.mark.parametrize( + "dtype,atol,rtol", + [ + (torch.float32, 1e-5, 1e-5), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 2e-2, 1.6e-2), + ], + ) + def test_forward_matches_fp32_reference(self, dtype, atol, rtol): + x, positions = _make_ascend_inputs(2, 8, 17, 128, dtype) + actual = self._op()(x, positions, theta=QWEN3_THETA) + expected = NativeRoPEOp().forward_fp32(x, positions, theta=QWEN3_THETA) + assert actual.shape == x.shape + assert actual.dtype == dtype + assert torch.allclose(actual.float(), expected, atol=atol, rtol=rtol) + + def test_per_batch_positions_match_reference(self): + x, positions = _make_ascend_inputs(3, 8, 11, 128, torch.bfloat16) + positions = torch.stack([positions + batch * 97 for batch in range(x.shape[0])]) + actual = self._op()(x, positions, theta=QWEN3_THETA) + expected = NativeRoPEOp().forward_fp32(x, positions, theta=QWEN3_THETA) + assert torch.allclose(actual.float(), expected, atol=2e-2, rtol=1.6e-2) + + def test_batch_invariance_is_bitwise(self): + x, positions = _make_ascend_inputs(4, 8, 13, 128, torch.bfloat16) + full = self._op()(x, positions, theta=QWEN3_THETA) + for batch in range(x.shape[0]): + single = self._op()(x[batch : batch + 1], positions, theta=QWEN3_THETA) + assert torch.equal(full[batch], single[0]) + + def test_backward_matches_transposed_reference_rotation(self): + x, positions = _make_ascend_inputs(2, 4, 9, 128, torch.float32) + grad_out = torch.randn_like(x) + + actual_x = x.detach().clone().requires_grad_(True) + self._op()(actual_x, positions, theta=QWEN3_THETA).backward(grad_out) + + expected_x = x.detach().clone().requires_grad_(True) + NativeRoPEOp().forward_fp32(expected_x, positions, theta=QWEN3_THETA).backward( + grad_out + ) + assert actual_x.grad is not None + assert expected_x.grad is not None + assert torch.allclose(actual_x.grad, expected_x.grad, atol=1e-5, rtol=1e-5) + + def test_empty_batch(self): + x = torch.empty(0, 8, 0, 128, device="npu", dtype=torch.bfloat16) + positions = torch.empty(0, device="npu", dtype=torch.long) + out = self._op()(x, positions) + assert out.shape == x.shape + assert out.dtype == x.dtype