diff --git a/xllm/compiler/tilelang/targets/ascend/kernels/embedding.py b/xllm/compiler/tilelang/targets/ascend/kernels/embedding.py new file mode 100644 index 0000000000..8b4a496d04 --- /dev/null +++ b/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 diff --git a/xllm/core/distributed_runtime/worker_service.cpp b/xllm/core/distributed_runtime/worker_service.cpp index 07e6e61806..14930970ca 100644 --- a/xllm/core/distributed_runtime/worker_service.cpp +++ b/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; diff --git a/xllm/core/kernels/npu/tilelang/CMakeLists.txt b/xllm/core/kernels/npu/tilelang/CMakeLists.txt index 3309b0737d..57304c5b57 100644 --- a/xllm/core/kernels/npu/tilelang/CMakeLists.txt +++ b/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 diff --git a/xllm/core/kernels/npu/tilelang/embedding_wrapper.cpp b/xllm/core/kernels/npu/tilelang/embedding_wrapper.cpp new file mode 100644 index 0000000000..1bc10987d9 --- /dev/null +++ b/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 +#include +#include + +#include + +#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(weight.size(0))}, + EmbeddingHiddenDim{static_cast(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(const_cast(weight.data_ptr())), + reinterpret_cast(const_cast(token_ids.data_ptr())), + reinterpret_cast(output.data_ptr()), + static_cast(token_ids.size(0)), + stream); + return output; +} + +} // namespace xllm::kernel::npu::tilelang diff --git a/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h b/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h index 99a0dc9f8d..a00fd11b31 100644 --- a/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h +++ b/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. diff --git a/xllm/core/layers/common/word_embedding_impl.cpp b/xllm/core/layers/common/word_embedding_impl.cpp index 844b02cce7..fbf1db469c 100644 --- a/xllm/core/layers/common/word_embedding_impl.cpp +++ b/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_); } diff --git a/xllm/core/runtime/forward_shared_memory_manager.cpp b/xllm/core/runtime/forward_shared_memory_manager.cpp index 694f953687..3d758fbb50 100644 --- a/xllm/core/runtime/forward_shared_memory_manager.cpp +++ b/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(payload_base), - {static_cast(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(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(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(input.input_host_buffer.data_ptr()); + } deserialize_forward_input_payload(data_ptr, total_size, input, diff --git a/xllm/core/runtime/forward_shared_memory_manager.h b/xllm/core/runtime/forward_shared_memory_manager.h index 17e9442610..d4656c9d4a 100644 --- a/xllm/core/runtime/forward_shared_memory_manager.h +++ b/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,