From 98001fb8d3c834a9fde0fbde55df8d0f3e718390 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 1 Sep 2026 11:16:05 +0800 Subject: [PATCH 1/4] feat(ascend): add batch-invariant embedding Ascend C operator Mirror the SM90 CUDA embedding forward (pure row gather) with an Ascend C kernel in csrc/ascend/embedding_ascend.asc: the copy is a bitwise byte move, so the Ascend output is bit-identical to the CUDA kernel for identical inputs; the fp32-output path upcasts afterwards (exact for bf16/fp16). Backward reuses the CUDA op's sorted-segment dweight formula, which is pure PyTorch and deterministic on NPU. Also ports the shared-module Ascend build from the rmsnorm_ascend branch (npu_module.cpp single pybind entry + setup.py bisheng build) and registers the op in the gtest operator specs and the NPU registry priority map. --- csrc/ascend/batch_invariant_logp_ascend.asc | 8 +- csrc/ascend/embedding_ascend.asc | 279 ++++++++++++++++++ csrc/ascend/npu_module.cpp | 29 ++ rl_engine/_C_npu.pyi | 6 + rl_engine/kernels/gtest/operator_specs.py | 1 + rl_engine/kernels/ops/ascend/__init__.py | 1 + .../kernels/ops/ascend/linear/__init__.py | 4 + .../kernels/ops/ascend/linear/embedding.py | 127 ++++++++ rl_engine/kernels/registry.py | 5 + rl_engine/tests/test_dispatch.py | 4 + setup.py | 116 +++++++- 11 files changed, 571 insertions(+), 9 deletions(-) create mode 100644 csrc/ascend/embedding_ascend.asc create mode 100644 csrc/ascend/npu_module.cpp create mode 100644 rl_engine/kernels/ops/ascend/linear/__init__.py create mode 100644 rl_engine/kernels/ops/ascend/linear/embedding.py 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/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 new file mode 100644 index 00000000..e8e6e856 --- /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 embedding_ascend_forward(torch::Tensor token_ids, + torch::Tensor weight, + bool output_fp32); + +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("embedding_ascend", + &embedding_ascend_forward, + "Batch-invariant token embedding (Ascend C forward)"); +} diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..729d6a4e 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,9 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... + +def 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 ca4a462e..1bc3d2ab 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -127,6 +127,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 ab85458d..800801d5 100644 --- a/rl_engine/kernels/ops/ascend/__init__.py +++ b/rl_engine/kernels/ops/ascend/__init__.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from . import linear # noqa: F401 from . import loss # 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 12ea9b21..d01f18f7 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_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" @@ -622,6 +623,10 @@ def __init__(self): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + 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 388c4aec..f2bcbd40 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -166,6 +166,10 @@ def fake_load_backend(backend): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + assert registry._priority_map["npu"]["embedding"] == [ + OpBackend.ASCEND_EMBEDDING, + OpBackend.PYTORCH_NATIVE_EMBEDDING, + ] def test_npu_available_handles_runtime_failure(monkeypatch): diff --git a/setup.py b/setup.py index 79f882d9..bc006932 100644 --- a/setup.py +++ b/setup.py @@ -3,10 +3,13 @@ import importlib.util import os +import platform +import subprocess +import sysconfig import warnings from pathlib import Path -from setuptools import find_packages, setup +from setuptools import Extension, find_packages, setup def _load_envs_module(): @@ -58,6 +61,110 @@ def _cuda_define_from_env(name: str, macro: str) -> list[str]: return [f"-D{macro}={parsed}"] +_ASCEND_EXTENSION_NAME = "rl_engine._C_npu" +_ASCEND_CPU_DIRS = {"aarch64": "aarch64-linux", "x86_64": "x86_64-linux"} + + +def _find_ascend_home() -> str: + """Locate the CANN toolkit root (must contain bin/bisheng).""" + candidates = [ + os.environ.get("ASCEND_HOME_PATH"), + os.environ.get("ASCEND_TOOLKIT_HOME"), + ] + candidates += [str(p) for p in sorted(Path.home().glob("Ascend/cann-*"), reverse=True)] + candidates.append("/usr/local/Ascend/ascend-toolkit/latest") + for cand in candidates: + if cand and (Path(cand) / "bin" / "bisheng").is_file(): + return cand + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 was requested but no CANN toolkit with bin/bisheng " + "was found. Set ASCEND_HOME_PATH to the toolkit root." + ) + + +def _ascend_extension_spec() -> Extension: + sources = ["csrc/ascend/npu_module.cpp"] + sources += sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + ext = Extension(name=_ASCEND_EXTENSION_NAME, sources=sources) + ext._rl_kernel_ascend = True # intercepted by the custom build_ext below + return ext + + +def _compile_ascend_extension(build_ext, ext) -> None: + """Compile the Ascend C extension with bisheng (torch's BuildExtension + does not know the .asc language, so we drive the compiler directly).""" + torch, _, _ = _load_torch_extension_tools() + try: + import torch_npu + except ModuleNotFoundError as exc: + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch_npu. Install a matching " + "torch_npu build first." + ) from exc + + ascend_home = _find_ascend_home() + cpu_dir = _ASCEND_CPU_DIRS.get(platform.machine()) + if cpu_dir is None: + raise RuntimeError(f"unsupported Ascend host architecture: {platform.machine()}") + arch = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-c220") + + bisheng = os.path.join(ascend_home, "bin", "bisheng") + torch_dir = os.path.dirname(torch.__file__) + tnpu_dir = os.path.dirname(torch_npu.__file__) + + includes = [ + f"-I{os.path.join(ascend_home, cpu_dir, 'asc', 'include')}", + f"-I{os.path.join(torch_dir, 'include')}", + f"-I{os.path.join(torch_dir, 'include', 'torch', 'csrc', 'api', 'include')}", + f"-I{os.path.join(tnpu_dir, 'include')}", + f"-I{sysconfig.get_paths()['include']}", + ] + defines = [f"-DTORCH_EXTENSION_NAME={_ASCEND_EXTENSION_NAME.rsplit('.', 1)[-1]}"] + + build_temp = os.path.join(build_ext.build_temp, "ascend") + os.makedirs(build_temp, exist_ok=True) + + objects = [] + for src in ext.sources: + obj = os.path.join(build_temp, Path(src).name + ".o") + cmd = [bisheng, "-std=c++17", "-O2", "-fPIC", "-c"] + if src.endswith(".asc"): + cmd += ["-x", "asc", f"--cce-aicore-arch={arch}"] + cmd += includes + defines + [src, "-o", obj] + subprocess.check_call(cmd) + objects.append(obj) + + out_path = build_ext.get_ext_fullpath(ext.name) + os.makedirs(os.path.dirname(out_path), exist_ok=True) + link = [bisheng, "-shared", *objects] + for lib_dir, libs in ( + (os.path.join(torch_dir, "lib"), ["torch", "torch_cpu", "torch_python", "c10"]), + (os.path.join(tnpu_dir, "lib"), ["torch_npu"]), + (os.path.join(ascend_home, "runtime", "lib64"), ["ascendcl"]), + (os.path.join(ascend_home, cpu_dir, "lib64"), ["runtime"]), + ): + link.append(f"-L{lib_dir}") + link += [f"-l{name}" for name in libs] + link += [ + f"-Wl,-rpath,{os.path.join(torch_dir, 'lib')}", + f"-Wl,-rpath,{os.path.join(tnpu_dir, 'lib')}", + "-o", + out_path, + ] + subprocess.check_call(link) + + +def _make_build_extension(BuildExtension): + class AscendAwareBuildExtension(BuildExtension): + def build_extension(self, ext): + if getattr(ext, "_rl_kernel_ascend", False): + _compile_ascend_extension(self, ext) + return + super().build_extension(ext) + + return AscendAwareBuildExtension + + _ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( "-Xfatbin", "-compress-all", @@ -231,7 +338,7 @@ def get_extensions(): if enable_sm90 and present_sm90: tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") + nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") if "-lcuda" not in extra_link_args: extra_link_args.append("-lcuda") @@ -266,6 +373,9 @@ def get_extensions(): ) ) + if envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): + extensions.append(_ascend_extension_spec()) + if _native_extension_required() and not extensions: raise RuntimeError( "rl_engine._C was requested but no CUDA/ROCm build environment is available. " @@ -280,7 +390,7 @@ def get_cmdclass(): _, BuildExtension, _ = _load_torch_extension_tools() if BuildExtension is None: return {} - return {"build_ext": BuildExtension} + return {"build_ext": _make_build_extension(BuildExtension)} setup( From 518225f9c1a0eeaa9ffcceb06b5378fe29273c31 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 1 Sep 2026 11:46:22 +0800 Subject: [PATCH 2/4] build(ascend): export CANN toolkit env vars from _find_ascend_home The bisheng driver and its Ascend C plugin resolve toolkit data (impl include dirs, stub JSON generation) through ASCEND_HOME_PATH; without it the plugin fails while compiling any kernel TU. Export the discovered toolkit root (plus ASCEND_TOOLKIT_HOME) for the compiler subprocesses so the build is self-sufficient when the caller's shell did not source the CANN env setup script. --- setup.py | 853 ++++++++++++++++++++++++++++--------------------------- 1 file changed, 430 insertions(+), 423 deletions(-) diff --git a/setup.py b/setup.py index bc006932..9c7cbd02 100644 --- a/setup.py +++ b/setup.py @@ -1,423 +1,430 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import platform -import subprocess -import sysconfig -import warnings -from pathlib import Path - -from setuptools import Extension, find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - 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", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - if not is_rocm: - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - 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}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - 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. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - return {"build_ext": _make_build_extension(BuildExtension)} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import platform +import subprocess +import sysconfig +import warnings +from pathlib import Path + +from setuptools import Extension, find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + 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(): + # The bisheng driver and its Ascend C plugin resolve toolkit data + # (impl include dirs, stub JSON generation) through these env + # vars. Without ASCEND_HOME_PATH the plugin crashes (segfault) + # while compiling any kernel TU, so export them for the compiler + # subprocesses once we know where the toolkit lives. + os.environ["ASCEND_HOME_PATH"] = cand + os.environ.setdefault("ASCEND_TOOLKIT_HOME", cand) + 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", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", + ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + 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}") + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + 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. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + return {"build_ext": _make_build_extension(BuildExtension)} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) From fa14cf119e0573a79723acaa09d7f0c1a4f8a13d Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 1 Sep 2026 12:36:06 +0800 Subject: [PATCH 3/4] test(ascend): add embedding ascend pytest suite, gtest npu runner, docs - tests/test_embedding_ascend.py: correctness vs the PyTorch reference (bitwise, all dtypes), fixed-order backward bitwise check, batch invariance across sizes/positions/block-striding, registry dispatch. - scripts/check_operator.py: --device npu support (auto-detect), ported from the batch-invariant-logp ascend PR. - docs/operators/embedding.md: Ascend backend row, NPU dispatch behavior, tests and implementation files. Verified on Ascend 910 / CANN 8.5.1: - gtest embedding ascend candidate fp32/bf16/fp16, output+gradient: max_abs=0.0 (bitwise) at 2x16x257x4096 and the 1x2x257x4096 smoke shape. - pytest tests/test_embedding_ascend.py: 28 passed. --- docs/operators/embedding.md | 23 ++- scripts/check_operator.py | 20 ++- tests/test_embedding_ascend.py | 280 +++++++++++++++++++++++++++++++++ 3 files changed, 320 insertions(+), 3 deletions(-) create mode 100644 tests/test_embedding_ascend.py 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/scripts/check_operator.py b/scripts/check_operator.py index 9dbca48d..ccf18a28 100644 --- a/scripts/check_operator.py +++ b/scripts/check_operator.py @@ -35,9 +35,24 @@ 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") + if torch.cuda.is_available(): + return torch.device("cuda") + if _npu_available(): + return torch.device("npu") + return torch.device("cpu") + if value == "npu" and not _npu_available(): + raise RuntimeError("--device npu was requested, but no Ascend NPU is available") 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") @@ -73,7 +88,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--candidate", default="pytorch", - help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton.", + help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton, " + "ascend.", ) parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") parser.add_argument("--device", default="auto") diff --git a/tests/test_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" From 42844510edb429923e333979d7e19a71f31a6c52 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 1 Sep 2026 12:47:35 +0800 Subject: [PATCH 4/4] style: fix black formatting for pre-commit CI --- rl_engine/_C_npu.pyi | 1 - 1 file changed, 1 deletion(-) diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index 729d6a4e..0403ee8a 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,7 +8,6 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... - def embedding_ascend( token_ids: torch.Tensor, weight: torch.Tensor,