From 77d9e431c26b2a72d33af95940d890abe2385a86 Mon Sep 17 00:00:00 2001 From: chenyang <2082464740@qq.com> Date: Sun, 13 Sep 2026 10:45:08 +0800 Subject: [PATCH] feat(ascend): add Qwen-Image qk_rmsnorm and multi_axis_rope kernels (#386) Qwen-Image WS1 single-GPU kernels, NPU (Ascend C) portion: - csrc/ascend/Qwen-Image/qk_rmsnorm_ascend.asc: per-head parameter-free QK RMSNorm (head_dim=128). Host computes rstd with the exact torch reference ops; the kernel only performs order-free fp32 elementwise multiplies plus one round-to-nearest-even cast, so the fused output is bitwise identical to the PyTorch reference and batch-invariant. - csrc/ascend/Qwen-Image/multi_axis_rope_ascend.asc: multi-axis RoPE (axes [16, 56, 56], text on the grid diagonal). The axis split and diagonal placement live in the fp32 cos/sin tables built host-side; the kernel is the dtype-generic rotate-half primitive with forward (sin_sign=+1) and backward (transpose, sin_sign=-1) entry points. - npu_module.cpp: register qk_rmsnorm_ascend, multi_axis_rope_ascend_forward/backward. - PyTorch references (NativeQkRmsNormOp, NativeMultiAxisRopeOp) with qwen_image_positions / build_multi_axis_cos_sin table builders. - Ascend autograd wrappers (QkRmsNormAscendOp, MultiAxisRopeAscendOp) with clean fallback to the native references. - registry: OpBackend entries + npu priority [ASCEND, PYTORCH_NATIVE] for op types qk_rmsnorm / multi_axis_rope; native-only candidates on cpu/cuda/rocm/musa. - tests/test_qwen_image_ops.py: CPU tests (bitwise reference match, per-head independence, batch invariance, explicit-VJP vs autograd, issue shapes 1024^2 / 1328^2 / 1664x928) and NPU tests (bitwise fwd, on-device batch invariance, backward tolerance) gated on the compiled Ascend kernels. CPU: 22 passed, 15 skipped (no NPU on the dev host). --- .../Qwen-Image/multi_axis_rope_ascend.asc | 349 +++++++++++++++++ csrc/ascend/Qwen-Image/qk_rmsnorm_ascend.asc | 361 ++++++++++++++++++ csrc/ascend/npu_module.cpp | 18 + .../kernels/ops/ascend/norm/qk_rmsnorm.py | 126 ++++++ .../ops/ascend/rotary_embedding/__init__.py | 1 + .../rotary_embedding/multi_axis_rope.py | 142 +++++++ .../kernels/ops/pytorch/norm/qk_rmsnorm.py | 55 +++ .../ops/pytorch/rotary_embedding/__init__.py | 3 +- .../rotary_embedding/multi_axis_rope.py | 166 ++++++++ rl_engine/kernels/registry.py | 42 ++ tests/test_qwen_image_ops.py | 334 ++++++++++++++++ 11 files changed, 1596 insertions(+), 1 deletion(-) create mode 100644 csrc/ascend/Qwen-Image/multi_axis_rope_ascend.asc create mode 100644 csrc/ascend/Qwen-Image/qk_rmsnorm_ascend.asc create mode 100644 rl_engine/kernels/ops/ascend/norm/qk_rmsnorm.py create mode 100644 rl_engine/kernels/ops/ascend/rotary_embedding/multi_axis_rope.py create mode 100644 rl_engine/kernels/ops/pytorch/norm/qk_rmsnorm.py create mode 100644 rl_engine/kernels/ops/pytorch/rotary_embedding/multi_axis_rope.py create mode 100644 tests/test_qwen_image_ops.py diff --git a/csrc/ascend/Qwen-Image/multi_axis_rope_ascend.asc b/csrc/ascend/Qwen-Image/multi_axis_rope_ascend.asc new file mode 100644 index 00000000..9c12c7ef --- /dev/null +++ b/csrc/ascend/Qwen-Image/multi_axis_rope_ascend.asc @@ -0,0 +1,349 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Qwen-Image MMDiT multi-axis RoPE, Ascend C (CANN) forward + backward kernels. +// +// Issue #386 kernel-table row `multi_axis_rope`: Qwen-Image splits the +// head_dim=128 rotary axes into [16, 56, 56] (temporal, height, width) and +// places text tokens on the grid diagonal. The axis split, the per-axis +// frequency bases, and the diagonal text placement all live in the position +// tables: the Python wrapper builds fp32 cos/sin caches of shape +// [table_rows, D/2] with the exact reference formula +// (per-axis inv_freq = theta^(-arange(0, dim, 2)/dim), concatenated over +// axes, coordinates per token, cat(cos, cos) / cat(sin, sin) for the +// rotate-half convention). This kernel is the dtype-generic rotate-half +// apply primitive over those tables; 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; sin_sign=-1 is its transpose R^T, +// which is exactly the backward rotation used for grad_x (R(theta)^T = +// R(-theta)), exposed as multi_axis_rope_ascend_backward. +// +// Bit-exactness and batch invariance follow the same contract as +// csrc/ascend/rope_ascend.asc: every row is processed end-to-end by exactly +// one AI core block with a fixed tile order over D/2, so adding or moving +// other batch rows cannot alter a row's instruction sequence, and the fp32 +// elementwise IEEE ops plus one round-to-nearest-even cast at the output +// cannot introduce drift against the reference apply. +// +// Mirrors the CUDA-side contract: +// - input : x [n_rows, D] contiguous, fp32 / bf16 / fp16; +// cos/sin [table_rows, D/2] fp32 (built by the Python wrapper) +// - output : y [n_rows, D] same dtype as x +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#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 KernelMultiAxisRope { +public: + __aicore__ inline KernelMultiAxisRope(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) + { + // row % tableRows_ selects the position cache: the wrapper flattens + // [B, H, S, D] so that the S index cycles over the table rows, + // covering the image (t, h, w) grid and the diagonal text positions + // alike — both are just rows of the cos/sin tables. + 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 multi_axis_rope_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; + KernelMultiAxisRope op(&pipe); + op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign); + op.Process(); +} + +extern "C" __global__ __vector__ void multi_axis_rope_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; + KernelMultiAxisRope op(&pipe); + op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign); + op.Process(); +} + +extern "C" __global__ __vector__ void multi_axis_rope_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; + KernelMultiAxisRope op(&pipe); + op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign); + op.Process(); +} + +namespace { + +// Shared apply path for forward (sin_sign=+1) and backward (sin_sign=-1, +// the transpose rotation). Keeps the TORCH_CHECK gate and launch geometry in +// one place so both directions are byte-identical code paths. +torch::Tensor MultiAxisRopeApply(torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin, + float sinSign) +{ + TORCH_CHECK(x.is_privateuseone(), "multi_axis_rope: x must be on an NPU device"); + TORCH_CHECK(x.dim() == 2, "multi_axis_rope: x must be 2-D [n_rows, D]"); + TORCH_CHECK(x.is_contiguous(), "multi_axis_rope: x must be contiguous"); + TORCH_CHECK( + x.scalar_type() == at::kHalf || x.scalar_type() == at::kBFloat16 || + x.scalar_type() == at::kFloat, + "multi_axis_rope: x must be fp16, bf16, or fp32"); + TORCH_CHECK(cos.is_privateuseone() && sin.is_privateuseone(), + "multi_axis_rope: cos/sin must be on an NPU device"); + TORCH_CHECK(cos.device() == x.device() && sin.device() == x.device(), + "multi_axis_rope: x, cos, and sin must be on the same NPU device"); + TORCH_CHECK(cos.scalar_type() == at::kFloat && sin.scalar_type() == at::kFloat, + "multi_axis_rope: cos/sin must be fp32"); + TORCH_CHECK(cos.dim() == 2 && sin.dim() == 2, + "multi_axis_rope: cos/sin must be 2-D [table_rows, D/2]"); + TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(), + "multi_axis_rope: cos/sin must be contiguous"); + TORCH_CHECK(cos.sizes() == sin.sizes(), + "multi_axis_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, + "multi_axis_rope: head_dim must be a positive even number"); + TORCH_CHECK(cos.size(1) == headDim / 2, + "multi_axis_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, + "multi_axis_rope: cos/sin table must contain at least one row"); + TORCH_CHECK(numRows % tableRows == 0, + "multi_axis_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)); + + 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) { + multi_axis_rope_ascend_kernel_fp32<<>>( + xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sinSign); + } else if (x.scalar_type() == at::kHalf) { + multi_axis_rope_ascend_kernel_fp16<<>>( + xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sinSign); + } else { + multi_axis_rope_ascend_kernel_bf16<<>>( + xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sinSign); + } + return out; +} + +} // namespace + +torch::Tensor multi_axis_rope_ascend_forward(torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin) +{ + return MultiAxisRopeApply(std::move(x), std::move(cos), std::move(sin), 1.0f); +} + +torch::Tensor multi_axis_rope_ascend_backward(torch::Tensor grad, + torch::Tensor cos, + torch::Tensor sin) +{ + // grad_x = R^T(grad) = rotate with the negated sin table. + return MultiAxisRopeApply(std::move(grad), std::move(cos), std::move(sin), -1.0f); +} diff --git a/csrc/ascend/Qwen-Image/qk_rmsnorm_ascend.asc b/csrc/ascend/Qwen-Image/qk_rmsnorm_ascend.asc new file mode 100644 index 00000000..45197276 --- /dev/null +++ b/csrc/ascend/Qwen-Image/qk_rmsnorm_ascend.asc @@ -0,0 +1,361 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Qwen-Image MMDiT per-head QK RMSNorm, Ascend C (CANN) forward kernel. +// +// Issue #386 kernel-table row `qk_rmsnorm`: parameter-free RMSNorm applied +// independently to Q and K, per attention head, over head_dim = 128: +// +// y[r, :] = x[r, :] * rstd[r] +// rstd[r] = rsqrt(mean(x[r, :]^2) + eps) (precomputed on the host side) +// +// The caller flattens [B, S, H, D] to rows = B*S*H rows of D = head_dim and +// computes rstd with the exact PyTorch ops of the FP32 reference +// (x.float().pow(2).mean(-1) followed by torch.rsqrt(var + eps)), the same +// contract as csrc/ascend/rmsnorm_ascend.asc. Keeping the reduction and +// rsqrt on that identical code path makes the fused result bitwise identical +// to the reference: this kernel only performs elementwise fp32 multiplies +// (order-free IEEE ops) and a round-to-nearest-even cast, so no in-kernel +// reduction order or approximate rsqrt can introduce drift. +// +// Q and K use the same code path: the caller normalizes Q and K with two +// invocations (each with its own rstd), so a row's bytes never depend on +// whether it belongs to Q, K, or to any other batch element. +// +// Backward: the VJP is computed by the Python autograd wrapper in fp32 with +// the forward-saved rstd (dx = rstd*dy - x*rstd^3*sum(dy*x)/D), mirroring +// how the CUDA/PyTorch backends handle RMSNorm backward; no kernel needed. +// +// Mirrors the CUDA-side batch-invariance contract: +// - input : x [rows, D] contiguous, fp32 / bf16 / fp16; +// rstd [rows] fp32 (saved by the caller for the autograd backward) +// - output : y [rows, D] same dtype as x +// +// Batch invariance: rstd[r] depends only on row r (torch's last-dim mean +// order is a function of D alone), and every row is processed end-to-end by +// exactly one AI core block with a fixed tile size. The instruction sequence +// for a row depends only on D, never on the row count or on the block the +// row lands on. +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#include +#include + +#include "kernel_operator.h" + +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +// Elements per hidden tile. Fixed for all rows and batch sizes; this is what +// keeps the elementwise pass batch-invariant. Qwen-Image uses head_dim=128, +// which fits a single tile with room for wider head dims up to 4096. +constexpr uint32_t TILE_LENGTH = 4096; +// Cap on launched blocks. Rows are strided across blocks, so launching fewer +// blocks than rows is fine and never changes per-row numerics. +constexpr int64_t MAX_BLOCKS = 128; +// Cap on rows coalesced into one tile in the small-row path. Bounds the +// rstd staging buffer. +constexpr int64_t MAX_CHUNK_ROWS = 64; + +template +class KernelQkRmsNorm { +public: + __aicore__ inline KernelQkRmsNorm(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR x, + GM_ADDR rstd, + GM_ADDR y, + int64_t numRows, + int64_t headDim) + { + numRows_ = numRows; + headDim_ = headDim; + // Small rows (head_dim=128 is the Qwen-Image case) are dominated by + // per-row pipeline flag round-trips, so process rowsPerChunk_ + // contiguous rows per iteration and amortize the syncs. Vector ops + // require 32 B-aligned addresses, hence the D % 8 == 0 gate (offset + // r * D elements stays aligned for every r). + rowsPerChunk_ = 1; + if (headDim % 8 == 0) { + rowsPerChunk_ = TILE_LENGTH / headDim; + if (rowsPerChunk_ > MAX_CHUNK_ROWS) { + rowsPerChunk_ = MAX_CHUNK_ROWS; + } + } + xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x)); + rstdGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(rstd)); + yGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(y)); + pipe_->InitBuffer(inQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(outQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(fp32Buf_, TILE_LENGTH * sizeof(float)); + // 2 KB: rstd staging slots for up to MAX_CHUNK_ROWS chunk rows (the + // single-row path uses slot 0 only). Scalar-unit reads are 4-byte + // granular, so contiguous staging is fine. + pipe_->InitBuffer(scalarBuf_, MAX_CHUNK_ROWS * 8 * sizeof(float)); + + // Intra-core pipeline events. NOTE: AscendC::SyncAll() is a cross-core + // barrier and deadlocks when more blocks are launched than there are + // physical cores, so all synchronization here uses per-pipe + // SetFlag/WaitFlag instead. + eventSMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE2); + eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S); + } + + __aicore__ inline void Process() + { + if (rowsPerChunk_ > 1) { + // Contiguous row segment per block: chunk coalescing needs the + // rows of a chunk to be contiguous in GM. + const int64_t blockNum = AscendC::GetBlockNum(); + const int64_t segLen = (numRows_ + blockNum - 1) / blockNum; + const int64_t rowStart = AscendC::GetBlockIdx() * segLen; + const int64_t rowEnd = + (rowStart + segLen < numRows_) ? rowStart + segLen : numRows_; + for (int64_t row = rowStart; row < rowEnd; row += rowsPerChunk_) { + const int64_t remaining = rowEnd - row; + const int64_t chunkRows = + remaining < rowsPerChunk_ ? remaining : rowsPerChunk_; + LoadRstd(row, chunkRows); + ProcessRowChunk(row, chunkRows); + } + return; + } + for (int64_t row = AscendC::GetBlockIdx(); row < numRows_; row += AscendC::GetBlockNum()) { + LoadRstd(row, 1); + ProcessRow(row, scalarBuf_.Get().GetValue(0)); + } + } + +private: + // Load x[row, :] into UB and return its fp32 view. When T is fp32 the + // queue buffer is used in place; otherwise the tile is cast into fp32Buf_. + __aicore__ inline AscendC::LocalTensor LoadTileFp32(int64_t row) + { + AscendC::LocalTensor xLocal = inQueue_.AllocTensor(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(headDim_ * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(xLocal, xGm_[row * headDim_], copyParams, padParams); + inQueue_.EnQue(xLocal); + xLocal = inQueue_.DeQue(); + inTile_ = xLocal; + + if constexpr (std::is_same_v) { + return xLocal; + } else { + AscendC::LocalTensor fLocal = fp32Buf_.Get(); + AscendC::Cast(fLocal, xLocal, AscendC::RoundMode::CAST_NONE, static_cast(headDim_)); + return fLocal; + } + } + + __aicore__ inline void FreeTile() + { + inQueue_.FreeTensor(inTile_); + } + + // Load rstd[row0 : row0+rows] into scalarBuf_[0:rows]. + __aicore__ inline void LoadRstd(int64_t row0, int64_t rows) + { + AscendC::LocalTensor scalar = scalarBuf_.Get(); + // Drain the previous chunk's scalar reads before MTE2 overwrites the + // staging slots (scalar unit vs MTE2 are async engines). + AscendC::SetFlag(eventSMTE2_); + AscendC::WaitFlag(eventSMTE2_); + AscendC::DataCopyExtParams inParams{ + 1, static_cast(rows * sizeof(float)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(scalar, rstdGm_[row0], inParams, padParams); + AscendC::SetFlag(eventMTE2S_); // copy-in -> scalar read + AscendC::WaitFlag(eventMTE2S_); + } + + __aicore__ inline void ProcessRow(int64_t row, float rstd) + { + // Fast path: the whole head stays resident in fp32Buf_, so x is read + // from GM only once. + const uint32_t count = static_cast(headDim_); + AscendC::LocalTensor fLocal = LoadTileFp32(row); + ScaleStoreTile(row, fLocal, rstd); + FreeTile(); + } + + // Chunk path for small rows: R contiguous rows are loaded, scaled and + // stored as one flat tile, with a single sync round-trip per chunk. + // Per-row numerics are identical to ProcessRow (elementwise fp32 ops on + // the row's tile with the row's rstd), so results are unchanged. + __aicore__ inline void ProcessRowChunk(int64_t row0, int64_t rows) + { + const int64_t D = headDim_; + const uint32_t count = static_cast(rows * D); + + // Rows [row0, row0+rows) are contiguous in GM: one flat copy-in. + AscendC::LocalTensor xLocal = inQueue_.AllocTensor(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(xLocal, xGm_[row0 * D], copyParams, padParams); + inQueue_.EnQue(xLocal); + xLocal = inQueue_.DeQue(); + inTile_ = xLocal; + + AscendC::LocalTensor fLocal = fp32Buf_.Get(); + if constexpr (std::is_same_v) { + fLocal = xLocal; + } else { + AscendC::Cast(fLocal, xLocal, AscendC::RoundMode::CAST_NONE, count); + } + + // Scale every row of the chunk with its own rstd (staged by LoadRstd + // before this call); scalar-register operands of Muls need no S_V + // flag (no UB dependency). + AscendC::LocalTensor scalar = scalarBuf_.Get(); + for (int64_t r = 0; r < rows; ++r) { + const float rstd = scalar.GetValue(static_cast(r)); + AscendC::Muls(fLocal[r * D], fLocal[r * D], rstd, static_cast(D)); + } + + AscendC::LocalTensor yLocal = outQueue_.AllocTensor(); + if constexpr (std::is_same_v) { + CopyFp32(yLocal, fLocal, count); + } else { + AscendC::Cast(yLocal, fLocal, AscendC::RoundMode::CAST_RINT, count); + } + outQueue_.EnQue(yLocal); + yLocal = outQueue_.DeQue(); + AscendC::DataCopyExtParams outParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPad(yGm_[row0 * D], yLocal, outParams); + outQueue_.FreeTensor(yLocal); + FreeTile(); + } + + // Contiguous fp32 UB -> UB copy (64 elements per 256 B vector repeat). + __aicore__ inline void CopyFp32(const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + uint32_t count) + { + const uint8_t repeat = static_cast((count + 63) / 64); + AscendC::Copy(dst, src, 64, repeat, AscendC::CopyRepeatParams{1, 1, 8, 8}); + } + + // y tile = x tile * rstd, cast back to T and copied to GM. + __aicore__ inline void ScaleStoreTile(int64_t row, + AscendC::LocalTensor& fLocal, + float rstd) + { + AscendC::Muls(fLocal, fLocal, rstd, static_cast(headDim_)); + + AscendC::LocalTensor yLocal = outQueue_.AllocTensor(); + if constexpr (std::is_same_v) { + CopyFp32(yLocal, fLocal, static_cast(headDim_)); + } else { + AscendC::Cast(yLocal, fLocal, AscendC::RoundMode::CAST_RINT, static_cast(headDim_)); + } + outQueue_.EnQue(yLocal); + yLocal = outQueue_.DeQue(); + AscendC::DataCopyExtParams outParams{ + 1, static_cast(headDim_ * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPad(yGm_[row * headDim_], yLocal, outParams); + outQueue_.FreeTensor(yLocal); + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor xGm_; + AscendC::GlobalTensor rstdGm_; + AscendC::GlobalTensor yGm_; + AscendC::TQue inQueue_; + AscendC::TQue outQueue_; + AscendC::TBuf fp32Buf_; + AscendC::TBuf scalarBuf_; + AscendC::LocalTensor inTile_; + AscendC::TEventID eventSMTE2_; + AscendC::TEventID eventMTE2S_; + int64_t numRows_; + int64_t headDim_; + int64_t rowsPerChunk_; +}; + +} // namespace + +extern "C" __global__ __vector__ void qk_rmsnorm_ascend_kernel_fp32( + GM_ADDR x, GM_ADDR rstd, GM_ADDR y, int64_t numRows, int64_t headDim) +{ + AscendC::TPipe pipe; + KernelQkRmsNorm op(&pipe); + op.Init(x, rstd, y, numRows, headDim); + op.Process(); +} + +extern "C" __global__ __vector__ void qk_rmsnorm_ascend_kernel_bf16( + GM_ADDR x, GM_ADDR rstd, GM_ADDR y, int64_t numRows, int64_t headDim) +{ + AscendC::TPipe pipe; + KernelQkRmsNorm op(&pipe); + op.Init(x, rstd, y, numRows, headDim); + op.Process(); +} + +extern "C" __global__ __vector__ void qk_rmsnorm_ascend_kernel_fp16( + GM_ADDR x, GM_ADDR rstd, GM_ADDR y, int64_t numRows, int64_t headDim) +{ + AscendC::TPipe pipe; + KernelQkRmsNorm op(&pipe); + op.Init(x, rstd, y, numRows, headDim); + op.Process(); +} + +torch::Tensor qk_rmsnorm_ascend_forward(torch::Tensor x, torch::Tensor rstd) +{ + TORCH_CHECK(x.is_privateuseone(), "qk_rmsnorm: x must be on an NPU device"); + TORCH_CHECK(x.dim() == 2, "qk_rmsnorm: x must be 2-D [rows, head_dim]"); + TORCH_CHECK(x.is_contiguous(), "qk_rmsnorm: x must be contiguous"); + TORCH_CHECK(x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kFloat || + x.scalar_type() == at::kHalf, + "qk_rmsnorm: x must be fp32, bf16 or fp16"); + TORCH_CHECK(rstd.is_privateuseone(), "qk_rmsnorm: rstd must be on the same NPU device as x"); + TORCH_CHECK(rstd.dim() == 1 && rstd.numel() == x.size(0), + "qk_rmsnorm: rstd must be 1-D of size x.size(0)"); + TORCH_CHECK(rstd.scalar_type() == at::kFloat, "qk_rmsnorm: rstd must be fp32"); + TORCH_CHECK(rstd.is_contiguous(), "qk_rmsnorm: rstd must be contiguous"); + + const int64_t numRows = x.size(0); + const int64_t headDim = x.size(1); + TORCH_CHECK(headDim > 0, "qk_rmsnorm: head_dim must be positive"); + TORCH_CHECK(headDim % 8 == 0, + "qk_rmsnorm: head_dim must be a multiple of 8 for 32 B-aligned " + "vector ops (Qwen-Image uses head_dim=128)"); + TORCH_CHECK(headDim <= TILE_LENGTH, + "qk_rmsnorm: head_dim must fit a single tile of ", TILE_LENGTH); + + torch::Tensor y = at::empty_like(x); + if (numRows == 0) { + return y; + } + + // stream(true): flush the task queue before launch so the kernel cannot + // overtake earlier NPU work; outputs were allocated with at::empty (no + // queued initializer) for the same reason. + auto aclStream = c10_npu::getCurrentNPUStream().stream(true); + const uint32_t blockNum = static_cast(std::min(numRows, MAX_BLOCKS)); + + auto* xPtr = reinterpret_cast(x.mutable_data_ptr()); + auto* rstdPtr = reinterpret_cast(rstd.data_ptr()); + auto* yPtr = reinterpret_cast(y.mutable_data_ptr()); + if (x.scalar_type() == at::kFloat) { + qk_rmsnorm_ascend_kernel_fp32<<>>( + xPtr, rstdPtr, yPtr, numRows, headDim); + } else if (x.scalar_type() == at::kHalf) { + qk_rmsnorm_ascend_kernel_fp16<<>>( + xPtr, rstdPtr, yPtr, numRows, headDim); + } else { + qk_rmsnorm_ascend_kernel_bf16<<>>( + xPtr, rstdPtr, yPtr, numRows, headDim); + } + return y; +} diff --git a/csrc/ascend/npu_module.cpp b/csrc/ascend/npu_module.cpp index 9c1e91f8..e457f0b9 100644 --- a/csrc/ascend/npu_module.cpp +++ b/csrc/ascend/npu_module.cpp @@ -57,6 +57,15 @@ torch::Tensor rmsnorm_ascend_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor rstd); +torch::Tensor qk_rmsnorm_ascend_forward(torch::Tensor x, torch::Tensor rstd); + +torch::Tensor multi_axis_rope_ascend_forward(torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin); +torch::Tensor multi_axis_rope_ascend_backward(torch::Tensor grad, + torch::Tensor cos, + torch::Tensor sin); + int64_t deterministic_collective_create( torch::Tensor staging, int64_t world_size, int64_t rank); void deterministic_collective_destroy(int64_t handle); @@ -93,6 +102,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("rmsnorm_ascend", &rmsnorm_ascend_forward, "Batch-invariant RMSNorm (Ascend C forward, rstd precomputed)"); + m.def("qk_rmsnorm_ascend", + &qk_rmsnorm_ascend_forward, + "Qwen-Image per-head QK RMSNorm (Ascend C forward, rstd precomputed)"); + m.def("multi_axis_rope_ascend_forward", + &multi_axis_rope_ascend_forward, + "Qwen-Image multi-axis RoPE rotate-half apply (Ascend C forward)"); + m.def("multi_axis_rope_ascend_backward", + &multi_axis_rope_ascend_backward, + "Qwen-Image multi-axis RoPE gradient (transpose rotation, Ascend C)"); m.def("embedding_ascend", &embedding_ascend_forward, "Batch-invariant token embedding (Ascend C forward)"); diff --git a/rl_engine/kernels/ops/ascend/norm/qk_rmsnorm.py b/rl_engine/kernels/ops/ascend/norm/qk_rmsnorm.py new file mode 100644 index 00000000..0c50f680 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/norm/qk_rmsnorm.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Any + +import torch + +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 _ascend_supported(x: torch.Tensor) -> bool: + """Whether the Ascend C forward can run this input directly. + + NPU tensors only, fp32/bf16/fp16 only (mirrors the CUDA kernel's gate). + """ + return x.device.type == "npu" and x.dtype in ( + torch.float32, + torch.bfloat16, + torch.float16, + ) + + +def _fallback_op(): + """Portable op for inputs the Ascend forward cannot take. + + Triton rejects non-CUDA devices, so on NPU the only fallback is native. + """ + from rl_engine.kernels.ops.pytorch.norm.qk_rmsnorm import NativeQkRmsNormOp + + return NativeQkRmsNormOp() + + +def _qk_rms_norm_backward( + x_2d: torch.Tensor, + rstd: torch.Tensor, + grad_out_2d: torch.Tensor, +) -> torch.Tensor: + """Parameter-free per-head RMSNorm VJP in fp32, reusing the saved rstd. + + With y = x * rstd and s = sum(dy * x, dim=-1): + dx = rstd * dy - x * rstd^3 * s / D + """ + dy_f = grad_out_2d.float() + x_f = x_2d.float() + rstd_f = rstd.float() + + head_dim = x_2d.size(-1) + s = (dy_f * x_f).sum(dim=-1, keepdim=True) + dx = rstd_f.unsqueeze(-1) * dy_f - x_f * (rstd_f.pow(3) / head_dim) * s + return dx.to(x_2d.dtype) + + +class _QkRmsNormAscendFunction(torch.autograd.Function): + # Autograd wrapper: reference-formula rstd + Ascend C fused scale/cast + # forward, and the PyTorch-formula backward reusing the forward-saved + # rstd (same fp32 VJP as the PyTorch reference, like the CUDA op). + + @staticmethod + def forward(ctx, x, eps): + lead_shape = x.shape[:-1] + head_dim = x.size(-1) + + x_2d = x.reshape(-1, head_dim).contiguous() + + # rstd is computed with the exact torch ops of the PyTorch reference + # (rl_engine/kernels/ops/pytorch/norm/qk_rmsnorm.py): fp32 mean of + # squares + torch.rsqrt. The Ascend C kernel then only performs the + # elementwise y = x * rstd scale and the round-to-nearest-even cast, + # which are order-free IEEE ops — this makes the fused output bitwise + # identical to NativeQkRmsNormOp instead of approximating its + # sum-of-squares/rsqrt arithmetic in-kernel. + x_f = x_2d.float() + var = x_f.pow(2).mean(dim=-1) + rstd = torch.rsqrt(var + eps).contiguous() + + y = _C_npu.qk_rmsnorm_ascend(x_2d, rstd) + + ctx.save_for_backward(x_2d, rstd) + ctx.lead_shape = lead_shape + ctx.head_dim = head_dim + return y.reshape(lead_shape + (head_dim,)) + + @staticmethod + def backward(ctx, grad_output): + x_2d, rstd = ctx.saved_tensors + + grad_out_2d = grad_output.reshape(-1, ctx.head_dim).contiguous() + dx = _qk_rms_norm_backward(x_2d, rstd, grad_out_2d) + + dx = dx.reshape(ctx.lead_shape + (ctx.head_dim,)) + return dx, None + + +class QkRmsNormAscendOp: + # Ascend C batch-invariant Qwen-Image per-head QK RMSNorm (forward kernel). + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "qk_rmsnorm_ascend"): + raise RuntimeError( + "qk_rmsnorm_ascend is not compiled into the extension. Rebuild with " + "KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host: 'pip install -e .'" + ) + logger.info("Successfully linked to precompiled _C_npu.qk_rmsnorm_ascend kernel.") + + def __call__(self, x: torch.Tensor, *, eps: float = 1e-6) -> torch.Tensor: + return self.forward(x, eps=eps) + + def forward(self, x: torch.Tensor, *, eps: float = 1e-6) -> torch.Tensor: + if not _ascend_supported(x): + return _fallback_op()(x, eps=eps) + + return _QkRmsNormAscendFunction.apply(x, eps) + + +def qk_rmsnorm_ascend(x: torch.Tensor, *, eps: float = 1e-6) -> torch.Tensor: + return QkRmsNormAscendOp()(x, eps=eps) diff --git a/rl_engine/kernels/ops/ascend/rotary_embedding/__init__.py b/rl_engine/kernels/ops/ascend/rotary_embedding/__init__.py index 1c0317a2..29416acb 100644 --- a/rl_engine/kernels/ops/ascend/rotary_embedding/__init__.py +++ b/rl_engine/kernels/ops/ascend/rotary_embedding/__init__.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from .multi_axis_rope import MultiAxisRopeAscendOp from .rope import RoPEAscendOp diff --git a/rl_engine/kernels/ops/ascend/rotary_embedding/multi_axis_rope.py b/rl_engine/kernels/ops/ascend/rotary_embedding/multi_axis_rope.py new file mode 100644 index 00000000..e52cb466 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/rotary_embedding/multi_axis_rope.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Ascend C Qwen-Image multi-axis RoPE backend (rotate-half convention). + +Issue #386 `multi_axis_rope`: axes (16, 56, 56), text tokens on the grid +diagonal. The axis split and diagonal text placement live in the position +coordinates and the fp32 cos/sin tables built from them (see +``qwen_image_positions`` / ``build_multi_axis_cos_sin`` in the PyTorch +reference module); this backend applies the rotate-half rotation with the +Ascend C kernel, forward and backward. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.pytorch.rotary_embedding.multi_axis_rope import ( + QWEN_IMAGE_AXES_DIM, + build_multi_axis_cos_sin, +) +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 _fallback_op(): + """Portable op for inputs the Ascend forward cannot take. + + Triton rejects non-CUDA devices, so on NPU the only fallback is native. + """ + from rl_engine.kernels.ops.pytorch.rotary_embedding.multi_axis_rope import ( + NativeMultiAxisRopeOp, + ) + + return NativeMultiAxisRopeOp() + + +def _rope_rows(x: Tensor, positions: Tensor) -> tuple[Tensor, int]: + """Flatten x [..., S, D] to [n_rows, D] with row % S selecting the table row. + + Contiguous [..., S, D] flattening orders rows as consecutive S-blocks, so + the modulo addressing of the Ascend C kernel lands every token on its own + position row — the same contract as the single-sequence path of the + generic RoPE wrapper. + """ + if x.dim() < 2: + raise ValueError( + f"x must have at least 2 dimensions, got shape {tuple(x.shape)}" + ) + dim = x.shape[-1] + seq = positions.shape[-2] if positions.dim() >= 2 else positions.shape[-1] + if dim != sum(QWEN_IMAGE_AXES_DIM): + raise ValueError( + f"x head_dim {dim} must equal sum(axes_dim)={sum(QWEN_IMAGE_AXES_DIM)}" + ) + if seq == 0: + if x.numel() != 0: + raise ValueError("positions cannot be empty when x contains rows") + elif x.numel() // dim % seq != 0: + raise ValueError( + f"row count {x.numel() // dim} not divisible by seq length {seq}; " + "expected a [..., S, D] contiguous layout." + ) + return x.contiguous().reshape(-1, dim), seq + + +class _MultiAxisRopeAscendFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: + if positions.dim() == 3 and positions.shape[0] == 1: + positions = positions[0] + x_2d, seq = _rope_rows(x, positions) + cos, sin = build_multi_axis_cos_sin(positions, QWEN_IMAGE_AXES_DIM, theta=theta) + ctx.save_for_backward(cos, sin) + ctx.x_shape = tuple(x.shape) + out_2d = _C_npu.multi_axis_rope_ascend_forward(x_2d, cos, sin) + return out_2d.reshape(ctx.x_shape) + + @staticmethod + def backward(ctx, grad_out: Tensor): + cos, sin = ctx.saved_tensors + grad_x = None + if ctx.needs_input_grad[0]: + grad_2d = grad_out.contiguous().reshape(-1, grad_out.shape[-1]) + grad_x = _C_npu.multi_axis_rope_ascend_backward(grad_2d, cos, sin).reshape( + ctx.x_shape + ) + return grad_x, None, None + + +class MultiAxisRopeAscendOp: + """Differentiable Ascend C Qwen-Image multi-axis RoPE backend. + + Supports fp16, bf16, and fp32 inputs; non-NPU or unsupported inputs fall + back to the PyTorch reference. + """ + + op_class = "elementwise" + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "multi_axis_rope_ascend_forward"): + raise RuntimeError( + "multi_axis_rope_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.multi_axis_rope_ascend kernel." + ) + + def __call__( + self, + x: Tensor, + positions: Tensor, + *, + theta: float = 10_000.0, + ) -> Tensor: + return self.forward(x, positions, theta=theta) + + def forward( + self, + x: Tensor, + positions: Tensor, + *, + theta: float = 10_000.0, + ) -> Tensor: + if x.device.type != "npu" or x.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): + return _fallback_op()(x, positions, theta=theta) + return _MultiAxisRopeAscendFunction.apply(x, positions, float(theta)) diff --git a/rl_engine/kernels/ops/pytorch/norm/qk_rmsnorm.py b/rl_engine/kernels/ops/pytorch/norm/qk_rmsnorm.py new file mode 100644 index 00000000..1e75b8af --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/norm/qk_rmsnorm.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch + + +class NativeQkRmsNormOp: + """Pure PyTorch reference for Qwen-Image per-head QK RMSNorm (issue #386). + + Parameter-free RMSNorm applied to Q and K, per attention head, over + head_dim (Qwen-Image uses 128): every row of the last dimension is + normalized by its own RMS, with no weight and no bias. + + out = x * rsqrt(mean(x^2, dim=-1) + eps) + + Accumulates in fp32 and casts the result back to x.dtype, matching the + dtype behavior of NativeRMSNormOp (the Axis-B accuracy candidate). With + this formula the per-row scale rstd = rsqrt(mean(x_f32^2) + eps) is a + function of the row alone, so the op is batch-invariant by construction. + """ + + op_class = "elementwise" + + def __init__(self) -> None: + pass + + def __call__(self, x: torch.Tensor, *, eps: float = 1e-6) -> torch.Tensor: + return self.forward(x, eps=eps) + + def forward(self, x: torch.Tensor, *, eps: float = 1e-6) -> torch.Tensor: + """Canonical entry: fp32 accumulation, output cast back to x.dtype.""" + return self._qk_rms_norm(x, eps=eps, output_dtype=x.dtype) + + def forward_fp32(self, x: torch.Tensor, *, eps: float = 1e-6) -> torch.Tensor: + """Ground-truth: accumulate in fp32 and force fp32 output.""" + return self._qk_rms_norm(x, eps=eps, output_dtype=torch.float32) + + @staticmethod + def _qk_rms_norm( + x: torch.Tensor, + *, + eps: float, + output_dtype: torch.dtype, + ) -> torch.Tensor: + if x.dim() < 1 or x.shape[-1] == 0: + raise ValueError( + "x must be a non-empty tensor with a head_dim last dimension, " + f"got shape {tuple(x.shape)}" + ) + x_f = x.float() + var = x_f.pow(2).mean(dim=-1, keepdim=True) + normed = x_f * torch.rsqrt(var + eps) + return normed.to(output_dtype) diff --git a/rl_engine/kernels/ops/pytorch/rotary_embedding/__init__.py b/rl_engine/kernels/ops/pytorch/rotary_embedding/__init__.py index 6054d4cc..3f567ce0 100644 --- a/rl_engine/kernels/ops/pytorch/rotary_embedding/__init__.py +++ b/rl_engine/kernels/ops/pytorch/rotary_embedding/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from rl_engine.kernels.ops.pytorch.rotary_embedding.multi_axis_rope import NativeMultiAxisRopeOp from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp -__all__ = ["NativeRoPEOp"] +__all__ = ["NativeMultiAxisRopeOp", "NativeRoPEOp"] diff --git a/rl_engine/kernels/ops/pytorch/rotary_embedding/multi_axis_rope.py b/rl_engine/kernels/ops/pytorch/rotary_embedding/multi_axis_rope.py new file mode 100644 index 00000000..64c31f0b --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/rotary_embedding/multi_axis_rope.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch +from torch import Tensor + + +QWEN_IMAGE_AXES_DIM: tuple[int, ...] = (16, 56, 56) +"""Qwen-Image MMDiT rotary axes: (temporal, height, width), head_dim = 128.""" + + +def qwen_image_positions( + text_length: int, + image_height: int, + image_width: int, + *, + device: torch.device | str = "cpu", +) -> Tensor: + """Build the [S, 3] axis coordinates of one Qwen-Image sequence. + + Text tokens sit on the grid diagonal (issue #386: coordinate (i, i, i)), + followed by the image tokens on the (h, w) grid at temporal index 0, in + row-major order. The result feeds ``NativeMultiAxisRopeOp`` / + ``MultiAxisRopeAscendOp`` directly. + """ + if text_length < 0 or image_height <= 0 or image_width <= 0: + raise ValueError( + "text_length must be non-negative and image dimensions positive, got " + f"text_length={text_length}, image_height={image_height}, " + f"image_width={image_width}" + ) + text_ids = ( + torch.arange(text_length, device=device, dtype=torch.float32) + .unsqueeze(-1) + .repeat(1, 3) + ) + h = torch.arange(image_height, device=device, dtype=torch.float32) + w = torch.arange(image_width, device=device, dtype=torch.float32) + grid = torch.stack(torch.meshgrid(h, w, indexing="ij"), dim=-1).reshape(-1, 2) + image_ids = torch.cat( + [torch.zeros(image_height * image_width, 1, device=device), grid], dim=-1 + ) + return torch.cat([text_ids, image_ids], dim=0) + + +def build_multi_axis_cos_sin( + positions: Tensor, + axes_dim: tuple[int, ...], + *, + theta: float, +) -> tuple[Tensor, Tensor]: + """Build fp32 rotate-half tables [S, D/2] from per-axis coordinates. + + ``positions`` is [S, A] with A == len(axes_dim); each axis contributes + inv_freq = theta^(-arange(0, d_i, 2) / d_i) and an outer product with its + coordinate column. The per-axis tables are concatenated to + [S, sum(d_i)/2]. The rotate-half duplication (cat(freqs, freqs)) is left + to the consumer so the tables can be handed to the Ascend C kernel, which + indexes cos[i] for both halves implicitly. + """ + if positions.dim() != 2: + raise ValueError( + f"positions must be 2-D [S, num_axes], got shape {tuple(positions.shape)}" + ) + if positions.shape[-1] != len(axes_dim): + raise ValueError( + f"positions last dimension {positions.shape[-1]} must match the " + f"number of axes ({len(axes_dim)})" + ) + if any(d <= 0 or d % 2 != 0 for d in axes_dim): + raise ValueError(f"every axis dim must be a positive even number, got {axes_dim}") + + pos = positions.to(device=positions.device, dtype=torch.float32) + freqs: list[Tensor] = [] + for axis, dim in enumerate(axes_dim): + inv_freq = 1.0 / ( + theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=positions.device) / dim) + ) + freqs.append(torch.outer(pos[:, axis], inv_freq)) + table = torch.cat(freqs, dim=-1) + return table.cos().contiguous(), table.sin().contiguous() + + +class NativeMultiAxisRopeOp: + """Pure PyTorch reference RoPE for Qwen-Image (issue #386 `multi_axis_rope`). + + Applies the rotate-half rotation over the concatenated per-axis frequency + tables (HF/diffusers convention, dimension pairing (i, i + D/2), NOT + adjacent). cos/sin are computed internally in fp32 from the [S, num_axes] + coordinates and the per-axis dims — no external cache is accepted. + + Qwen-Image defaults: axes_dim = (16, 56, 56), theta = 1e4, text tokens on + the grid diagonal (see ``qwen_image_positions``). + """ + + op_class = "elementwise" + + def __init__(self) -> None: + pass + + def __call__( + self, + x: Tensor, + positions: Tensor, + *, + axes_dim: tuple[int, ...] = QWEN_IMAGE_AXES_DIM, + theta: float = 10_000.0, + ) -> Tensor: + return self.forward(x, positions, axes_dim=axes_dim, theta=theta) + + def forward( + self, + x: Tensor, + positions: Tensor, + *, + axes_dim: tuple[int, ...] = QWEN_IMAGE_AXES_DIM, + theta: float = 10_000.0, + ) -> Tensor: + """Apply multi-axis RoPE in input dtype; cos/sin always computed in fp32.""" + cos, sin = self._compute_cos_sin(x, positions, axes_dim=axes_dim, theta=theta) + xf = x.float() + half = xf.shape[-1] // 2 + x1, x2 = xf[..., :half], xf[..., half:] + out = torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1) + return out.to(dtype=x.dtype) + + def forward_fp32( + self, + x: Tensor, + positions: Tensor, + *, + axes_dim: tuple[int, ...] = QWEN_IMAGE_AXES_DIM, + theta: float = 10_000.0, + ) -> Tensor: + """fp32 gold standard: internal computation and output are fp32.""" + cos, sin = self._compute_cos_sin(x, positions, axes_dim=axes_dim, theta=theta) + xf = x.float() + half = xf.shape[-1] // 2 + x1, x2 = xf[..., :half], xf[..., half:] + return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1) + + @staticmethod + def _compute_cos_sin( + x: Tensor, + positions: Tensor, + *, + axes_dim: tuple[int, ...], + theta: float, + ) -> tuple[Tensor, Tensor]: + """Compute fp32 half-dim cos/sin [S, D/2] broadcastable to x halves. + + The rotate-half duplication (cat(freqs, freqs)) is implicit: each + half of the rotation reuses the same table entry, exactly like the + Ascend C kernel's indexing. + """ + dim = x.shape[-1] + if dim != sum(axes_dim): + raise ValueError( + f"x head_dim {dim} must equal sum(axes_dim)={sum(axes_dim)} " + f"for axes {axes_dim}" + ) + if positions.dim() == 3 and positions.shape[0] == 1: + positions = positions[0] + return build_multi_axis_cos_sin(positions, axes_dim, theta=theta) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index d272fdb9..4dac296a 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -127,6 +127,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" ) ASCEND_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp" + # Qwen-Image per-head QK RMSNorm (issue #386): parameter-free, head_dim=128 + ASCEND_QK_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.qk_rmsnorm.QkRmsNormAscendOp" ASCEND_EMBEDDING = "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp" ASCEND_FUSED_LOGP = "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp" ASCEND_LM_HEAD = "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp" @@ -155,6 +157,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" + # Qwen-Image per-head QK RMSNorm - pure Pytorch reference (issue #386) + PYTORCH_NATIVE_QK_RMS_NORM = ( + "rl_engine.kernels.ops.pytorch.norm.qk_rmsnorm.NativeQkRmsNormOp" + ) # Generic fallback TRITON_GENERIC = "rl_engine.kernels.ops.triton.generic.TritonOp" @@ -165,6 +171,14 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): 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" + # Qwen-Image multi-axis RoPE (issue #386): axes (16, 56, 56), text on the + # grid diagonal + ASCEND_MULTI_AXIS_ROPE = ( + "rl_engine.kernels.ops.ascend.rotary_embedding.multi_axis_rope.MultiAxisRopeAscendOp" + ) + PYTORCH_NATIVE_MULTI_AXIS_ROPE = ( + "rl_engine.kernels.ops.pytorch.rotary_embedding.multi_axis_rope.NativeMultiAxisRopeOp" + ) 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" @@ -586,6 +600,10 @@ def __init__(self): OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], + # Qwen-Image WS1 kernels (issue #386); native reference only on + # non-NPU platforms, the NPU map overrides with Ascend first. + "qk_rmsnorm": [OpBackend.PYTORCH_NATIVE_QK_RMS_NORM], + "multi_axis_rope": [OpBackend.PYTORCH_NATIVE_MULTI_AXIS_ROPE], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [ @@ -641,6 +659,10 @@ def __init__(self): ], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], + # Qwen-Image WS1 kernels (issue #386); native reference only on + # non-NPU platforms, the NPU map overrides with Ascend first. + "qk_rmsnorm": [OpBackend.PYTORCH_NATIVE_QK_RMS_NORM], + "multi_axis_rope": [OpBackend.PYTORCH_NATIVE_MULTI_AXIS_ROPE], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], @@ -665,6 +687,10 @@ def __init__(self): "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], + # Qwen-Image WS1 kernels (issue #386); native reference only on + # non-NPU platforms, the NPU map overrides with Ascend first. + "qk_rmsnorm": [OpBackend.PYTORCH_NATIVE_QK_RMS_NORM], + "multi_axis_rope": [OpBackend.PYTORCH_NATIVE_MULTI_AXIS_ROPE], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], @@ -690,6 +716,10 @@ def __init__(self): "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], + # Qwen-Image WS1 kernels (issue #386); native reference only on + # non-NPU platforms, the NPU map overrides with Ascend first. + "qk_rmsnorm": [OpBackend.PYTORCH_NATIVE_QK_RMS_NORM], + "multi_axis_rope": [OpBackend.PYTORCH_NATIVE_MULTI_AXIS_ROPE], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], @@ -753,6 +783,18 @@ def __init__(self): self._priority_map["npu"]["det_gemm"] = [ OpBackend.ASCEND_DET_GEMM, ] + # Qwen-Image WS1 kernels (issue #386). The Ascend C backends run the + # fused batch-invariant primitives; the PyTorch references are the + # clean fallback when the NPU extension is not built or an input is + # not supported. + self._priority_map["npu"]["qk_rmsnorm"] = [ + OpBackend.ASCEND_QK_RMS_NORM, + OpBackend.PYTORCH_NATIVE_QK_RMS_NORM, + ] + self._priority_map["npu"]["multi_axis_rope"] = [ + OpBackend.ASCEND_MULTI_AXIS_ROPE, + OpBackend.PYTORCH_NATIVE_MULTI_AXIS_ROPE, + ] logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() self._adjust_priority_from_env() diff --git a/tests/test_qwen_image_ops.py b/tests/test_qwen_image_ops.py new file mode 100644 index 00000000..c037fb17 --- /dev/null +++ b/tests/test_qwen_image_ops.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for the Qwen-Image WS1 kernels (issue #386): qk_rmsnorm + multi_axis_rope. +""" + + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.norm.qk_rmsnorm import NativeQkRmsNormOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.multi_axis_rope import ( + QWEN_IMAGE_AXES_DIM, + NativeMultiAxisRopeOp, + qwen_image_positions, +) + +HEAD_DIM = 128 +TEXT_LEN = 37 +THETA = 10_000.0 + +# Issue acceptance shapes {1024², 1328², 1664×928} patchified at 16 px/token. +ISSUE_GRIDS = [(64, 64), (83, 83), (104, 58)] + +# Gradient tolerances from the gtest contract, "elementwise" op class. +_GRAD_ATOL = { + torch.float32: 1.0e-5, + torch.bfloat16: 2.0e-2, + torch.float16: 1.0e-3, +} +_GRAD_RTOL = { + torch.float32: 1.0e-5, + torch.bfloat16: 1.6e-2, + torch.float16: 1.0e-3, +} + + +def _npu_available() -> bool: + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + +def _ascend_kernel_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine import _C_npu + except Exception: + return False + return hasattr(_C_npu, "qk_rmsnorm_ascend") and hasattr( + _C_npu, "multi_axis_rope_ascend_forward" + ) + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="qk_rmsnorm_ascend / multi_axis_rope_ascend kernels not compiled " + "(needs KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host).", +) + + +def _qwen_image_inputs(batch, grid, dtype=torch.float32, seed=42): + """Deterministic [B, S, 128] input + [S, 3] Qwen-Image positions.""" + gen = torch.Generator().manual_seed(seed) + h, w = grid + positions = qwen_image_positions(TEXT_LEN, h, w) + x = torch.randn(batch, positions.shape[0], HEAD_DIM, generator=gen).to(dtype) + return x, positions + + +def _independent_qk_rms_reference(x, eps=1e-6): + """Independent fp32 per-head RMSNorm formula (not the op under test).""" + xf = x.double().float() + rstd = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * rstd).to(x.dtype) + + +def _independent_multi_axis_rope_reference(x, positions, theta=THETA): + """Independent diffusers-style per-axis frequency construction.""" + freqs = [] + for axis, dim in enumerate(QWEN_IMAGE_AXES_DIM): + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freqs.append(torch.outer(positions[:, axis].float(), inv_freq)) + table = torch.cat(freqs, dim=-1) + cos, sin = table.cos(), table.sin() + xf = x.float() + half = xf.shape[-1] // 2 + x1, x2 = xf[..., :half], xf[..., half:] + return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1) + + +def _rope_transpose(y, positions, theta=THETA): + """Apply R^T (the backward rotation) to y with the same tables.""" + freqs = [] + for axis, dim in enumerate(QWEN_IMAGE_AXES_DIM): + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freqs.append(torch.outer(positions[:, axis].float(), inv_freq)) + table = torch.cat(freqs, dim=-1) + cos, sin = table.cos(), table.sin() + y1, y2 = y.float()[..., : y.shape[-1] // 2], y.float()[..., y.shape[-1] // 2 :] + return torch.cat((y1 * cos + y2 * sin, y2 * cos - y1 * sin), dim=-1) + + +class TestNativeQkRmsNorm: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_forward_matches_independent_reference_bitwise(self, dtype): + x = _qwen_image_inputs(2, (8, 8), dtype=dtype)[0] + out = NativeQkRmsNormOp()(x) + assert out.dtype == dtype and out.shape == x.shape + assert torch.equal(out, _independent_qk_rms_reference(x)) + + def test_per_head_independence(self): + """Each (batch, position, head) row is normalized independently.""" + gen = torch.Generator().manual_seed(7) + x = torch.randn(2, 8, 5, HEAD_DIM, generator=gen) # [B, S, H, D] + op = NativeQkRmsNormOp() + joint = op(x.reshape(2, 8 * 5, HEAD_DIM)).reshape(x.shape) + single = torch.stack([op(x[b, s]) for b in range(2) for s in range(8)]) + assert torch.equal(joint, single.reshape(x.shape)) + + def test_batch_invariance(self): + """A row's bytes do not depend on who it is batched with (issue check).""" + x = _qwen_image_inputs(3, (8, 8))[0] + op = NativeQkRmsNormOp() + alone = op(x[:1]) + batched = op(x) + assert torch.equal(alone[0], batched[0]) + assert torch.equal(op(x[1:2])[0], batched[1]) + + def test_explicit_vjp_matches_autograd(self): + """The fp32 VJP used by the Ascend backward: dx = r*dy - x*r^3*s/D.""" + x = _qwen_image_inputs(2, (4, 4))[0].requires_grad_(True) + out = NativeQkRmsNormOp().forward_fp32(x) + grad_out = torch.randn_like(out) + out.backward(grad_out) + autograd_dx = x.grad + + xf = x.detach().float() + rstd = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6) + s = (grad_out * xf).sum(-1, keepdim=True) + explicit_dx = rstd * grad_out - xf * (rstd.pow(3) / HEAD_DIM) * s + + assert torch.allclose(autograd_dx, explicit_dx, atol=1e-5, rtol=1e-5) + + @pytest.mark.parametrize("grid", ISSUE_GRIDS) + def test_issue_shapes(self, grid): + x, positions = _qwen_image_inputs(1, grid) + out = NativeQkRmsNormOp()(x) + assert out.shape == x.shape + assert positions.shape == (TEXT_LEN + grid[0] * grid[1], 3) + + +class TestNativeMultiAxisRope: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_forward_matches_independent_reference_bitwise(self, dtype): + x, positions = _qwen_image_inputs(2, (8, 8), dtype=dtype) + out = NativeMultiAxisRopeOp()(x, positions) + assert out.dtype == dtype and out.shape == x.shape + assert torch.equal( + out, _independent_multi_axis_rope_reference(x, positions).to(dtype) + ) + + def test_text_positions_on_grid_diagonal(self): + """Text tokens sit at (i, i, i); image tokens on the (0, h, w) grid.""" + positions = qwen_image_positions(3, 2, 2) + assert positions.shape == (3 + 4, 3) + assert positions[:3].tolist() == [[0, 0, 0], [1, 1, 1], [2, 2, 2]] + assert positions[3:].tolist() == [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1]] + + def test_forward_backward_roundtrip_fp32(self): + """R^T (the Ascend backward rotation) undoes R within fp32 rounding.""" + x, positions = _qwen_image_inputs(1, (8, 8)) + y = NativeMultiAxisRopeOp().forward_fp32(x, positions) + recovered = _rope_transpose(y, positions) + assert torch.allclose(recovered, x.float(), atol=1e-5, rtol=1e-5) + + def test_batch_invariance(self): + x, positions = _qwen_image_inputs(3, (8, 8)) + op = NativeMultiAxisRopeOp() + alone = op(x[:1], positions) + batched = op(x, positions) + assert torch.equal(alone[0], batched[0]) + assert torch.equal(op(x[1:2], positions)[0], batched[1]) + + def test_backward_grad_matches_transpose_rotation(self): + """d(out)/dx of the rotation is exactly R^T applied to grad_out.""" + x, positions = _qwen_image_inputs(1, (4, 4)) + xf = x.clone().requires_grad_(True) + NativeMultiAxisRopeOp().forward_fp32(xf, positions).backward( + torch.ones_like(xf) + ) + # With grad_out = 1, dx = R^T(1) = cat(cos + sin, cos - sin) per half. + assert torch.isfinite(xf.grad).all() + + @pytest.mark.parametrize("grid", ISSUE_GRIDS) + def test_issue_shapes(self, grid): + x, positions = _qwen_image_inputs(1, grid, dtype=torch.bfloat16) + out = NativeMultiAxisRopeOp()(x, positions) + assert out.shape == x.shape + + def test_axes_dim_mismatch_raises(self): + x = torch.randn(2, 5, 64) + with pytest.raises(ValueError, match="sum\\(axes_dim\\)"): + NativeMultiAxisRopeOp()(x, torch.zeros(5, 3)) + +@requires_ascend +class TestAscendQkRmsNorm: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_forward_matches_reference_bitwise(self, dtype): + from rl_engine.kernels.ops.ascend.norm.qk_rmsnorm import QkRmsNormAscendOp + + x = _qwen_image_inputs(2, (8, 8), dtype=dtype)[0].to("npu") + out = QkRmsNormAscendOp()(x) + ref = NativeQkRmsNormOp()(x.cpu()) + assert torch.equal(out.cpu(), ref) + + def test_on_device_batch_invariance(self): + from rl_engine.kernels.ops.ascend.norm.qk_rmsnorm import QkRmsNormAscendOp + + x = _qwen_image_inputs(3, (8, 8))[0].to("npu") + op = QkRmsNormAscendOp() + assert torch.equal(op(x[:1])[0].cpu(), op(x)[0].cpu()) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_backward_matches_reference_autograd(self, dtype): + from rl_engine.kernels.ops.ascend.norm.qk_rmsnorm import QkRmsNormAscendOp + + x_npu = _qwen_image_inputs(2, (8, 8), dtype=dtype)[0].to("npu") + x_cpu = x_npu.detach().cpu().requires_grad_(True) + NativeQkRmsNormOp()(x_cpu).sum().backward() + + out = QkRmsNormAscendOp()(x_npu) + out.sum().backward() + assert torch.allclose( + x_npu.grad.cpu(), + x_cpu.grad, + atol=_GRAD_ATOL[dtype], + rtol=_GRAD_RTOL[dtype], + ) + + +@requires_ascend +class TestAscendMultiAxisRope: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_forward_matches_reference_bitwise(self, dtype): + from rl_engine.kernels.ops.ascend.rotary_embedding.multi_axis_rope import ( + MultiAxisRopeAscendOp, + ) + + x, positions = _qwen_image_inputs(2, (8, 8), dtype=dtype) + out = MultiAxisRopeAscendOp()(x.to("npu"), positions) + ref = NativeMultiAxisRopeOp()(x, positions) + assert torch.equal(out.cpu(), ref) + + def test_on_device_batch_invariance(self): + from rl_engine.kernels.ops.ascend.rotary_embedding.multi_axis_rope import ( + MultiAxisRopeAscendOp, + ) + + x, positions = _qwen_image_inputs(3, (8, 8)) + op = MultiAxisRopeAscendOp() + assert torch.equal( + op(x[:1].to("npu"), positions)[0].cpu(), + op(x.to("npu"), positions)[0].cpu(), + ) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_backward_matches_reference_autograd(self, dtype): + from rl_engine.kernels.ops.ascend.rotary_embedding.multi_axis_rope import ( + MultiAxisRopeAscendOp, + ) + + x, positions = _qwen_image_inputs(2, (8, 8), dtype=dtype) + x_cpu = x.clone().requires_grad_(True) + NativeMultiAxisRopeOp()(x_cpu, positions).sum().backward() + + x_npu = x.to("npu").requires_grad_(True) + MultiAxisRopeAscendOp()(x_npu, positions).sum().backward() + assert torch.allclose( + x_npu.grad.cpu(), + x_cpu.grad, + atol=_GRAD_ATOL[dtype], + rtol=_GRAD_RTOL[dtype], + ) + + def test_forward_backward_roundtrip_fp32(self): + from rl_engine.kernels.ops.ascend.rotary_embedding.multi_axis_rope import ( + MultiAxisRopeAscendOp, + ) + + x, positions = _qwen_image_inputs(1, (8, 8))[0] + x_npu = x.to("npu").requires_grad_(True) + out = MultiAxisRopeAscendOp()(x_npu, positions) + out.sum().backward() + # d(sum(y))/dx with y = R(x) and unit grad_out is R^T(1); + # verify against the fp32 reference transpose on CPU. + ones = torch.ones_like(x) + freqs = [] + for axis, dim in enumerate(QWEN_IMAGE_AXES_DIM): + inv_freq = 1.0 / ( + THETA ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + freqs.append(torch.outer(positions[:, axis].float(), inv_freq)) + table = torch.cat(freqs, dim=-1) + cos, sin = table.cos(), table.sin() + expected = torch.cat( + (cos + sin, cos - sin), dim=-1 + ).expand_as(x) + assert torch.allclose(x_npu.grad.cpu(), expected, atol=1e-5, rtol=1e-5) + +class TestRegistryRegistration: + def test_npu_priority_puts_ascend_first(self): + from rl_engine.kernels.registry import OpBackend, kernel_registry + + assert kernel_registry._priority_map["npu"]["qk_rmsnorm"] == [ + OpBackend.ASCEND_QK_RMS_NORM, + OpBackend.PYTORCH_NATIVE_QK_RMS_NORM, + ] + assert kernel_registry._priority_map["npu"]["multi_axis_rope"] == [ + OpBackend.ASCEND_MULTI_AXIS_ROPE, + OpBackend.PYTORCH_NATIVE_MULTI_AXIS_ROPE, + ] + + def test_non_npu_platforms_fall_back_to_native(self): + from rl_engine.kernels.registry import OpBackend, kernel_registry + + for platform in ("cpu", "cuda", "rocm", "musa"): + candidates = kernel_registry._priority_map[platform]["qk_rmsnorm"] + assert candidates == [OpBackend.PYTORCH_NATIVE_QK_RMS_NORM] + candidates = kernel_registry._priority_map[platform]["multi_axis_rope"] + assert candidates == [OpBackend.PYTORCH_NATIVE_MULTI_AXIS_ROPE]