Skip to content

feat(ascend): add RoPE kernel #378

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
erfgss wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(ascend): add RoPE kernel #378

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9476a64
feat(ascend): add RoPE kernel
erfgss Sep 2, 2026
ef74179
refactor(ascend): rename ops_ascend.cpp to npu_module.cpp
erfgss Sep 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 0 additions & 7 deletions csrc/ascend/batch_invariant_logp_ascend.asc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -307,10 +307,3 @@ std::vector<torch::Tensor> 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)");
}
25 changes: 25 additions & 0 deletions csrc/ascend/npu_module.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 RL-Kernel Contributors

#include <torch/extension.h>

#include <vector>

std::vector<torch::Tensor> 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)
Comment thread
zhangj1an marked this conversation as resolved.
{
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)");
}
296 changes: 296 additions & 0 deletions csrc/ascend/rope_ascend.asc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 <algorithm>
#include <type_traits>

#include "kernel_operator.h"

#include <torch/extension.h>

#include "torch_npu/csrc/core/npu/NPUStream.h"

namespace {

constexpr uint32_t ROPE_TILE_HALF = 4096;
constexpr int64_t ROPE_MAX_BLOCKS = 128;

template <typename T>
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<AscendC::HardEvent::V_MTE2>(0);
AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(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<AscendC::HardEvent::V_MTE2>(0);
AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(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<uint32_t>(
remaining < ROPE_TILE_HALF ? remaining : ROPE_TILE_HALF);

// Previous vector reads are complete before MTE2 reuses input/cache buffers.
AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(0);
CopyIn(row, tableRow, start, count);
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(0);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(0);

// Previous MTE3 reads are complete before V reuses output buffers.
AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(0);
Compute(count);

// The next MTE2 tile may reuse its buffers after all vector reads finish.
AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(0);
AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0);
AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0);
CopyOut(row, start, count);
AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(0);
}
}

__aicore__ inline void CopyIn(int64_t row,
int64_t tableRow,
int64_t start,
uint32_t count)
{
AscendC::DataCopyExtParams xParams{
1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
AscendC::DataCopyExtParams fpParams{
1, static_cast<uint32_t>(count * sizeof(float)), 0, 0, 0};
AscendC::DataCopyPadExtParams<T> xPad{false, 0, 0, 0};
AscendC::DataCopyPadExtParams<float> fpPad{false, 0, 0, 0};

const int64_t xBase = row * headDim_ + start;
if constexpr (std::is_same_v<T, float>) {
AscendC::DataCopyPad(x1FpBuf_.Get<float>(), xGm_[xBase], xParams, xPad);
AscendC::DataCopyPad(
x2FpBuf_.Get<float>(), xGm_[xBase + halfDim_], xParams, xPad);
} else {
AscendC::DataCopyPad(x1InBuf_.Get<T>(), xGm_[xBase], xParams, xPad);
AscendC::DataCopyPad(x2InBuf_.Get<T>(), xGm_[xBase + halfDim_], xParams, xPad);
}

const int64_t cacheBase = tableRow * halfDim_ + start;
AscendC::DataCopyPad(cosBuf_.Get<float>(), cosGm_[cacheBase], fpParams, fpPad);
AscendC::DataCopyPad(sinBuf_.Get<float>(), sinGm_[cacheBase], fpParams, fpPad);
}

__aicore__ inline void Compute(uint32_t count)
{
AscendC::LocalTensor<float> x1 = x1FpBuf_.Get<float>();
AscendC::LocalTensor<float> x2 = x2FpBuf_.Get<float>();
if constexpr (!std::is_same_v<T, float>) {
AscendC::Cast(x1, x1InBuf_.Get<T>(), AscendC::RoundMode::CAST_NONE, count);
AscendC::Cast(x2, x2InBuf_.Get<T>(), AscendC::RoundMode::CAST_NONE, count);
}

AscendC::LocalTensor<float> cos = cosBuf_.Get<float>();
AscendC::LocalTensor<float> sin = sinBuf_.Get<float>();
AscendC::LocalTensor<float> out1 = out1FpBuf_.Get<float>();
AscendC::LocalTensor<float> out2 = out2FpBuf_.Get<float>();
AscendC::LocalTensor<float> tmp = tmpFpBuf_.Get<float>();

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<T, float>) {
AscendC::Cast(
out1Buf_.Get<T>(), out1, AscendC::RoundMode::CAST_RINT, count);
AscendC::Cast(
out2Buf_.Get<T>(), 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<uint32_t>(count * sizeof(T)), 0, 0, 0};
const int64_t outBase = row * headDim_ + start;
if constexpr (std::is_same_v<T, float>) {
AscendC::DataCopyPad(outGm_[outBase], out1FpBuf_.Get<float>(), outParams);
AscendC::DataCopyPad(
outGm_[outBase + halfDim_], out2FpBuf_.Get<float>(), outParams);
} else {
AscendC::DataCopyPad(outGm_[outBase], out1Buf_.Get<T>(), outParams);
AscendC::DataCopyPad(outGm_[outBase + halfDim_], out2Buf_.Get<T>(), outParams);
}
}

AscendC::TPipe* pipe_;
AscendC::GlobalTensor<T> xGm_;
AscendC::GlobalTensor<float> cosGm_;
AscendC::GlobalTensor<float> sinGm_;
AscendC::GlobalTensor<T> outGm_;
AscendC::TBuf<AscendC::TPosition::VECCALC> x1InBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> x2InBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> x1FpBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> x2FpBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> cosBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> sinBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> out1FpBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> out2FpBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> tmpFpBuf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> out1Buf_;
AscendC::TBuf<AscendC::TPosition::VECCALC> 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<float> 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<half> 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<bfloat16_t> 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);

