bugfix: order SHM H2D with graph replay on NPU.

xllm/compiler/tilelang/targets/ascend/kernels/embedding.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+
+import tilelang
+import tilelang.language as T
+
+from .utils import DEFAULT_ASCEND_PASS_CONFIGS, detect_vec_core_num
+from ....common.spec import DispatchField, TilelangKernel, register_kernel
+
+MAX_TOKENS = 192
+VOCAB_SIZE = 248320
+HIDDEN_DIMS = (2560, 5120)
+DTYPE = "bfloat16"
+VEC_NUM = 2
+
+
+def build_embedding_kernel(vocab_size: int, hidden_dim: int, vec_core_num: int):
+    if vec_core_num <= 0 or vec_core_num % VEC_NUM != 0:
+        raise ValueError("vec_core_num must be positive and divisible by 2")
+
+    @T.prim_func
+    def embedding_kernel(
+        weight: T.Tensor((vocab_size, hidden_dim), DTYPE),
+        token_ids: T.Tensor((MAX_TOKENS,), "int32"),
+        output: T.Tensor((MAX_TOKENS, hidden_dim), DTYPE),
+        num_tokens: T.int32,
+    ):
+        with T.Kernel(vec_core_num // VEC_NUM, is_npu=True) as (cid, vid):
+            task_id = cid * VEC_NUM + vid
+            row_ub = T.alloc_ub((hidden_dim,), DTYPE)
+            rows_per_task = (num_tokens + vec_core_num - 1) // vec_core_num
+            for row_local in T.serial(rows_per_task):
+                row = task_id + row_local * vec_core_num
+                if row < num_tokens:
+                    token_id = token_ids[row]
+                    if token_id >= 0:
+                        if token_id < vocab_size:
+                            T.copy(weight[token_id, 0], row_ub)
+                        else:
+                            T.tile.fill(row_ub, 0)
+                    else:
+                        T.tile.fill(row_ub, 0)
+                    T.copy(row_ub, output[row, 0])
+
+    return embedding_kernel
+
+
+@register_kernel
+class EmbeddingKernel(TilelangKernel):
+    DISPATCH_SCHEMA = [
+        DispatchField("vocab_size", "int32"),
+        DispatchField("hidden_dim", "int32"),
+        DispatchField("dtype", "dtype"),
+    ]
+    SPECIALIZATIONS = [
+        {
+            "variant_key": f"v248320_h{hidden_dim}_bf16",
+            "vocab_size": VOCAB_SIZE,
+            "hidden_dim": hidden_dim,
+            "dtype": DTYPE,
+        }
+        for hidden_dim in HIDDEN_DIMS
+    ]
+
+    @staticmethod
+    def generate_source(vocab_size: int, hidden_dim: int, dtype: str) -> str:
+        if dtype != DTYPE:
+            raise ValueError(f"embedding only supports {DTYPE}, got {dtype}")
+        tilelang.disable_cache()
+        kernel = build_embedding_kernel(
+            vocab_size=vocab_size,
+            hidden_dim=hidden_dim,
+            vec_core_num=detect_vec_core_num(),
+        )
+        with tilelang.tvm.transform.PassContext(
+            opt_level=3, config=DEFAULT_ASCEND_PASS_CONFIGS
+        ):
+            return tilelang.engine.lower(kernel).kernel_source

xllm/core/distributed_runtime/worker_service.cpp

@@ -315,7 +315,13 @@ void WorkerService::create_polling_shm_thread(
         Timer timer;
         while (true) {
           ForwardInput fwd_input;
-          input_shm_manager->input_read(fwd_input, device_);
+          // NPU graph task updates cannot safely overlap an H2D enqueue from
+          // the SHM polling thread. Keep scheduler overlap, but defer device
+          // materialization to WorkerImpl's ordered prepare stream.
+          const bool defer_graph_overlap_h2d =
+              options_.enable_schedule_overlap() && options_.enable_graph();
+          input_shm_manager->input_read(
+              fwd_input, device_, !defer_graph_overlap_h2d);
           timer.reset();
           // model output variables
           torch::Tensor next_tokens;

xllm/core/kernels/npu/tilelang/CMakeLists.txt

@@ -127,6 +127,11 @@ tilelang_register_runtime_kernel(
   WRAPPER_SRCS rope_wrapper.cpp
 )
 
+tilelang_register_runtime_kernel(
+  NAME embedding
+  WRAPPER_SRCS embedding_wrapper.cpp
+)
+
 tilelang_register_runtime_kernel(
   NAME fused_gdn_gating
   WRAPPER_SRCS fused_gdn_gating_wrapper.cpp

xllm/core/kernels/npu/tilelang/embedding_wrapper.cpp

@@ -0,0 +1,75 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <c10/core/DeviceType.h>
+#include <glog/logging.h>
+#include <torch_npu/csrc/core/npu/NPUStream.h>
+
+#include <cstdint>
+
+#include "acl/acl.h"
+#include "dispatch_registry.h"
+#include "tilelang_ops_api.h"
+
+#ifndef XLLM_TL_EMBEDDING_REGISTRY_INC
+#error "XLLM_TL_EMBEDDING_REGISTRY_INC is not defined"
+#endif
+
+namespace xllm::kernel::npu::tilelang {
+namespace {
+
+#include XLLM_TL_EMBEDDING_REGISTRY_INC
+
+EmbeddingSpecialization build_runtime_specialization(
+    const torch::Tensor& weight) {
+  return make_embedding_specialization(
+      EmbeddingVocabSize{static_cast<int32_t>(weight.size(0))},
+      EmbeddingHiddenDim{static_cast<int32_t>(weight.size(1))},
+      EmbeddingDType{to_tilelang_dtype(weight.scalar_type())});
+}
+
+}  // namespace
+
+torch::Tensor embedding(const torch::Tensor& weight,
+                        const torch::Tensor& token_ids) {
+  CHECK(weight.defined() && token_ids.defined());
+  CHECK(weight.device().type() == c10::DeviceType::PrivateUse1 &&
+        token_ids.device().type() == c10::DeviceType::PrivateUse1);
+  CHECK_EQ(weight.dim(), 2);
+  CHECK_EQ(token_ids.dim(), 1);
+  CHECK_EQ(weight.scalar_type(), c10::ScalarType::BFloat16);
+  CHECK_EQ(token_ids.scalar_type(), c10::ScalarType::Int);
+  CHECK(weight.is_contiguous());
+  CHECK(token_ids.is_contiguous());
+  CHECK_GT(token_ids.size(0), 0);
+  CHECK_LE(token_ids.size(0), 192);
+
+  const auto* entry =
+      find_embedding_kernel_entry(build_runtime_specialization(weight));
+  CHECK(entry != nullptr) << "TileLang embedding: no compiled variant";
+
+  torch::Tensor output =
+      torch::empty({token_ids.size(0), weight.size(1)}, weight.options());
+  aclrtStream stream = c10_npu::getCurrentNPUStream(weight.device().index())
+                           .stream(/*need_empty=*/true);
+  entry->fn(reinterpret_cast<uint8_t*>(const_cast<void*>(weight.data_ptr())),
+            reinterpret_cast<uint8_t*>(const_cast<void*>(token_ids.data_ptr())),
+            reinterpret_cast<uint8_t*>(output.data_ptr()),
+            static_cast<int32_t>(token_ids.size(0)),
+            stream);
+  return output;
+}
+
+}  // namespace xllm::kernel::npu::tilelang

xllm/core/kernels/npu/tilelang/tilelang_ops_api.h

@@ -24,6 +24,9 @@ limitations under the License.
 
 namespace xllm::kernel::npu::tilelang {
 
+torch::Tensor embedding(const torch::Tensor& weight,
+                        const torch::Tensor& token_ids);
+
 // Public TileLang kernel APIs exported to the xLLM NPU runtime.
 //
 // Apply TileLang RoPE kernel in-place on a single input tensor.

xllm/core/layers/common/word_embedding_impl.cpp

@@ -15,6 +15,10 @@ limitations under the License.
 
 #include "word_embedding_impl.h"
 
+#if defined(USE_NPU)
+#include "core/kernels/npu/tilelang/tilelang_ops_api.h"
+#endif
+
 namespace xllm {
 namespace layer {
 
@@ -47,8 +51,16 @@ WordEmbeddingImpl::WordEmbeddingImpl(int64_t num_embeddings,
 // The input to the module is a list of indices, and the output is the
 // corresponding word embeddings.
 torch::Tensor WordEmbeddingImpl::forward(torch::Tensor input) {
-  namespace F = torch::nn::functional;
-  auto output = F::embedding(input, weight_);
+#if defined(USE_NPU)
+  const bool use_tilelang =
+      input.numel() <= 192 && weight_.size(0) == 248320 &&
+      (weight_.size(1) == 2560 || weight_.size(1) == 5120);
+  auto output = use_tilelang
+                    ? xllm::kernel::npu::tilelang::embedding(weight_, input)
+                    : torch::nn::functional::embedding(input, weight_);
+#else
+  auto output = torch::nn::functional::embedding(input, weight_);
+#endif
   if (world_size_ > 1) {
     output = xllm::parallel_state::gather(output, parallel_args_.tp_group_);
   }

xllm/core/runtime/forward_shared_memory_manager.cpp

@@ -2226,13 +2226,18 @@ inline void initialize_device_buffer_session(ReadContext& context,
 #endif
 
   auto& session = *context.device_session;
-  torch::Tensor host_input_buffer =
-      torch::from_blob(const_cast<char*>(payload_base),
-                       {static_cast<int64_t>(payload_size)},
-                       torch::TensorOptions()
-                           .dtype(torch::kUInt8)
-                           .device(torch::kCPU)
-                           .pinned_memory(/*pinned_memory=*/true));
+  // POSIX shared-memory pages are not pinned merely because TensorOptions says
+  // so. Own a genuinely pinned staging copy and keep it alive with ForwardInput
+  // until the asynchronous H2D has completed.
+  forward_input.input_host_buffer =
+      torch::empty({static_cast<int64_t>(payload_size)},
+                   torch::TensorOptions()
+                       .dtype(torch::kUInt8)
+                       .device(torch::kCPU)
+                       .pinned_memory(/*pinned_memory=*/true));
+  std::memcpy(
+      forward_input.input_host_buffer.data_ptr(), payload_base, payload_size);
+  const torch::Tensor& host_input_buffer = forward_input.input_host_buffer;
 
   auto device_options =
       torch::TensorOptions().dtype(torch::kUInt8).device(device);
@@ -2369,7 +2374,8 @@ inline void deserialize_forward_input_payload(
   read_linear_state_cache_ops(context, input_params.linear_state_cache_ops);
   normalize_linear_state_ids(input_params.embedding.linear_state_ids,
                              input_params.meta.num_sequences);
-  if (!input_params.embedding.linear_state_ids.empty()) {
+  if (materialize_device_buffer &&
+      !input_params.embedding.linear_state_ids.empty()) {
     input_params.embedding.linear_state_indices =
         torch::tensor(input_params.embedding.linear_state_ids, torch::kInt)
             .to(device, /*non_blocking=*/true);
@@ -3143,8 +3149,10 @@ bool ForwardSharedMemoryManager::input_write(const ForwardInput& input) {
   return true;
 }
 
-void ForwardSharedMemoryManager::input_read(ForwardInput& input,
-                                            const torch::Device& device) {
+void ForwardSharedMemoryManager::input_read(
+    ForwardInput& input,
+    const torch::Device& device,
+    bool materialize_device_buffer_on_read) {
   while (true) {
     if (control_ptr_->version != last_version_) {
       last_version_ = control_ptr_->version;
@@ -3159,13 +3167,25 @@ void ForwardSharedMemoryManager::input_read(ForwardInput& input,
   uint64_t total_size;
   read_data(data_ptr, total_size);
   bool materialize_device_buffer = false;
+  bool own_pinned_host_payload = false;
 #if defined(USE_NPU)
-  materialize_device_buffer = true;
+  materialize_device_buffer = materialize_device_buffer_on_read;
+  own_pinned_host_payload = !materialize_device_buffer;
 #elif defined(USE_CUDA)
   materialize_device_buffer = device.type() == torch::kCUDA;
 #elif defined(USE_MLU)
   materialize_device_buffer = device.type() == torch::kPrivateUse1;
 #endif
+  if (own_pinned_host_payload) {
+    input.input_host_buffer =
+        torch::empty({static_cast<int64_t>(total_size)},
+                     torch::TensorOptions()
+                         .dtype(torch::kUInt8)
+                         .device(torch::kCPU)
+                         .pinned_memory(/*pinned_memory=*/true));
+    std::memcpy(input.input_host_buffer.data_ptr(), data_ptr, total_size);
+    data_ptr = static_cast<const char*>(input.input_host_buffer.data_ptr());
+  }
   deserialize_forward_input_payload(data_ptr,
                                     total_size,
                                     input,

xllm/core/runtime/forward_shared_memory_manager.h

@@ -104,7 +104,9 @@ class ForwardSharedMemoryManager : public SharedMemoryManager {
   };
 
   bool input_write(const ForwardInput& input);
-  void input_read(ForwardInput& input, const torch::Device& device);
+  void input_read(ForwardInput& input,
+                  const torch::Device& device,
+                  bool materialize_device_buffer_on_read = true);
   bool raw_output_write(
       const torch::Tensor& next_tokens,
       const torch::Tensor& logprobs,