perf: eliminate validate-to-draft bubbles with asynchronous device-side handoff.

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)

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 <gtest/gtest.h>
+#include <torch/torch.h>
+#include <torch_npu/csrc/libs/init_npu.h>
+#include <torch_npu/torch_npu.h>
+
+#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

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

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 <gtest/gtest.h>
+#include <torch/torch.h>
+
+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

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;

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<torch::Tensor>& 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);
 }
 

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<torch::Tensor>& seq_lens_host,
                   torch::Tensor& output);
 
 std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(

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