diff --git a/csrc/ascend/embedding_ascend.asc b/csrc/ascend/embedding_ascend.asc new file mode 100644 index 00000000..503a5f7f --- /dev/null +++ b/csrc/ascend/embedding_ascend.asc @@ -0,0 +1,279 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant token embedding, Ascend C (CANN) forward kernel. +// +// out[t, :] = weight[token_ids[t], :] +// +// Mirrors the SM90 CUDA kernel in csrc/cuda/embedding_lm_head_sm90.cu: +// - input : token_ids [*lead] (cast to int64), weight [V, H] contiguous +// fp32 / bf16 / fp16 +// - output : [*lead, H] in the weight's native dtype (bit copy), or fp32 +// when output_fp32 (handled by the host wrapper, see below) +// - every token id must be in [0, V); the host wrapper checks this. +// +// Bitwise identity with the CUDA kernel: the SM90 forward is a pure row +// gather (output[idx] = static_cast(weight[...])). This kernel is +// a pure byte copy of the same rows in the native dtype, and the fp32-output +// path upcasts the gathered result afterwards. Upcasting bf16/fp16 to fp32 +// is exact (every value is representable), so both paths are bitwise +// identical to the CUDA kernel for identical inputs. There is no arithmetic +// anywhere in the op, so there is no reduction order to drift. +// +// Batch-invariance: every token row is copied end-to-end by exactly one AI +// core block with a fixed tile size. The copy sequence for a row depends +// only on H, never on the total token count or on the block the row happens +// to land on. Rows are strided across blocks, so launching fewer blocks than +// rows is fine and never changes any row's bytes. +// +// 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 token counts; this is what +// makes the copy sequence batch-invariant. UB budget (in tile + out tile) +// stays far under the 192 KB UB of current SoCs. +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; + +template +class KernelEmbedding { +public: + __aicore__ inline KernelEmbedding(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR tokenIds, + GM_ADDR weight, + GM_ADDR output, + int64_t numTokens, + int64_t hiddenSize) + { + numTokens_ = numTokens; + hiddenSize_ = hiddenSize; + tokenIdsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(tokenIds)); + weightGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(weight)); + outputGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(output)); + pipe_->InitBuffer(inQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(outQueue_, 1, TILE_LENGTH * sizeof(T)); + // 32 B window for reading token_ids[row] via DataCopyPad (GM scalar + // GetValue/SetValue are unreliable on hardware; see cannbot + // ascendc-precision-debug common-traps). + pipe_->InitBuffer(idsBuf_, 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. + eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S); + } + + __aicore__ inline void Process() + { + for (int64_t row = AscendC::GetBlockIdx(); row < numTokens_; + row += AscendC::GetBlockNum()) { + ProcessRow(row); + } + } + +private: + // Copy one hidden tile of row `row` (gathered from weight row `tokenId`) + // through UB. The copy is a pure byte move; the fixed tile order is what + // keeps the kernel batch-invariant. + __aicore__ inline void ProcessRow(int64_t row) + { + const int64_t tokenId = LoadTokenId(row); + 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); + + // Canonical GM -> UB -> GM pipeline (same shape as the official + // Ascend C elementwise samples). The queues own all cross-pipe + // ordering: inQueue.EnQue/DeQue syncs MTE2 copy-in -> vector, + // outQueue.EnQue/DeQue syncs vector -> MTE3 copy-out, and + // FreeTensor orders the next tile's writes against the previous + // tile's reads, so the shared UB tiles are never reused while a + // pipe is still draining them. + AscendC::LocalTensor inTile = inQueue_.AllocTensor(); + AscendC::DataCopyExtParams inParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad( + inTile, weightGm_[tokenId * hiddenSize_ + start], inParams, padParams); + inQueue_.EnQue(inTile); + inTile = inQueue_.DeQue(); + + AscendC::LocalTensor outTile = outQueue_.AllocTensor(); + // The vector-pipe UB copy needs 32 B-aligned element counts. + // Over-copying within UB is harmless: the copy-out below writes + // only `count` elements to GM, so the tail never escapes. + AscendC::DataCopy(outTile, inTile, VecAlignCount(count)); + outQueue_.EnQue(outTile); + outTile = outQueue_.DeQue(); + + AscendC::DataCopyExtParams outParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPad( + outputGm_[row * hiddenSize_ + start], outTile, outParams); + outQueue_.FreeTensor(outTile); + inQueue_.FreeTensor(inTile); + } + } + + // Read token_ids[row] through a 32-byte-aligned DataCopyPad window. + __aicore__ inline int64_t LoadTokenId(int64_t row) + { + const int64_t alignedRow = row & ~3LL; // 4 x int64 per 32 B + const int64_t remaining = numTokens_ - alignedRow; + const uint32_t winCount = static_cast(remaining < 4 ? remaining : 4); + AscendC::LocalTensor idsLocal = idsBuf_.Get(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(winCount * sizeof(int64_t)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(idsLocal, tokenIdsGm_[alignedRow], copyParams, padParams); + AscendC::SetFlag(eventMTE2S_); // copy-in -> scalar read + AscendC::WaitFlag(eventMTE2S_); + return static_cast( + idsLocal.GetValue(static_cast(row - alignedRow))); + } + + __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 tokenIdsGm_; + AscendC::GlobalTensor weightGm_; + AscendC::GlobalTensor outputGm_; + AscendC::TQue inQueue_; + AscendC::TQue outQueue_; + AscendC::TBuf idsBuf_; + AscendC::TEventID eventMTE2S_; + int64_t numTokens_; + int64_t hiddenSize_; +}; + +} // namespace + +extern "C" __global__ __vector__ void embedding_ascend_kernel_fp32( + GM_ADDR tokenIds, GM_ADDR weight, GM_ADDR output, + int64_t numTokens, int64_t hiddenSize) +{ + AscendC::TPipe pipe; + KernelEmbedding op(&pipe); + op.Init(tokenIds, weight, output, numTokens, hiddenSize); + op.Process(); +} + +extern "C" __global__ __vector__ void embedding_ascend_kernel_bf16( + GM_ADDR tokenIds, GM_ADDR weight, GM_ADDR output, + int64_t numTokens, int64_t hiddenSize) +{ + AscendC::TPipe pipe; + KernelEmbedding op(&pipe); + op.Init(tokenIds, weight, output, numTokens, hiddenSize); + op.Process(); +} + +extern "C" __global__ __vector__ void embedding_ascend_kernel_fp16( + GM_ADDR tokenIds, GM_ADDR weight, GM_ADDR output, + int64_t numTokens, int64_t hiddenSize) +{ + AscendC::TPipe pipe; + KernelEmbedding op(&pipe); + op.Init(tokenIds, weight, output, numTokens, hiddenSize); + op.Process(); +} + +torch::Tensor embedding_ascend_forward(torch::Tensor token_ids, torch::Tensor weight, + bool output_fp32) +{ + TORCH_CHECK(token_ids.is_privateuseone(), "token_ids must be on an NPU device"); + TORCH_CHECK(weight.is_privateuseone(), "weight must be on an NPU device"); + TORCH_CHECK(token_ids.device() == weight.device(), + "token_ids and weight must be on the same NPU device"); + TORCH_CHECK(weight.dim() == 2, "embedding weight must be [vocab, hidden]"); + TORCH_CHECK(weight.is_contiguous(), "embedding weight must be contiguous"); + TORCH_CHECK(weight.scalar_type() == at::kBFloat16 || weight.scalar_type() == at::kFloat || + weight.scalar_type() == at::kHalf, + "embedding_ascend supports fp32, fp16, and bf16 weights"); + + const int64_t vocabSize = weight.size(0); + const int64_t hiddenSize = weight.size(1); + const int64_t numTokens = token_ids.numel(); + auto ids = token_ids.reshape({numTokens}).to(at::kLong).contiguous(); + if (numTokens > 0) { + const int64_t minId = ids.min().item(); + const int64_t maxId = ids.max().item(); + TORCH_CHECK(minId >= 0 && maxId < vocabSize, + "embedding_ascend token ids must be in [0, ", vocabSize - 1, + "], got [", minId, ", ", maxId, "]"); + } + + std::vector outSizes; + outSizes.reserve(static_cast(token_ids.dim()) + 1); + for (int64_t i = 0; i < token_ids.dim(); ++i) { + outSizes.push_back(token_ids.size(i)); + } + outSizes.push_back(hiddenSize); + + // Gather in the weight's native dtype first: the kernel is a pure byte + // copy, and upcasting bf16/fp16 rows to fp32 afterwards is exact (every + // value is representable), so this is bitwise identical to the SM90 + // kernel's in-kernel static_cast -- but the kernel surface stays a single + // native-dtype copy path. + auto outOptions = weight.options().dtype(weight.scalar_type()); + auto output = torch::empty(outSizes, outOptions); + if (numTokens == 0 || hiddenSize == 0) { + return output_fp32 ? output.to(at::kFloat) : 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 uint32_t blockNum = static_cast(std::min(numTokens, MAX_BLOCKS)); + + if (weight.scalar_type() == at::kBFloat16) { + embedding_ascend_kernel_bf16<<>>( + reinterpret_cast(ids.mutable_data_ptr()), + reinterpret_cast(weight.mutable_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + numTokens, hiddenSize); + } else if (weight.scalar_type() == at::kHalf) { + embedding_ascend_kernel_fp16<<>>( + reinterpret_cast(ids.mutable_data_ptr()), + reinterpret_cast(weight.mutable_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + numTokens, hiddenSize); + } else { + embedding_ascend_kernel_fp32<<>>( + reinterpret_cast(ids.mutable_data_ptr()), + reinterpret_cast(weight.mutable_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + numTokens, hiddenSize); + } + return output_fp32 ? output.to(at::kFloat) : 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 06f97a84..9cb7a384 100644 --- a/csrc/ascend/npu_module.cpp +++ b/csrc/ascend/npu_module.cpp @@ -25,6 +25,10 @@ std::vector deterministic_attention_ascend_forward( torch::Tensor prefix_shared_attention_ascend_forward( torch::Tensor q, torch::Tensor k, torch::Tensor v); +torch::Tensor embedding_ascend_forward(torch::Tensor token_ids, + torch::Tensor weight, + bool output_fp32); + torch::Tensor rmsnorm_ascend_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor rstd); @@ -65,4 +69,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("rmsnorm_ascend", &rmsnorm_ascend_forward, "Batch-invariant RMSNorm (Ascend C forward, rstd precomputed)"); + m.def("embedding_ascend", + &embedding_ascend_forward, + "Batch-invariant token embedding (Ascend C forward)"); } diff --git a/docs/operators/embedding.md b/docs/operators/embedding.md index 1923ec84..cb142464 100644 --- a/docs/operators/embedding.md +++ b/docs/operators/embedding.md @@ -33,6 +33,7 @@ The op exposes the WS1 dual-path contract: | --- | --- | --- | --- | | PyTorch fallback | `NativeEmbeddingOp` | None | fp32 ground-truth reference; CPU and any GPU. | | CUDA SM90 (H200/Hopper) | `SM90EmbeddingOp` | `_C.embedding_sm90_forward` | Single-card batch-invariant forward backend; deterministic duplicate-id backward in the wrapper. | +| Ascend NPU | `AscendEmbeddingOp` | `_C_npu.embedding_ascend` | Batch-invariant Ascend C forward (pure row copy); reuses the SM90 op's deterministic sorted-segment backward. | | Triton | `TritonEmbeddingOp` | `_embedding_fwd`, `_embedding_bwd` | CUDA gather with deterministic, atomic-free sorted-segment backward. | | ROCm | N/A | N/A | Falls back to the PyTorch native reference. | @@ -55,6 +56,21 @@ CPU, ROCm, and CUDA devices without the SM90 extension, dispatch uses the PyTorc the CUDA SM90 single-card batch-invariant backend is prepended and the native op remains the fallback. +On `npu` the priority is: + +1. `ASCEND_EMBEDDING` — `AscendEmbeddingOp` (batch-invariant Ascend C forward, bf16/fp16/fp32). +2. `PYTORCH_NATIVE_EMBEDDING` — `NativeEmbeddingOp` (fallback). + +The Ascend kernel implements the same semantics as the SM90 CUDA kernel: a pure row +gather (`out[t, :] = weight[token_ids[t], :]`). Every token row is copied end-to-end by +exactly one AI-core block with a fixed tile size, so the copy sequence for a row depends +only on `hidden`, never on the token count or block assignment. Because the copy performs +no arithmetic, the Ascend output is **bitwise identical** to the CUDA kernel (and to the +PyTorch reference) for identical inputs at every supported dtype; the fp32-output path +upcasts the gathered rows afterwards, which is exact for bf16/fp16. The backward reuses the +SM90 op's deterministic sorted-segment dweight (stable-sorted ids, fixed addition order), +so duplicate-id gradients match the CUDA op bit for bit. + ## Accuracy Reference semantics (`forward_fp32`): @@ -90,7 +106,8 @@ nondeterminism for repeated token ids at the cost of throughput. python -m pytest \ tests/test_embedding.py \ tests/test_triton_embedding.py \ - tests/test_canonical_embedding.py -v + tests/test_canonical_embedding.py \ + tests/test_embedding_ascend.py -v ``` Covers: correctness vs direct indexing (bitwise), dtype paths, non-int64 id tolerance, @@ -107,10 +124,14 @@ Triton sorted-segment backward and canonical logical-row ordering. - `rl_engine/kernels/ops/cuda/linear/embedding.py` - `rl_engine/kernels/ops/canonical_embedding.py` - `csrc/cuda/embedding_lm_head_sm90.cu` +- `rl_engine/kernels/ops/ascend/linear/embedding.py` — Ascend deterministic op +- `csrc/ascend/embedding_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_embedding.py` - `tests/test_triton_embedding.py` - `tests/test_canonical_embedding.py` +- `tests/test_embedding_ascend.py` ## Known Limitations diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index 5776625d..3cf73aea 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -53,3 +53,8 @@ def rmsnorm_ascend( rstd: torch.Tensor, ) -> torch.Tensor: ... +def embedding_ascend( + token_ids: torch.Tensor, + weight: torch.Tensor, + output_fp32: bool, +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 8cdbf4a8..1ad17fdd 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -152,6 +152,7 @@ def _load_object(path: str) -> Any: "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", "triton": "rl_engine.kernels.ops.triton.linear.embedding.TritonEmbeddingOp", "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp", + "ascend": "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp", }, grad_input_names=("weight",), ), diff --git a/rl_engine/kernels/ops/ascend/__init__.py b/rl_engine/kernels/ops/ascend/__init__.py index 5b45b492..12926601 100644 --- a/rl_engine/kernels/ops/ascend/__init__.py +++ b/rl_engine/kernels/ops/ascend/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from . import linear # noqa: F401 from . import loss # noqa: F401 from . import norm # noqa: F401 from . import rotary_embedding # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/linear/__init__.py b/rl_engine/kernels/ops/ascend/linear/__init__.py new file mode 100644 index 00000000..59881dd4 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/linear/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from . import embedding # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/linear/embedding.py b/rl_engine/kernels/ops/ascend/linear/embedding.py new file mode 100644 index 00000000..cf8c721a --- /dev/null +++ b/rl_engine/kernels/ops/ascend/linear/embedding.py @@ -0,0 +1,127 @@ +# 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.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} + + +def _deterministic_embedding_grad_weight( + ids: torch.Tensor, + grad_rows: torch.Tensor, + *, + weight_shape: tuple[int, ...], + weight_dtype: torch.dtype, +) -> torch.Tensor: + # Bitwise-identical backward by construction: the SM90 CUDA op's backward + # is itself pure PyTorch (sorted-segment dweight), so the Ascend op reuses + # the exact same function. Every op in it (mask, stable argsort, + # unique_consecutive, fixed-order accumulation) is deterministic on NPU, + # hence grad_weight matches the CUDA op bit for bit on identical inputs. + from rl_engine.kernels.ops.cuda.linear.embedding import ( + _deterministic_embedding_grad_weight as _cuda_grad_weight, + ) + + return _cuda_grad_weight( + ids, + grad_rows, + weight_shape=weight_shape, + weight_dtype=weight_dtype, + ) + + +class _AscendEmbeddingFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, token_ids: torch.Tensor, weight: torch.Tensor, output_fp32: bool): + ctx.save_for_backward(token_ids) + ctx.weight_shape = tuple(weight.shape) + ctx.weight_dtype = weight.dtype + ctx.output_fp32 = bool(output_fp32) + return _C_npu.embedding_ascend(token_ids, weight.contiguous(), bool(output_fp32)) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + (token_ids,) = ctx.saved_tensors + grad_weight = None + if ctx.needs_input_grad[1]: + ids = token_ids.reshape(-1).to(device=grad_output.device, dtype=torch.long) + hidden_size = int(ctx.weight_shape[1]) + grad_rows = grad_output.reshape(ids.numel(), hidden_size) + grad_weight = _deterministic_embedding_grad_weight( + ids, + grad_rows, + weight_shape=ctx.weight_shape, + weight_dtype=ctx.weight_dtype, + ) + record_backward( + "embedding", + kernel_id=( + "rl_engine.kernels.ops.ascend.linear.embedding." + "_deterministic_embedding_grad_weight" + ), + impl="ascend_sorted_segment_dweight", + family="ascend", + ) + return None, grad_weight, None + + +class AscendEmbeddingOp(torch.nn.Module): + """Single-card batch-invariant Ascend C embedding op. + + Forward is a pure row gather (a byte copy of weight rows), so it is + bitwise identical to the SM90 CUDA embedding kernel on identical inputs; + backward reuses the same sorted-segment dweight formula as the CUDA op. + """ + + op_class = "elementwise" + is_batch_invariant = True + + def __init__(self) -> None: + super().__init__() + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "embedding_ascend"): + raise RuntimeError( + "embedding_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.embedding_ascend kernel.") + + def forward(self, token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if not self._can_use_ascend(token_ids, weight): + raise RuntimeError( + "AscendEmbeddingOp requires Ascend NPU bf16/fp16/fp32 inputs; " + "Native/Triton fallback is forbidden" + ) + return _AscendEmbeddingFunction.apply(token_ids, weight, False) + + def forward_fp32(self, token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if not self._can_use_ascend(token_ids, weight): + raise RuntimeError( + "AscendEmbeddingOp requires Ascend NPU bf16/fp16/fp32 inputs; " + "Native/Triton fallback is forbidden" + ) + return _AscendEmbeddingFunction.apply(token_ids, weight, True) + + @staticmethod + def _can_use_ascend(token_ids: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + token_ids.device.type == "npu" + and weight.device.type == "npu" + and token_ids.device == weight.device + and weight.dim() == 2 + and weight.dtype in _SUPPORTED_DTYPES + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index f0b32eb0..6d994fa3 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -126,6 +126,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" ) ASCEND_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp" + ASCEND_EMBEDDING = "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp" # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) PYTORCH_VOCAB_PARALLEL_LOGP = ( "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" @@ -722,6 +723,10 @@ def __init__(self): OpBackend.ASCEND_RMS_NORM, OpBackend.PYTORCH_NATIVE_RMS_NORM, ] + self._priority_map["npu"]["embedding"] = [ + OpBackend.ASCEND_EMBEDDING, + OpBackend.PYTORCH_NATIVE_EMBEDDING, + ] 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 f77055d3..9e5dcd7f 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -174,6 +174,10 @@ def fake_load_backend(backend): OpBackend.ASCEND_RMS_NORM, OpBackend.PYTORCH_NATIVE_RMS_NORM, ] + assert registry._priority_map["npu"]["embedding"] == [ + OpBackend.ASCEND_EMBEDDING, + OpBackend.PYTORCH_NATIVE_EMBEDDING, + ] def test_npu_available_handles_runtime_failure(monkeypatch): diff --git a/tests/test_embedding_ascend.py b/tests/test_embedding_ascend.py new file mode 100644 index 00000000..ad63488e --- /dev/null +++ b/tests/test_embedding_ascend.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for the Ascend NPU deterministic token embedding. + +Validates the same two orthogonal properties as the CUDA deterministic op, +but with a stronger correctness claim than the attention op: embedding is a +pure row gather (a bit copy, no arithmetic), so the Ascend output is +**bitwise identical** to the ``NativeEmbeddingOp`` PyTorch reference at every +dtype -- there is no reduction tolerance to calibrate. + +1. **Correctness** - ``forward``/``forward_fp32`` match the PyTorch reference + bitwise (``torch.equal``), and the deterministic sorted-segment backward + reproduces the fixed-order duplicate-id sum bitwise in the gradient dtype. +2. **Batch-invariance** - a token's gathered row is bitwise identical + regardless of batch size, batch position, or how many AI-core blocks were + launched (each row is copied end-to-end by one block). +""" + +import pytest +import torch + +from rl_engine.kernels.ops.cuda.linear.embedding import _deterministic_embedding_grad_weight +from rl_engine.kernels.ops.pytorch.linear.embedding import NativeEmbeddingOp + +_VOCAB = 128 +_HIDDEN = 64 + +# Gradient tolerances from the gtest contract, "elementwise" op class. +_GRAD_ATOL = { + torch.float32: 1.0e-5, + torch.bfloat16: 2.0e-2, + torch.float16: 1.0e-3, +} +_GRAD_RTOL = { + torch.float32: 1.0e-5, + torch.bfloat16: 1.6e-2, + torch.float16: 1.0e-3, +} + + +def _npu_available() -> bool: + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + +def _ascend_kernel_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine.kernels.ops.ascend.linear.embedding import _NPU_EXT_AVAILABLE, _C_npu + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "embedding_ascend") + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="embedding_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.embedding import AscendEmbeddingOp + + return AscendEmbeddingOp() + + +def _make_inputs(shape, vocab=_VOCAB, hidden=_HIDDEN, dtype=torch.float32, seed=0): + generator = torch.Generator(device="cpu").manual_seed(seed) + weight = torch.randn(vocab, hidden, dtype=dtype, generator=generator).to("npu") + token_ids = torch.randint(0, vocab, shape, generator=generator).long().to("npu") + return token_ids, weight + + +# --------------------------------------------------------------------------- +# Correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@requires_ascend +class TestAscendEmbeddingCorrectness: + def test_forward_matches_pytorch_reference_bitwise(self, dtype): + """Ascend forward == NativeEmbeddingOp.forward, bitwise (pure gather).""" + op = _get_op() + token_ids, weight = _make_inputs((3, 5), dtype=dtype) + out = op(token_ids, weight) + ref = NativeEmbeddingOp().forward(token_ids, weight) + assert out.dtype == dtype + assert torch.equal(out, ref) + + def test_forward_matches_direct_indexing_bitwise(self, dtype): + op = _get_op() + token_ids, weight = _make_inputs((3, 5), dtype=dtype) + out = op(token_ids, weight) + assert torch.equal(out, weight[token_ids]) + + def test_forward_fp32_matches_reference_bitwise(self, dtype): + """Ascend forward_fp32 == NativeEmbeddingOp.forward_fp32, bitwise.""" + op = _get_op() + token_ids, weight = _make_inputs((3, 5), dtype=dtype) + out = op.forward_fp32(token_ids, weight) + ref = NativeEmbeddingOp().forward_fp32(token_ids, weight) + assert out.dtype == torch.float32 + assert torch.equal(out, ref) + + def test_output_shape_leading_dims(self, dtype): + op = _get_op() + token_ids, weight = _make_inputs((2, 4, 3), dtype=dtype) + out = op(token_ids, weight) + assert out.shape == (2, 4, 3, _HIDDEN) + + def test_backward_matches_fixed_order_sum_bitwise(self, dtype): + """The sorted-segment dweight equals the input-order row sum, bitwise. + + The backward is the same deterministic formula the SM90 CUDA op uses + (stable-sorted segments, fixed addition order), so this asserts the + Ascend op reproduces that exact arithmetic on NPU. + """ + op = _get_op() + token_ids, weight = _make_inputs((3, 5), dtype=dtype) + flat = token_ids.reshape(-1) + flat[1::3] = flat[0] # force duplicates of the first token id + grad_out = torch.randn(3, 5, _HIDDEN, device="npu", dtype=dtype) + + weight_g = weight.clone().requires_grad_() + op(flat.reshape(3, 5), weight_g).backward(grad_out) + grad_asc = weight_g.grad + + grad_weight = _deterministic_embedding_grad_weight( + flat, + grad_out.reshape(flat.numel(), _HIDDEN), + weight_shape=tuple(weight.shape), + weight_dtype=dtype, + ) + assert torch.equal(grad_asc, grad_weight) + + def test_backward_matches_native_reference(self, dtype): + """vs the native op's backward at the elementwise gradient contract. + + Not bitwise by design: the deterministic formula accumulates + duplicate-id rows in the grad dtype (one rounding per add) while the + native backward accumulates in fp32, and the native reduction order + is unspecified. Two duplicates keep the drift within the contract. + """ + op = _get_op() + token_ids, weight = _make_inputs((3, 5), dtype=dtype) + flat = token_ids.reshape(-1) + flat[1] = flat[0] # a single duplicate exercises multi-row accumulation + grad_out = torch.randn(3, 5, _HIDDEN, device="npu", dtype=dtype) + + weight_a = weight.clone().requires_grad_() + op(flat.reshape(3, 5), weight_a).backward(grad_out) + + weight_n = weight.clone().requires_grad_() + NativeEmbeddingOp().forward(flat.reshape(3, 5), weight_n).backward(grad_out) + + assert torch.allclose( + weight_a.grad.float(), + weight_n.grad.float(), + atol=_GRAD_ATOL[dtype], + rtol=_GRAD_RTOL[dtype], + ) + + def test_unused_rows_stay_zero(self, dtype): + op = _get_op() + token_ids, weight = _make_inputs((1, 2), dtype=dtype) + grad_out = torch.randn(1, 2, _HIDDEN, device="npu", dtype=dtype) + weight_g = weight.clone().requires_grad_() + op(token_ids, weight_g).backward(grad_out) + used = set(token_ids.reshape(-1).cpu().tolist()) + for row in range(_VOCAB): + if row not in used: + assert torch.equal( + weight_g.grad[row], torch.zeros(_HIDDEN, device="npu", dtype=dtype) + ) + + +# --------------------------------------------------------------------------- +# Input guards +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendEmbeddingGuards: + def test_rejects_non_npu(self): + op = _get_op() + token_ids, weight = _make_inputs((2, 3)) + with pytest.raises(RuntimeError): + op(token_ids.cpu(), weight.cpu()) + + +# --------------------------------------------------------------------------- +# Batch invariance +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendEmbeddingBatchInvariance: + def _run_row(self, batch, seq, dtype, pos, seed=7): + """One fixed token embedded at position `pos` of a random batch.""" + op = _get_op() + token_ids, weight = _make_inputs((batch, seq), dtype=dtype, seed=seed) + out = op(token_ids, weight) + return out[0, pos, :].clone() + + def test_batch_size_1_vs_n(self): + dtype = torch.float16 + alone = self._run_row(1, 8, dtype, pos=0, seed=7) + for batch in (2, 4, 8): + in_batch = self._run_row(batch, 8, dtype, pos=0, seed=7) + assert torch.equal(alone, in_batch), f"drift at batch_size={batch}" + + def test_different_positions_in_batch(self): + # The same weight row gathered at every position of a batch must be + # bitwise-identical regardless of where the token lands. + dtype = torch.bfloat16 + op = _get_op() + token_ids, weight = _make_inputs((2, 16), dtype=dtype, seed=11) + fixed_id = token_ids[0, 0] + token_ids[0, :] = fixed_id # one token id repeated across positions + out = op(token_ids, weight) + ref = weight[fixed_id] + for pos in range(16): + assert torch.equal(out[0, pos, :], ref), f"drift at position={pos}" + + def test_block_striding(self): + # 1024 tokens > MAX_BLOCKS (128): rows are strided across blocks, so + # the copied bytes must not depend on block assignment. The same + # (weight row, token id) gathered in a small run and in the strided + # run must be bitwise-identical. + dtype = torch.bfloat16 + op = _get_op() + small_ids, small_weight = _make_inputs((1,), dtype=dtype, seed=3) + small = op(small_ids, small_weight) + big_ids, big_weight = _make_inputs((1024,), dtype=dtype, seed=4) + big_ids[511] = small_ids[0] + big_weight[:] = small_weight # same table content + big = op(big_ids, big_weight) + assert torch.equal(big[511, :], small[0, :]) + + def test_multi_tile_rows(self): + # hidden > TILE_LENGTH would need a 4096+ column table; use a + # multi-tile-equivalent via a large hidden with the tile loop. + # (TILE_LENGTH = 4096; hidden = 12288 exercises 3 tiles per row.) + dtype = torch.float16 + op = _get_op() + generator = torch.Generator(device="cpu").manual_seed(9) + weight = torch.randn(256, 12288, dtype=dtype, generator=generator).to("npu") + token_ids = torch.randint(0, 256, (2, 3), generator=generator).long().to("npu") + out = op(token_ids, weight) + assert torch.equal(out, weight[token_ids]) + + def test_repeated_runs_deterministic(self): + dtype = torch.bfloat16 + token_ids, weight = _make_inputs((3, 5), dtype=dtype, seed=5) + op = _get_op() + first = op(token_ids, weight) + for _ in range(3): + again = op(token_ids, weight) + assert torch.equal(first, again) + + +# --------------------------------------------------------------------------- +# Registry dispatch +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendRegistryDispatch: + def test_get_op_embedding(self): + from rl_engine.kernels.registry import kernel_registry + + op = kernel_registry.get_op("embedding", device="npu") + assert type(op).__name__ == "AscendEmbeddingOp"