diff --git a/benchmarks/benchmark_rmsnorm.py b/benchmarks/benchmark_rmsnorm.py index 81325281..5227729c 100644 --- a/benchmarks/benchmark_rmsnorm.py +++ b/benchmarks/benchmark_rmsnorm.py @@ -14,14 +14,22 @@ except ImportError: HAS_CUDA_EXT = False +try: + from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp + + HAS_ASCEND_EXT = True +except (ImportError, OSError, RuntimeError): + HAS_ASCEND_EXT = False + def bench(fn, x, w, dy, warmup=20, iters=100): + 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) - torch.cuda.synchronize() + sync() start = time.time() for _ in range(iters): @@ -29,7 +37,7 @@ def bench(fn, x, w, dy, warmup=20, iters=100): w.grad = None y = fn(x, w) y.backward(dy) - torch.cuda.synchronize() + sync() return (time.time() - start) * 1000.0 / iters @@ -41,7 +49,7 @@ def main(): args = parser.parse_args() dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16 - device = "cuda" + device = "npu" if torch.npu.is_available() else "cuda" T, H = args.T, args.H torch.manual_seed(0) @@ -72,6 +80,15 @@ def make_inputs(): else: 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) + 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__": 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..1bfb12aa --- /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); + +torch::Tensor rmsnorm_ascend_forward(torch::Tensor x, + torch::Tensor weight, + torch::Tensor rstd); + +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, rstd precomputed)"); +} diff --git a/csrc/ascend/rmsnorm_ascend.asc b/csrc/ascend/rmsnorm_ascend.asc new file mode 100644 index 00000000..b88341ef --- /dev/null +++ b/csrc/ascend/rmsnorm_ascend.asc @@ -0,0 +1,425 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant RMSNorm, Ascend C (CANN) forward kernel. +// +// 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; +// rstd [N] fp32 (saved by the caller for the autograd backward) +// - output : y [N, H] same dtype as x +// +// 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. + +#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. 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 +// rstd staging buffer. +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 rstd, + GM_ADDR y, + int64_t numRows, + int64_t hiddenSize) + { + numRows_ = numRows; + hiddenSize_ = hiddenSize; + 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)); + 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)); + // 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 (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; + 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, 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}); + } + + // 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) + { + if (singleTile_) { + // Fast path: the whole row stays resident in fp32Buf_ across the + // scaling, so x is read from GM only once. + const uint32_t count = static_cast(hiddenSize_); + AscendC::LocalTensor fLocal = LoadTileFp32(row, 0, count); + 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; + 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 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_; + 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); + } + + // 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(); + AscendC::LocalTensor wFp32 = wFp32Buf_.Get(); + for (int64_t r = 0; r < rows; ++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)); + } + + 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(); + } + + // 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); + } + + __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 rstdGm_; + AscendC::GlobalTensor yGm_; + AscendC::TQue inQueue_; + AscendC::TQue wQueue_; + AscendC::TQue outQueue_; + AscendC::TBuf fp32Buf_; + AscendC::TBuf wFp32Buf_; + AscendC::TBuf scalarBuf_; + AscendC::LocalTensor inTile_; + AscendC::TEventID eventSMTE2_; + AscendC::TEventID eventMTE2S_; + int64_t numRows_; + int64_t hiddenSize_; + bool singleTile_; + int64_t rowsPerChunk_; +}; + +} // namespace + +extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_fp32( + 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, rstd, y, numRows, hiddenSize); + op.Process(); +} + +extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_bf16( + 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, rstd, y, numRows, hiddenSize); + op.Process(); +} + +extern "C" __global__ __vector__ void rmsnorm_ascend_kernel_fp16( + 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, rstd, y, numRows, hiddenSize); + op.Process(); +} + +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]"); + 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"); + 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); + if (numRows == 0) { + return y; + } + + 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)); + + 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()), + 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()), + 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()), + numRows, hiddenSize); + } + return y; +} diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..5b18296b 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,8 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... +def rmsnorm_ascend( + x: torch.Tensor, + weight: 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/__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..80cd30fe --- /dev/null +++ b/rl_engine/kernels/ops/ascend/norm/rmsnorm.py @@ -0,0 +1,155 @@ +# 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: 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): + lead_shape = x.shape[:-1] + hidden = x.size(-1) + + x_2d = x.reshape(-1, hidden).contiguous() + + # 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 + 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/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/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..14e89322 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,118 @@ 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 _NPU_EXT_AVAILABLE, _C_npu + 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, 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]) +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)