From 8b39ad4f1f3be41720056e6098a1c595574f4964 Mon Sep 17 00:00:00 2001 From: chenyang <2082464740@qq.com> Date: Sun, 30 Aug 2026 16:35:21 +0800 Subject: [PATCH 1/4] feat(ascend): add batch-invariant RMSNorm Ascend C operator - csrc/ascend/rmsnorm_ascend.asc: fixed-tile (4096) fp32-accumulation forward kernel for fp32/bf16/fp16, outputs y and rstd; Rsqrt refined with 2 Newton-Raphson iterations (fp32 y err ~5e-7); small-row chunk coalescing amortizes pipeline syncs (H%8==0, up to 64 rows/iter) - csrc/ascend/npu_module.cpp: single PYBIND11_MODULE for rl_engine._C_npu binding all Ascend ops (moved out of batch_invariant_logp_ascend.asc) - setup.py: Ascend build path (bisheng -x asc, dav-c220 default) gated by KERNEL_ALIGN_FORCE_ASCEND=1 - rl_engine/kernels/ops/ascend/norm/rmsnorm.py: RMSNormAscendOp with fp32 VJP backward reusing forward-saved rstd, native fallback - registry: ASCEND_RMS_NORM, NPU priority [ascend, native] - tests: 24 Ascend cases (accuracy / fwd+bwd / bitwise batch-invariance), dispatch assertions updated - benchmarks: NPU device dispatch + ascend branch Measured on Ascend910_9362 (fwd+bwd, bf16) vs PyTorch native: T=1024/H=4096 1.24x, T=8192/H=4096 1.57x, T=1024/H=8192 1.49x, T=4096/H=5120 1.46x, T=32768/H=128 1.20x, T=16384/H=256 1.21x --- benchmarks/benchmark_rmsnorm.py | 72 ++- csrc/ascend/batch_invariant_logp_ascend.asc | 8 +- csrc/ascend/npu_module.cpp | 29 + csrc/ascend/rmsnorm_ascend.asc | 529 ++++++++++++++++++ rl_engine/_C_npu.pyi | 6 + rl_engine/kernels/ops/ascend/__init__.py | 1 + rl_engine/kernels/ops/ascend/norm/__init__.py | 2 + rl_engine/kernels/ops/ascend/norm/rmsnorm.py | 145 +++++ rl_engine/kernels/registry.py | 5 + rl_engine/tests/test_dispatch.py | 6 +- setup.py | 116 +++- tests/test_rms_norm.py | 93 ++- 12 files changed, 981 insertions(+), 31 deletions(-) create mode 100644 csrc/ascend/npu_module.cpp create mode 100644 csrc/ascend/rmsnorm_ascend.asc create mode 100644 rl_engine/kernels/ops/ascend/norm/__init__.py create mode 100644 rl_engine/kernels/ops/ascend/norm/rmsnorm.py diff --git a/benchmarks/benchmark_rmsnorm.py b/benchmarks/benchmark_rmsnorm.py index 81325281..31b37c26 100644 --- a/benchmarks/benchmark_rmsnorm.py +++ b/benchmarks/benchmark_rmsnorm.py @@ -4,24 +4,47 @@ import torch from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp -from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton +from rl_engine.platforms.device import device_ctx -try: - from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda +if device_ctx.device_type != "npu": + from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton - HAS_CUDA_EXT = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward") -except ImportError: - HAS_CUDA_EXT = False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda + + HAS_CUDA_EXT = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward") + except ImportError: + HAS_CUDA_EXT = False + + +def _accel(): + """The active accelerator module (torch.npu on Ascend, torch.cuda otherwise).""" + if device_ctx.device_type == "npu": + return torch.npu + return torch.cuda + + +def _maybe_ascend_op(): + """The Ascend C op, or None when unavailable (no NPU / not built).""" + if device_ctx.device_type != "npu": + return None + try: + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp + + return RMSNormAscendOp() + except (ImportError, OSError, RuntimeError): + return None def bench(fn, x, w, dy, warmup=20, iters=100): + accel = _accel() for _ in range(warmup): x.grad = None w.grad = None y = fn(x, w) y.backward(dy) - torch.cuda.synchronize() + accel.synchronize() start = time.time() for _ in range(iters): @@ -29,7 +52,7 @@ def bench(fn, x, w, dy, warmup=20, iters=100): w.grad = None y = fn(x, w) y.backward(dy) - torch.cuda.synchronize() + accel.synchronize() return (time.time() - start) * 1000.0 / iters @@ -37,11 +60,11 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--T", type=int, default=1024) parser.add_argument("--H", type=int, default=4096) - parser.add_argument("--dtype", choices=["fp16", "bf16"], default="bf16") + parser.add_argument("--dtype", choices=["fp16", "bf16", "fp32"], default="bf16") args = parser.parse_args() - dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16 - device = "cuda" + dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[args.dtype] + device = "npu" if device_ctx.device_type == "npu" else "cuda" T, H = args.T, args.H torch.manual_seed(0) @@ -61,16 +84,25 @@ def make_inputs(): t_ref = bench(lambda a, b: native.forward(a, b), x, w, dy) print(f"pytorch ref : {t_ref:.4f} ms") - x, w = make_inputs() - t_tri = bench(lambda a, b: rmsnorm_triton(a, b), x, w, dy) - print(f"triton : {t_tri:.4f} ms | speedup vs ref: {t_ref / t_tri:.2f}x") - - if HAS_CUDA_EXT: + if device == "cuda": x, w = make_inputs() - t_cuda = bench(lambda a, b: rmsnorm_cuda(a, b), x, w, dy) - print(f"cuda : {t_cuda:.4f} ms | speedup vs ref: {t_ref / t_cuda:.2f}x") + t_tri = bench(lambda a, b: rmsnorm_triton(a, b), x, w, dy) + print(f"triton : {t_tri:.4f} ms | speedup vs ref: {t_ref / t_tri:.2f}x") + + if HAS_CUDA_EXT: + x, w = make_inputs() + t_cuda = bench(lambda a, b: rmsnorm_cuda(a, b), x, w, dy) + print(f"cuda : {t_cuda:.4f} ms | speedup vs ref: {t_ref / t_cuda:.2f}x") + else: + print("cuda : skipped, extension is not built") else: - print("cuda : skipped, extension is not built") + ascend_op = _maybe_ascend_op() + if ascend_op is not None: + x, w = make_inputs() + t_asc = bench(lambda a, b: ascend_op(a, b), x, w, dy) + print(f"ascend : {t_asc:.4f} ms | speedup vs ref: {t_ref / t_asc:.2f}x") + else: + print("ascend : skipped, extension is not built") if __name__ == "__main__": diff --git a/csrc/ascend/batch_invariant_logp_ascend.asc b/csrc/ascend/batch_invariant_logp_ascend.asc index dead4cbe..3b7e46b9 100644 --- a/csrc/ascend/batch_invariant_logp_ascend.asc +++ b/csrc/ascend/batch_invariant_logp_ascend.asc @@ -308,9 +308,5 @@ 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)"); -} +// The PYBIND11_MODULE for rl_engine._C_npu lives in npu_module.cpp so that +// every Ascend op shares one compiled module. diff --git a/csrc/ascend/npu_module.cpp b/csrc/ascend/npu_module.cpp new file mode 100644 index 00000000..f94dac4c --- /dev/null +++ b/csrc/ascend/npu_module.cpp @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Pybind entry point for the rl_engine._C_npu extension. The Ascend C kernels +// and their torch host wrappers live in the sibling *.asc files; this TU only +// declares and binds them so every Ascend op shares one compiled module. +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#include + +std::vector batch_invariant_logp_ascend_forward(torch::Tensor logits, + torch::Tensor target, + int64_t ignore_index); + +std::vector rmsnorm_ascend_forward(torch::Tensor x, + torch::Tensor weight, + double eps); + +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("rmsnorm_ascend", + &rmsnorm_ascend_forward, + "Batch-invariant RMSNorm (Ascend C forward)"); +} diff --git a/csrc/ascend/rmsnorm_ascend.asc b/csrc/ascend/rmsnorm_ascend.asc new file mode 100644 index 00000000..a87660f4 --- /dev/null +++ b/csrc/ascend/rmsnorm_ascend.asc @@ -0,0 +1,529 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant RMSNorm, Ascend C (CANN) forward kernel. +// +// y[n, :] = x[n, :] * rsqrt(mean(x[n, :]^2) + eps) * weight[:] +// +// Mirrors the CUDA kernel in csrc/cuda/rmsnorm.cu: +// - input : x [N, H] contiguous, fp32 / bf16 / fp16; weight [H] same dtype +// - output : y [N, H] same dtype as x, rstd [N] fp32 +// (rstd = rsqrt(mean(x^2) + eps), saved for the autograd backward) +// +// Batch-invariance: every row is processed end-to-end by exactly one AI core +// block with a fixed tile size and a fixed reduction order (fp32 sum of +// squares over fixed-order tiles). The instruction sequence for a row depends +// only on H, never on N or on the block the row happens to land on. +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#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 +// makes the reduction order batch-invariant. UB budget (in/out/weight tiles + +// fp32 tile + fp32 weight tile + square tile + reduce scratch) stays under +// the 192 KB UB of DAV_2201 SoCs even for fp32 in/out tiles. +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 +// staging buffers (aBuf/rBuf) and the per-row reduce scratch slots. +constexpr int64_t MAX_CHUNK_ROWS = 64; + +template +class KernelRmsNorm { +public: + __aicore__ inline KernelRmsNorm(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR x, + GM_ADDR weight, + GM_ADDR y, + GM_ADDR rstd, + int64_t numRows, + int64_t hiddenSize, + float eps) + { + numRows_ = numRows; + hiddenSize_ = hiddenSize; + eps_ = eps; + singleTile_ = hiddenSize <= static_cast(TILE_LENGTH); + // Small rows 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 + // H % 8 == 0 gate (offset r * H floats stays aligned for every r). + rowsPerChunk_ = 1; + if (singleTile_ && hiddenSize % 8 == 0) { + rowsPerChunk_ = TILE_LENGTH / hiddenSize; + if (rowsPerChunk_ > MAX_CHUNK_ROWS) { + rowsPerChunk_ = MAX_CHUNK_ROWS; + } + } + xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x)); + weightGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(weight)); + yGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(y)); + rstdGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(rstd)); + pipe_->InitBuffer(inQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(outQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(wQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(fp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(wFp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(sqBuf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(reduceBuf_, TILE_LENGTH * sizeof(float)); + // 2 KB: per-row reduce scratch at scalar[8*r] (each slot 32-byte + // aligned) for r in [0, MAX_CHUNK_ROWS); the single-row path only + // uses floats [0,8) for scratch and [8,16) for rstd staging. + pipe_->InitBuffer(scalarBuf_, MAX_CHUNK_ROWS * 8 * sizeof(float)); + // 512 B: aBuf [0,64) holds meanSq+eps per chunk row, rBuf [64,128) + // holds the refined rstd per chunk row (flushed to GM in one burst). + pipe_->InitBuffer(stageBuf_, MAX_CHUNK_ROWS * 2 * 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. + eventVS_ = pipe_->FetchEventID(AscendC::HardEvent::V_S); + eventSV_ = pipe_->FetchEventID(AscendC::HardEvent::S_V); + eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3); + eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + } + + __aicore__ inline void Process() + { + if (singleTile_) { + // Cache the fp32 weight tile once; every row reuses it. + LoadWeightFp32(0, static_cast(hiddenSize_)); + } + 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; + ProcessRowChunk(row, remaining < rowsPerChunk_ ? remaining : rowsPerChunk_); + } + return; + } + for (int64_t row = AscendC::GetBlockIdx(); row < numRows_; row += AscendC::GetBlockNum()) { + ProcessRow(row); + } + } + +private: + // Load x[row, start:start+count] 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, + int64_t start, + uint32_t count) + { + 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_[row * hiddenSize_ + start], 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, count); + return fLocal; + } + } + + __aicore__ inline void FreeTile() + { + inQueue_.FreeTensor(inTile_); + } + + // Load weight[start:start+count] and cast it into wFp32Buf_[0:count]. + __aicore__ inline void LoadWeightFp32(int64_t start, uint32_t count) + { + AscendC::LocalTensor wLocal = wQueue_.AllocTensor(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(wLocal, weightGm_[start], copyParams, padParams); + wQueue_.EnQue(wLocal); + wLocal = wQueue_.DeQue(); + AscendC::LocalTensor wFp32 = wFp32Buf_.Get(); + if constexpr (std::is_same_v) { + CopyFp32(wFp32, wLocal, count); + } else { + AscendC::Cast(wFp32, wLocal, AscendC::RoundMode::CAST_NONE, count); + } + wQueue_.FreeTensor(wLocal); + } + + // 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}); + } + + __aicore__ inline float TileSumSquare(const AscendC::LocalTensor& fLocal, + uint32_t count) + { + // Square in place (destroys the tile's x values), then reduce with a + // fixed order. src and work tensors are distinct buffers. + AscendC::Mul(fLocal, fLocal, fLocal, count); + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceSum(scalar, fLocal, rTmp, static_cast(count)); + WaitVector(); // vector -> scalar read + return scalar.GetValue(0); + } + + __aicore__ inline void ProcessRow(int64_t row) + { + float sumSq = 0.0f; + + if (singleTile_) { + // Fast path: the whole row stays resident in fp32Buf_ across the + // sum-of-squares and the scaling, so x is read from GM only once. + const uint32_t count = static_cast(hiddenSize_); + AscendC::LocalTensor fLocal = LoadTileFp32(row, 0, count); + // Square into sqBuf_ so the row values survive in fLocal; reduce + // scratch lives in reduceBuf_ (src/work must not alias). + AscendC::LocalTensor sq = sqBuf_.Get(); + AscendC::Mul(sq, fLocal, fLocal, count); + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceSum(scalar, sq, rTmp, static_cast(count)); + WaitVector(); // vector -> scalar read + sumSq = scalar.GetValue(0); + + const float rstd = ComputeRstd(sumSq, row); + ScaleStoreTile(row, 0, count, fLocal, rstd); + FreeTile(); + return; + } + + const int64_t tileCount = (hiddenSize_ + TILE_LENGTH - 1) / TILE_LENGTH; + + // Pass 1: sum of squares with a fixed tile order. + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_LENGTH; + const uint32_t count = TileCount(start); + AscendC::LocalTensor fLocal = LoadTileFp32(row, start, count); + sumSq += TileSumSquare(fLocal, count); + FreeTile(); + } + + const float rstd = ComputeRstd(sumSq, row); + + // Pass 2: y = x * rstd * w with the same fixed tile order. + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_LENGTH; + const uint32_t count = TileCount(start); + AscendC::LocalTensor fLocal = LoadTileFp32(row, start, count); + LoadWeightFp32(start, count); + ScaleStoreTile(row, start, count, 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 reduction order is identical to ProcessRow (fp32 sum over the + // row's single tile), so numerics are unchanged. + __aicore__ inline void ProcessRowChunk(int64_t row0, int64_t rows) + { + const int64_t H = hiddenSize_; + const uint32_t count = static_cast(rows * H); + + // 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 * H], 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); + } + + // Squares of the whole chunk, then one aligned ReduceSum per row. + // All reduces complete before the first scalar read (single V_S). + AscendC::LocalTensor sq = sqBuf_.Get(); + AscendC::Mul(sq, fLocal, fLocal, count); + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + for (int64_t r = 0; r < rows; ++r) { + AscendC::ReduceSum(scalar[8 * r], sq[r * H], rTmp, + static_cast(H)); + } + WaitVector(); // vector -> scalar reads + + // a[r] = mean(x_r^2) + eps staged for the vector Rsqrt. + AscendC::LocalTensor stage = stageBuf_.Get(); + AscendC::LocalTensor aBuf = stage; + AscendC::LocalTensor rBuf = stage[MAX_CHUNK_ROWS]; + for (int64_t r = 0; r < rows; ++r) { + aBuf.SetValue(r, scalar.GetValue(static_cast(8 * r)) / + static_cast(H) + + eps_); + } + AscendC::SetFlag(eventSV_); // scalar write -> vector op + AscendC::WaitFlag(eventSV_); + const uint32_t rPad = static_cast((rows + 7) / 8 * 8); + AscendC::Rsqrt(rBuf, aBuf, rPad); + WaitVector(); // vector -> scalar reads + // Newton-Raphson refinement on the scalar unit (see ComputeRstd). + for (int64_t r = 0; r < rows; ++r) { + const float a = aBuf.GetValue(static_cast(r)); + float rstd = rBuf.GetValue(static_cast(r)); + rstd = rstd * (1.5f - 0.5f * a * rstd * rstd); + rstd = rstd * (1.5f - 0.5f * a * rstd * rstd); + rBuf.SetValue(static_cast(r), rstd); + } + + // Scale every row of the chunk; scalar-register operands of Muls/Mul + // need no S_V flag (no UB dependency). + AscendC::LocalTensor wFp32 = wFp32Buf_.Get(); + for (int64_t r = 0; r < rows; ++r) { + const float rstd = rBuf.GetValue(static_cast(r)); + AscendC::Muls(fLocal[r * H], fLocal[r * H], rstd, static_cast(H)); + AscendC::Mul(fLocal[r * H], fLocal[r * H], wFp32, static_cast(H)); + } + + 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 * H], yLocal, outParams); + outQueue_.FreeTensor(yLocal); + FreeTile(); + + // Flush the chunk's rstd values in one contiguous burst. + AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out + AscendC::WaitFlag(eventSMTE3_); + AscendC::DataCopyExtParams rstdParams{ + 1, static_cast(rows * sizeof(float)), 0, 0, 0}; + AscendC::DataCopyPad(rstdGm_[row0], rBuf, rstdParams); + // Drain MTE3 before the next chunk stages new values into rBuf. + AscendC::SetFlag(eventMTE3S_); + AscendC::WaitFlag(eventMTE3S_); + } + + // rstd = rsqrt(mean(x^2) + eps), staged to rstdGm_[row] for the backward. + __aicore__ inline float ComputeRstd(float sumSq, int64_t row) + { + // The scalar unit has no rsqrt, so run a 1-element vector Rsqrt + // (count padded to 8; scalarBuf_ is 32 B aligned). The Rsqrt + // instruction is only a ~2^-9 relative approximation, so refine it + // with Newton-Raphson on the scalar unit: r = r*(1.5 - 0.5*a*r*r) + // converges to rsqrt(a) at ~1 ulp after two iterations. + const float meanSqEps = sumSq / static_cast(hiddenSize_) + eps_; + AscendC::LocalTensor scalar = scalarBuf_.Get(); + scalar.SetValue(0, meanSqEps); + AscendC::SetFlag(eventSV_); // scalar write -> vector op + AscendC::WaitFlag(eventSV_); + AscendC::Rsqrt(scalar, scalar, 8); + WaitVector(); // vector -> scalar read + float rstd = scalar.GetValue(0); + rstd = rstd * (1.5f - 0.5f * meanSqEps * rstd * rstd); + rstd = rstd * (1.5f - 0.5f * meanSqEps * rstd * rstd); + + // GlobalTensor.SetValue is unreliable on hardware (cannbot + // ascendc-precision-debug common-traps), so stage rstd in UB and + // DataCopyPad it to GM instead of a scalar GM store. + scalar.SetValue(8, rstd); + AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out + AscendC::WaitFlag(eventSMTE3_); + AscendC::DataCopyExtParams outParams{1, sizeof(float), 0, 0, 0}; + AscendC::DataCopyPad(rstdGm_[row], scalar[8], outParams); + // Drain MTE3 before the next row stages a new value into scalarBuf_. + AscendC::SetFlag(eventMTE3S_); + AscendC::WaitFlag(eventMTE3S_); + return rstd; + } + + // y tile = x tile * rstd * w tile, cast back to T and copied to GM. + __aicore__ inline void ScaleStoreTile(int64_t row, + int64_t start, + uint32_t count, + const AscendC::LocalTensor& fLocal, + float rstd) + { + AscendC::LocalTensor wFp32 = wFp32Buf_.Get(); + // Vector pipe is in-order: the LoadWeightFp32 cast into wFp32Buf_ + // needs no extra flag before these reads. + AscendC::Muls(fLocal, fLocal, rstd, count); + AscendC::Mul(fLocal, fLocal, wFp32, count); + + 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_[row * hiddenSize_ + start], yLocal, outParams); + outQueue_.FreeTensor(yLocal); + } + + // Wait until all outstanding vector-pipe results are readable as scalars. + __aicore__ inline void WaitVector() + { + AscendC::SetFlag(eventVS_); + AscendC::WaitFlag(eventVS_); + } + + __aicore__ inline uint32_t TileCount(int64_t start) const + { + const int64_t remaining = hiddenSize_ - start; + return static_cast(remaining < TILE_LENGTH ? remaining : TILE_LENGTH); + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor xGm_; + AscendC::GlobalTensor weightGm_; + AscendC::GlobalTensor yGm_; + AscendC::GlobalTensor rstdGm_; + AscendC::TQue inQueue_; + AscendC::TQue wQueue_; + AscendC::TQue outQueue_; + AscendC::TBuf fp32Buf_; + AscendC::TBuf wFp32Buf_; + AscendC::TBuf sqBuf_; + AscendC::TBuf reduceBuf_; + AscendC::TBuf scalarBuf_; + AscendC::TBuf stageBuf_; + AscendC::LocalTensor inTile_; + AscendC::TEventID eventVS_; + AscendC::TEventID eventSV_; + AscendC::TEventID eventSMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t numRows_; + int64_t hiddenSize_; + float eps_; + bool singleTile_; + int64_t rowsPerChunk_; +}; + +} // namespace + +extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_fp32( + GM_ADDR x, GM_ADDR weight, GM_ADDR y, GM_ADDR rstd, + int64_t numRows, int64_t hiddenSize, float eps) +{ + AscendC::TPipe pipe; + KernelRmsNorm op(&pipe); + op.Init(x, weight, y, rstd, numRows, hiddenSize, eps); + op.Process(); +} + +extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_bf16( + GM_ADDR x, GM_ADDR weight, GM_ADDR y, GM_ADDR rstd, + int64_t numRows, int64_t hiddenSize, float eps) +{ + AscendC::TPipe pipe; + KernelRmsNorm op(&pipe); + op.Init(x, weight, y, rstd, numRows, hiddenSize, eps); + op.Process(); +} + +extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_fp16( + GM_ADDR x, GM_ADDR weight, GM_ADDR y, GM_ADDR rstd, + int64_t numRows, int64_t hiddenSize, float eps) +{ + AscendC::TPipe pipe; + KernelRmsNorm op(&pipe); + op.Init(x, weight, y, rstd, numRows, hiddenSize, eps); + op.Process(); +} + +std::vector rmsnorm_ascend_forward(torch::Tensor x, + torch::Tensor weight, + double eps) +{ + TORCH_CHECK(x.is_privateuseone(), "x must be on an NPU device"); + TORCH_CHECK(x.dim() == 2, "x must be 2-D [N, H]"); + TORCH_CHECK(x.is_contiguous(), "x must be contiguous"); + TORCH_CHECK(x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kFloat || + x.scalar_type() == at::kHalf, + "x must be fp32, bf16 or fp16"); + TORCH_CHECK(x.size(-1) > 0, "hidden size must be positive"); + TORCH_CHECK(weight.is_privateuseone(), "weight must be on the same NPU device as x"); + TORCH_CHECK(weight.dim() == 1 && weight.numel() == x.size(-1), + "weight must be 1-D of size x.size(-1)"); + TORCH_CHECK(weight.scalar_type() == x.scalar_type(), "weight dtype must match x dtype"); + + const int64_t numRows = x.size(0); + const int64_t hiddenSize = x.size(1); + + torch::Tensor y = at::empty_like(x); + torch::Tensor rstd = at::empty({numRows}, x.options().dtype(at::kFloat)); + if (numRows == 0) { + return {y, rstd}; + } + + torch::Tensor weightContig = weight.contiguous(); + + // 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)); + const float epsF = static_cast(eps); + + if (x.scalar_type() == at::kBFloat16) { + rmsnorm_ascend_kernel_bf16<<>>( + reinterpret_cast(x.mutable_data_ptr()), + reinterpret_cast(weightContig.mutable_data_ptr()), + reinterpret_cast(y.mutable_data_ptr()), + reinterpret_cast(rstd.mutable_data_ptr()), + numRows, hiddenSize, epsF); + } else if (x.scalar_type() == at::kHalf) { + rmsnorm_ascend_kernel_fp16<<>>( + reinterpret_cast(x.mutable_data_ptr()), + reinterpret_cast(weightContig.mutable_data_ptr()), + reinterpret_cast(y.mutable_data_ptr()), + reinterpret_cast(rstd.mutable_data_ptr()), + numRows, hiddenSize, epsF); + } else { + rmsnorm_ascend_kernel_fp32<<>>( + reinterpret_cast(x.mutable_data_ptr()), + reinterpret_cast(weightContig.mutable_data_ptr()), + reinterpret_cast(y.mutable_data_ptr()), + reinterpret_cast(rstd.mutable_data_ptr()), + numRows, hiddenSize, epsF); + } + return {y, rstd}; +} diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..15216d9d 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,9 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... + +def rmsnorm_ascend( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/ops/ascend/__init__.py b/rl_engine/kernels/ops/ascend/__init__.py index ab85458d..e89f9acc 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 norm # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/norm/__init__.py b/rl_engine/kernels/ops/ascend/norm/__init__.py new file mode 100644 index 00000000..86cf4c9d --- /dev/null +++ b/rl_engine/kernels/ops/ascend/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/kernels/ops/ascend/norm/rmsnorm.py b/rl_engine/kernels/ops/ascend/norm/rmsnorm.py new file mode 100644 index 00000000..a2fd9461 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/norm/rmsnorm.py @@ -0,0 +1,145 @@ +# 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.rms_norm import NativeRMSNormOp + + return NativeRMSNormOp() + + +def _rms_norm_backward( + x_2d: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + grad_out_2d: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """RMSNorm VJP in fp32, reusing the forward-saved rstd. + + With y = x * rstd * w and s = sum(dy * w * x, dim=-1): + dx = rstd * (dy * w) - x * rstd^3 * s / H + dw = sum_rows(dy * x * rstd) + """ + dy_f = grad_out_2d.float() + x_f = x_2d.float() + w_f = weight.float() + rstd_f = rstd.float() + + dyw = dy_f * w_f + s = (dyw * x_f).sum(dim=-1) + hidden = x_2d.size(-1) + dx = rstd_f.unsqueeze(-1) * dyw - x_f * (rstd_f.pow(3) / hidden).unsqueeze(-1) * s.unsqueeze(-1) + dw = (dy_f * x_f * rstd_f.unsqueeze(-1)).sum(dim=0) + return dx.to(x_2d.dtype), dw.to(weight.dtype) + + +class _RMSNormAscendFunction(torch.autograd.Function): + # Autograd wrapper: Ascend C forward + PyTorch-formula backward. + # The CUDA op reuses a dedicated dw/dx kernel; here the backward uses the + # same fp32 VJP formula as the PyTorch reference, reusing the + # forward-saved rstd. + + @staticmethod + def forward(ctx, x, weight, eps): + lead_shape = x.shape[:-1] + hidden = x.size(-1) + + x_2d = x.reshape(-1, hidden).contiguous() + + y, rstd = _C_npu.rmsnorm_ascend(x_2d, weight, float(eps)) + + ctx.save_for_backward(x_2d, weight, rstd) + ctx.eps = eps + ctx.lead_shape = lead_shape + return y.reshape(lead_shape + (hidden,)) + + @staticmethod + def backward(ctx, grad_output): + x_2d, weight, rstd = ctx.saved_tensors + hidden = x_2d.size(-1) + + grad_out_2d = grad_output.reshape(-1, hidden).contiguous() + dx, dw = _rms_norm_backward(x_2d, weight, rstd, grad_out_2d) + + dx = dx.reshape(ctx.lead_shape + (hidden,)) + return dx, dw, None + + +class RMSNormAscendOp: + # Ascend C batch-invariant RMSNorm (forward kernel). + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "rmsnorm_ascend"): + raise RuntimeError( + "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.rmsnorm_ascend kernel.") + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + *, + eps: float = 1e-6, + ) -> torch.Tensor: + return self.forward(x, weight, eps=eps) + + def forward( + self, + x: torch.Tensor, + weight: torch.Tensor, + *, + eps: float = 1e-6, + ) -> torch.Tensor: + if weight.dim() != 1 or weight.shape[0] != x.shape[-1]: + raise ValueError( + f"weight must be 1-D of size x.shape[-1]={x.shape[-1]}, " + f"got tuple(weight.shape)={tuple(weight.shape)}" + ) + + if not _ascend_supported(x) or weight.dtype != x.dtype: + return _fallback_op()(x, weight, eps=eps) + + return _RMSNormAscendFunction.apply(x, weight, eps) + + +def rmsnorm_ascend( + x: torch.Tensor, + weight: torch.Tensor, + *, + eps: float = 1e-6, +) -> torch.Tensor: + return RMSNormAscendOp()(x, weight, eps=eps) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 12ea9b21..ad416d6b 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -107,6 +107,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ASCEND_BATCH_INVARIANT_LOGP = ( "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" ) + ASCEND_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp" # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) PYTORCH_VOCAB_PARALLEL_LOGP = ( "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" @@ -622,6 +623,10 @@ def __init__(self): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + self._priority_map["npu"]["rms_norm"] = [ + OpBackend.ASCEND_RMS_NORM, + OpBackend.PYTORCH_NATIVE_RMS_NORM, + ] 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..ef67fcbb 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -161,11 +161,15 @@ def fake_load_backend(backend): registry.get_op("rms_norm") assert registry._priority_map["npu"].keys() == registry._priority_map["cpu"].keys() - assert loaded[0] == OpBackend.PYTORCH_NATIVE_RMS_NORM + assert loaded[0] == OpBackend.ASCEND_RMS_NORM assert registry._priority_map["npu"]["batch_invariant_logp"] == [ OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + assert registry._priority_map["npu"]["rms_norm"] == [ + OpBackend.ASCEND_RMS_NORM, + OpBackend.PYTORCH_NATIVE_RMS_NORM, + ] def test_npu_available_handles_runtime_failure(monkeypatch): diff --git a/setup.py b/setup.py index 79f882d9..bc006932 100644 --- a/setup.py +++ b/setup.py @@ -3,10 +3,13 @@ import importlib.util import os +import platform +import subprocess +import sysconfig import warnings from pathlib import Path -from setuptools import find_packages, setup +from setuptools import Extension, find_packages, setup def _load_envs_module(): @@ -58,6 +61,110 @@ def _cuda_define_from_env(name: str, macro: str) -> list[str]: return [f"-D{macro}={parsed}"] +_ASCEND_EXTENSION_NAME = "rl_engine._C_npu" +_ASCEND_CPU_DIRS = {"aarch64": "aarch64-linux", "x86_64": "x86_64-linux"} + + +def _find_ascend_home() -> str: + """Locate the CANN toolkit root (must contain bin/bisheng).""" + candidates = [ + os.environ.get("ASCEND_HOME_PATH"), + os.environ.get("ASCEND_TOOLKIT_HOME"), + ] + candidates += [str(p) for p in sorted(Path.home().glob("Ascend/cann-*"), reverse=True)] + candidates.append("/usr/local/Ascend/ascend-toolkit/latest") + for cand in candidates: + if cand and (Path(cand) / "bin" / "bisheng").is_file(): + return cand + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 was requested but no CANN toolkit with bin/bisheng " + "was found. Set ASCEND_HOME_PATH to the toolkit root." + ) + + +def _ascend_extension_spec() -> Extension: + sources = ["csrc/ascend/npu_module.cpp"] + sources += sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + ext = Extension(name=_ASCEND_EXTENSION_NAME, sources=sources) + ext._rl_kernel_ascend = True # intercepted by the custom build_ext below + return ext + + +def _compile_ascend_extension(build_ext, ext) -> None: + """Compile the Ascend C extension with bisheng (torch's BuildExtension + does not know the .asc language, so we drive the compiler directly).""" + torch, _, _ = _load_torch_extension_tools() + try: + import torch_npu + except ModuleNotFoundError as exc: + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch_npu. Install a matching " + "torch_npu build first." + ) from exc + + ascend_home = _find_ascend_home() + cpu_dir = _ASCEND_CPU_DIRS.get(platform.machine()) + if cpu_dir is None: + raise RuntimeError(f"unsupported Ascend host architecture: {platform.machine()}") + arch = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-c220") + + bisheng = os.path.join(ascend_home, "bin", "bisheng") + torch_dir = os.path.dirname(torch.__file__) + tnpu_dir = os.path.dirname(torch_npu.__file__) + + includes = [ + f"-I{os.path.join(ascend_home, cpu_dir, 'asc', 'include')}", + f"-I{os.path.join(torch_dir, 'include')}", + f"-I{os.path.join(torch_dir, 'include', 'torch', 'csrc', 'api', 'include')}", + f"-I{os.path.join(tnpu_dir, 'include')}", + f"-I{sysconfig.get_paths()['include']}", + ] + defines = [f"-DTORCH_EXTENSION_NAME={_ASCEND_EXTENSION_NAME.rsplit('.', 1)[-1]}"] + + build_temp = os.path.join(build_ext.build_temp, "ascend") + os.makedirs(build_temp, exist_ok=True) + + objects = [] + for src in ext.sources: + obj = os.path.join(build_temp, Path(src).name + ".o") + cmd = [bisheng, "-std=c++17", "-O2", "-fPIC", "-c"] + if src.endswith(".asc"): + cmd += ["-x", "asc", f"--cce-aicore-arch={arch}"] + cmd += includes + defines + [src, "-o", obj] + subprocess.check_call(cmd) + objects.append(obj) + + out_path = build_ext.get_ext_fullpath(ext.name) + os.makedirs(os.path.dirname(out_path), exist_ok=True) + link = [bisheng, "-shared", *objects] + for lib_dir, libs in ( + (os.path.join(torch_dir, "lib"), ["torch", "torch_cpu", "torch_python", "c10"]), + (os.path.join(tnpu_dir, "lib"), ["torch_npu"]), + (os.path.join(ascend_home, "runtime", "lib64"), ["ascendcl"]), + (os.path.join(ascend_home, cpu_dir, "lib64"), ["runtime"]), + ): + link.append(f"-L{lib_dir}") + link += [f"-l{name}" for name in libs] + link += [ + f"-Wl,-rpath,{os.path.join(torch_dir, 'lib')}", + f"-Wl,-rpath,{os.path.join(tnpu_dir, 'lib')}", + "-o", + out_path, + ] + subprocess.check_call(link) + + +def _make_build_extension(BuildExtension): + class AscendAwareBuildExtension(BuildExtension): + def build_extension(self, ext): + if getattr(ext, "_rl_kernel_ascend", False): + _compile_ascend_extension(self, ext) + return + super().build_extension(ext) + + return AscendAwareBuildExtension + + _ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( "-Xfatbin", "-compress-all", @@ -231,7 +338,7 @@ def get_extensions(): if enable_sm90 and present_sm90: tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") + nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") if "-lcuda" not in extra_link_args: extra_link_args.append("-lcuda") @@ -266,6 +373,9 @@ def get_extensions(): ) ) + if envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): + extensions.append(_ascend_extension_spec()) + if _native_extension_required() and not extensions: raise RuntimeError( "rl_engine._C was requested but no CUDA/ROCm build environment is available. " @@ -280,7 +390,7 @@ def get_cmdclass(): _, BuildExtension, _ = _load_torch_extension_tools() if BuildExtension is None: return {} - return {"build_ext": BuildExtension} + return {"build_ext": _make_build_extension(BuildExtension)} setup( diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 6572603e..5a5a8a99 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -234,7 +234,7 @@ def test_backward_batch_invariance_slice(): assert torch.equal(x_slice.grad, grad_x_full_sliced) -# 10. Registry dispatch resolves to the native op +# 10. Registry dispatch resolves to the hardware op when available def test_registry_dispatches_rms_norm(): from rl_engine.kernels.registry import kernel_registry @@ -242,6 +242,11 @@ def test_registry_dispatches_rms_norm(): if torch.cuda.is_available() and _HAS_CUDA_RMSNORM: assert isinstance(op, RMSNormCudaOp) assert hasattr(op, "forward") + elif _ascend_rmsnorm_available(): + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp + + assert isinstance(op, RMSNormAscendOp) + assert hasattr(op, "forward") else: assert isinstance(op, NativeRMSNormOp) assert hasattr(op, "forward") and hasattr(op, "forward_fp32") @@ -426,3 +431,89 @@ def test_cuda_rms_norm_masked_dw_layout_invariance(): torch.testing.assert_close(dw1.float(), ref_dw1.float(), atol=atol, rtol=rtol) torch.testing.assert_close(dw2.float(), ref_dw2.float(), atol=atol, rtol=rtol) torch.testing.assert_close(dw3.float(), ref_dw3.float(), atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# Ascend C kernel (rl_engine._C_npu.rmsnorm_ascend) +# --------------------------------------------------------------------------- + +from rl_engine.platforms.device import _npu_available # noqa: E402 + + +def _ascend_rmsnorm_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine.kernels.ops.ascend.norm.rmsnorm import _C_npu, _NPU_EXT_AVAILABLE + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "rmsnorm_ascend") + + +requires_ascend_rmsnorm = pytest.mark.skipif( + not _ascend_rmsnorm_available(), + reason="rmsnorm_ascend kernel not compiled " + "(needs KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host).", +) + + +@requires_ascend_rmsnorm +@pytest.mark.parametrize("hidden", [_HIDDEN, _HEAD_DIM, 1000, 12288]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_ascend_rms_norm_forward_matches_manual_reference(hidden, dtype): + """Ascend forward vs the hand-written fp32 reference (tolerance-based).""" + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp + + torch.manual_seed(0) + x = torch.randn((32, hidden), device="npu", dtype=torch.float32).to(dtype) + w = torch.randn((hidden,), device="npu", dtype=torch.float32).to(dtype) + + y = RMSNormAscendOp()(x, w, eps=_EPS) + ref = _manual_rms_norm(x, w) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(y.float(), ref, atol=atol, rtol=rtol) + + +@requires_ascend_rmsnorm +@pytest.mark.parametrize("hidden", [_HIDDEN, _HEAD_DIM, 12288]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_ascend_rms_norm_backward_matches_native(hidden, dtype): + """Ascend forward + VJP backward vs the native op's forward/backward.""" + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp + + torch.manual_seed(0) + op = RMSNormAscendOp() + native = NativeRMSNormOp() + x = torch.randn((64, hidden), device="npu", dtype=torch.float32).to(dtype) + w = torch.randn((hidden,), device="npu", dtype=torch.float32).to(dtype) + dy = torch.randn((64, hidden), device="npu", dtype=torch.float32).to(dtype) + + y_a, dx_a, dw_a = _run_forward_backward(lambda a, b: op(a, b, eps=_EPS), x, w, dy) + y_n, dx_n, dw_n = _run_forward_backward(lambda a, b: native(a, b, eps=_EPS), x, w, dy) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(y_a.float(), y_n.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(dx_a.float(), dx_n.float(), atol=atol, rtol=rtol) + # dw accumulates over rows, so allow the tolerance to grow with sqrt(rows). + torch.testing.assert_close(dw_a.float(), dw_n.float(), atol=8 * atol, rtol=rtol) + + +@requires_ascend_rmsnorm +@pytest.mark.parametrize("hidden", [_HIDDEN, _HEAD_DIM, 12288]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_ascend_rms_norm_batch_invariance_bitwise(hidden, dtype): + """A row's output must not depend on how many rows share the batch.""" + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp + + torch.manual_seed(0) + op = RMSNormAscendOp() + x_full = torch.randn((64, hidden), device="npu", dtype=torch.float32).to(dtype) + w = torch.randn((hidden,), device="npu", dtype=torch.float32).to(dtype) + + y_full = op(x_full, w, eps=_EPS) + y_slice = op(x_full[:7].clone(), w, eps=_EPS) + y_single = op(x_full[13:14].clone(), w, eps=_EPS) + + assert torch.equal(y_full[:7], y_slice) + assert torch.equal(y_full[13:14], y_single) From 7e05e9a040cff068b1672390f3d49fbde6b2a72b Mon Sep 17 00:00:00 2001 From: chenyang <2082464740@qq.com> Date: Sun, 30 Aug 2026 17:18:59 +0800 Subject: [PATCH 2/4] style: fix black and isort formatting for pre-commit CI --- rl_engine/_C_npu.pyi | 1 - tests/test_rms_norm.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index 15216d9d..5fff0d23 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,7 +8,6 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... - def rmsnorm_ascend( x: torch.Tensor, weight: torch.Tensor, diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 5a5a8a99..b7d7932e 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -444,7 +444,7 @@ def _ascend_rmsnorm_available() -> bool: if not _npu_available(): return False try: - from rl_engine.kernels.ops.ascend.norm.rmsnorm import _C_npu, _NPU_EXT_AVAILABLE + from rl_engine.kernels.ops.ascend.norm.rmsnorm import _NPU_EXT_AVAILABLE, _C_npu except Exception: return False return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "rmsnorm_ascend") From f87bb61778f7461059069cf0c207ed221e4d7bd7 Mon Sep 17 00:00:00 2001 From: chenyang <2082464740@qq.com> Date: Sun, 30 Aug 2026 17:23:59 +0800 Subject: [PATCH 3/4] refactor(bench): restore cuda/triton benchmark, add ascend c as additive --- benchmarks/benchmark_rmsnorm.py | 73 +++++++++++++-------------------- 1 file changed, 29 insertions(+), 44 deletions(-) diff --git a/benchmarks/benchmark_rmsnorm.py b/benchmarks/benchmark_rmsnorm.py index 31b37c26..5227729c 100644 --- a/benchmarks/benchmark_rmsnorm.py +++ b/benchmarks/benchmark_rmsnorm.py @@ -4,47 +4,32 @@ import torch from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp -from rl_engine.platforms.device import device_ctx +from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton -if device_ctx.device_type != "npu": - from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton +try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda - try: - from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda + HAS_CUDA_EXT = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward") +except ImportError: + HAS_CUDA_EXT = False - HAS_CUDA_EXT = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward") - except ImportError: - HAS_CUDA_EXT = False +try: + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp - -def _accel(): - """The active accelerator module (torch.npu on Ascend, torch.cuda otherwise).""" - if device_ctx.device_type == "npu": - return torch.npu - return torch.cuda - - -def _maybe_ascend_op(): - """The Ascend C op, or None when unavailable (no NPU / not built).""" - if device_ctx.device_type != "npu": - return None - try: - from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp - - return RMSNormAscendOp() - except (ImportError, OSError, RuntimeError): - return None + HAS_ASCEND_EXT = True +except (ImportError, OSError, RuntimeError): + HAS_ASCEND_EXT = False def bench(fn, x, w, dy, warmup=20, iters=100): - accel = _accel() + sync = torch.npu.synchronize if x.device.type == "npu" else torch.cuda.synchronize for _ in range(warmup): x.grad = None w.grad = None y = fn(x, w) y.backward(dy) - accel.synchronize() + sync() start = time.time() for _ in range(iters): @@ -52,7 +37,7 @@ def bench(fn, x, w, dy, warmup=20, iters=100): w.grad = None y = fn(x, w) y.backward(dy) - accel.synchronize() + sync() return (time.time() - start) * 1000.0 / iters @@ -60,11 +45,11 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--T", type=int, default=1024) parser.add_argument("--H", type=int, default=4096) - parser.add_argument("--dtype", choices=["fp16", "bf16", "fp32"], default="bf16") + parser.add_argument("--dtype", choices=["fp16", "bf16"], default="bf16") args = parser.parse_args() - dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[args.dtype] - device = "npu" if device_ctx.device_type == "npu" else "cuda" + dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16 + device = "npu" if torch.npu.is_available() else "cuda" T, H = args.T, args.H torch.manual_seed(0) @@ -84,19 +69,19 @@ def make_inputs(): t_ref = bench(lambda a, b: native.forward(a, b), x, w, dy) print(f"pytorch ref : {t_ref:.4f} ms") - if device == "cuda": - x, w = make_inputs() - t_tri = bench(lambda a, b: rmsnorm_triton(a, b), x, w, dy) - print(f"triton : {t_tri:.4f} ms | speedup vs ref: {t_ref / t_tri:.2f}x") + x, w = make_inputs() + t_tri = bench(lambda a, b: rmsnorm_triton(a, b), x, w, dy) + print(f"triton : {t_tri:.4f} ms | speedup vs ref: {t_ref / t_tri:.2f}x") - if HAS_CUDA_EXT: - x, w = make_inputs() - t_cuda = bench(lambda a, b: rmsnorm_cuda(a, b), x, w, dy) - print(f"cuda : {t_cuda:.4f} ms | speedup vs ref: {t_ref / t_cuda:.2f}x") - else: - print("cuda : skipped, extension is not built") + if HAS_CUDA_EXT: + x, w = make_inputs() + t_cuda = bench(lambda a, b: rmsnorm_cuda(a, b), x, w, dy) + print(f"cuda : {t_cuda:.4f} ms | speedup vs ref: {t_ref / t_cuda:.2f}x") else: - ascend_op = _maybe_ascend_op() + print("cuda : skipped, extension is not built") + + if device == "npu": + ascend_op = RMSNormAscendOp() if HAS_ASCEND_EXT else None if ascend_op is not None: x, w = make_inputs() t_asc = bench(lambda a, b: ascend_op(a, b), x, w, dy) From 3cd6e749705829d521080b6456a1afdf70fe8b4e Mon Sep 17 00:00:00 2001 From: chenyang <2082464740@qq.com> Date: Tue, 1 Sep 2026 14:51:21 +0800 Subject: [PATCH 4/4] fix(ascend): register rms_norm gtest candidate, make kernel bitwise identical to reference --- csrc/ascend/npu_module.cpp | 8 +- csrc/ascend/rmsnorm_ascend.asc | 272 ++++++------------- rl_engine/_C_npu.pyi | 4 +- rl_engine/kernels/gtest/operator_specs.py | 1 + rl_engine/kernels/ops/ascend/norm/rmsnorm.py | 20 +- scripts/check_operator.py | 14 +- tests/test_rms_norm.py | 29 ++ 7 files changed, 148 insertions(+), 200 deletions(-) diff --git a/csrc/ascend/npu_module.cpp b/csrc/ascend/npu_module.cpp index f94dac4c..1bfb12aa 100644 --- a/csrc/ascend/npu_module.cpp +++ b/csrc/ascend/npu_module.cpp @@ -14,9 +14,9 @@ std::vector batch_invariant_logp_ascend_forward(torch::Tensor log torch::Tensor target, int64_t ignore_index); -std::vector rmsnorm_ascend_forward(torch::Tensor x, - torch::Tensor weight, - double eps); +torch::Tensor rmsnorm_ascend_forward(torch::Tensor x, + torch::Tensor weight, + torch::Tensor rstd); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -25,5 +25,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) "Batch-invariant selected-token log-probability (Ascend C forward)"); m.def("rmsnorm_ascend", &rmsnorm_ascend_forward, - "Batch-invariant RMSNorm (Ascend C forward)"); + "Batch-invariant RMSNorm (Ascend C forward, rstd precomputed)"); } diff --git a/csrc/ascend/rmsnorm_ascend.asc b/csrc/ascend/rmsnorm_ascend.asc index a87660f4..b88341ef 100644 --- a/csrc/ascend/rmsnorm_ascend.asc +++ b/csrc/ascend/rmsnorm_ascend.asc @@ -3,17 +3,26 @@ // // Batch-invariant RMSNorm, Ascend C (CANN) forward kernel. // -// y[n, :] = x[n, :] * rsqrt(mean(x[n, :]^2) + eps) * weight[:] +// y[n, :] = x[n, :] * rstd[n] * weight[:] +// rstd[n] = rsqrt(mean(x[n, :]^2) + eps) (precomputed on the host side) +// +// The row scale rstd is computed by the caller with the exact PyTorch ops of +// the reference (rl_engine/kernels/ops/pytorch/norm/rms_norm.py: +// x.float().pow(2).mean(-1) followed by torch.rsqrt(var + eps)). 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. // // Mirrors the CUDA kernel in csrc/cuda/rmsnorm.cu: -// - input : x [N, H] contiguous, fp32 / bf16 / fp16; weight [H] same dtype -// - output : y [N, H] same dtype as x, rstd [N] fp32 -// (rstd = rsqrt(mean(x^2) + eps), saved for the autograd backward) +// - input : x [N, H] contiguous, fp32 / bf16 / fp16; weight [H] same dtype; +// rstd [N] fp32 (saved by the caller for the autograd backward) +// - output : y [N, H] same dtype as x // -// Batch-invariance: every row is processed end-to-end by exactly one AI core -// block with a fixed tile size and a fixed reduction order (fp32 sum of -// squares over fixed-order tiles). The instruction sequence for a row depends -// only on H, never on N or on the block the row happens to land on. +// Batch-invariance: rstd[n] depends only on row n (torch's last-dim mean +// order is a function of H 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 H, never on N 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. @@ -29,15 +38,15 @@ namespace { // Elements per hidden tile. Fixed for all rows and batch sizes; this is what -// makes the reduction order batch-invariant. UB budget (in/out/weight tiles + -// fp32 tile + fp32 weight tile + square tile + reduce scratch) stays under -// the 192 KB UB of DAV_2201 SoCs even for fp32 in/out tiles. +// keeps the elementwise pass batch-invariant. UB budget (in/out/weight tiles +// + fp32 tile + fp32 weight tile + rstd staging) stays well under the +// 192 KB UB of DAV_2201 SoCs even for fp32 in/out tiles. 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 -// staging buffers (aBuf/rBuf) and the per-row reduce scratch slots. +// rstd staging buffer. constexpr int64_t MAX_CHUNK_ROWS = 64; template @@ -47,15 +56,13 @@ public: __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, - GM_ADDR y, GM_ADDR rstd, + GM_ADDR y, int64_t numRows, - int64_t hiddenSize, - float eps) + int64_t hiddenSize) { numRows_ = numRows; hiddenSize_ = hiddenSize; - eps_ = eps; singleTile_ = hiddenSize <= static_cast(TILE_LENGTH); // Small rows are dominated by per-row pipeline flag round-trips, so // process rowsPerChunk_ contiguous rows per iteration and amortize @@ -70,31 +77,24 @@ public: } xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x)); weightGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(weight)); - yGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(y)); 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(wQueue_, 1, TILE_LENGTH * sizeof(T)); pipe_->InitBuffer(fp32Buf_, TILE_LENGTH * sizeof(float)); pipe_->InitBuffer(wFp32Buf_, TILE_LENGTH * sizeof(float)); - pipe_->InitBuffer(sqBuf_, TILE_LENGTH * sizeof(float)); - pipe_->InitBuffer(reduceBuf_, TILE_LENGTH * sizeof(float)); - // 2 KB: per-row reduce scratch at scalar[8*r] (each slot 32-byte - // aligned) for r in [0, MAX_CHUNK_ROWS); the single-row path only - // uses floats [0,8) for scratch and [8,16) for rstd staging. + // 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)); - // 512 B: aBuf [0,64) holds meanSq+eps per chunk row, rBuf [64,128) - // holds the refined rstd per chunk row (flushed to GM in one burst). - pipe_->InitBuffer(stageBuf_, MAX_CHUNK_ROWS * 2 * 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. - eventVS_ = pipe_->FetchEventID(AscendC::HardEvent::V_S); - eventSV_ = pipe_->FetchEventID(AscendC::HardEvent::S_V); - eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3); - eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + eventSMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE2); + eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S); } __aicore__ inline void Process() @@ -113,12 +113,16 @@ public: (rowStart + segLen < numRows_) ? rowStart + segLen : numRows_; for (int64_t row = rowStart; row < rowEnd; row += rowsPerChunk_) { const int64_t remaining = rowEnd - row; - ProcessRowChunk(row, remaining < rowsPerChunk_ ? remaining : rowsPerChunk_); + 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()) { - ProcessRow(row); + LoadRstd(row, 1); + ProcessRow(row, scalarBuf_.Get().GetValue(0)); } } @@ -181,58 +185,36 @@ private: AscendC::Copy(dst, src, 64, repeat, AscendC::CopyRepeatParams{1, 1, 8, 8}); } - __aicore__ inline float TileSumSquare(const AscendC::LocalTensor& fLocal, - uint32_t count) + // Load rstd[row0 : row0+rows] into scalarBuf_[0:rows]. + __aicore__ inline void LoadRstd(int64_t row0, int64_t rows) { - // Square in place (destroys the tile's x values), then reduce with a - // fixed order. src and work tensors are distinct buffers. - AscendC::Mul(fLocal, fLocal, fLocal, count); - AscendC::LocalTensor rTmp = reduceBuf_.Get(); AscendC::LocalTensor scalar = scalarBuf_.Get(); - AscendC::ReduceSum(scalar, fLocal, rTmp, static_cast(count)); - WaitVector(); // vector -> scalar read - return scalar.GetValue(0); + // 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) + __aicore__ inline void ProcessRow(int64_t row, float rstd) { - float sumSq = 0.0f; - if (singleTile_) { // Fast path: the whole row stays resident in fp32Buf_ across the - // sum-of-squares and the scaling, so x is read from GM only once. + // scaling, so x is read from GM only once. const uint32_t count = static_cast(hiddenSize_); AscendC::LocalTensor fLocal = LoadTileFp32(row, 0, count); - // Square into sqBuf_ so the row values survive in fLocal; reduce - // scratch lives in reduceBuf_ (src/work must not alias). - AscendC::LocalTensor sq = sqBuf_.Get(); - AscendC::Mul(sq, fLocal, fLocal, count); - AscendC::LocalTensor rTmp = reduceBuf_.Get(); - AscendC::LocalTensor scalar = scalarBuf_.Get(); - AscendC::ReduceSum(scalar, sq, rTmp, static_cast(count)); - WaitVector(); // vector -> scalar read - sumSq = scalar.GetValue(0); - - const float rstd = ComputeRstd(sumSq, row); ScaleStoreTile(row, 0, count, fLocal, rstd); FreeTile(); return; } + // Fixed tile order over the row; rstd is the same for every tile. const int64_t tileCount = (hiddenSize_ + TILE_LENGTH - 1) / TILE_LENGTH; - - // Pass 1: sum of squares with a fixed tile order. - for (int64_t tile = 0; tile < tileCount; ++tile) { - const int64_t start = tile * TILE_LENGTH; - const uint32_t count = TileCount(start); - AscendC::LocalTensor fLocal = LoadTileFp32(row, start, count); - sumSq += TileSumSquare(fLocal, count); - FreeTile(); - } - - const float rstd = ComputeRstd(sumSq, row); - - // Pass 2: y = x * rstd * w with the same fixed tile order. for (int64_t tile = 0; tile < tileCount; ++tile) { const int64_t start = tile * TILE_LENGTH; const uint32_t count = TileCount(start); @@ -245,8 +227,8 @@ private: // 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 reduction order is identical to ProcessRow (fp32 sum over the - // row's single tile), so numerics are unchanged. + // 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 H = hiddenSize_; @@ -269,46 +251,13 @@ private: AscendC::Cast(fLocal, xLocal, AscendC::RoundMode::CAST_NONE, count); } - // Squares of the whole chunk, then one aligned ReduceSum per row. - // All reduces complete before the first scalar read (single V_S). - AscendC::LocalTensor sq = sqBuf_.Get(); - AscendC::Mul(sq, fLocal, fLocal, count); - AscendC::LocalTensor rTmp = reduceBuf_.Get(); + // Scale every row of the chunk with its own rstd (staged by LoadRstd + // before this call); scalar-register operands of Muls/Mul need no + // S_V flag (no UB dependency). AscendC::LocalTensor scalar = scalarBuf_.Get(); - for (int64_t r = 0; r < rows; ++r) { - AscendC::ReduceSum(scalar[8 * r], sq[r * H], rTmp, - static_cast(H)); - } - WaitVector(); // vector -> scalar reads - - // a[r] = mean(x_r^2) + eps staged for the vector Rsqrt. - AscendC::LocalTensor stage = stageBuf_.Get(); - AscendC::LocalTensor aBuf = stage; - AscendC::LocalTensor rBuf = stage[MAX_CHUNK_ROWS]; - for (int64_t r = 0; r < rows; ++r) { - aBuf.SetValue(r, scalar.GetValue(static_cast(8 * r)) / - static_cast(H) + - eps_); - } - AscendC::SetFlag(eventSV_); // scalar write -> vector op - AscendC::WaitFlag(eventSV_); - const uint32_t rPad = static_cast((rows + 7) / 8 * 8); - AscendC::Rsqrt(rBuf, aBuf, rPad); - WaitVector(); // vector -> scalar reads - // Newton-Raphson refinement on the scalar unit (see ComputeRstd). - for (int64_t r = 0; r < rows; ++r) { - const float a = aBuf.GetValue(static_cast(r)); - float rstd = rBuf.GetValue(static_cast(r)); - rstd = rstd * (1.5f - 0.5f * a * rstd * rstd); - rstd = rstd * (1.5f - 0.5f * a * rstd * rstd); - rBuf.SetValue(static_cast(r), rstd); - } - - // Scale every row of the chunk; scalar-register operands of Muls/Mul - // need no S_V flag (no UB dependency). AscendC::LocalTensor wFp32 = wFp32Buf_.Get(); for (int64_t r = 0; r < rows; ++r) { - const float rstd = rBuf.GetValue(static_cast(r)); + const float rstd = scalar.GetValue(static_cast(r)); AscendC::Muls(fLocal[r * H], fLocal[r * H], rstd, static_cast(H)); AscendC::Mul(fLocal[r * H], fLocal[r * H], wFp32, static_cast(H)); } @@ -326,49 +275,6 @@ private: AscendC::DataCopyPad(yGm_[row0 * H], yLocal, outParams); outQueue_.FreeTensor(yLocal); FreeTile(); - - // Flush the chunk's rstd values in one contiguous burst. - AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out - AscendC::WaitFlag(eventSMTE3_); - AscendC::DataCopyExtParams rstdParams{ - 1, static_cast(rows * sizeof(float)), 0, 0, 0}; - AscendC::DataCopyPad(rstdGm_[row0], rBuf, rstdParams); - // Drain MTE3 before the next chunk stages new values into rBuf. - AscendC::SetFlag(eventMTE3S_); - AscendC::WaitFlag(eventMTE3S_); - } - - // rstd = rsqrt(mean(x^2) + eps), staged to rstdGm_[row] for the backward. - __aicore__ inline float ComputeRstd(float sumSq, int64_t row) - { - // The scalar unit has no rsqrt, so run a 1-element vector Rsqrt - // (count padded to 8; scalarBuf_ is 32 B aligned). The Rsqrt - // instruction is only a ~2^-9 relative approximation, so refine it - // with Newton-Raphson on the scalar unit: r = r*(1.5 - 0.5*a*r*r) - // converges to rsqrt(a) at ~1 ulp after two iterations. - const float meanSqEps = sumSq / static_cast(hiddenSize_) + eps_; - AscendC::LocalTensor scalar = scalarBuf_.Get(); - scalar.SetValue(0, meanSqEps); - AscendC::SetFlag(eventSV_); // scalar write -> vector op - AscendC::WaitFlag(eventSV_); - AscendC::Rsqrt(scalar, scalar, 8); - WaitVector(); // vector -> scalar read - float rstd = scalar.GetValue(0); - rstd = rstd * (1.5f - 0.5f * meanSqEps * rstd * rstd); - rstd = rstd * (1.5f - 0.5f * meanSqEps * rstd * rstd); - - // GlobalTensor.SetValue is unreliable on hardware (cannbot - // ascendc-precision-debug common-traps), so stage rstd in UB and - // DataCopyPad it to GM instead of a scalar GM store. - scalar.SetValue(8, rstd); - AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out - AscendC::WaitFlag(eventSMTE3_); - AscendC::DataCopyExtParams outParams{1, sizeof(float), 0, 0, 0}; - AscendC::DataCopyPad(rstdGm_[row], scalar[8], outParams); - // Drain MTE3 before the next row stages a new value into scalarBuf_. - AscendC::SetFlag(eventMTE3S_); - AscendC::WaitFlag(eventMTE3S_); - return rstd; } // y tile = x tile * rstd * w tile, cast back to T and copied to GM. @@ -398,13 +304,6 @@ private: outQueue_.FreeTensor(yLocal); } - // Wait until all outstanding vector-pipe results are readable as scalars. - __aicore__ inline void WaitVector() - { - AscendC::SetFlag(eventVS_); - AscendC::WaitFlag(eventVS_); - } - __aicore__ inline uint32_t TileCount(int64_t start) const { const int64_t remaining = hiddenSize_ - start; @@ -414,25 +313,19 @@ private: AscendC::TPipe* pipe_; AscendC::GlobalTensor xGm_; AscendC::GlobalTensor weightGm_; - AscendC::GlobalTensor yGm_; AscendC::GlobalTensor rstdGm_; + AscendC::GlobalTensor yGm_; AscendC::TQue inQueue_; AscendC::TQue wQueue_; AscendC::TQue outQueue_; AscendC::TBuf fp32Buf_; AscendC::TBuf wFp32Buf_; - AscendC::TBuf sqBuf_; - AscendC::TBuf reduceBuf_; AscendC::TBuf scalarBuf_; - AscendC::TBuf stageBuf_; AscendC::LocalTensor inTile_; - AscendC::TEventID eventVS_; - AscendC::TEventID eventSV_; - AscendC::TEventID eventSMTE3_; - AscendC::TEventID eventMTE3S_; + AscendC::TEventID eventSMTE2_; + AscendC::TEventID eventMTE2S_; int64_t numRows_; int64_t hiddenSize_; - float eps_; bool singleTile_; int64_t rowsPerChunk_; }; @@ -440,38 +333,38 @@ private: } // namespace extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_fp32( - GM_ADDR x, GM_ADDR weight, GM_ADDR y, GM_ADDR rstd, - int64_t numRows, int64_t hiddenSize, float eps) + GM_ADDR x, GM_ADDR weight, GM_ADDR rstd, GM_ADDR y, + int64_t numRows, int64_t hiddenSize) { AscendC::TPipe pipe; KernelRmsNorm op(&pipe); - op.Init(x, weight, y, rstd, numRows, hiddenSize, eps); + op.Init(x, weight, rstd, y, numRows, hiddenSize); op.Process(); } extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_bf16( - GM_ADDR x, GM_ADDR weight, GM_ADDR y, GM_ADDR rstd, - int64_t numRows, int64_t hiddenSize, float eps) + GM_ADDR x, GM_ADDR weight, GM_ADDR rstd, GM_ADDR y, + int64_t numRows, int64_t hiddenSize) { AscendC::TPipe pipe; KernelRmsNorm op(&pipe); - op.Init(x, weight, y, rstd, numRows, hiddenSize, eps); + op.Init(x, weight, rstd, y, numRows, hiddenSize); op.Process(); } extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_fp16( - GM_ADDR x, GM_ADDR weight, GM_ADDR y, GM_ADDR rstd, - int64_t numRows, int64_t hiddenSize, float eps) + GM_ADDR x, GM_ADDR weight, GM_ADDR rstd, GM_ADDR y, + int64_t numRows, int64_t hiddenSize) { AscendC::TPipe pipe; KernelRmsNorm op(&pipe); - op.Init(x, weight, y, rstd, numRows, hiddenSize, eps); + op.Init(x, weight, rstd, y, numRows, hiddenSize); op.Process(); } -std::vector rmsnorm_ascend_forward(torch::Tensor x, - torch::Tensor weight, - double eps) +torch::Tensor rmsnorm_ascend_forward(torch::Tensor x, + torch::Tensor weight, + torch::Tensor rstd) { TORCH_CHECK(x.is_privateuseone(), "x must be on an NPU device"); TORCH_CHECK(x.dim() == 2, "x must be 2-D [N, H]"); @@ -484,14 +377,18 @@ std::vector rmsnorm_ascend_forward(torch::Tensor x, TORCH_CHECK(weight.dim() == 1 && weight.numel() == x.size(-1), "weight must be 1-D of size x.size(-1)"); TORCH_CHECK(weight.scalar_type() == x.scalar_type(), "weight dtype must match x dtype"); + TORCH_CHECK(rstd.is_privateuseone(), "rstd must be on the same NPU device as x"); + TORCH_CHECK(rstd.dim() == 1 && rstd.numel() == x.size(0), + "rstd must be 1-D of size x.size(0)"); + TORCH_CHECK(rstd.scalar_type() == at::kFloat, "rstd must be fp32"); + TORCH_CHECK(rstd.is_contiguous(), "rstd must be contiguous"); const int64_t numRows = x.size(0); const int64_t hiddenSize = x.size(1); torch::Tensor y = at::empty_like(x); - torch::Tensor rstd = at::empty({numRows}, x.options().dtype(at::kFloat)); if (numRows == 0) { - return {y, rstd}; + return y; } torch::Tensor weightContig = weight.contiguous(); @@ -501,29 +398,28 @@ std::vector rmsnorm_ascend_forward(torch::Tensor x, // queued initializer) for the same reason. auto aclStream = c10_npu::getCurrentNPUStream().stream(true); const uint32_t blockNum = static_cast(std::min(numRows, MAX_BLOCKS)); - const float epsF = static_cast(eps); if (x.scalar_type() == at::kBFloat16) { rmsnorm_ascend_kernel_bf16<<>>( reinterpret_cast(x.mutable_data_ptr()), reinterpret_cast(weightContig.mutable_data_ptr()), + reinterpret_cast(rstd.data_ptr()), reinterpret_cast(y.mutable_data_ptr()), - reinterpret_cast(rstd.mutable_data_ptr()), - numRows, hiddenSize, epsF); + numRows, hiddenSize); } else if (x.scalar_type() == at::kHalf) { rmsnorm_ascend_kernel_fp16<<>>( reinterpret_cast(x.mutable_data_ptr()), reinterpret_cast(weightContig.mutable_data_ptr()), + reinterpret_cast(rstd.data_ptr()), reinterpret_cast(y.mutable_data_ptr()), - reinterpret_cast(rstd.mutable_data_ptr()), - numRows, hiddenSize, epsF); + numRows, hiddenSize); } else { rmsnorm_ascend_kernel_fp32<<>>( reinterpret_cast(x.mutable_data_ptr()), reinterpret_cast(weightContig.mutable_data_ptr()), + reinterpret_cast(rstd.data_ptr()), reinterpret_cast(y.mutable_data_ptr()), - reinterpret_cast(rstd.mutable_data_ptr()), - numRows, hiddenSize, epsF); + numRows, hiddenSize); } - return {y, rstd}; + return y; } diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index 5fff0d23..5b18296b 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -11,5 +11,5 @@ def batch_invariant_logp_ascend( def rmsnorm_ascend( x: torch.Tensor, weight: torch.Tensor, - eps: float, -) -> list[torch.Tensor]: ... + rstd: torch.Tensor, +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index ca4a462e..7486f67a 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -42,6 +42,7 @@ def _load_object(path: str) -> Any: "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp", "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", + "ascend": "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp", }, grad_input_names=("x", "weight"), ), diff --git a/rl_engine/kernels/ops/ascend/norm/rmsnorm.py b/rl_engine/kernels/ops/ascend/norm/rmsnorm.py index a2fd9461..80cd30fe 100644 --- a/rl_engine/kernels/ops/ascend/norm/rmsnorm.py +++ b/rl_engine/kernels/ops/ascend/norm/rmsnorm.py @@ -66,10 +66,9 @@ def _rms_norm_backward( class _RMSNormAscendFunction(torch.autograd.Function): - # Autograd wrapper: Ascend C forward + PyTorch-formula backward. - # The CUDA op reuses a dedicated dw/dx kernel; here the backward uses the - # same fp32 VJP formula as the PyTorch reference, reusing the - # forward-saved rstd. + # 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, weight, eps): @@ -78,7 +77,18 @@ def forward(ctx, x, weight, eps): x_2d = x.reshape(-1, hidden).contiguous() - y, rstd = _C_npu.rmsnorm_ascend(x_2d, weight, float(eps)) + # rstd is computed with the exact torch ops of the PyTorch reference + # (rl_engine/kernels/ops/pytorch/norm/rms_norm.py): fp32 mean of + # squares + torch.rsqrt. The Ascend C kernel then only performs the + # elementwise y = x * rstd * w scale and the round-to-nearest-even + # cast, which are order-free IEEE ops — this makes the fused output + # bitwise identical to NativeRMSNormOp 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.rmsnorm_ascend(x_2d, weight, rstd) ctx.save_for_backward(x_2d, weight, rstd) ctx.eps = eps diff --git a/scripts/check_operator.py b/scripts/check_operator.py index 9dbca48d..4bc33164 100644 --- a/scripts/check_operator.py +++ b/scripts/check_operator.py @@ -35,12 +35,23 @@ def _parse_dtype(value: str) -> torch.dtype: raise ValueError(f"unsupported dtype: {value}") +def _npu_available() -> bool: + """torch.npu only exists after torch_npu is imported; probe defensively.""" + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + def _select_device(value: str) -> torch.device: if value == "auto": return torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device(value) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("--device cuda was requested, but CUDA is not available") + if device.type == "npu" and not _npu_available(): + raise RuntimeError("--device npu was requested, but no Ascend NPU is available") return device @@ -73,7 +84,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--candidate", default="pytorch", - help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton.", + help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton, " + "ascend.", ) parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") parser.add_argument("--device", default="auto") diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index b7d7932e..14e89322 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -475,6 +475,35 @@ def test_ascend_rms_norm_forward_matches_manual_reference(hidden, dtype): torch.testing.assert_close(y.float(), ref, atol=atol, rtol=rtol) +@requires_ascend_rmsnorm +@pytest.mark.parametrize("hidden", [_HIDDEN, _HEAD_DIM, 1000, 12288]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("rows", [1, 7, 64, 257]) +def test_ascend_rms_norm_bitwise_identical_to_native(hidden, dtype, rows): + """Ascend forward must be bitwise identical to the PyTorch reference. + + The row scale rstd is computed with the exact reference torch ops + (fp32 mean of squares + torch.rsqrt), so the fused kernel — which only + performs order-free elementwise multiplies and an RNE cast — must match + NativeRMSNormOp bit-for-bit on every dtype, including the H > tile and + H % 8 != 0 paths. + """ + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp + + torch.manual_seed(0) + op = RMSNormAscendOp() + native = NativeRMSNormOp() + + x = torch.randn((rows, hidden), device="npu", dtype=torch.float32).to(dtype) + w = torch.randn((hidden,), device="npu", dtype=torch.float32).to(dtype) + + y_ascend = op(x, w, eps=_EPS) + y_native = native(x, w, eps=_EPS) + + assert y_ascend.dtype == y_native.dtype == dtype + assert torch.equal(y_ascend, y_native) + + @requires_ascend_rmsnorm @pytest.mark.parametrize("hidden", [_HIDDEN, _HEAD_DIM, 12288]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])