[WS1][Ascend] [Qwen3-8b] Embedding ops

csrc/ascend/batch_invariant_logp_ascend.asc

@@ -308,9 +308,5 @@ std::vector<torch::Tensor> 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.

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<output_t>(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 <type_traits>
+
+#include "kernel_operator.h"
+
+#include <torch/extension.h>
+
+#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 <typename T>
+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<T> inTile = inQueue_.AllocTensor<T>();
+            AscendC::DataCopyExtParams inParams{
+                1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
+            AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
+            AscendC::DataCopyPad(
+                inTile, weightGm_[tokenId * hiddenSize_ + start], inParams, padParams);
+            inQueue_.EnQue(inTile);
+            inTile = inQueue_.DeQue<T>();
+
+            AscendC::LocalTensor<T> outTile = outQueue_.AllocTensor<T>();
+            // 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<T>();
+
+            AscendC::DataCopyExtParams outParams{
+                1, static_cast<uint32_t>(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<uint32_t>(remaining < 4 ? remaining : 4);
+        AscendC::LocalTensor<int64_t> idsLocal = idsBuf_.Get<int64_t>();
+        AscendC::DataCopyExtParams copyParams{
+            1, static_cast<uint32_t>(winCount * sizeof(int64_t)), 0, 0, 0};
+        AscendC::DataCopyPadExtParams<int64_t> padParams{false, 0, 0, 0};
+        AscendC::DataCopyPad(idsLocal, tokenIdsGm_[alignedRow], copyParams, padParams);
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_S>(eventMTE2S_);  // copy-in -> scalar read
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_S>(eventMTE2S_);
+        return static_cast<int64_t>(
+            idsLocal.GetValue(static_cast<uint32_t>(row - alignedRow)));
+    }
+
+    __aicore__ inline uint32_t TileCount(int64_t start) const
+    {
+        const int64_t remaining = hiddenSize_ - start;
+        return static_cast<uint32_t>(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<int64_t> tokenIdsGm_;
+    AscendC::GlobalTensor<T> weightGm_;
+    AscendC::GlobalTensor<T> outputGm_;
+    AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueue_;
+    AscendC::TQue<AscendC::TPosition::VECOUT, 1> outQueue_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> 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<float> 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<bfloat16_t> 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<half> 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<int64_t>();
+        const int64_t maxId = ids.max().item<int64_t>();
+        TORCH_CHECK(minId >= 0 && maxId < vocabSize,
+                    "embedding_ascend token ids must be in [0, ", vocabSize - 1,
+                    "], got [", minId, ", ", maxId, "]");
+    }
+
+    std::vector<int64_t> outSizes;
+    outSizes.reserve(static_cast<size_t>(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<uint32_t>(std::min(numTokens, MAX_BLOCKS));
+
+    if (weight.scalar_type() == at::kBFloat16) {
+        embedding_ascend_kernel_bf16<<<blockNum, nullptr, aclStream>>>(
+            reinterpret_cast<uint8_t*>(ids.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(weight.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(output.mutable_data_ptr()),
+            numTokens, hiddenSize);
+    } else if (weight.scalar_type() == at::kHalf) {
+        embedding_ascend_kernel_fp16<<<blockNum, nullptr, aclStream>>>(
+            reinterpret_cast<uint8_t*>(ids.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(weight.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(output.mutable_data_ptr()),
+            numTokens, hiddenSize);
+    } else {
+        embedding_ascend_kernel_fp32<<<blockNum, nullptr, aclStream>>>(
+            reinterpret_cast<uint8_t*>(ids.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(weight.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(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.

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 <torch/extension.h>
+
+std::vector<torch::Tensor> 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)");
+}

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
 

rl_engine/_C_npu.pyi

@@ -8,3 +8,8 @@ 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: ...

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",),
     ),

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