From 6a3e0baf0b2f345964a9b71344e122d3578f2c6f Mon Sep 17 00:00:00 2001 From: xumaojun Date: Thu, 23 Jul 2026 16:57:32 +0800 Subject: [PATCH] perf(mtp): optimize asynchronous draft scheduling --- tests/core/kernels/npu/CMakeLists.txt | 19 + .../npu/mtp_prepare_next_draft_test.cpp | 105 +++ tests/core/runtime/CMakeLists.txt | 16 + tests/core/runtime/mtp_async_state_test.cpp | 107 ++++ third_party/torch_npu_ops | 2 +- third_party/xllm_ops | 2 +- .../core/framework/model/model_input_params.h | 5 + xllm/core/kernels/npu/attention.cpp | 2 + xllm/core/kernels/npu/npu_ops_api.h | 1 + xllm/core/kernels/npu/xllm_ops/CMakeLists.txt | 1 + .../npu/xllm_ops/mtp_prepare_next_draft.cpp | 174 +++++ xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h | 17 + xllm/core/layers/common/attention_metadata.h | 1 + .../common/attention_metadata_builder.cpp | 2 + .../npu_deepseek_v32_decoder_layer_impl.cpp | 9 +- xllm/core/layers/npu_torch/attention.cpp | 21 +- xllm/core/platform/stream.cpp | 21 + xllm/core/platform/stream_event.h | 5 + xllm/core/runtime/CMakeLists.txt | 4 + xllm/core/runtime/mtp_async_input_builder.cpp | 144 +++++ xllm/core/runtime/mtp_async_input_builder.h | 43 ++ xllm/core/runtime/mtp_async_state.cpp | 169 +++++ xllm/core/runtime/mtp_async_state.h | 76 +++ xllm/core/runtime/mtp_worker_impl.cpp | 606 +++++++++++++++--- xllm/core/runtime/mtp_worker_impl.h | 67 +- xllm/core/runtime/worker_impl.cpp | 47 +- xllm/core/runtime/worker_impl.h | 3 +- 27 files changed, 1543 insertions(+), 126 deletions(-) create mode 100644 tests/core/kernels/npu/mtp_prepare_next_draft_test.cpp create mode 100644 tests/core/runtime/mtp_async_state_test.cpp create mode 100644 xllm/core/kernels/npu/xllm_ops/mtp_prepare_next_draft.cpp create mode 100644 xllm/core/runtime/mtp_async_input_builder.cpp create mode 100644 xllm/core/runtime/mtp_async_input_builder.h create mode 100644 xllm/core/runtime/mtp_async_state.cpp create mode 100644 xllm/core/runtime/mtp_async_state.h diff --git a/tests/core/kernels/npu/CMakeLists.txt b/tests/core/kernels/npu/CMakeLists.txt index 97ad595459..192732277c 100644 --- a/tests/core/kernels/npu/CMakeLists.txt +++ b/tests/core/kernels/npu/CMakeLists.txt @@ -16,4 +16,23 @@ include(cc_test) # glog::glog # ) +cc_test( + NAME + mtp_prepare_next_draft_test + SRCS + mtp_prepare_next_draft_test.cpp + DEPS + :xllm_ops + ascendcl + nnopbase + torch + torch_npu + GTest::gtest_main + glog::glog +) +if(EXISTS "$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") + target_link_options(mtp_prepare_next_draft_test PRIVATE + "-Wl,-rpath-link,$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") +endif() + add_subdirectory(tilelang) diff --git a/tests/core/kernels/npu/mtp_prepare_next_draft_test.cpp b/tests/core/kernels/npu/mtp_prepare_next_draft_test.cpp new file mode 100644 index 0000000000..91d6e8d76d --- /dev/null +++ b/tests/core/kernels/npu/mtp_prepare_next_draft_test.cpp @@ -0,0 +1,105 @@ +/* Copyright 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 "core/kernels/npu/xllm_ops/xllm_ops_api.h" + +namespace xllm::kernel::npu { +namespace { + +class MtpPrepareNextDraftTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); } + + static void TearDownTestSuite() { torch_npu::finalize_npu(); } +}; + +TEST_F(MtpPrepareNextDraftTest, ProducesExpectedOutputsForMixedAcceptance) { + constexpr int32_t kBlockSize = 4; + constexpr int64_t kHiddenSize = 16; + const torch::Tensor accepted_tokens_cpu = torch::tensor( + {{10, 11, 12, 13}, {20, 21, -1, -1}, {30, -1, -1, -1}}, torch::kLong); + const torch::Tensor accepted_embeddings_cpu = + torch::arange(3 * 4 * kHiddenSize, torch::kFloat) + .reshape({3, 4, kHiddenSize}) + .to(torch::kBFloat16); + const torch::Tensor placeholder_cpu = + torch::full({kHiddenSize}, -7.0, torch::kBFloat16); + const torch::Tensor base_positions_cpu = + torch::tensor({4, 8, 12}, torch::kInt); + const torch::Tensor base_kv_seq_lens_cpu = + torch::tensor({5, 9, 13}, torch::kInt); + const torch::Tensor block_tables_cpu = torch::tensor( + {{10, 11, 12, 13, 14}, {20, 21, 22, 23, 24}, {30, 31, 32, 33, 34}}, + torch::kInt); + const torch::Device npu_device("npu:0"); + + const auto output = + try_mtp_prepare_next_draft(accepted_tokens_cpu.to(npu_device), + accepted_embeddings_cpu.to(npu_device), + placeholder_cpu.to(npu_device), + base_positions_cpu.to(npu_device), + base_kv_seq_lens_cpu.to(npu_device), + block_tables_cpu.to(npu_device), + kBlockSize); + ASSERT_TRUE(output.has_value()); + + const torch::Tensor expected_tokens = + torch::tensor({12, 13, 20, 21, 30, 30}, torch::kInt); + const torch::Tensor expected_embeddings = + torch::stack({accepted_embeddings_cpu[0][2], + accepted_embeddings_cpu[0][3], + accepted_embeddings_cpu[1][0], + accepted_embeddings_cpu[1][1], + placeholder_cpu, + accepted_embeddings_cpu[2][0]}); + const torch::Tensor expected_positions = + torch::tensor({7, 8, 9, 10, 12, 13}, torch::kInt); + const torch::Tensor expected_kv_seq_lens = + torch::tensor({9, 11, 14}, torch::kInt); + const torch::Tensor expected_slots = + torch::tensor({47, 48, 91, 90, 134, 133}, torch::kInt); + + EXPECT_TRUE(torch::equal(output->token_ids.cpu(), expected_tokens)); + EXPECT_TRUE(torch::equal(output->embeddings.cpu(), expected_embeddings)); + EXPECT_TRUE(torch::equal(output->positions.cpu(), expected_positions)); + EXPECT_TRUE(torch::equal(output->kv_seq_lens.cpu(), expected_kv_seq_lens)); + EXPECT_TRUE(torch::equal(output->cache_slots.cpu(), expected_slots)); +} + +TEST_F(MtpPrepareNextDraftTest, RejectsUnsupportedHostInputs) { + const torch::Tensor tokens = torch::tensor({{1, 2}}, torch::kLong); + const torch::Tensor embeddings = torch::zeros({1, 2, 16}, torch::kBFloat16); + const torch::Tensor placeholder = torch::zeros({16}, torch::kBFloat16); + const torch::Tensor positions = torch::tensor({1}, torch::kInt); + const torch::Tensor kv_seq_lens = torch::tensor({2}, torch::kInt); + const torch::Tensor block_tables = torch::tensor({{0, 1}}, torch::kInt); + + EXPECT_FALSE(try_mtp_prepare_next_draft(tokens, + embeddings, + placeholder, + positions, + kv_seq_lens, + block_tables, + /*block_size=*/4) + .has_value()); +} + +} // namespace +} // namespace xllm::kernel::npu diff --git a/tests/core/runtime/CMakeLists.txt b/tests/core/runtime/CMakeLists.txt index cae555ea01..c242538656 100644 --- a/tests/core/runtime/CMakeLists.txt +++ b/tests/core/runtime/CMakeLists.txt @@ -77,6 +77,22 @@ cc_test( GTest::gtest_main ) +cc_test( + NAME + mtp_async_state_test + SRCS + mtp_async_state_test.cpp + "${PROJECT_SOURCE_DIR}/xllm/core/runtime/mtp_async_state.cpp" + DEPS + torch + glog::glog + GTest::gtest_main +) +if(EXISTS "$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") + target_link_options(mtp_async_state_test PRIVATE + "-Wl,-rpath-link,$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") +endif() + cc_test( NAME cp_input_partition_test diff --git a/tests/core/runtime/mtp_async_state_test.cpp b/tests/core/runtime/mtp_async_state_test.cpp new file mode 100644 index 0000000000..65f042f82e --- /dev/null +++ b/tests/core/runtime/mtp_async_state_test.cpp @@ -0,0 +1,107 @@ +/* Copyright 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 "core/runtime/mtp_async_state.h" + +#include +#include + +namespace xllm::mtp_async { +namespace { + +TEST(MtpAsyncStateTest, ClassifiesSupportedCombinedDraftExecutionPaths) { + EXPECT_EQ(classify_combined_draft_execution_path("qwen3_5_mtp"), + CombinedDraftExecutionPath::QWEN3_5_PAGED_ATTENTION); + EXPECT_EQ(classify_combined_draft_execution_path("qwen3_5_moe_mtp"), + CombinedDraftExecutionPath::QWEN3_5_PAGED_ATTENTION); + EXPECT_EQ(classify_combined_draft_execution_path("glm_moe_dsa_mtp"), + CombinedDraftExecutionPath::GLM5_DSA); + EXPECT_EQ(classify_combined_draft_execution_path("mimo_mtp"), + CombinedDraftExecutionPath::UNSUPPORTED); +} + +TEST(MtpAsyncStateTest, BuildsMixedAcceptanceStateWithoutHostRoundTrip) { + const torch::Tensor accepted_tokens = torch::tensor( + {{10, 11, 12, 13}, {20, 21, -1, -1}, {30, -1, -1, -1}}, torch::kLong); + const torch::Tensor accepted_embeddings = + torch::arange(24, torch::kFloat).reshape({3, 4, 2}); + const torch::Tensor placeholder = torch::tensor({-100.0, -101.0}); + const torch::Tensor base_positions = torch::tensor({100, 200, 300}); + const torch::Tensor base_kv_seq_lens = torch::tensor({101, 201, 301}); + + const AcceptedState state = build_accepted_state(accepted_tokens, + accepted_embeddings, + placeholder, + base_positions, + base_kv_seq_lens); + + EXPECT_TRUE(torch::equal(state.accepted_lengths, + torch::tensor({4, 2, 1}, torch::kLong))); + EXPECT_TRUE(torch::equal(state.all_draft_accepted, + torch::tensor({true, false, false}))); + EXPECT_TRUE(torch::equal(state.last_tokens, torch::tensor({13, 21, 30}))); + EXPECT_TRUE(torch::equal(state.previous_tokens, torch::tensor({12, 20, 30}))); + EXPECT_TRUE(torch::equal(state.last_embeddings, + torch::stack({accepted_embeddings[0][3], + accepted_embeddings[1][1], + accepted_embeddings[2][0]}))); + EXPECT_TRUE(torch::equal(state.previous_embeddings, + torch::stack({accepted_embeddings[0][2], + accepted_embeddings[1][0], + placeholder}))); + EXPECT_TRUE(torch::equal(state.base_positions, + torch::tensor({104, 202, 301}, torch::kLong))); + EXPECT_TRUE(torch::equal(state.base_kv_seq_lens, + torch::tensor({105, 203, 302}, torch::kLong))); +} + +TEST(MtpAsyncStateTest, BuildsRowMetadataForChunkedAndDecodeLayouts) { + AcceptedState state; + state.base_positions = torch::tensor({104, 202, 301}, torch::kLong); + state.base_kv_seq_lens = torch::tensor({105, 203, 302}, torch::kLong); + const torch::Tensor offsets = torch::tensor({-1, 0}, torch::kLong); + + EXPECT_TRUE(torch::equal( + make_row_positions(state, offsets), + torch::tensor({{103, 104}, {201, 202}, {300, 301}}, torch::kLong))); + EXPECT_TRUE(torch::equal( + make_kv_seq_lens(state, offsets, /*use_chunked_prefill=*/true), + state.base_kv_seq_lens)); + EXPECT_TRUE(torch::equal( + make_kv_seq_lens(state, offsets, /*use_chunked_prefill=*/false), + torch::tensor({104, 105, 202, 203, 301, 302}, torch::kLong))); +} + +TEST(MtpAsyncStateTest, RedirectsUnusedRepairRowsToScratchPositions) { + AcceptedState state; + state.base_positions = torch::tensor({104, 202, 301}, torch::kLong); + state.all_draft_accepted = torch::tensor({true, false, false}); + + EXPECT_TRUE(torch::equal(make_repair_cache_positions(state), + torch::tensor({103, 203, 302}, torch::kLong))); +} + +TEST(MtpAsyncStateTest, MapsPositionsAcrossCacheBlockBoundaries) { + const torch::Tensor block_tables = + torch::tensor({{10, 11, 12}, {20, 21, 22}}, torch::kInt); + const torch::Tensor positions = torch::tensor({{3, 4}, {7, 8}}, torch::kLong); + + EXPECT_TRUE(torch::equal( + map_positions_to_cache_slots(block_tables, positions, /*block_size=*/4), + torch::tensor({43, 44, 87, 88}, torch::kInt))); +} + +} // namespace +} // namespace xllm::mtp_async diff --git a/third_party/torch_npu_ops b/third_party/torch_npu_ops index adc29dc9d2..ed3de5f5b7 160000 --- a/third_party/torch_npu_ops +++ b/third_party/torch_npu_ops @@ -1 +1 @@ -Subproject commit adc29dc9d2e69875524b8a54bfc251463275d66e +Subproject commit ed3de5f5b7b83c4743e59ba8aa63ea77bdaedbcc diff --git a/third_party/xllm_ops b/third_party/xllm_ops index 9e436ee34a..7c05f17775 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit 9e436ee34a69d25a9806d537c2a8de1f525b1331 +Subproject commit 7c05f1777518551493fc1a2f79f6affdff0794cd diff --git a/xllm/core/framework/model/model_input_params.h b/xllm/core/framework/model/model_input_params.h index 1ed4263e2a..bceaacd538 100644 --- a/xllm/core/framework/model/model_input_params.h +++ b/xllm/core/framework/model/model_input_params.h @@ -412,6 +412,10 @@ struct AttentionDeviceInput { struct AttentionInput { AttentionHostInput host; AttentionDeviceInput device; + // Some asynchronous decode paths produce exact sequence lengths on device + // after host metadata preparation has finished. Keep the host lengths for + // shape planning, but make PagedAttention consume the device values. + bool use_device_kv_seq_lens = false; torch::Tensor attention_host_buffer; torch::Tensor attention_device_buffer; uint64_t attention_buffer_bytes = 0; @@ -422,6 +426,7 @@ struct AttentionInput { AttentionInput out; out.host = host; out.device = device.to(target_device); + out.use_device_kv_seq_lens = use_device_kv_seq_lens; out.attention_host_buffer = attention_host_buffer; out.attention_device_buffer = attention_device_buffer; out.attention_buffer_bytes = attention_buffer_bytes; diff --git a/xllm/core/kernels/npu/attention.cpp b/xllm/core/kernels/npu/attention.cpp index cd6a0e9f13..bd1d357dc9 100644 --- a/xllm/core/kernels/npu/attention.cpp +++ b/xllm/core/kernels/npu/attention.cpp @@ -73,6 +73,7 @@ void batch_decode(const torch::Tensor& query, float scale, const torch::Tensor& block_table, const torch::Tensor& seq_lens, + const std::optional& seq_lens_host, torch::Tensor& output) { int64_t head_size = query.size(-1); int64_t num_heads = query.size(-2); @@ -87,6 +88,7 @@ void batch_decode(const torch::Tensor& query, scale, block_table, seq_lens, + seq_lens_host, o); } diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index dea67faf54..d7259e9a31 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -45,6 +45,7 @@ void batch_decode(const torch::Tensor& query, float scale, const torch::Tensor& block_table, const torch::Tensor& seq_lens, + const std::optional& seq_lens_host, torch::Tensor& output); std::tuple npu_fused_infer_attention( diff --git a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt index 459f53b85a..e7df9f1705 100644 --- a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt +++ b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt @@ -9,6 +9,7 @@ cc_library( SRCS replace_token.cpp laser_attention.cpp + mtp_prepare_next_draft.cpp top_k_top_p.cpp quant_matmul.cpp quantize.cpp diff --git a/xllm/core/kernels/npu/xllm_ops/mtp_prepare_next_draft.cpp b/xllm/core/kernels/npu/xllm_ops/mtp_prepare_next_draft.cpp new file mode 100644 index 0000000000..8eb5e1769a --- /dev/null +++ b/xllm/core/kernels/npu/xllm_ops/mtp_prepare_next_draft.cpp @@ -0,0 +1,174 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#include +#include +#include +#include + +#include + +#include "acl/acl.h" +#include "aclnn_mtp_prepare_next_draft.h" +#include "core/common/macros.h" +#include "core/kernels/npu/utils.h" +#include "xllm_ops_api.h" + +namespace xllm::kernel::npu { +namespace { + +bool is_npu_tensor(const torch::Tensor& tensor) { + return tensor.defined() && tensor.numel() > 0 && + tensor.device().type() == c10::DeviceType::PrivateUse1; +} + +bool is_supported_mtp_prepare_input(const torch::Tensor& accepted_tokens, + const torch::Tensor& accepted_embeddings, + const torch::Tensor& embedding_placeholder, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens, + const torch::Tensor& block_tables, + int64_t block_size) { + if (!is_npu_tensor(accepted_tokens) || !is_npu_tensor(accepted_embeddings) || + !is_npu_tensor(embedding_placeholder) || !is_npu_tensor(base_positions) || + !is_npu_tensor(base_kv_seq_lens) || !is_npu_tensor(block_tables) || + block_size <= 0 || accepted_tokens.dim() != 2 || + accepted_embeddings.dim() != 3 || block_tables.dim() != 2) { + return false; + } + + const int64_t batch_size = accepted_tokens.size(0); + const int64_t speculative_width = accepted_tokens.size(1); + const int64_t hidden_size = accepted_embeddings.size(2); + const torch::Device device = accepted_tokens.device(); + const torch::ScalarType embedding_type = accepted_embeddings.scalar_type(); + return batch_size > 0 && speculative_width > 0 && hidden_size > 0 && + accepted_tokens.scalar_type() == torch::kLong && + (embedding_type == torch::kFloat16 || + embedding_type == torch::kBFloat16) && + embedding_placeholder.scalar_type() == embedding_type && + base_positions.scalar_type() == torch::kInt && + base_kv_seq_lens.scalar_type() == torch::kInt && + block_tables.scalar_type() == torch::kInt && + accepted_embeddings.size(0) == batch_size && + accepted_embeddings.size(1) == speculative_width && + embedding_placeholder.numel() == hidden_size && + base_positions.numel() >= batch_size && + base_kv_seq_lens.numel() >= batch_size && + block_tables.size(0) >= batch_size && + (hidden_size * accepted_embeddings.element_size()) % 32 == 0 && + accepted_tokens.is_contiguous() && + accepted_embeddings.is_contiguous() && + embedding_placeholder.is_contiguous() && + base_positions.is_contiguous() && base_kv_seq_lens.is_contiguous() && + block_tables.is_contiguous() && + accepted_embeddings.device() == device && + embedding_placeholder.device() == device && + base_positions.device() == device && + base_kv_seq_lens.device() == device && block_tables.device() == device; +} + +} // namespace + +std::optional try_mtp_prepare_next_draft( + const torch::Tensor& accepted_tokens, + const torch::Tensor& accepted_embeddings, + const torch::Tensor& embedding_placeholder, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens, + const torch::Tensor& block_tables, + int64_t block_size) { + if (!is_supported_mtp_prepare_input(accepted_tokens, + accepted_embeddings, + embedding_placeholder, + base_positions, + base_kv_seq_lens, + block_tables, + block_size)) { + return std::nullopt; + } + + const int64_t batch_size = accepted_tokens.size(0); + const int64_t hidden_size = accepted_embeddings.size(2); + const torch::Tensor position_rows = + base_positions.flatten().slice(0, 0, batch_size); + const torch::Tensor kv_seq_len_rows = + base_kv_seq_lens.flatten().slice(0, 0, batch_size); + + MtpPrepareNextDraftOutput output; + output.token_ids = torch::empty({batch_size * 2}, + accepted_tokens.options().dtype(torch::kInt)); + output.embeddings = torch::empty({batch_size * 2, hidden_size}, + accepted_embeddings.options()); + output.positions = torch::empty({batch_size * 2}, base_positions.options()); + output.kv_seq_lens = torch::empty({batch_size}, base_kv_seq_lens.options()); + output.cache_slots = torch::empty({batch_size * 2}, base_positions.options()); + + aclTensor* accepted_tokens_acl = nullptr; + aclTensor* accepted_embeddings_acl = nullptr; + aclTensor* embedding_placeholder_acl = nullptr; + aclTensor* base_positions_acl = nullptr; + aclTensor* base_kv_seq_lens_acl = nullptr; + aclTensor* block_tables_acl = nullptr; + aclTensor* token_ids_acl = nullptr; + aclTensor* embeddings_acl = nullptr; + aclTensor* positions_acl = nullptr; + aclTensor* kv_seq_lens_acl = nullptr; + aclTensor* cache_slots_acl = nullptr; + create_acltensor(&accepted_tokens_acl, accepted_tokens); + create_acltensor(&accepted_embeddings_acl, accepted_embeddings); + create_acltensor(&embedding_placeholder_acl, embedding_placeholder); + create_acltensor(&base_positions_acl, position_rows); + create_acltensor(&base_kv_seq_lens_acl, kv_seq_len_rows); + create_acltensor(&block_tables_acl, block_tables); + create_acltensor(&token_ids_acl, output.token_ids); + create_acltensor(&embeddings_acl, output.embeddings); + create_acltensor(&positions_acl, output.positions); + create_acltensor(&kv_seq_lens_acl, output.kv_seq_lens); + create_acltensor(&cache_slots_acl, output.cache_slots); + + const int32_t device_id = accepted_tokens.device().index(); + aclrtStream stream = c10_npu::getCurrentNPUStream(device_id).stream(); + uint64_t workspace_size = 0; + aclOpExecutor* executor = nullptr; + CHECK_ACL_SUCCESS( + aclnnMtpPrepareNextDraftGetWorkspaceSize(accepted_tokens_acl, + accepted_embeddings_acl, + embedding_placeholder_acl, + base_positions_acl, + base_kv_seq_lens_acl, + block_tables_acl, + block_size, + token_ids_acl, + embeddings_acl, + positions_acl, + kv_seq_lens_acl, + cache_slots_acl, + &workspace_size, + &executor), + "mtp_prepare_next_draft: failed to get workspace size"); + torch::Tensor workspace; + void* workspace_addr = nullptr; + if (workspace_size > 0) { + workspace = torch::empty({static_cast(workspace_size)}, + accepted_tokens.options().dtype(torch::kUInt8)); + workspace_addr = workspace.data_ptr(); + } + CHECK_ACL_SUCCESS(aclnnMtpPrepareNextDraft( + workspace_addr, workspace_size, executor, stream), + "mtp_prepare_next_draft: kernel launch failed"); + + aclDestroyTensor(accepted_tokens_acl); + aclDestroyTensor(accepted_embeddings_acl); + aclDestroyTensor(embedding_placeholder_acl); + aclDestroyTensor(base_positions_acl); + aclDestroyTensor(base_kv_seq_lens_acl); + aclDestroyTensor(block_tables_acl); + aclDestroyTensor(token_ids_acl); + aclDestroyTensor(embeddings_acl); + aclDestroyTensor(positions_acl); + aclDestroyTensor(kv_seq_lens_acl); + aclDestroyTensor(cache_slots_acl); + return output; +} + +} // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h b/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h index 52e076cc4a..753ec21ea7 100644 --- a/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h +++ b/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h @@ -70,6 +70,23 @@ torch::Tensor laser_attention(const torch::Tensor& q_bnsd, double scale_value, int64_t head_num); +struct MtpPrepareNextDraftOutput { + torch::Tensor token_ids; + torch::Tensor embeddings; + torch::Tensor positions; + torch::Tensor kv_seq_lens; + torch::Tensor cache_slots; +}; + +std::optional try_mtp_prepare_next_draft( + const torch::Tensor& accepted_tokens, + const torch::Tensor& accepted_embeddings, + const torch::Tensor& embedding_placeholder, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens, + const torch::Tensor& block_tables, + int64_t block_size); + void beam_search_rec(const torch::Tensor& logprobs, const torch::Tensor& top_tokens, const torch::Tensor& top_logprobs, diff --git a/xllm/core/layers/common/attention_metadata.h b/xllm/core/layers/common/attention_metadata.h index 21a33ba2b1..626f515b9f 100644 --- a/xllm/core/layers/common/attention_metadata.h +++ b/xllm/core/layers/common/attention_metadata.h @@ -174,6 +174,7 @@ struct AttentionMetadata { #if defined(USE_NPU) // for npu bool is_spec_verify = false; + bool use_device_kv_seq_lens = false; torch::Tensor q_seq_lens_host; torch::Tensor kv_seq_lens_host; // For ACL graph execution - fixed-address device tiling data for diff --git a/xllm/core/layers/common/attention_metadata_builder.cpp b/xllm/core/layers/common/attention_metadata_builder.cpp index ee47f881b3..24dd036861 100644 --- a/xllm/core/layers/common/attention_metadata_builder.cpp +++ b/xllm/core/layers/common/attention_metadata_builder.cpp @@ -104,6 +104,8 @@ AttentionMetadata build_attention_metadata( #if defined(USE_NPU) attn_metadata.is_spec_verify = params.is_spec_verify; + attn_metadata.use_device_kv_seq_lens = + params.attention.use_device_kv_seq_lens; attn_metadata.use_expanded_decode_for_spec_verify_attention = params.graph.use_expanded_decode_for_spec_verify_attention; if (attn_metadata.use_expanded_decode_for_spec_verify_attention) { diff --git a/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp b/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp index 7d8cd78d30..b223533061 100644 --- a/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp +++ b/xllm/core/layers/npu/npu_deepseek_v32_decoder_layer_impl.cpp @@ -1215,8 +1215,13 @@ void NpuDeepseekV32DecoderLayerImpl::build_node_variant_pack( node.variantPack.inTensors.at(WEIGHT_COUNT_PER_LAYER + 11) = atb_speed::Utils::AtTensor2Tensor( input_params.attention.device.kv_seq_lens); - node.variantPack.inTensors.at(WEIGHT_COUNT_PER_LAYER + 11).hostData = - const_cast(input_params.attention.host.kv_seq_lens.data()); + if (!input_params.attention.use_device_kv_seq_lens) { + node.variantPack.inTensors.at(WEIGHT_COUNT_PER_LAYER + 11).hostData = + const_cast(input_params.attention.host.kv_seq_lens.data()); + } else { + node.variantPack.inTensors.at(WEIGHT_COUNT_PER_LAYER + 11).hostData = + nullptr; + } } node.variantPack.inTensors.at(WEIGHT_COUNT_PER_LAYER + 12) = diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index df42b89b44..f53b0956e4 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -153,7 +153,8 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, } else { kv_seq_lens = attn_metadata.expanded_kv_seq_lens; } - } else if (attn_metadata.kv_seq_lens_host.defined()) { + } else if (!attn_metadata.use_device_kv_seq_lens && + attn_metadata.kv_seq_lens_host.defined()) { kv_seq_lens = attn_metadata.kv_seq_lens_host; } else { // Fallback if host tensor isn't prepared. @@ -173,13 +174,17 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, output); } else { // Standard PagedAttention path - xllm::kernel::npu::batch_decode(query, - k_cache, - v_cache.value_or(torch::Tensor()), - scale_, - block_table, - kv_seq_lens, - output); + xllm::kernel::npu::batch_decode( + query, + k_cache, + v_cache.value_or(torch::Tensor()), + scale_, + block_table, + kv_seq_lens, + attn_metadata.use_device_kv_seq_lens + ? std::make_optional(attn_metadata.kv_seq_lens_host) + : std::nullopt, + output); } } diff --git a/xllm/core/platform/stream.cpp b/xllm/core/platform/stream.cpp index 382a6ebbb2..965b05283d 100644 --- a/xllm/core/platform/stream.cpp +++ b/xllm/core/platform/stream.cpp @@ -45,6 +45,27 @@ PlatformStream get_stream_from_pool() { return c10::hip::getStreamFromPool(); } } // namespace +bool StreamEvent::synchronize() { +#if defined(USE_NPU) + const aclError ret = aclrtSynchronizeEvent(npu_event_); + if (ret != ACL_SUCCESS) { + LOG(ERROR) << "Failed to synchronize NPU event: " << ret; + return false; + } + return true; +#else + try { + c10_event_.synchronize(); + return true; + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to synchronize event: " << e.what(); + } catch (...) { + LOG(ERROR) << "Failed to synchronize event: unknown exception"; + } + return false; +#endif +} + Stream::Stream(const int32_t timeout) : stream_(get_stream_from_pool()), timeout_(timeout) {} diff --git a/xllm/core/platform/stream_event.h b/xllm/core/platform/stream_event.h index dc5bca71d8..12ed3cd06d 100644 --- a/xllm/core/platform/stream_event.h +++ b/xllm/core/platform/stream_event.h @@ -45,6 +45,11 @@ class StreamEvent final { c10::Event& c10_event() { return c10_event_; } #endif + // Block the calling host thread until the recorded device work completes. + // This is intentionally different from Stream::wait_event(), which only + // installs a device-stream dependency and returns immediately on the host. + bool synchronize(); + DISALLOW_COPY_AND_ASSIGN(StreamEvent); private: diff --git a/xllm/core/runtime/CMakeLists.txt b/xllm/core/runtime/CMakeLists.txt index 18ac236cc3..54334bbac1 100644 --- a/xllm/core/runtime/CMakeLists.txt +++ b/xllm/core/runtime/CMakeLists.txt @@ -34,7 +34,9 @@ cc_library( worker_client.h xservice_client.h speculative_worker_impl.h + mtp_async_input_builder.h mtp_worker_impl.h + mtp_async_state.h suffix_worker_impl.h eagle3_worker_impl.h dflash_worker_impl.h @@ -68,7 +70,9 @@ cc_library( xservice_client.cpp params_utils.cpp speculative_worker_impl.cpp + mtp_async_input_builder.cpp mtp_worker_impl.cpp + mtp_async_state.cpp suffix_worker_impl.cpp eagle3_worker_impl.cpp dflash_worker_impl.cpp diff --git a/xllm/core/runtime/mtp_async_input_builder.cpp b/xllm/core/runtime/mtp_async_input_builder.cpp new file mode 100644 index 0000000000..131ba4ad72 --- /dev/null +++ b/xllm/core/runtime/mtp_async_input_builder.cpp @@ -0,0 +1,144 @@ +/* Copyright 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 "core/runtime/mtp_async_input_builder.h" + +#include + +#if defined(USE_NPU) +#include "kernels/npu/xllm_ops/xllm_ops_api.h" +#endif + +#include "core/runtime/forward_params.h" +#include "core/runtime/mtp_async_state.h" + +namespace xllm::mtp_async { +namespace { + +torch::Tensor build_device_cache_slots(const ForwardInput& input, + const torch::Tensor& positions, + int32_t block_size) { + CHECK_EQ(positions.dim(), 2); + if (!input.input_params.multi_block_tables.empty()) { + return torch::zeros_like(positions, positions.options().dtype(torch::kInt)) + .flatten(); + } + return map_positions_to_cache_slots( + input.input_params.attention.device.block_tables, positions, block_size); +} + +void apply_device_row_metadata(ForwardInput& input, + const ForwardInput& block_table_source, + const AcceptedState& state, + const torch::Tensor& offsets, + int32_t block_size, + bool use_chunked_prefill) { + torch::Tensor row_positions = make_row_positions(state, offsets); + input.positions = row_positions.flatten().to(input.positions.options()); + input.input_params.attention.device.new_cache_slots = + build_device_cache_slots(block_table_source, row_positions, block_size); + torch::Tensor kv_seq_lens = + make_kv_seq_lens(state, offsets, use_chunked_prefill); + input.input_params.attention.device.kv_seq_lens = + kv_seq_lens.to(input.input_params.attention.device.kv_seq_lens.options()); +} + +#if defined(USE_NPU) +void apply_mtp_prepare_output( + ForwardInput& draft_input, + const kernel::npu::MtpPrepareNextDraftOutput& output, + bool use_chunked_prefill) { + draft_input.token_ids = output.token_ids; + draft_input.input_params.embedding.input_embedding = output.embeddings; + draft_input.positions = output.positions; + if (use_chunked_prefill) { + draft_input.input_params.attention.device.kv_seq_lens = output.kv_seq_lens; + } else { + draft_input.input_params.attention.device.kv_seq_lens = + torch::stack({output.kv_seq_lens - 1, output.kv_seq_lens}, /*dim=*/1) + .flatten(); + } + draft_input.input_params.attention.device.new_cache_slots = + output.cache_slots; +} +#endif + +} // namespace + +void prepare_next_draft_from_accepted_state( + ForwardInput& draft_input, + const ForwardInput& block_table_source, + const torch::Tensor& accepted_tokens, + const torch::Tensor& accepted_embeddings, + const torch::Tensor& embedding_placeholder, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens, + bool use_chunked_prefill, + int32_t block_size) { +#if defined(USE_NPU) + if (block_table_source.input_params.multi_block_tables.empty()) { + const auto output = kernel::npu::try_mtp_prepare_next_draft( + accepted_tokens, + accepted_embeddings, + embedding_placeholder, + base_positions, + base_kv_seq_lens, + block_table_source.input_params.attention.device.block_tables, + block_size); + if (output.has_value()) { + apply_mtp_prepare_output(draft_input, *output, use_chunked_prefill); + return; + } + } +#endif + + AcceptedState state = build_accepted_state(accepted_tokens, + accepted_embeddings, + embedding_placeholder, + base_positions, + base_kv_seq_lens); + // Generate offsets on device to avoid a synchronizing host-to-device copy. + torch::Tensor extend_offsets = torch::arange( + /*start=*/-1, + /*end=*/1, + torch::TensorOptions() + .dtype(torch::kLong) + .device(accepted_tokens.device())); + apply_device_row_metadata(draft_input, + block_table_source, + state, + extend_offsets, + block_size, + use_chunked_prefill); + + // On rejection, redirect the shape-stabilizing repair row to a future + // scratch position so it cannot overwrite valid draft KV state. + torch::Tensor previous_cache_positions = make_repair_cache_positions(state); + torch::Tensor cache_positions = + torch::stack({previous_cache_positions, state.base_positions}, + /*dim=*/1); + draft_input.input_params.attention.device.new_cache_slots = + build_device_cache_slots(block_table_source, cache_positions, block_size); + draft_input.token_ids = + torch::stack({state.previous_tokens, state.last_tokens}, /*dim=*/1) + .flatten() + .to(draft_input.token_ids.options()); + draft_input.input_params.embedding.input_embedding = + torch::stack({state.previous_embeddings, state.last_embeddings}, + /*dim=*/1) + .flatten(/*start_dim=*/0, /*end_dim=*/1); +} + +} // namespace xllm::mtp_async diff --git a/xllm/core/runtime/mtp_async_input_builder.h b/xllm/core/runtime/mtp_async_input_builder.h new file mode 100644 index 0000000000..e08ff768e4 --- /dev/null +++ b/xllm/core/runtime/mtp_async_input_builder.h @@ -0,0 +1,43 @@ +/* Copyright 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. +==============================================================================*/ + +#pragma once + +#include + +#include + +namespace xllm { + +struct ForwardInput; + +namespace mtp_async { + +// Applies accepted target state to the fixed [repair, current] draft layout. +// The NPU path uses one fused preparation kernel when possible and otherwise +// falls back to equivalent Torch tensor operations. +void prepare_next_draft_from_accepted_state( + ForwardInput& draft_input, + const ForwardInput& block_table_source, + const torch::Tensor& accepted_tokens, + const torch::Tensor& accepted_embeddings, + const torch::Tensor& embedding_placeholder, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens, + bool use_chunked_prefill, + int32_t block_size); + +} // namespace mtp_async +} // namespace xllm diff --git a/xllm/core/runtime/mtp_async_state.cpp b/xllm/core/runtime/mtp_async_state.cpp new file mode 100644 index 0000000000..18ba8cd5e3 --- /dev/null +++ b/xllm/core/runtime/mtp_async_state.cpp @@ -0,0 +1,169 @@ +/* Copyright 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 "core/runtime/mtp_async_state.h" + +#include + +#include + +namespace xllm::mtp_async { +namespace { + +torch::Tensor gather_sequence_rows(const torch::Tensor& values, + const torch::Tensor& indices) { + CHECK_GE(values.dim(), 2); + CHECK_EQ(values.size(0), indices.numel()); + torch::Tensor gather_index = + indices.to(torch::dtype(torch::kLong).device(indices.device())) + .view({-1, 1}); + for (int64_t dim = 2; dim < values.dim(); ++dim) { + gather_index = gather_index.unsqueeze(-1); + } + std::vector expanded_shape = values.sizes().vec(); + expanded_shape[1] = 1; + gather_index = gather_index.expand(expanded_shape); + return values.gather(/*dim=*/1, gather_index).squeeze(/*dim=*/1); +} + +} // namespace + +CombinedDraftExecutionPath classify_combined_draft_execution_path( + std::string_view model_type) { + if (model_type == "qwen3_5_mtp" || model_type == "qwen3_5_moe_mtp") { + return CombinedDraftExecutionPath::QWEN3_5_PAGED_ATTENTION; + } + if (model_type == "glm_moe_dsa_mtp") { + return CombinedDraftExecutionPath::GLM5_DSA; + } + return CombinedDraftExecutionPath::UNSUPPORTED; +} + +bool can_use_combined_decode(bool enable_schedule_overlap, + bool execution_path_supported) { + return enable_schedule_overlap && execution_path_supported; +} + +AcceptedState build_accepted_state(const torch::Tensor& accepted_tokens, + const torch::Tensor& accepted_embeddings, + const torch::Tensor& embedding_placeholder, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens) { + CHECK_EQ(accepted_tokens.dim(), 2); + CHECK_EQ(accepted_embeddings.dim(), 3); + const int64_t batch_size = accepted_tokens.size(0); + CHECK_EQ(accepted_embeddings.size(0), batch_size); + CHECK_GE(base_positions.numel(), batch_size); + CHECK_GE(base_kv_seq_lens.numel(), batch_size); + + AcceptedState state; + state.accepted_lengths = + accepted_tokens.ge(0).sum(/*dim=*/1).to(torch::kLong); + state.all_draft_accepted = + state.accepted_lengths.eq(accepted_tokens.size(/*dim=*/1)); + torch::Tensor last_indices = (state.accepted_lengths - 1).clamp_min(0); + torch::Tensor previous_indices = (state.accepted_lengths - 2).clamp_min(0); + state.last_tokens = gather_sequence_rows(accepted_tokens, last_indices); + torch::Tensor gathered_previous_tokens = + gather_sequence_rows(accepted_tokens, previous_indices); + torch::Tensor has_previous = state.accepted_lengths.gt(1); + state.previous_tokens = + torch::where(has_previous, gathered_previous_tokens, state.last_tokens); + state.last_embeddings = + gather_sequence_rows(accepted_embeddings, last_indices); + torch::Tensor gathered_previous_embeddings = + gather_sequence_rows(accepted_embeddings, previous_indices); + torch::Tensor placeholder = embedding_placeholder; + if (placeholder.dim() == 1) { + placeholder = placeholder.unsqueeze(0); + } + placeholder = placeholder.expand_as(gathered_previous_embeddings); + state.previous_embeddings = torch::where(has_previous.view({batch_size, 1}), + gathered_previous_embeddings, + placeholder); + + state.base_positions = + base_positions.flatten().slice(0, 0, batch_size).to(torch::kLong) + + state.accepted_lengths; + state.base_kv_seq_lens = + base_kv_seq_lens.flatten().slice(0, 0, batch_size).to(torch::kLong) + + state.accepted_lengths; + return state; +} + +torch::Tensor make_row_positions(const AcceptedState& state, + const torch::Tensor& offsets) { + return state.base_positions.unsqueeze(1) + + offsets.to(state.base_positions.options()).unsqueeze(0); +} + +torch::Tensor make_kv_seq_lens(const AcceptedState& state, + const torch::Tensor& offsets, + bool use_chunked_prefill) { + if (use_chunked_prefill) { + return state.base_kv_seq_lens; + } + return (state.base_kv_seq_lens.unsqueeze(1) + + offsets.to(state.base_kv_seq_lens.options()).unsqueeze(0)) + .flatten(); +} + +torch::Tensor make_repair_cache_positions(const AcceptedState& state) { + return torch::where(state.all_draft_accepted, + state.base_positions - 1, + state.base_positions + 1); +} + +torch::Tensor map_positions_to_cache_slots(const torch::Tensor& block_tables, + const torch::Tensor& positions, + int32_t block_size) { + CHECK_EQ(positions.dim(), 2); + CHECK(block_tables.defined()); + CHECK_GT(block_size, 0); + const int64_t batch_size = positions.size(0); + torch::Tensor position_long = + positions.to(torch::dtype(torch::kLong).device(positions.device())); + torch::Tensor block_indices = + torch::floor_divide(position_long, block_size) + .to(torch::dtype(torch::kLong).device(position_long.device())); + torch::Tensor block_ids = block_tables.slice(/*dim=*/0, 0, batch_size) + .to(torch::kLong) + .gather(/*dim=*/1, block_indices); + return (block_ids * block_size + position_long.remainder(block_size)) + .to(torch::kInt) + .flatten(); +} + +torch::Tensor select_rows(const torch::Tensor& values, + const torch::Tensor& selected_rows) { + if (!values.defined() || !selected_rows.defined() || + selected_rows.numel() == 0) { + return values; + } + CHECK_GE(values.dim(), 1); + if (values.size(0) == selected_rows.numel()) { + return values; + } + if (selected_rows.device().is_cpu()) { + CHECK_GT(values.size(0), selected_rows.max().item()) + << "MTP selected row index exceeds available rows."; + } + torch::Tensor index = + selected_rows.to(torch::dtype(torch::kLong).device(values.device())) + .contiguous(); + return values.index_select(/*dim=*/0, index); +} + +} // namespace xllm::mtp_async diff --git a/xllm/core/runtime/mtp_async_state.h b/xllm/core/runtime/mtp_async_state.h new file mode 100644 index 0000000000..8b72d2cc27 --- /dev/null +++ b/xllm/core/runtime/mtp_async_state.h @@ -0,0 +1,76 @@ +/* Copyright 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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include + +namespace xllm::mtp_async { + +enum class CombinedDraftExecutionPath { + UNSUPPORTED, + QWEN3_5_PAGED_ATTENTION, + GLM5_DSA, +}; + +CombinedDraftExecutionPath classify_combined_draft_execution_path( + std::string_view model_type); + +bool can_use_combined_decode(bool enable_schedule_overlap, + bool execution_path_supported); + +// Device-resident state derived from target verification. base_positions and +// base_kv_seq_lens point at the logical position immediately after the accepted +// prefix and are therefore the base of the next draft iteration. +struct AcceptedState { + torch::Tensor accepted_lengths; + torch::Tensor all_draft_accepted; + torch::Tensor last_tokens; + torch::Tensor previous_tokens; + torch::Tensor last_embeddings; + torch::Tensor previous_embeddings; + torch::Tensor base_positions; + torch::Tensor base_kv_seq_lens; +}; + +AcceptedState build_accepted_state(const torch::Tensor& accepted_tokens, + const torch::Tensor& accepted_embeddings, + const torch::Tensor& embedding_placeholder, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens); + +torch::Tensor make_row_positions(const AcceptedState& state, + const torch::Tensor& offsets); + +torch::Tensor make_kv_seq_lens(const AcceptedState& state, + const torch::Tensor& offsets, + bool use_chunked_prefill); + +// The repair row is useful only when all draft tokens were accepted. On a +// rejection it is redirected to a future scratch position so it cannot +// overwrite valid draft KV state. +torch::Tensor make_repair_cache_positions(const AcceptedState& state); + +torch::Tensor map_positions_to_cache_slots(const torch::Tensor& block_tables, + const torch::Tensor& positions, + int32_t block_size); + +torch::Tensor select_rows(const torch::Tensor& values, + const torch::Tensor& selected_rows); + +} // namespace xllm::mtp_async diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index 82812c1893..872eeb4415 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include "common/metrics.h" @@ -31,6 +32,8 @@ limitations under the License. #include "core/framework/config/speculative_config.h" #include "core/framework/kv_cache/kv_cache_estimation.h" #include "core/framework/multimodal/mm_data.h" +#include "core/runtime/mtp_async_input_builder.h" +#include "core/runtime/mtp_async_state.h" #include "spec_input_builder.h" #include "util/pretty_print.h" #include "util/slice.h" @@ -264,31 +267,6 @@ bool should_reuse_mtp_topk_indices(const ModelArgs& model_args, return reuse_enabled; } -torch::Tensor select_mtp_topk_indices_for_next_step( - const torch::Tensor& topk_indices, - const SamplingParameters& sampling_params) { - if (!topk_indices.defined() || - !sampling_params.selected_token_idxes.defined()) { - return topk_indices; - } - const torch::Tensor& selected_idxes = sampling_params.selected_token_idxes; - if (selected_idxes.numel() == 0 || - topk_indices.size(0) == selected_idxes.numel()) { - return topk_indices; - } - CHECK_GE(topk_indices.dim(), 1) - << "MTP DSA top-k indices must have at least one dimension."; - if (selected_idxes.device().is_cpu()) { - CHECK_GT(topk_indices.size(0), selected_idxes.max().item()) - << "MTP selected top-k index exceeds top-k rows."; - } - torch::Tensor index = - selected_idxes - .to(torch::dtype(torch::kLong).device(topk_indices.device())) - .contiguous(); - return topk_indices.index_select(/*dim=*/0, index); -} - std::optional run_llm_no_sync_impl( LLMWorkerImpl& worker, const ForwardInput& input, @@ -296,7 +274,10 @@ std::optional run_llm_no_sync_impl( Stream& compute_stream, ForwardInput& processed_input) { worker.prepare_work_before_execute_on_stream( - input, processed_input, prepare_stream); + input, + processed_input, + prepare_stream, + /*record_ready_event=*/&prepare_stream != &compute_stream); return worker.execute_no_sync_on_stream( processed_input, compute_stream, /*record_ready_event=*/false); } @@ -566,7 +547,8 @@ bool is_qwen3_5_target_model_type(const std::string& model_type) { } bool is_qwen3_5_draft_model_type(const std::string& model_type) { - return model_type == "qwen3_5_mtp" || model_type == "qwen3_5_moe_mtp"; + return mtp_async::classify_combined_draft_execution_path(model_type) == + mtp_async::CombinedDraftExecutionPath::QWEN3_5_PAGED_ATTENTION; } bool is_mimo_target_model_type(const std::string& model_type) { @@ -853,6 +835,18 @@ void MTPWorkerImpl::prepare_work_before_execute(const ForwardInput& input, std::optional MTPWorkerImpl::step_empty( const ForwardInput& input) { + if (pending_draft_context_.output.has_value()) { + // The preceding validation may have speculatively submitted draft1 before + // the scheduler learned that the batch had finished. Keep its graph/input + // buffers alive until the queued work completes, then discard the result. + // This is a batch-exit slow path and is never taken in steady decode. + const int32_t ret = compute_stream_->synchronize(); + CHECK_EQ(ret, 0) << "failed to drain final MTP draft prelaunch, ret=" + << ret; + pending_draft_context_ = PendingDraftContext(); + } + flush_pending_target_context(); + if (!input.input_params.meta.batch_forward_type.is_decode()) { ForwardInput target_prepared; ForwardInput draft_prepared; @@ -922,6 +916,8 @@ std::optional MTPWorkerImpl::step_empty( std::optional MTPWorkerImpl::step_prefill( const ForwardInput& input) { + flush_pending_target_context(); + Timer timer; ForwardInput target_prepared; ForwardInput draft_prepared; @@ -1083,12 +1079,48 @@ std::optional MTPWorkerImpl::step_decode( stabilize_decode_host_tensors(input); } const int32_t num_speculative_tokens = options_.num_speculative_tokens(); + // Reuse draft-0 prelaunched for this same batch. + const bool use_prelaunched_first_draft = + can_use_combined_first_draft() && pending_draft_context_matches(input); + const bool matching_device_target_context = + pending_target_context_matches(input); + // Consume this batch's pending target context directly on device. + const bool use_device_target_context = + can_use_combined_first_draft() && matching_device_target_context && + device_target_context_ready_for_batch(input); + if (pending_draft_context_.output.has_value() && + !use_prelaunched_first_draft) { + // A batch transition invalidates the speculative prelaunch. Drain it + // before releasing its graph/input buffers; this slow path is outside + // steady decode and preserves cache/buffer lifetime correctness. + const int32_t ret = compute_stream_->synchronize(); + CHECK_EQ(ret, 0) << "failed to drain stale MTP draft prelaunch, ret=" + << ret; + pending_draft_context_ = PendingDraftContext(); + } + if (!use_device_target_context) { + // Batch transitions are uncommon in steady decode. Materialize the most + // recent target state only for that fallback; the normal path below never + // synchronizes the worker thread with the NPU. + flush_pending_target_context(); + if (matching_device_target_context) { + // The first target handoff for a new batch establishes the scheduler's + // corrected position/KV base. Subsequent handoffs can derive that base + // fully on device without waiting for the scheduler. + device_context_ready_embedding_ids_ = + input.input_params.embedding.embedding_ids; + device_context_ready_request_ids_ = + input.input_params.embedding.request_ids; + } else if (!device_target_context_ready_for_batch(input)) { + device_context_ready_embedding_ids_.clear(); + device_context_ready_request_ids_.clear(); + } + } std::vector draft_outputs; ForwardInput current_draft_input, validate_input, next_step_input; std::vector draft_prepared(num_speculative_tokens); Timer timer; - CHECK(embedding_cache_ != nullptr) << "MTP embedding cache is not allocated"; const auto& embedding = input.input_params.embedding; @@ -1129,20 +1161,81 @@ std::optional MTPWorkerImpl::step_decode( } } - // Get decode state of last step - std::vector last_states = - embedding_cache_->read_decode_states( - input.input_params.embedding.embedding_ids, - input.input_params.embedding.request_ids); - CHECK_EQ(last_states.size(), - input.input_params.embedding.embedding_ids.size()) - << "decode target state count mismatch"; - check_mtp_decode_states(last_states, - input.input_params.embedding.request_ids, - input.token_ids_host, - enable_schedule_overlap()); - update_decode_step_input(input, last_states); - prepare_draft_extend_inputs(input, last_states, current_draft_input); + ForwardInput metadata_template = input; + if (use_prelaunched_first_draft) { + // The first draft was fully prepared and submitted by the preceding + // run_validate(). Host metadata is corrected below after the accepted + // token handoff; it is only needed for draft steps 1..N-1 and target. + } else if (use_device_target_context) { + c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); + + // Build fixed-shape host metadata immediately while target verification is + // still running. Use the maximum accepted draft offset for conservative + // graph planning; actual values replace every device tensor below. + int32_t* template_positions = + metadata_template.positions_host.data_ptr(); + int32_t* template_tokens = + metadata_template.token_ids_host.data_ptr(); + auto& template_kv_lens = + metadata_template.input_params.attention.host.kv_seq_lens; + for (int32_t seq_id = 0; + seq_id < metadata_template.input_params.meta.num_sequences; + ++seq_id) { + template_positions[seq_id] += num_speculative_tokens; + template_kv_lens[seq_id] += num_speculative_tokens; + if (template_tokens[seq_id] < 0) { + template_tokens[seq_id] = 0; + } + } + + std::vector template_states( + metadata_template.input_params.meta.num_sequences); + const torch::Tensor& placeholder = + embedding_cache_->embedding_placeholder(); + for (int32_t seq_id = 0; + seq_id < metadata_template.input_params.meta.num_sequences; + ++seq_id) { + template_states[seq_id].valid = true; + template_states[seq_id].request_id = + metadata_template.input_params.embedding.request_ids[seq_id]; + template_states[seq_id].token_id = template_tokens[seq_id]; + template_states[seq_id].embedding = placeholder; + } + prepare_draft_extend_inputs(metadata_template, + template_states, + current_draft_input, + /*force_two_rows=*/true); + current_draft_input.input_params.attention.use_device_kv_seq_lens = true; + wait_metadata_ready_event(current_draft_input, *compute_stream_); + clear_ready_events(current_draft_input); + + mtp_async::prepare_next_draft_from_accepted_state( + current_draft_input, + input, + pending_target_context_.accepted_tokens, + pending_target_context_.accepted_embeddings, + embedding_cache_->embedding_placeholder(), + pending_target_context_.base_positions, + pending_target_context_.base_kv_seq_lens, + /*use_chunked_prefill=*/false, + options_.block_size()); + } else { + // First decode after prefill and batch transitions use the host cache. + std::vector last_states = + embedding_cache_->read_decode_states( + input.input_params.embedding.embedding_ids, + input.input_params.embedding.request_ids); + CHECK_EQ(last_states.size(), + input.input_params.embedding.embedding_ids.size()) + << "decode target state count mismatch"; + check_mtp_decode_states(last_states, + input.input_params.embedding.request_ids, + input.token_ids_host, + enable_schedule_overlap()); + update_decode_step_input(input, last_states); + metadata_template = input; + prepare_draft_extend_inputs(input, last_states, current_draft_input); + } draft_outputs.reserve(num_speculative_tokens); const bool reuse_mtp_topk_indices = should_reuse_mtp_topk_indices( draft_impl_->context_.get_model_args(), enable_schedule_overlap()); @@ -1151,18 +1244,50 @@ std::optional MTPWorkerImpl::step_decode( if (reuse_mtp_topk_indices) { current_draft_input.input_params.dsa_topk_indices = mtp_topk_indices; } - std::optional draft_output_opt = - run_llm_no_sync_impl(*draft_impl_, - current_draft_input, - *prepare_stream_, - *compute_stream_, - draft_prepared[draft_idx]); + std::optional draft_output_opt; + if (use_prelaunched_first_draft && draft_idx == 0) { + draft_output_opt = std::move(pending_draft_context_.output); + draft_prepared[draft_idx] = + std::move(pending_draft_context_.prepared_input); + pending_draft_context_ = PendingDraftContext(); + } else { + draft_output_opt = run_llm_no_sync_impl(*draft_impl_, + current_draft_input, + *compute_stream_, + *compute_stream_, + draft_prepared[draft_idx]); + } + + if ((use_device_target_context || use_prelaunched_first_draft) && + draft_idx == 0) { + // The next draft forward is already queued behind target validation. + // It can start immediately when rejection sampling finishes while the + // worker materializes the accepted state for later draft/target metadata + // on CPU. This synchronization is therefore outside the NPU critical + // path rather than sitting between target and draft launches. + flush_pending_target_context(); + std::vector resolved_states = + embedding_cache_->read_decode_states( + input.input_params.embedding.embedding_ids, + input.input_params.embedding.request_ids); + // The scheduler input still contains the conservative placeholder used + // to launch the first draft. Materialize the accepted host state before + // comparing token ids; otherwise this check compares the new target + // result with the intentionally stale overlap template. + input.token_ids_host = torch::full_like(input.token_ids_host, -1); + update_decode_step_input(input, resolved_states); + check_mtp_decode_states(resolved_states, + input.input_params.embedding.request_ids, + input.token_ids_host, + /*allow_overlap_fake_token=*/false); + metadata_template = input; + } // Overlap next-step input preparation with async draft forward. if (draft_idx == num_speculative_tokens - 1) { - prepare_validate_inputs(input, validate_input); + prepare_validate_inputs(metadata_template, validate_input); } else { - prepare_draft_inputs(input, next_step_input, draft_idx + 1); + prepare_draft_inputs(metadata_template, next_step_input, draft_idx + 1); } CHECK(draft_output_opt.has_value()) @@ -1172,9 +1297,9 @@ std::optional MTPWorkerImpl::step_decode( { c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); if (reuse_mtp_topk_indices) { - mtp_topk_indices = select_mtp_topk_indices_for_next_step( + mtp_topk_indices = mtp_async::select_rows( draft_outputs.back().dsa_topk_indices, - current_draft_input.sampling_params); + draft_prepared[draft_idx].sampling_params.selected_token_idxes); } // Unify this step's draft next_tokens across the consensus group before // process_draft_sample_output() compresses the still-full [batch, vocab] @@ -1201,8 +1326,9 @@ std::optional MTPWorkerImpl::step_decode( if (last_output.embeddings.defined()) { current_draft_input.input_params.embedding.input_embedding = last_output.embeddings; + // input_embedding is produced and consumed on compute_stream_; FIFO + // ordering replaces the same-stream EventRecord/EventWait pair. } - record_current_metadata_ready_event(current_draft_input, *compute_stream_); } COUNTER_ADD(speculative_execution_latency_seconds_draft, timer.elapsed_seconds()); @@ -1260,13 +1386,25 @@ std::optional MTPWorkerImpl::run_validate( draft_outputs, validate_input, *compute_stream_); ForwardOutput target_output = run_llm_no_sync_impl(*impl_, validate_input, - *prepare_stream_, + *compute_stream_, *compute_stream_, target_prepared) .value(); COUNTER_ADD(speculative_execution_latency_seconds_target, timer.elapsed_seconds()); + const bool prelaunch_next_first_draft = + can_use_combined_first_draft() && + device_target_context_ready_for_batch(input); + ForwardInput next_first_draft_input; + if (prelaunch_next_first_draft) { + // This input is independent of the accepted token. Prepare it on the + // auxiliary stream while target verification is still executing; the + // compute stream consumes it through a device-side event after rejection + // sampling, with no host synchronization. + prepare_next_first_draft_template(input, next_first_draft_input); + } + // verify the proposals with target and update the batch timer.reset(); SampleOutput val_output; @@ -1277,21 +1415,84 @@ std::optional MTPWorkerImpl::run_validate( COUNTER_ADD(speculative_execution_latency_seconds_validation, timer.elapsed_seconds()); - // Catch-all for cross-rank RNG divergence: unify the accepted next_tokens to - // the consensus group's rank 0 so all_draft_accepted and the next - // draft-extend row layout agree across ranks. - if (get_optimization_config().enable_spec_token_broadcast && - enable_schedule_overlap()) { + const int64_t batch_size = val_output.next_tokens.size(0); + const int64_t num_val_tokens = options_.num_speculative_tokens() + 1; + CHECK_EQ(validate_input.positions.numel(), batch_size * num_val_tokens) + << "validate positions must contain one row per speculative token"; + const torch::Tensor& validate_kv_seq_lens = + validate_input.input_params.attention.device.kv_seq_lens; + CHECK_GE(validate_kv_seq_lens.numel(), batch_size) + << "validate KV lengths must be sequence-scoped"; + torch::Tensor accepted_tokens_host = + acquire_accepted_tokens_host_buffer(val_output.next_tokens); + torch::Tensor accepted_tokens_cpu_result = accepted_tokens_host; + torch::Tensor base_positions; + torch::Tensor base_kv_seq_lens; + StreamEventPtr handoff_ready_event; + { c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); - broadcast_spec_tokens(val_output.next_tokens, - spec_broadcast_group(parallel_args_)); + + // Catch-all for cross-rank RNG divergence: unify accepted tokens before + // deriving any device-resident state used by the next draft iteration. + if (get_optimization_config().enable_spec_token_broadcast && + enable_schedule_overlap()) { + broadcast_spec_tokens(val_output.next_tokens, + spec_broadcast_group(parallel_args_)); + } + + base_positions = validate_input.positions.view({batch_size, num_val_tokens}) + .select(/*dim=*/1, /*index=*/0) + .contiguous(); + base_kv_seq_lens = validate_kv_seq_lens.flatten().slice(0, 0, batch_size) - + options_.num_speculative_tokens(); + + accepted_tokens_host.copy_(val_output.next_tokens, + /*non_blocking=*/true); + // The event covers consensus, base-state derivation, and the D2H handoff. + handoff_ready_event = compute_stream_->record_event(); + } + if (handoff_ready_event == nullptr) { + const int32_t ret = compute_stream_->synchronize(); + CHECK_EQ(ret, 0) << "failed to synchronize MTP handoff, ret=" << ret; + } + stage_target_context_write(input, + val_output, + base_positions, + base_kv_seq_lens, + handoff_ready_event, + std::move(accepted_tokens_host)); + if (prelaunch_next_first_draft) { + // Submit the next iteration's first draft before returning to the + // scheduler. This is the actual asynchronous boundary: scheduler/host + // accepted-state work can no longer sit between target validation and the + // next draft launch. + enqueue_next_first_draft(input, + val_output, + base_positions, + base_kv_seq_lens, + std::move(next_first_draft_input)); + } + target_output.ready_event = handoff_ready_event; + + // Target validation consumes all draft outputs on the same compute stream. + // Keep their prepared inputs alive until the target handoff event completes. + for (const ForwardOutput& draft_output : draft_outputs) { + if (draft_output.retained_input != nullptr) { + target_output.retained_input_dependencies.emplace_back( + draft_output.retained_input); + } + target_output.retained_input_dependencies.insert( + target_output.retained_input_dependencies.end(), + draft_output.retained_input_dependencies.begin(), + draft_output.retained_input_dependencies.end()); } - const int32_t ret = compute_stream_->synchronize(); - CHECK_EQ(ret, 0) << "failed to synchronize MTP compute stream, ret=" << ret; - release_retained_inputs(target_output); - val_output.next_tokens = val_output.next_tokens.to(torch::kCPU); - write_target_context_to_cache(input, val_output); + if (!enable_schedule_overlap()) { + flush_pending_target_context(); + release_retained_inputs(target_output); + target_output.ready_event.reset(); + val_output.next_tokens = std::move(accepted_tokens_cpu_result); + } if (!enable_schedule_overlap() && !driver_ && !dp_driver_) { return std::nullopt; @@ -1302,19 +1503,225 @@ std::optional MTPWorkerImpl::run_validate( return target_output; } -void MTPWorkerImpl::write_target_context_to_cache( +void MTPWorkerImpl::stage_target_context_write( const ForwardInput& input, - const SampleOutput& validate_output) { + const SampleOutput& validate_output, + torch::Tensor base_positions, + torch::Tensor base_kv_seq_lens, + StreamEventPtr ready_event, + torch::Tensor accepted_tokens_host) { + CHECK(!pending_target_context_.accepted_tokens.defined()) + << "previous MTP target context must be flushed before staging another"; + pending_target_context_.embedding_ids = + input.input_params.embedding.embedding_ids; + pending_target_context_.request_ids = + input.input_params.embedding.request_ids; + pending_target_context_.accepted_tokens = validate_output.next_tokens; + pending_target_context_.accepted_tokens_host = + std::move(accepted_tokens_host); + pending_target_context_.accepted_embeddings = validate_output.embeddings; + pending_target_context_.base_positions = std::move(base_positions); + pending_target_context_.base_kv_seq_lens = std::move(base_kv_seq_lens); + pending_target_context_.ready_event = std::move(ready_event); +} + +torch::Tensor MTPWorkerImpl::acquire_accepted_tokens_host_buffer( + const torch::Tensor& accepted_tokens) { + CHECK(accepted_tokens.defined()) << "accepted tokens must be defined"; + CHECK_GT(accepted_tokens.numel(), 0) << "accepted tokens must not be empty"; + CHECK(!pending_target_context_.accepted_tokens.defined()) + << "accepted-token host buffer is still in use"; + + const int64_t required_capacity = accepted_tokens.numel(); + const int64_t configured_capacity = + static_cast(options_.max_seqs_per_batch()) * + (static_cast(options_.num_speculative_tokens()) + 1); + const bool needs_allocation = + !accepted_tokens_host_buffer_.defined() || + accepted_tokens_host_buffer_.scalar_type() != + accepted_tokens.scalar_type() || + accepted_tokens_host_buffer_.numel() < required_capacity; + if (needs_allocation) { + const int64_t capacity = std::max(required_capacity, configured_capacity); + accepted_tokens_host_buffer_ = torch::empty( + {capacity}, + accepted_tokens.options().device(torch::kCPU).pinned_memory(true)); + } + + return accepted_tokens_host_buffer_ + .narrow(/*dim=*/0, /*start=*/0, required_capacity) + .view(accepted_tokens.sizes()); +} + +bool MTPWorkerImpl::pending_target_context_matches( + const ForwardInput& input) const { + return pending_target_context_.accepted_tokens.defined() && + pending_target_context_.embedding_ids == + input.input_params.embedding.embedding_ids && + pending_target_context_.request_ids == + input.input_params.embedding.request_ids; +} + +bool MTPWorkerImpl::device_target_context_ready_for_batch( + const ForwardInput& input) const { + return device_context_ready_embedding_ids_ == + input.input_params.embedding.embedding_ids && + device_context_ready_request_ids_ == + input.input_params.embedding.request_ids; +} + +void MTPWorkerImpl::flush_pending_target_context() { + if (!pending_target_context_.accepted_tokens.defined()) { + return; + } + CHECK(pending_target_context_.ready_event == nullptr || + pending_target_context_.ready_event->synchronize()) + << "failed to wait for pending MTP target context"; CHECK(embedding_cache_ != nullptr) << "embedding_cache_ must be initialized before target cache write"; - CHECK(!input.input_params.embedding.embedding_ids.empty()) - << "target context cache write requires embedding ids"; embedding_cache_->write_target_context( - input.input_params.embedding.embedding_ids, - input.input_params.embedding.request_ids, + pending_target_context_.embedding_ids, + pending_target_context_.request_ids, + pending_target_context_.accepted_tokens_host, + pending_target_context_.accepted_embeddings, + options_.num_speculative_tokens()); + pending_target_context_ = PendingTargetContext(); +} + +bool MTPWorkerImpl::supports_combined_first_draft_execution() const { +#if defined(USE_NPU) + if (draft_impl_ == nullptr || + draft_impl_->get_status() == WorkerImpl::Status::UNINITIALIZED) { + return false; + } + + // The prelaunch builds an eager two-row DECODE input. The ATB speculative + // path expects CHUNKED_PREFILL metadata instead, and the current device + // handoff is rank-local rather than DP-aware. + if (::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel() || + parallel_args_.dp_size() > 1) { + return false; + } + + const ModelArgs& draft_args = draft_impl_->context_.get_model_args(); + switch (mtp_async::classify_combined_draft_execution_path( + draft_args.model_type())) { + case mtp_async::CombinedDraftExecutionPath::GLM5_DSA: + return true; + case mtp_async::CombinedDraftExecutionPath::QWEN3_5_PAGED_ATTENTION: + return device_.unwrap().is_privateuseone() && + ::xllm::KernelConfig::get_instance().npu_kernel_backend() == + "TORCH"; + case mtp_async::CombinedDraftExecutionPath::UNSUPPORTED: + return false; + } + return false; +#else + return false; +#endif +} + +bool MTPWorkerImpl::can_use_combined_first_draft() const { + return mtp_async::can_use_combined_decode( + enable_schedule_overlap(), supports_combined_first_draft_execution()); +} + +void MTPWorkerImpl::prepare_next_first_draft_template( + const ForwardInput& input, + ForwardInput& combined_input) { + CHECK(embedding_cache_ != nullptr); + + ForwardInput metadata_template = input; + const int32_t num_speculative_tokens = options_.num_speculative_tokens(); + int32_t* template_positions = + metadata_template.positions_host.data_ptr(); + int32_t* template_tokens = + metadata_template.token_ids_host.data_ptr(); + auto& template_kv_lens = + metadata_template.input_params.attention.host.kv_seq_lens; + for (int32_t seq_id = 0; + seq_id < metadata_template.input_params.meta.num_sequences; + ++seq_id) { + template_positions[seq_id] += num_speculative_tokens; + template_kv_lens[seq_id] += num_speculative_tokens; + if (template_tokens[seq_id] < 0) { + template_tokens[seq_id] = 0; + } + } + + std::vector template_states( + metadata_template.input_params.meta.num_sequences); + const torch::Tensor& placeholder = embedding_cache_->embedding_placeholder(); + for (int32_t seq_id = 0; + seq_id < metadata_template.input_params.meta.num_sequences; + ++seq_id) { + template_states[seq_id].valid = true; + template_states[seq_id].request_id = + metadata_template.input_params.embedding.request_ids[seq_id]; + template_states[seq_id].token_id = template_tokens[seq_id]; + template_states[seq_id].embedding = placeholder; + } + + prepare_draft_extend_inputs(metadata_template, + template_states, + combined_input, + /*force_two_rows=*/true, + /*wait_for_compute_stream=*/false); + combined_input.skip_sampling_for_logits_only = false; + combined_input.input_params.attention.use_device_kv_seq_lens = true; +} + +void MTPWorkerImpl::enqueue_next_first_draft( + const ForwardInput& input, + const SampleOutput& validate_output, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens, + ForwardInput combined_input) { + CHECK(!pending_draft_context_.output.has_value()) + << "MTP first-draft prelaunch was not consumed"; + CHECK(validate_output.next_tokens.defined()); + CHECK(validate_output.embeddings.defined()); + CHECK(embedding_cache_ != nullptr); + + c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); + wait_metadata_ready_event(combined_input, *compute_stream_); + clear_ready_events(combined_input); + + // Interleave [repair, current] rows in one decode batch. Every transformer + // layer projects both rows, writes both KV rows, and only then launches + // PagedAttention. Same-stream ordering therefore makes repair KV visible to + // the current row without a host wait or a separate repair forward. + mtp_async::prepare_next_draft_from_accepted_state( + combined_input, + input, validate_output.next_tokens, validate_output.embeddings, - options_.num_speculative_tokens()); + embedding_cache_->embedding_placeholder(), + base_positions, + base_kv_seq_lens, + /*use_chunked_prefill=*/false, + options_.block_size()); + + pending_draft_context_.embedding_ids = + input.input_params.embedding.embedding_ids; + pending_draft_context_.request_ids = input.input_params.embedding.request_ids; + pending_draft_context_.output = + run_llm_no_sync_impl(*draft_impl_, + combined_input, + *compute_stream_, + *compute_stream_, + pending_draft_context_.prepared_input); + CHECK(pending_draft_context_.output.has_value()) + << "failed to prelaunch next MTP first draft"; +} + +bool MTPWorkerImpl::pending_draft_context_matches( + const ForwardInput& input) const { + return pending_draft_context_.output.has_value() && + pending_draft_context_.embedding_ids == + input.input_params.embedding.embedding_ids && + pending_draft_context_.request_ids == + input.input_params.embedding.request_ids; } void MTPWorkerImpl::process_draft_sample_output(SampleOutput& sample_output) { @@ -1593,14 +2000,23 @@ void MTPWorkerImpl::prepare_validate_inputs(const ForwardInput& input, } #endif validate_input.device_tensors_ready = true; + // This metadata is independent of the in-flight final draft. Keep it on the + // auxiliary stream and hand it to the compute stream with a device event. finish_metadata_prepare(*prepare_stream_, validate_input); } void MTPWorkerImpl::prepare_draft_extend_inputs( const ForwardInput& base_input, const std::vector& last_states, - ForwardInput& extend_input) { + ForwardInput& extend_input, + bool force_two_rows, + bool wait_for_compute_stream) { c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); + // Regular draft preparation may consume tensors produced by the previous + // compute. The placeholder-only first-draft prelaunch has no such dependency. + if (wait_for_compute_stream) { + prepare_stream_->wait_stream(*compute_stream_); + } extend_input = base_input; extend_input.sampling_params.return_probs = !extend_input.sampling_params.all_greedy_sample; @@ -1696,7 +2112,8 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( selected_row_idx.emplace_back(2 * seq_id + 1); continue; } - const bool use_two_rows = dp_enabled || state.all_draft_accepted; + const bool use_two_rows = + force_two_rows || dp_enabled || state.all_draft_accepted; if (use_two_rows) { int32_t prev_token_id = state.prev_token_id; int32_t prev_position_offset = -1; @@ -1767,12 +2184,6 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( } input_params.attention.rebuild_device_buffer(device_); - // Establish cross-stream ordering before stacking the draft-extend - // embeddings. The stack runs on prepare_stream_, but the embedding rows it - // reads are produced/finalized by work enqueued on compute_stream_. Make - // prepare_stream_ wait on compute_stream_ via an event. - prepare_stream_->wait_stream(*compute_stream_); - input_params.embedding.input_embedding = torch::stack(expanded_embeddings); if (!input_params.parallel.dp_global_token_nums.empty()) { @@ -1796,21 +2207,28 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( params.selected_token_idxes.defined() ? params.selected_token_idxes.options() : torch::dtype(torch::kInt).device(device_); - params.selected_token_idxes = - safe_to(specBuilder::make_cpu_int_tensor(selected_row_idx), - idx_options, - /*non_blocking=*/true); - if (!params.sample_idxes.defined()) { - std::vector sample_idxes_vec; - sample_idxes_vec.reserve(num_sequences); - for (int32_t i = 0; i < num_sequences; ++i) { - sample_idxes_vec.emplace_back(i); - } - params.sample_idxes = - safe_to(specBuilder::make_cpu_int_tensor(sample_idxes_vec), + if (use_chunked_prefill || dp_enabled || force_two_rows) { + // These layouts always append two rows per sequence and select the second + // row. Build the tiny control tensor directly on device; copying a + // temporary pinned CPU tensor forces its allocator to synchronize before + // the asynchronous H2D has completed. + params.selected_token_idxes = torch::arange( + /*start=*/1, + /*end=*/2 * num_sequences, + /*step=*/2, + idx_options); + } else { + params.selected_token_idxes = + safe_to(specBuilder::make_cpu_int_tensor(selected_row_idx), idx_options, /*non_blocking=*/true); } + if (!params.sample_idxes.defined()) { + // This control tensor is always the identity mapping. Generate it directly + // on device instead of allocating a short-lived pinned H2D source. + params.sample_idxes = torch::arange( + /*start=*/0, /*end=*/num_sequences, idx_options); + } extend_input.device_tensors_ready = true; finish_metadata_prepare(*prepare_stream_, extend_input); } @@ -1872,6 +2290,8 @@ void MTPWorkerImpl::prepare_draft_inputs(const ForwardInput& input, // token_ids is intentionally filled later from the previous draft output. draft_input.device_tensors_ready = false; + // Positions/KV metadata do not depend on the in-flight draft result. Prepare + // them concurrently; token ids and embeddings are filled on compute_stream. finish_metadata_prepare(*prepare_stream_, draft_input); } diff --git a/xllm/core/runtime/mtp_worker_impl.h b/xllm/core/runtime/mtp_worker_impl.h index e961019cee..47c8e5c452 100644 --- a/xllm/core/runtime/mtp_worker_impl.h +++ b/xllm/core/runtime/mtp_worker_impl.h @@ -125,10 +125,53 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { void prepare_draft_extend_inputs( const ForwardInput& base_input, const std::vector& last_states, - ForwardInput& extend_input); - - void write_target_context_to_cache(const ForwardInput& input, - const SampleOutput& validate_output); + ForwardInput& extend_input, + bool force_two_rows = false, + bool wait_for_compute_stream = true); + + struct PendingTargetContext { + std::vector embedding_ids; + std::vector request_ids; + // Both tensors stay on device. A steady-state overlap step consumes them + // by queueing gather/update ops behind rejection sampling on the same + // stream. They are materialized on CPU only when the batch shape/order + // changes and the host cache fallback is required. + torch::Tensor accepted_tokens; + torch::Tensor accepted_tokens_host; + torch::Tensor accepted_embeddings; + torch::Tensor base_positions; + torch::Tensor base_kv_seq_lens; + StreamEventPtr ready_event; + }; + + struct PendingDraftContext { + std::vector embedding_ids; + std::vector request_ids; + std::optional output; + ForwardInput prepared_input; + }; + + void stage_target_context_write(const ForwardInput& input, + const SampleOutput& validate_output, + torch::Tensor base_positions, + torch::Tensor base_kv_seq_lens, + StreamEventPtr ready_event, + torch::Tensor accepted_tokens_host); + torch::Tensor acquire_accepted_tokens_host_buffer( + const torch::Tensor& accepted_tokens); + bool pending_target_context_matches(const ForwardInput& input) const; + bool device_target_context_ready_for_batch(const ForwardInput& input) const; + void flush_pending_target_context(); + bool supports_combined_first_draft_execution() const; + bool can_use_combined_first_draft() const; + void prepare_next_first_draft_template(const ForwardInput& input, + ForwardInput& combined_input); + void enqueue_next_first_draft(const ForwardInput& input, + const SampleOutput& validate_output, + const torch::Tensor& base_positions, + const torch::Tensor& base_kv_seq_lens, + ForwardInput combined_input); + bool pending_draft_context_matches(const ForwardInput& input) const; protected: // Draft model worker @@ -137,6 +180,22 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { // Embedding cache for speculative decoding std::shared_ptr embedding_cache_; + // Rejection sampling produces accepted state on the compute stream. Keep + // that state device-resident so the next overlap task can be fully enqueued + // without waiting for target verification to finish. + PendingTargetContext pending_target_context_; + std::vector device_context_ready_embedding_ids_; + std::vector device_context_ready_request_ids_; + // A single persistent pinned destination is sufficient for accepted-token + // D2H: the preceding pending target context is always flushed before the + // next validation can submit another copy. The pending context holds a view + // into this storage until the copy event is synchronized and CPU consumers + // have finished reading it. + torch::Tensor accepted_tokens_host_buffer_; + // Draft step 0 is submitted at the tail of the preceding target validation, + // before control returns to the scheduler. The following scheduler turn + // consumes this output and only submits draft steps 1..N-1. + PendingDraftContext pending_draft_context_; // Whether validation directly uses selected-only draft_probs [B, S]. // If false, selected-only cache values are restored to dense [B, S, V]. bool enable_opt_validate_probs_ = false; diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 060875fce8..1979f7291c 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -204,13 +204,28 @@ void prepare_input_params_for_linear_attention(ModelInputParams& input_params) { input_params.parallel.query_start_loc[i] + seq_len; } - torch::Tensor has_initial_state_tensor = - input_params.attention.device.kv_cache_tokens_nums > 0; - torch::Tensor has_initial_state_int64 = has_initial_state_tensor.contiguous() + const std::vector& host_kv_cache_tokens_nums = + input_params.attention.host.kv_cache_tokens_nums; + std::vector has_initial_state; + if (!host_kv_cache_tokens_nums.empty()) { + has_initial_state.reserve(host_kv_cache_tokens_nums.size()); + for (int32_t num_tokens : host_kv_cache_tokens_nums) { + has_initial_state.emplace_back(num_tokens > 0 ? 1 : 0); + } + } else { + // Compatibility fallback for inputs that only carry the device view. + torch::Tensor has_initial_state_tensor = + input_params.attention.device.kv_cache_tokens_nums > 0; + torch::Tensor has_initial_state_cpu = has_initial_state_tensor.contiguous() .view({-1}) .to(torch::kCPU) .to(torch::kInt64); - const int64_t has_initial_state_size = has_initial_state_int64.size(0); + has_initial_state.assign(has_initial_state_cpu.data_ptr(), + has_initial_state_cpu.data_ptr() + + has_initial_state_cpu.numel()); + } + const int64_t has_initial_state_size = + static_cast(has_initial_state.size()); CHECK_GT(has_initial_state_size, 0) << "kv_cache_tokens_nums must not be empty for linear attention"; CHECK(batch_size == has_initial_state_size || @@ -219,21 +234,16 @@ void prepare_input_params_for_linear_attention(ModelInputParams& input_params) { << "size, kv_cache_tokens_nums_size=" << has_initial_state_size << ", batch_size=" << batch_size; if (batch_size == has_initial_state_size) { - input_params.parallel.has_initial_state = std::vector( - has_initial_state_int64.data_ptr(), - has_initial_state_int64.data_ptr() + batch_size); + input_params.parallel.has_initial_state = std::move(has_initial_state); return; } const int64_t repeat_count = batch_size / has_initial_state_size; input_params.parallel.has_initial_state.clear(); input_params.parallel.has_initial_state.reserve(batch_size); - const int64_t* has_initial_state_ptr = - has_initial_state_int64.data_ptr(); for (int64_t i = 0; i < has_initial_state_size; ++i) { for (int64_t repeat_idx = 0; repeat_idx < repeat_count; ++repeat_idx) { - input_params.parallel.has_initial_state.push_back( - has_initial_state_ptr[i]); + input_params.parallel.has_initial_state.push_back(has_initial_state[i]); } } } @@ -760,7 +770,8 @@ void WorkerImpl::prepare_work_before_execute(const ForwardInput& input, void WorkerImpl::prepare_work_before_execute_on_stream( const ForwardInput& input, ForwardInput& processed_input, - Stream& prepare_stream) { + Stream& prepare_stream, + bool record_ready_event) { #if defined(USE_NPU) // Without device_capture_lock, ACL graph capture will be interrupted by the // synchronization H2D of data update streams asynchronously scheduled by @@ -1008,11 +1019,15 @@ void WorkerImpl::prepare_work_before_execute_on_stream( prepare_device_on_stream(); - StreamEventPtr event = prepare_stream.record_event(); - if (event == nullptr) { - prepare_stream.synchronize(); + if (record_ready_event) { + StreamEventPtr event = prepare_stream.record_event(); + if (event == nullptr) { + prepare_stream.synchronize(); + } + processed_input.metadata_ready_event = event; + } else { + processed_input.metadata_ready_event.reset(); } - processed_input.metadata_ready_event = event; } void WorkerImpl::apply_kv_block_swaps(const ModelInputParams& input_params) { diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index 50fae27196..979d11c76c 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -114,7 +114,8 @@ class WorkerImpl { ForwardInput& processed_inputs); void prepare_work_before_execute_on_stream(const ForwardInput& input, ForwardInput& processed_input, - Stream& prepare_stream); + Stream& prepare_stream, + bool record_ready_event = true); // Internal helper shared by worker pipelines before model execution. virtual void apply_kv_block_swaps(const ModelInputParams& input_params);