@coderabbitai coderabbitai Bot Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🔎 Supported by static analysis

🏁 Script executed:

#!/bin/bash
# Run on a host with two NPUs and the PR extension installed.
python - <<'PY'
import sys
import torch
import torch_npu

from rl_engine.kernels.ops.ascend.rotary_embedding.rope import RoPEAscendOp
from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp

if torch.npu.device_count() < 2:
    sys.exit(0)

torch.npu.set_device(0)
x = torch.randn(1, 2, 7, 128, device="npu:1", dtype=torch.float16)
positions = torch.arange(7, device="npu:1", dtype=torch.long)

actual = RoPEAscendOp().forward(x, positions)
expected = NativeRoPEOp().forward(x, positions)
torch.npu.synchronize(1)
torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3)
PY

Repository: RL-Align/RL-Kernel

Length of output: 274

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- target source ---'
sed -n '235,305p' csrc/ascend/rope_ascend.asc

printf '%s\n' '--- nearby device/stream symbols ---'
rg -n -C 3 'getCurrentNPUStream|NPUGuard|DeviceGuard|set_device|x\.device|device\(' csrc/ascend csrc include 2>/dev/null | head -240

printf '%s\n' '--- repository files related to the launcher ---'
git ls-files | rg '(^|/)(rope_ascend\.asc|rope|npu|ascend)' | head -160

Repository: RL-Align/RL-Kernel

Length of output: 23243

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- Ascend source includes and declarations ---'
sed -n '1,90p' csrc/ascend/rope_ascend.asc

printf '%s\n' '--- Ascend operator binding and call path ---'
rg -n -C 5 'rope_apply_ascend_forward|rope_ascend_forward|rotary|rope' csrc/ascend/ops_ascend.cpp rl_engine/kernels/ops/ascend/rotary_embedding/rope.py

printf '%s\n' '--- NPU guard and stream API references in repository/configuration ---'
rg -n -C 4 'c10_npu|NPUGuard|NPU.*Guard|Optional.*NPU|DeviceGuard|getCurrentNPUStream' . -g '!*.md' -g '!*.pyc' | head -260

Repository: RL-Align/RL-Kernel

Length of output: 14455

🌐 Web query:

torch_npu c10_npu getCurrentNPUStream NPUGuardImpl current device source

💡 Result:

