diff --git a/csrc/ascend/lm_head_ascend.asc b/csrc/ascend/lm_head_ascend.asc new file mode 100644 index 00000000..4d71aca1 --- /dev/null +++ b/csrc/ascend/lm_head_ascend.asc @@ -0,0 +1,390 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant LM-head projection, Ascend C (CANN) forward kernel. +// +// out[n, v] = dot(hidden[n, :], weight[v, :]) (+ bias[v]) +// +// Mirrors the SM90 CUDA kernel in csrc/cuda/embedding_lm_head_sm90.cu: +// - input : hidden [N, H] contiguous, fp32 / bf16 / fp16; weight [V, H] +// contiguous, same dtype; bias [V] optional (cast to fp32) +// - output : [N, V] in the compute dtype (hidden dtype, or fp32 when +// output_fp32 -- the host wrapper pre-casts the inputs, exactly +// like the CUDA wrapper) +// - every output element owns its full hidden-dimension reduction inside +// one block: no Split-K, no second-pass merge, so the reduction order +// depends only on H, never on N or the block the element lands on. +// +// The math follows the CUDA kernel's structure: fp32 accumulation over a +// fixed tile order (products -> per-tile sum -> sequential scalar +// accumulation), bias added in fp32, final cast to the output dtype with +// round-to-nearest. The per-tile sums use the Ascend vector unit's fixed +// hardware reduction tree (fixed per tile size) instead of CUDA's +// warp-shuffle tree, so cross-platform bitwise parity with the CUDA kernel +// is not claimed -- the guarantee is the same one the CUDA kernel provides +// on its platform: batch-invariant determinism. +// +// Batch-invariance: one output element per block iteration, fixed tile +// order over H, rows/elements strided across blocks (MAX_BLOCKS cap), so a +// row's logits are bitwise identical across batch sizes, row positions, and +// block assignments on the NPU. +// +// 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 elements and batch sizes; this is +// what makes the reduction order batch-invariant. UB budget (two native +// tiles + four fp32 tiles) stays well under the 192 KB UB of current SoCs. +constexpr uint32_t TILE_LENGTH = 4096; +// Cap on launched blocks. Elements are strided across blocks, so launching +// fewer blocks than elements is fine and never changes per-element numerics. +constexpr int64_t MAX_BLOCKS = 128; + +template +class KernelLmHead { +public: + __aicore__ inline KernelLmHead(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR hidden, + GM_ADDR weight, + GM_ADDR bias, + GM_ADDR output, + int64_t numRows, + int64_t vocabSize, + int64_t hiddenSize, + bool hasBias) + { + numRows_ = numRows; + vocabSize_ = vocabSize; + hiddenSize_ = hiddenSize; + hasBias_ = hasBias; + hiddenGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(hidden)); + weightGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(weight)); + biasGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(bias)); + outputGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(output)); + pipe_->InitBuffer(hQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(wQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(hFp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(wFp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(prodBuf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(reduceBuf_, TILE_LENGTH * sizeof(float)); + // 64 B: floats [0,8) for reduce/log scratch, floats [8,16) for the + // output staging slots (both halves 32-byte aligned for DataCopyPad). + pipe_->InitBuffer(scalarBuf_, 64); + // 32 B window for reading bias[v] via DataCopyPad (GM scalar + // GetValue/SetValue are unreliable on hardware; see cannbot + // ascendc-precision-debug common-traps). + pipe_->InitBuffer(biasBuf_, 32); + + // Intra-core pipeline events. NOTE: AscendC::SyncAll() is a cross-core + // barrier and deadlocks when more blocks are launched than there are + // physical cores (resident blocks wait for unscheduled ones), so all + // synchronization here uses per-pipe SetFlag/WaitFlag instead. + eventVS_ = pipe_->FetchEventID(AscendC::HardEvent::V_S); + eventSV_ = pipe_->FetchEventID(AscendC::HardEvent::S_V); + eventMTE2V_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_V); + eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S); + eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3); + eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + } + + __aicore__ inline void Process() + { + const int64_t total = numRows_ * vocabSize_; + for (int64_t idx = AscendC::GetBlockIdx(); idx < total; + idx += AscendC::GetBlockNum()) { + const int64_t row = idx / vocabSize_; + const int64_t col = idx - row * vocabSize_; + ProcessElement(row, col); + } + } + +private: + // Load one hidden/weight tile pair into fp32 UB buffers. + __aicore__ inline void LoadTilesFp32(int64_t row, + int64_t col, + int64_t start, + uint32_t count) + { + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + + AscendC::LocalTensor hTile = hQueue_.AllocTensor(); + AscendC::DataCopyPad(hTile, hiddenGm_[row * hiddenSize_ + start], copyParams, padParams); + hQueue_.EnQue(hTile); + hTile = hQueue_.DeQue(); + + AscendC::LocalTensor wTile = wQueue_.AllocTensor(); + AscendC::DataCopyPad(wTile, weightGm_[col * hiddenSize_ + start], copyParams, padParams); + wQueue_.EnQue(wTile); + wTile = wQueue_.DeQue(); + + AscendC::LocalTensor hFp32 = hFp32Buf_.Get(); + AscendC::LocalTensor wFp32 = wFp32Buf_.Get(); + if constexpr (std::is_same_v) { + // No cast needed: copy the native tiles into the fp32 buffers. + // (The queue tiles ARE fp32 views; staging keeps all later vector + // work on the dedicated buffers with a single ordering path.) + AscendC::SetFlag(eventMTE2V_); + AscendC::WaitFlag(eventMTE2V_); + AscendC::DataCopy(hFp32, hTile, VecAlignCount(count)); + AscendC::DataCopy(wFp32, wTile, VecAlignCount(count)); + } else { + AscendC::SetFlag(eventMTE2V_); + AscendC::WaitFlag(eventMTE2V_); + AscendC::Cast(hFp32, hTile, AscendC::RoundMode::CAST_NONE, count); + AscendC::Cast(wFp32, wTile, AscendC::RoundMode::CAST_NONE, count); + } + hQueue_.FreeTensor(hTile); + wQueue_.FreeTensor(wTile); + } + + // Wait until all outstanding vector-pipe results are readable as scalars. + __aicore__ inline void WaitVector() + { + AscendC::SetFlag(eventVS_); + AscendC::WaitFlag(eventVS_); + } + + __aicore__ inline void ProcessElement(int64_t row, int64_t col) + { + const int64_t tileCount = (hiddenSize_ + TILE_LENGTH - 1) / TILE_LENGTH; + float acc = 0.0f; + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_LENGTH; + const uint32_t count = TileCount(start); + LoadTilesFp32(row, col, start, count); + + AscendC::LocalTensor hFp32 = hFp32Buf_.Get(); + AscendC::LocalTensor wFp32 = wFp32Buf_.Get(); + AscendC::LocalTensor prod = prodBuf_.Get(); + AscendC::Mul(prod, hFp32, wFp32, count); + + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceSum(scalar, prod, rTmp, static_cast(count)); + WaitVector(); // vector -> scalar read + // Sequential accumulation over tiles in fixed order. + acc += scalar.GetValue(0); + } + + if (hasBias_) { + acc += LoadBias(col); + } + + // Stage the output in UB, then DataCopyPad to GM. GlobalTensor.SetValue + // is unreliable on hardware (cannbot ascendc-precision-debug + // common-traps), so scalar GM stores are not used. + AscendC::LocalTensor scalar = scalarBuf_.Get(); + scalar.SetValue(0, acc); + AscendC::SetFlag(eventSV_); // scalar write -> vector op + AscendC::WaitFlag(eventSV_); + if constexpr (std::is_same_v) { + AscendC::SetFlag(eventSMTE3_); + AscendC::WaitFlag(eventSMTE3_); + AscendC::DataCopyExtParams outParams{1, sizeof(float), 0, 0, 0}; + AscendC::DataCopyPad(outputGm_[row * vocabSize_ + col], scalar[0], outParams); + } else { + // Cast the fp32 scalar to the output dtype (round-to-nearest), + // padded to a 32 B vector-pipe minimum. + AscendC::LocalTensor tScalar = scalarBuf_.Get(); + const uint32_t castCount = 32 / sizeof(T); + AscendC::Cast(tScalar, scalar, AscendC::RoundMode::CAST_RINT, castCount); + WaitVector(); // vector cast -> scalar read + AscendC::SetFlag(eventSMTE3_); + AscendC::WaitFlag(eventSMTE3_); + AscendC::DataCopyExtParams outParams{1, sizeof(T), 0, 0, 0}; + AscendC::DataCopyPad(outputGm_[row * vocabSize_ + col], tScalar[0], outParams); + } + // Drain MTE3 before the next element stages new values into scalarBuf_. + AscendC::SetFlag(eventMTE3S_); + AscendC::WaitFlag(eventMTE3S_); + } + + // Read bias[col] through a 32-byte-aligned DataCopyPad window. + __aicore__ inline float LoadBias(int64_t col) + { + const int64_t alignedCol = col & ~7LL; // 8 x fp32 per 32 B + const int64_t remaining = vocabSize_ - alignedCol; + const uint32_t winCount = static_cast(remaining < 8 ? remaining : 8); + AscendC::LocalTensor bLocal = biasBuf_.Get(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(winCount * sizeof(float)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(bLocal, biasGm_[alignedCol], copyParams, padParams); + AscendC::SetFlag(eventMTE2S_); // copy-in -> scalar read + AscendC::WaitFlag(eventMTE2S_); + return bLocal.GetValue(static_cast(col - alignedCol)); + } + + __aicore__ inline uint32_t TileCount(int64_t start) const + { + const int64_t remaining = hiddenSize_ - start; + return static_cast(remaining < TILE_LENGTH ? remaining : TILE_LENGTH); + } + + // Round an element count up to a 32 B boundary (vector-pipe minimum). + __aicore__ inline uint32_t VecAlignCount(uint32_t count) const + { + constexpr uint32_t elemsPer32B = 32 / sizeof(T); + return (count + elemsPer32B - 1) / elemsPer32B * elemsPer32B; + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor hiddenGm_; + AscendC::GlobalTensor weightGm_; + AscendC::GlobalTensor biasGm_; + AscendC::GlobalTensor outputGm_; + AscendC::TQue hQueue_; + AscendC::TQue wQueue_; + AscendC::TBuf hFp32Buf_; + AscendC::TBuf wFp32Buf_; + AscendC::TBuf prodBuf_; + AscendC::TBuf reduceBuf_; + AscendC::TBuf scalarBuf_; + AscendC::TBuf biasBuf_; + AscendC::TEventID eventVS_; + AscendC::TEventID eventSV_; + AscendC::TEventID eventMTE2V_; + AscendC::TEventID eventMTE2S_; + AscendC::TEventID eventSMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t numRows_; + int64_t vocabSize_; + int64_t hiddenSize_; + bool hasBias_; +}; + +} // namespace + +extern "C" __global__ __vector__ void lm_head_ascend_kernel_fp32( + GM_ADDR hidden, GM_ADDR weight, GM_ADDR bias, GM_ADDR output, + int64_t numRows, int64_t vocabSize, int64_t hiddenSize, bool hasBias) +{ + AscendC::TPipe pipe; + KernelLmHead op(&pipe); + op.Init(hidden, weight, bias, output, numRows, vocabSize, hiddenSize, hasBias); + op.Process(); +} + +extern "C" __global__ __vector__ void lm_head_ascend_kernel_bf16( + GM_ADDR hidden, GM_ADDR weight, GM_ADDR bias, GM_ADDR output, + int64_t numRows, int64_t vocabSize, int64_t hiddenSize, bool hasBias) +{ + AscendC::TPipe pipe; + KernelLmHead op(&pipe); + op.Init(hidden, weight, bias, output, numRows, vocabSize, hiddenSize, hasBias); + op.Process(); +} + +extern "C" __global__ __vector__ void lm_head_ascend_kernel_fp16( + GM_ADDR hidden, GM_ADDR weight, GM_ADDR bias, GM_ADDR output, + int64_t numRows, int64_t vocabSize, int64_t hiddenSize, bool hasBias) +{ + AscendC::TPipe pipe; + KernelLmHead op(&pipe); + op.Init(hidden, weight, bias, output, numRows, vocabSize, hiddenSize, hasBias); + op.Process(); +} + +torch::Tensor lm_head_ascend_forward(torch::Tensor hidden, + torch::Tensor weight, + torch::optional bias, + bool output_fp32) +{ + TORCH_CHECK(hidden.is_privateuseone(), "hidden must be on an NPU device"); + TORCH_CHECK(weight.is_privateuseone(), "weight must be on an NPU device"); + TORCH_CHECK(hidden.device() == weight.device(), + "hidden and weight must be on the same NPU device"); + TORCH_CHECK(hidden.dim() >= 2, "hidden must have shape [..., hidden]"); + TORCH_CHECK(weight.dim() == 2, "lm_head weight must be [vocab, hidden]"); + TORCH_CHECK(hidden.size(-1) == weight.size(1), "hidden/weight hidden-dim mismatch"); + TORCH_CHECK(hidden.scalar_type() == at::kFloat || hidden.scalar_type() == at::kHalf || + hidden.scalar_type() == at::kBFloat16, + "lm_head_ascend supports fp32, fp16, and bf16 hidden states"); + TORCH_CHECK(weight.scalar_type() == hidden.scalar_type(), + "lm_head_ascend requires weight to match the hidden dtype"); + + const int64_t hiddenSize = hidden.size(-1); + TORCH_CHECK(hiddenSize > 0, "lm_head hidden dimension must be non-zero"); + const int64_t vocabSize = weight.size(0); + const int64_t numRows = hidden.numel() / hiddenSize; + // Mirror the CUDA wrapper: pre-cast the inputs to the compute dtype. + const at::ScalarType computeDtype = output_fp32 ? at::kFloat : hidden.scalar_type(); + + auto hidden2d = hidden.reshape({numRows, hiddenSize}).to(computeDtype).contiguous(); + auto weight2d = weight.to(computeDtype).contiguous(); + torch::Tensor biasF; + uint8_t* biasPtr = nullptr; + bool hasBias = false; + if (bias.has_value()) { + TORCH_CHECK(bias->is_privateuseone(), "lm_head bias must be on an NPU device"); + TORCH_CHECK(bias->device() == hidden.device(), + "lm_head bias must be on the same NPU device as hidden"); + TORCH_CHECK(bias->dim() == 1, "lm_head bias must be 1-D [vocab]"); + TORCH_CHECK(bias->numel() == vocabSize, "lm_head bias must have vocab elements"); + TORCH_CHECK(bias->scalar_type() == at::kFloat || bias->scalar_type() == at::kHalf || + bias->scalar_type() == at::kBFloat16, + "lm_head_ascend supports fp32, fp16, and bf16 bias"); + biasF = bias->reshape({vocabSize}).to(at::kFloat).contiguous(); + biasPtr = reinterpret_cast(biasF.mutable_data_ptr()); + hasBias = true; + } + + auto outOptions = hidden.options().dtype(computeDtype); + std::vector outSizes; + outSizes.reserve(static_cast(hidden.dim())); + for (int64_t i = 0; i < hidden.dim() - 1; ++i) { + outSizes.push_back(hidden.size(i)); + } + outSizes.push_back(vocabSize); + auto output = torch::empty(outSizes, outOptions); + if (numRows == 0 || vocabSize == 0) { + return output; + } + + // 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 int64_t total = numRows * vocabSize; + const uint32_t blockNum = static_cast(std::min(total, MAX_BLOCKS)); + + if (computeDtype == at::kBFloat16) { + lm_head_ascend_kernel_bf16<<>>( + reinterpret_cast(hidden2d.mutable_data_ptr()), + reinterpret_cast(weight2d.mutable_data_ptr()), + biasPtr, + reinterpret_cast(output.mutable_data_ptr()), + numRows, vocabSize, hiddenSize, hasBias); + } else if (computeDtype == at::kHalf) { + lm_head_ascend_kernel_fp16<<>>( + reinterpret_cast(hidden2d.mutable_data_ptr()), + reinterpret_cast(weight2d.mutable_data_ptr()), + biasPtr, + reinterpret_cast(output.mutable_data_ptr()), + numRows, vocabSize, hiddenSize, hasBias); + } else { + lm_head_ascend_kernel_fp32<<>>( + reinterpret_cast(hidden2d.mutable_data_ptr()), + reinterpret_cast(weight2d.mutable_data_ptr()), + biasPtr, + reinterpret_cast(output.mutable_data_ptr()), + numRows, vocabSize, hiddenSize, hasBias); + } + return output; +} + +// 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 index 5cd38b99..6d05ae0b 100644 --- a/csrc/ascend/npu_module.cpp +++ b/csrc/ascend/npu_module.cpp @@ -32,6 +32,11 @@ torch::Tensor embedding_ascend_forward(torch::Tensor token_ids, torch::Tensor fused_logp_ascend_forward(torch::Tensor logits, torch::Tensor target); +torch::Tensor lm_head_ascend_forward(torch::Tensor hidden, + torch::Tensor weight, + torch::optional bias, + bool output_fp32); + torch::Tensor rmsnorm_ascend_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor rstd); @@ -78,4 +83,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("fused_logp_ascend", &fused_logp_ascend_forward, "Batch-invariant fused selected-token log-probability (Ascend C forward)"); + m.def("lm_head_ascend", + &lm_head_ascend_forward, + "Batch-invariant LM-head projection (Ascend C forward)"); } diff --git a/docs/operators/lm_head.md b/docs/operators/lm_head.md index 901d1759..f8a16c24 100644 --- a/docs/operators/lm_head.md +++ b/docs/operators/lm_head.md @@ -37,6 +37,7 @@ The op exposes the WS1 dual-path contract: | --- | --- | --- | --- | | PyTorch fallback | `NativeLMHeadOp` | None | fp32 ground-truth reference; CPU and any GPU. | | CUDA SM90 (H200/Hopper) | `SM90LMHeadOp` | `_C.lm_head_sm90_forward` | Single-card batch-invariant forward backend; no Split-K; bf16 backward uses deterministic GEMM. | +| Ascend NPU | `AscendLMHeadOp` | `_C_npu.lm_head_ascend` | Single-card batch-invariant forward backend: one output element per block, full K reduction in fp32 over a fixed tile order; fp32-formula VJP backward. | | ROCm / Triton | N/A | N/A | Falls back to the PyTorch native reference. | ## Tensor Contract @@ -89,6 +90,21 @@ hidden row and vocab weight row independently, so large-vocab projections are ex to be memory-bandwidth bound compared with a tiled GEMM. This path exists to preserve a fixed hidden-dimension accumulation order for the WS1/H200 correctness gate. +On `npu` the priority is: + +1. `ASCEND_LM_HEAD` — `AscendLMHeadOp` (batch-invariant Ascend C forward; fp32/bf16/fp16). +2. `PYTORCH_NATIVE_LM_HEAD` — `NativeLMHeadOp` (fallback). + +The Ascend kernel implements the same structure as the SM90 CUDA kernel: one output +element per block iteration, the full hidden-dimension reduction inside that block over +a fixed tile order (products -> per-tile sum -> sequential scalar accumulation), bias +added in fp32, final cast to the output dtype. There is no Split-K, so a row's logits +depend only on H — never on N or block assignment — and are bitwise identical across +batch sizes, row positions, and block assignments on the NPU. The per-tile sums use the +Ascend vector unit's fixed hardware tree instead of CUDA's warp-shuffle tree, so the +comparison against the PyTorch reference (torch.mv) is tolerance-based per the +reduction contract, not bitwise. + For bf16 H200 training, `SM90LMHeadOp.backward` routes `dhidden` through `_C.det_gemm_da` and `dweight` through `_C.det_gemm_db` (`hidden.T @ dlogits`, transposed back to the HF `[vocab, hidden]` layout). The wrapper fails fast if those deterministic @@ -97,18 +113,25 @@ GEMM symbols are missing instead of silently falling back to cuBLAS for bf16 gra ## Tests ```bash -python -m pytest tests/test_lm_head.py -v +python -m pytest tests/test_lm_head.py tests/test_lm_head_ascend.py -v ``` Covers fp32 correctness vs the fixed-K reference, precision-context safety, bf16/fp16 accuracy, output shape, bias semantics, Axis-A batch invariance, input purity, gradient flow to `hidden` and `weight`, registry dispatch, and a GPU-only smoke test at the real -Qwen3-8B dimensions. +Qwen3-8B dimensions. The Ascend suite adds: contract-tolerance correctness vs the +reference (fp32/bf16/fp16, with and without bias), fp32-formula VJP backward, bitwise +batch invariance (batch sizes 1 vs {2,4,16,300}, row positions, multi-tile H=10000, +repeated runs), and NPU registry dispatch. ## Implementation Files - `rl_engine/kernels/ops/pytorch/linear/lm_head.py` - `rl_engine/kernels/ops/cuda/linear/lm_head.py` - `csrc/cuda/embedding_lm_head_sm90.cu` +- `rl_engine/kernels/ops/ascend/linear/lm_head.py` — Ascend deterministic op +- `csrc/ascend/lm_head_ascend.asc` — Ascend C forward kernel +- `csrc/ascend/npu_module.cpp` — shared pybind entry for `rl_engine._C_npu` - `rl_engine/kernels/registry.py` - `tests/test_lm_head.py` +- `tests/test_lm_head_ascend.py` diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index 0286293c..dccab1bb 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -63,3 +63,9 @@ def fused_logp_ascend( logits: torch.Tensor, target: torch.Tensor, ) -> torch.Tensor: ... +def lm_head_ascend( + hidden: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + output_fp32: bool, +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index c7a0e0e6..dae7e986 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -166,6 +166,7 @@ def _load_object(path: str) -> Any: "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", "triton": "rl_engine.kernels.ops.triton.linear.lm_head.TritonLMHeadOp", "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp", + "ascend": "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp", }, grad_input_names=("hidden", "weight"), ), diff --git a/rl_engine/kernels/ops/ascend/linear/__init__.py b/rl_engine/kernels/ops/ascend/linear/__init__.py index 59881dd4..dd74a18f 100644 --- a/rl_engine/kernels/ops/ascend/linear/__init__.py +++ b/rl_engine/kernels/ops/ascend/linear/__init__.py @@ -2,3 +2,4 @@ # Copyright (c) 2026 RL-Kernel Contributors from . import embedding # noqa: F401 +from . import lm_head # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/linear/lm_head.py b/rl_engine/kernels/ops/ascend/linear/lm_head.py new file mode 100644 index 00000000..722fe892 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/linear/lm_head.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +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 + +_SUPPORTED_DTYPES = {torch.float32, torch.float16, torch.bfloat16} + + +class _AscendLMHeadFunction(torch.autograd.Function): + """Autograd bridge for the Ascend batch-invariant LM-head forward. + + The VJP is the standard linear formula computed in fp32 on the NPU + (``grad_hidden = grad @ W``, ``grad_weight = grad^T @ H``, bias = the + fixed-order row sum), then cast back to the input dtypes. The CUDA op + uses its declared deterministic GEMM for the backward; on NPU the + plain torch matmuls are the deterministic-in-practice equivalent, and + the gtest compares gradients at the reduction contract tolerance. + """ + + @staticmethod + def forward(ctx, hidden, weight, bias, output_fp32: bool): + bias_to_save = ( + bias if bias is not None else torch.empty(0, device=hidden.device, dtype=hidden.dtype) + ) + ctx.save_for_backward(hidden, weight, bias_to_save) + ctx.has_bias = bias is not None + ctx.output_fp32 = bool(output_fp32) + return _C_npu.lm_head_ascend(hidden, weight.contiguous(), bias, bool(output_fp32)) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + hidden, weight, bias = ctx.saved_tensors + grad_2d = grad_output.reshape(-1, weight.size(0)).contiguous().float() + hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous().float() + weight_f = weight.contiguous().float() + grad_hidden = grad_weight = grad_bias = None + if ctx.needs_input_grad[0]: + grad_hidden = grad_2d @ weight_f + grad_hidden = grad_hidden.reshape_as(hidden).to(hidden.dtype) + if ctx.needs_input_grad[1]: + grad_weight = (grad_2d.t() @ hidden_2d).to(weight.dtype) + if ctx.has_bias and ctx.needs_input_grad[2]: + rows = grad_output.reshape(-1, weight.size(0)).float() + acc = torch.zeros((rows.shape[1],), device=rows.device, dtype=torch.float32) + for index in range(rows.shape[0]): + acc = acc + rows[index] + grad_bias = acc.to(bias.dtype) + record_backward( + "lm_head", + kernel_id="rl_engine.kernels.ops.ascend.linear.lm_head._AscendLMHeadFunction", + impl="ascend_fp32_matmul_vjp", + family="ascend", + ) + return grad_hidden, grad_weight, grad_bias, None + + +class AscendLMHeadOp: + """Single-card batch-invariant Ascend LM-head op. + + The Ascend C forward mirrors the SM90 CUDA kernel's structure: one + output element per block iteration, full hidden-dimension fp32 reduction + inside that block over a fixed tile order, bias added in fp32, final + cast to the output dtype. There is no Split-K and no algorithm + selection, so a row's logits do not depend on batch layout. + """ + + op_class = "reduction" + is_batch_invariant = True + backward_impl = "ascend_fp32_matmul_vjp" + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "lm_head_ascend"): + raise RuntimeError( + "lm_head_ascend is not compiled into the extension. " + "Rebuild on an Ascend NPU host with KERNEL_ALIGN_FORCE_ASCEND=1." + ) + logger.info("Successfully linked to precompiled _C_npu.lm_head_ascend kernel.") + + def __call__( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.forward(hidden, weight, bias=bias) + + def forward( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not self._can_use_ascend(hidden, weight, bias): + raise RuntimeError( + "AscendLMHeadOp requires Ascend NPU bf16/fp16/fp32 inputs; " + "Native/Triton fallback is forbidden" + ) + return _AscendLMHeadFunction.apply(hidden, weight, bias, False) + + def forward_fp32( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not self._can_use_ascend(hidden, weight, bias): + raise RuntimeError( + "AscendLMHeadOp requires Ascend NPU bf16/fp16/fp32 inputs; " + "Native/Triton fallback is forbidden" + ) + return _AscendLMHeadFunction.apply(hidden, weight, bias, True) + + def parameter_vjp_contributions_fp32(self, *, hidden, weight, grad_output, bias=None): + del weight, bias + rows_h = hidden.reshape(-1, hidden.size(-1)).float() + rows_g = grad_output.reshape(-1, grad_output.size(-1)).float() + return {"weight": rows_g[:, :, None] * rows_h[:, None, :]} + + @staticmethod + def _can_use_ascend( + hidden: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + ) -> bool: + bias_ok = bias is None or ( + bias.device.type == "npu" + and bias.device == hidden.device + and bias.dim() == 1 + and bias.dtype in _SUPPORTED_DTYPES + ) + return ( + hidden.device.type == "npu" + and weight.device.type == "npu" + and hidden.device == weight.device + and hidden.dim() >= 2 + and weight.dim() == 2 + and hidden.size(-1) == weight.size(1) + and hidden.dtype in _SUPPORTED_DTYPES + and weight.dtype in _SUPPORTED_DTYPES + and bias_ok + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index cae4fdac..bd61db59 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -128,6 +128,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ASCEND_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp" ASCEND_EMBEDDING = "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp" ASCEND_FUSED_LOGP = "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp" + ASCEND_LM_HEAD = "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp" # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) PYTORCH_VOCAB_PARALLEL_LOGP = ( "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" @@ -732,6 +733,10 @@ def __init__(self): OpBackend.ASCEND_FUSED_LOGP, OpBackend.PYTORCH_NATIVE, ] + self._priority_map["npu"]["lm_head"] = [ + OpBackend.ASCEND_LM_HEAD, + OpBackend.PYTORCH_NATIVE_LM_HEAD, + ] 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 bad5e32f..11b12f05 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -182,6 +182,10 @@ def fake_load_backend(backend): OpBackend.ASCEND_FUSED_LOGP, OpBackend.PYTORCH_NATIVE, ] + assert registry._priority_map["npu"]["lm_head"] == [ + OpBackend.ASCEND_LM_HEAD, + OpBackend.PYTORCH_NATIVE_LM_HEAD, + ] def test_npu_available_handles_runtime_failure(monkeypatch): diff --git a/tests/test_lm_head_ascend.py b/tests/test_lm_head_ascend.py new file mode 100644 index 00000000..43cf353c --- /dev/null +++ b/tests/test_lm_head_ascend.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for the Ascend NPU batch-invariant LM-head projection. + +Validates the same two orthogonal properties as the CUDA deterministic op: + +1. **Correctness** - output matches the ``NativeLMHeadOp.forward_fp32`` ground + truth within the reduction contract tolerance. The Ascend C kernel mirrors + the SM90 CUDA kernel's structure (one output element per block, full + hidden-dimension fp32 reduction over a fixed tile order); the hardware + reduction trees differ from CUDA's (and from torch.mv's), so the + comparison is tolerance-based. +2. **Batch-invariance** - a row's logits are bitwise identical regardless of + batch size, batch position, or how many AI-core blocks were launched + (each element is reduced end-to-end by one block; no Split-K merge). +""" + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.linear.lm_head import NativeLMHeadOp + +# Accuracy tolerances from the gtest contract, "reduction" op class. +_ATOL = { + torch.float32: 1.0e-4, + torch.bfloat16: 5.0e-2, + torch.float16: 1.0e-3, +} +_RTOL = { + torch.float32: 1.0e-4, + torch.bfloat16: 2.0e-2, + torch.float16: 1.0e-3, +} +# Gradient tolerances from the gtest contract, "gradient_accuracy" reduction row. +_GRAD_ATOL = { + torch.float32: 1.0e-4, + torch.bfloat16: 1.0e-1, + torch.float16: 1.0e-3, +} +_GRAD_RTOL = { + torch.float32: 1.0e-4, + torch.bfloat16: 2.0e-2, + torch.float16: 1.0e-3, +} + + +def _npu_available() -> bool: + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + +def _ascend_kernel_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine.kernels.ops.ascend.linear.lm_head import _NPU_EXT_AVAILABLE, _C_npu + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "lm_head_ascend") + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="lm_head_ascend kernel not compiled " + "(needs KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host).", +) + + +def _get_op(): + from rl_engine.kernels.ops.ascend.linear.lm_head import AscendLMHeadOp + + return AscendLMHeadOp() + + +def _make_inputs(shape, vocab, hidden, dtype, seed=0, with_bias=False): + generator = torch.Generator(device="cpu").manual_seed(seed) + hidden_t = torch.randn(*shape, hidden, dtype=dtype, generator=generator).to("npu") + weight = torch.randn(vocab, hidden, dtype=dtype, generator=generator).to("npu") + bias = torch.randn(vocab, dtype=dtype, generator=generator).to("npu") if with_bias else None + return hidden_t, weight, bias + + +# --------------------------------------------------------------------------- +# Correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@requires_ascend +class TestAscendLMHeadCorrectness: + def test_forward_matches_pytorch_reference(self, dtype): + op = _get_op() + hidden, weight, _ = _make_inputs((3, 5), 129, 1000, dtype) + out = op(hidden, weight) + ref = NativeLMHeadOp().forward_fp32(hidden, weight) + assert out.dtype == dtype + assert out.shape == (3, 5, 129) + assert torch.allclose(out.float(), ref.float(), atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_forward_fp32_matches_reference(self, dtype): + op = _get_op() + hidden, weight, _ = _make_inputs((3, 5), 129, 1000, dtype) + out = op.forward_fp32(hidden, weight) + ref = NativeLMHeadOp().forward_fp32(hidden, weight) + assert out.dtype == torch.float32 + assert torch.allclose(out, ref, atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_forward_with_bias(self, dtype): + op = _get_op() + hidden, weight, bias = _make_inputs((3, 5), 129, 1000, dtype, with_bias=True) + out = op(hidden, weight, bias=bias) + ref = NativeLMHeadOp().forward_fp32(hidden, weight, bias=bias) + assert torch.allclose(out.float(), ref.float(), atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_output_shape_leading_dims(self, dtype): + op = _get_op() + hidden, weight, _ = _make_inputs((2, 4, 3), 17, 64, dtype) + out = op(hidden, weight) + assert out.shape == (2, 4, 3, 17) + + def test_backward_matches_native_reference(self, dtype): + # The native reference runs the fp32 path: torch.mv on NPU rejects + # bf16, and the fp32 path keeps both VJPs in the same accumulation + # dtype so the comparison isolates the matmul-tree drift. + op = _get_op() + hidden, weight, _ = _make_inputs((3, 5), 129, 1000, dtype) + grad_out = torch.randn(3, 5, 129, device="npu", dtype=dtype) + + h_a = hidden.clone().requires_grad_() + w_a = weight.clone().requires_grad_() + op(h_a, w_a).backward(grad_out) + + h_n = hidden.clone().requires_grad_() + w_n = weight.clone().requires_grad_() + NativeLMHeadOp().forward_fp32(h_n, w_n).backward(grad_out) + + assert torch.allclose( + h_a.grad.float(), h_n.grad.float(), atol=_GRAD_ATOL[dtype], rtol=_GRAD_RTOL[dtype] + ) + assert torch.allclose( + w_a.grad.float(), w_n.grad.float(), atol=_GRAD_ATOL[dtype], rtol=_GRAD_RTOL[dtype] + ) + + def test_backward_with_bias(self, dtype): + op = _get_op() + hidden, weight, bias = _make_inputs((3, 5), 129, 1000, dtype, with_bias=True) + grad_out = torch.randn(3, 5, 129, device="npu", dtype=dtype) + + h_a = hidden.clone().requires_grad_() + w_a = weight.clone().requires_grad_() + b_a = bias.clone().requires_grad_() + op(h_a, w_a, bias=b_a).backward(grad_out) + assert b_a.grad is not None + assert b_a.grad.shape == (129,) + assert torch.isfinite(b_a.grad).all() + + b_n = bias.clone().requires_grad_() + NativeLMHeadOp().forward_fp32(hidden, weight, bias=b_n).backward(grad_out) + assert torch.allclose( + b_a.grad.float(), b_n.grad.float(), atol=_GRAD_ATOL[dtype], rtol=_GRAD_RTOL[dtype] + ) + + +# --------------------------------------------------------------------------- +# Batch invariance +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendLMHeadBatchInvariance: + def test_batch_size_1_vs_n(self): + # One fixed hidden row embedded in batches of growing size: its logits + # must be bitwise identical regardless of batch size. + dtype = torch.bfloat16 + op = _get_op() + alone_hidden, weight, _ = _make_inputs((1,), 129, 1000, dtype, seed=7) + alone = op(alone_hidden, weight)[0] + for batch in (2, 4, 16, 300): # 300 rows -> > MAX_BLOCKS strided blocks + hidden, _, _ = _make_inputs((batch,), 129, 1000, dtype, seed=7) + hidden[0].copy_(alone_hidden[0]) + in_batch = op(hidden, weight)[0] # same weight table throughout + assert torch.equal(alone, in_batch), f"drift at batch_size={batch}" + + def test_different_positions_in_batch(self): + # The same row content copied to every position projects bitwise + # identically regardless of where it lands. + dtype = torch.float16 + op = _get_op() + hidden, weight, _ = _make_inputs((8,), 129, 1000, dtype, seed=11) + base = hidden[0].clone() + for pos in range(1, 8): + hidden[pos].copy_(base) + out = op(hidden, weight) + for pos in range(1, 8): + assert torch.equal(out[pos], out[0]), f"drift at position={pos}" + + def test_multi_tile_rows(self): + # hidden > TILE_LENGTH (4096): the reduction spans multiple tiles. + dtype = torch.float32 + op = _get_op() + hidden, weight, _ = _make_inputs((4,), 129, 10000, dtype, seed=5) + out = op(hidden, weight) + assert torch.equal(out, op(hidden, weight)) + + def test_repeated_runs_deterministic(self): + dtype = torch.bfloat16 + hidden, weight, _ = _make_inputs((3, 5), 129, 1000, dtype, seed=5) + op = _get_op() + first = op(hidden, weight) + for _ in range(3): + again = op(hidden, weight) + assert torch.equal(first, again) + + +# --------------------------------------------------------------------------- +# Registry dispatch +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendRegistryDispatch: + def test_get_op_lm_head(self): + from rl_engine.kernels.registry import kernel_registry + + op = kernel_registry.get_op("lm_head", device="npu") + assert type(op).__name__ == "AscendLMHeadOp"