In the torch_npu architecture, NPUGuardImpl is an implementation of the DeviceGuardImplInterface that manages NPU device and stream contexts [1][2]. It is central to how PyTorch interacts with Ascend NPU hardware by providing the necessary logic to handle device switching and stream management [3][4]. Key technical relationships include: 1. NPUGuardImpl and Stream Access: The NPUGuardImpl class implements the getStream method by calling c10_npu::getCurrentNPUStream [5][6]. Specifically, it retrieves the current stream for a given device index: c10::Stream NPUGuardImpl::getStream(c10::Device d) const noexcept { return c10_npu::getCurrentNPUStream(d.index).unwrap; } 2. getCurrentNPUStream Implementation: This function is defined in torch_npu/csrc/core/npu/NPUStream.cpp [7][8]. It initializes NPU streams if necessary, determines the device index (defaulting to the current device if -1 is passed), and returns the current stream associated with that device from internal storage [7][8]. 3. Role of NPUGuard: NPUGuard provides RAII-style management for NPU device contexts, utilizing the NPUGuardImpl class to perform operations like setting the device or managing streams efficiently [9]. These components are part of the torch_npu internal core, which facilitates the integration of Ascend NPU backends into the broader PyTorch C10 infrastructure [3][4].

Citations:

🌐 Web query:

site:github.com/Ascend/pytorch "NPUGuard device_guard" "NPUStream.h"

💡 Result:

In the context of the Ascend PyTorch (torch_npu) codebase, NPUGuard and NPUStream.h are integral components for managing NPU (Neural Processing Unit) device contexts and asynchronous stream operations [1][2]. NPUGuard is a RAII (Resource Acquisition Is Initialization) utility class—similar in pattern to c10::DeviceGuard—used to automatically set the active NPU device for the duration of a scope [1][3]. When an NPUGuard object is instantiated with a specific device index (e.g., NPUGuard device_guard{device_index}), it ensures that subsequent operations within that scope are executed on the intended NPU device [1][2]. This is commonly used in initialization functions, stream management, and synchronization routines to prevent accidental execution on the wrong device [1][3][2]. NPUStream.h defines the interface and management logic for NPU streams, which are used to manage asynchronous execution queues on the NPU [1]. These streams allow for parallelized operations and are managed via pools within the torch_npu library [1][2]. Developers typically use these components together by: 1. Including NPUStream.h and NPUGuard.h where NPU-specific operations are performed [1][3][2]. 2. Using NPUGuard to switch to the target device context before creating, initializing, or synchronizing streams [1][2]. 3. Interfacing with NPUStream objects to handle command queues and stream synchronization [1][2]. Example usage pattern: NPUGuard device_guard{device_index}; // Operations here are bound to device_index auto stream = c10_npu::getCurrentNPUStream(device_index); These utilities are core to the torch_npu framework's ability to maintain compatibility with PyTorch's device management abstractions while supporting the specific requirements of Ascend NPUs [1][2].

Citations:

Guard x.device() before selecting the launch stream.

getCurrentNPUStream() uses the current NPU when called without a device argument. If the current NPU differs from x.device(), this launch can use a stream on the wrong NPU for the tensor pointers. Add an RAII NPUGuard for x.device() before line 277, or use the device-specific stream API.

const uint32_t blockNum = static_cast<uint32_t>(std::min(numRows, ROPE_MAX_BLOCKS));
const float sign = static_cast<float>(sin_sign);

auto* xPtr = reinterpret_cast<uint8_t*>(x.mutable_data_ptr());
auto* cosPtr = reinterpret_cast<uint8_t*>(cos.mutable_data_ptr());
auto* sinPtr = reinterpret_cast<uint8_t*>(sin.mutable_data_ptr());
auto* outPtr = reinterpret_cast<uint8_t*>(out.mutable_data_ptr());
if (x.scalar_type() == at::kFloat) {
rope_apply_ascend_kernel_fp32<<<blockNum, nullptr, aclStream>>>(
xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sign);
} else if (x.scalar_type() == at::kHalf) {
rope_apply_ascend_kernel_fp16<<<blockNum, nullptr, aclStream>>>(
xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sign);
} else {
rope_apply_ascend_kernel_bf16<<<blockNum, nullptr, aclStream>>>(
xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sign);
}
return out;
}
18 changes: 11 additions & 7 deletions docs/operators/rope.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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:
Expand All @@ -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

Expand Down Expand Up @@ -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`
Loading
Loading