diff --git a/tests/core/framework/batch/batch_test.cpp b/tests/core/framework/batch/batch_test.cpp index 3fc5ea8a8d..7b8d6061a2 100644 --- a/tests/core/framework/batch/batch_test.cpp +++ b/tests/core/framework/batch/batch_test.cpp @@ -2056,7 +2056,7 @@ TEST(BatchTest, OverlapMTPReplacementKeepsCompositeKvBlocks) { .max_seqs_per_batch(max_seqs_per_batch) .manager_types({1, 0, 0}) .compress_ratios({0, 4, 128}); - CompositeBlockManager manager(options); + CompositeBlockManager manager(build_composite_leaves(options)); RequestSamplingParam sampling_param; StoppingChecker stopping_checker; @@ -2064,8 +2064,8 @@ TEST(BatchTest, OverlapMTPReplacementKeepsCompositeKvBlocks) { Sequence seq = make_overlap_sequence( {1, 10, 11}, /*seq_capacity=*/128, &sampling_param, &stopping_checker); - ASSERT_TRUE(manager.allocate_for_sequence(&seq, seq.num_prompt_tokens())); - ASSERT_EQ(seq.kv_state().num_kv_blocks(), 0u); + ASSERT_TRUE(manager.allocate_sequence(&seq, seq.num_prompt_tokens())); + ASSERT_EQ(seq.kv_state().num_blocks(BlockType::KV), 0u); ASSERT_GT(seq.kv_state().current_max_tokens_capacity(), 0u); seq.kv_state().incr_kv_cache_tokens_num(seq.num_prompt_tokens() - 1); diff --git a/tests/core/kernels/npu/tilelang/CMakeLists.txt b/tests/core/kernels/npu/tilelang/CMakeLists.txt index ede23639b0..f7048c9e0f 100644 --- a/tests/core/kernels/npu/tilelang/CMakeLists.txt +++ b/tests/core/kernels/npu/tilelang/CMakeLists.txt @@ -13,6 +13,45 @@ cc_test( unified_dlog ) +cc_test( + NAME + spec_verify_token_update_wrapper_test + SRCS + spec_verify_token_update_wrapper_test.cpp + DEPS + :tilelang_kernels + torch + GTest::gtest_main + glog::glog + unified_dlog +) + +cc_test( + NAME + spec_verify_attention_tiling_update_wrapper_test + SRCS + spec_verify_attention_tiling_update_wrapper_test.cpp + DEPS + :tilelang_kernels + torch + GTest::gtest_main + glog::glog + unified_dlog +) + +cc_test( + NAME + spec_verify_metadata_update_wrapper_test + SRCS + spec_verify_metadata_update_wrapper_test.cpp + DEPS + :tilelang_kernels + torch + GTest::gtest_main + glog::glog + unified_dlog +) + cc_test( NAME fused_gdn_gating_wrapper_test diff --git a/tests/core/kernels/npu/tilelang/spec_verify_attention_tiling_update_wrapper_test.cpp b/tests/core/kernels/npu/tilelang/spec_verify_attention_tiling_update_wrapper_test.cpp new file mode 100644 index 0000000000..a27df9ccae --- /dev/null +++ b/tests/core/kernels/npu/tilelang/spec_verify_attention_tiling_update_wrapper_test.cpp @@ -0,0 +1,74 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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://www.apache.org/licenses/LICENSE-2.0 + +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 "core/kernels/npu/tilelang/tilelang_ops_api.h" + +namespace xllm::kernel::npu::tilelang { +namespace { + +class TileLangSpecVerifyAttentionTilingUpdateTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); } + static void TearDownTestSuite() { torch_npu::finalize_npu(); } +}; + +TEST_F(TileLangSpecVerifyAttentionTilingUpdateTest, UpdatesDynamicKvLengths) { + const torch::Device device("npu:0"); + const torch::TensorOptions i32 = + torch::TensorOptions().dtype(torch::kInt32).device(device); + const torch::Tensor kv_seq_lens = + torch::tensor({125, 126, 127, 128, 129, 190}, i32); + torch::Tensor tiling_data = torch::full({262144}, -1, i32); + + spec_verify_attention_tiling_update( + kv_seq_lens, tiling_data, /*block_size=*/128); + + const torch::Tensor result = tiling_data.cpu(); + EXPECT_EQ(result[22].item(), 190); + EXPECT_EQ(result[23].item(), 256); + for (int64_t row = 0; row < 6; ++row) { + const int64_t offset = 45 + row * 17; + EXPECT_EQ(result[offset].item(), 125 + row + (row == 5 ? 60 : 0)); + } + EXPECT_EQ(result[21].item(), -1); + EXPECT_EQ(result[24].item(), -1); +} + +TEST_F(TileLangSpecVerifyAttentionTilingUpdateTest, SupportsMtp3AndMtp4Widths) { + const torch::Device device("npu:0"); + const auto i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + for (const int64_t spec_width : {4, 5}) { + const auto kv_seq_lens = torch::arange(125, 125 + spec_width, i32); + auto tiling_data = torch::full({262144}, -1, i32); + + spec_verify_attention_tiling_update( + kv_seq_lens, tiling_data, /*block_size=*/128); + + const auto result = tiling_data.cpu(); + const int32_t max_kv = 124 + static_cast(spec_width); + EXPECT_EQ(result[22].item(), max_kv); + EXPECT_EQ(result[23].item(), ((max_kv + 127) / 128) * 128); + for (int64_t row = 0; row < spec_width; ++row) { + EXPECT_EQ(result[45 + row * 17].item(), 125 + row); + } + } +} + +} // namespace +} // namespace xllm::kernel::npu::tilelang diff --git a/tests/core/kernels/npu/tilelang/spec_verify_metadata_update_wrapper_test.cpp b/tests/core/kernels/npu/tilelang/spec_verify_metadata_update_wrapper_test.cpp new file mode 100644 index 0000000000..4d5700de9f --- /dev/null +++ b/tests/core/kernels/npu/tilelang/spec_verify_metadata_update_wrapper_test.cpp @@ -0,0 +1,178 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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://www.apache.org/licenses/LICENSE-2.0 + +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/tilelang/tilelang_ops_api.h" + +namespace xllm::kernel::npu::tilelang { +namespace { + +class TileLangSpecVerifyMetadataUpdateTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); } + static void TearDownTestSuite() { torch_npu::finalize_npu(); } +}; + +TEST_F(TileLangSpecVerifyMetadataUpdateTest, ReportsCompiledLayouts) { + EXPECT_TRUE(has_spec_verify_metadata_update_specialization(4, 64)); + EXPECT_TRUE(has_spec_verify_metadata_update_specialization(5, 64)); + EXPECT_TRUE(has_spec_verify_metadata_update_specialization(6, 64)); + EXPECT_FALSE(has_spec_verify_metadata_update_specialization(3, 64)); + EXPECT_FALSE(has_spec_verify_metadata_update_specialization(6, 128)); +} + +TEST_F(TileLangSpecVerifyMetadataUpdateTest, WritesCompletePersistentState) { + const auto device = torch::Device("npu:0"); + const auto i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + auto positions = torch::arange(125, 131, i32); + auto kv_seq_lens = torch::tensor({130}, i32); + auto slots = torch::arange(1000, 1006, i32); + auto block_table = torch::arange(2000, 2064, i32); + auto linear_state = torch::tensor({17}, i32); + auto num_accepted = torch::tensor({4}, i32); + + auto persistent_positions = torch::full({3, 8}, -1, i32); + auto persistent_q = torch::full({1}, -1, i32); + auto persistent_kv = torch::full({1}, -1, i32); + auto persistent_slots = torch::full({8}, -1, i32); + auto persistent_block = torch::full({64}, -1, i32); + auto persistent_linear = torch::full({1}, -1, i32); + auto persistent_accepted = torch::full({1}, -1, i32); + auto persistent_q_cu = torch::full({2}, -1, i32); + auto persistent_expanded_kv = torch::full({8}, -1, i32); + auto persistent_expanded_block = torch::full({6, 64}, -1, i32); + + std::vector position_rows; + std::vector expanded_block_rows; + for (int64_t i = 0; i < 3; ++i) { + position_rows.emplace_back(persistent_positions.select(0, i)); + } + for (int64_t i = 0; i < 6; ++i) { + expanded_block_rows.emplace_back(persistent_expanded_block.select(0, i)); + } + + spec_verify_metadata_update(positions, + kv_seq_lens, + slots, + block_table, + linear_state, + num_accepted, + position_rows, + persistent_q, + persistent_kv, + persistent_slots, + persistent_block, + persistent_linear, + persistent_accepted, + persistent_q_cu, + persistent_expanded_kv, + expanded_block_rows); + for (int64_t i = 0; i < 3; ++i) { + EXPECT_TRUE(torch::equal( + persistent_positions.cpu()[i], + torch::tensor({125, 126, 127, 128, 129, 130, 0, 0}, torch::kInt32))); + } + EXPECT_EQ(persistent_q.cpu()[0].item(), 6); + EXPECT_EQ(persistent_kv.cpu()[0].item(), 130); + EXPECT_TRUE( + torch::equal(persistent_slots.cpu(), + torch::tensor({1000, 1001, 1002, 1003, 1004, 1005, 0, 0}, + torch::kInt32))); + EXPECT_TRUE(torch::equal(persistent_block.cpu(), block_table.cpu())); + EXPECT_EQ(persistent_linear.cpu()[0].item(), 17); + EXPECT_EQ(persistent_accepted.cpu()[0].item(), 4); + EXPECT_TRUE(torch::equal(persistent_q_cu.cpu(), + torch::tensor({0, 6}, torch::kInt32))); + EXPECT_TRUE(torch::equal( + persistent_expanded_kv.cpu(), + torch::tensor({125, 126, 127, 128, 129, 130, 0, 0}, torch::kInt32))); + for (int64_t i = 0; i < 6; ++i) { + EXPECT_TRUE( + torch::equal(persistent_expanded_block.cpu()[i], block_table.cpu())); + } +} + +TEST_F(TileLangSpecVerifyMetadataUpdateTest, SupportsMtp3AndMtp4Widths) { + const auto device = torch::Device("npu:0"); + const auto i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + for (const int64_t spec_width : {4, 5}) { + auto positions = torch::arange(125, 125 + spec_width, i32); + auto kv_seq_lens = torch::tensor({130}, i32); + auto slots = torch::arange(1000, 1000 + spec_width, i32); + auto block_table = torch::arange(2000, 2064, i32); + auto linear_state = torch::tensor({17}, i32); + auto num_accepted = torch::tensor({2}, i32); + auto persistent_positions = torch::full({3, 8}, -1, i32); + auto persistent_q = torch::full({1}, -1, i32); + auto persistent_kv = torch::full({1}, -1, i32); + auto persistent_slots = torch::full({8}, -1, i32); + auto persistent_block = torch::full({64}, -1, i32); + auto persistent_linear = torch::full({1}, -1, i32); + auto persistent_accepted = torch::full({1}, -1, i32); + auto persistent_q_cu = torch::full({2}, -1, i32); + auto persistent_expanded_kv = torch::full({8}, -1, i32); + auto persistent_expanded_block = torch::full({6, 64}, -1, i32); + std::vector position_rows; + std::vector expanded_block_rows; + for (int64_t row = 0; row < 3; ++row) { + position_rows.emplace_back(persistent_positions.select(0, row)); + } + for (int64_t row = 0; row < 6; ++row) { + expanded_block_rows.emplace_back( + persistent_expanded_block.select(0, row)); + } + + spec_verify_metadata_update(positions, + kv_seq_lens, + slots, + block_table, + linear_state, + num_accepted, + position_rows, + persistent_q, + persistent_kv, + persistent_slots, + persistent_block, + persistent_linear, + persistent_accepted, + persistent_q_cu, + persistent_expanded_kv, + expanded_block_rows); + + auto expected_positions = torch::zeros({8}, torch::kInt32); + expected_positions.narrow(0, 0, spec_width).copy_(positions.cpu()); + auto expected_slots = torch::zeros({8}, torch::kInt32); + expected_slots.narrow(0, 0, spec_width).copy_(slots.cpu()); + for (int64_t row = 0; row < 3; ++row) { + EXPECT_TRUE( + torch::equal(persistent_positions.cpu()[row], expected_positions)); + } + EXPECT_TRUE(torch::equal(persistent_slots.cpu(), expected_slots)); + EXPECT_EQ(persistent_q.cpu()[0].item(), spec_width); + EXPECT_EQ(persistent_q_cu.cpu()[1].item(), spec_width); + for (int64_t row = 0; row < 6; ++row) { + EXPECT_TRUE(torch::equal(persistent_expanded_block.cpu()[row], + block_table.cpu())); + } + } +} + +} // namespace +} // namespace xllm::kernel::npu::tilelang diff --git a/tests/core/kernels/npu/tilelang/spec_verify_token_update_wrapper_test.cpp b/tests/core/kernels/npu/tilelang/spec_verify_token_update_wrapper_test.cpp new file mode 100644 index 0000000000..60a89f2fb9 --- /dev/null +++ b/tests/core/kernels/npu/tilelang/spec_verify_token_update_wrapper_test.cpp @@ -0,0 +1,102 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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://www.apache.org/licenses/LICENSE-2.0 + +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/tilelang/tilelang_ops_api.h" + +namespace xllm::kernel::npu::tilelang { +namespace { + +class TileLangSpecVerifyTokenUpdateTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); } + static void TearDownTestSuite() { torch_npu::finalize_npu(); } +}; + +TEST_F(TileLangSpecVerifyTokenUpdateTest, ReportsCompiledWidths) { + EXPECT_TRUE(has_spec_verify_token_update_specialization(4)); + EXPECT_TRUE(has_spec_verify_token_update_specialization(5)); + EXPECT_TRUE(has_spec_verify_token_update_specialization(6)); + EXPECT_FALSE(has_spec_verify_token_update_specialization(3)); + EXPECT_FALSE(has_spec_verify_token_update_specialization(8)); +} + +TEST_F(TileLangSpecVerifyTokenUpdateTest, PacksBaseAndDraftTokens) { + const torch::Device device("npu:0"); + const torch::TensorOptions i32 = + torch::TensorOptions().dtype(torch::kInt32).device(device); + const torch::TensorOptions i64 = + torch::TensorOptions().dtype(torch::kInt64).device(device); + const torch::Tensor base_token = torch::tensor({42}, i32); + std::vector draft_tokens; + draft_tokens.reserve(5); + for (int64_t token = 1; token <= 5; ++token) { + draft_tokens.emplace_back(torch::tensor({token}, i64)); + } + torch::Tensor persistent_tokens = torch::full({8}, -1, i32); + + spec_verify_token_update(base_token, draft_tokens, persistent_tokens); + + EXPECT_TRUE( + torch::equal(persistent_tokens.cpu(), + torch::tensor({42, 1, 2, 3, 4, 5, 0, 0}, torch::kInt32))); +} + +TEST_F(TileLangSpecVerifyTokenUpdateTest, PacksMtp3AndZerosTail) { + const torch::Device device("npu:0"); + const auto i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + const auto i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto base_token = torch::tensor({42}, i32); + std::vector draft_tokens = { + torch::tensor({1}, i64), + torch::tensor({2}, i64), + torch::tensor({3}, i64), + }; + auto persistent_tokens = torch::full({8}, -1, i32); + + spec_verify_token_update(base_token, draft_tokens, persistent_tokens); + + EXPECT_TRUE( + torch::equal(persistent_tokens.cpu(), + torch::tensor({42, 1, 2, 3, 0, 0, 0, 0}, torch::kInt32))); +} + +TEST_F(TileLangSpecVerifyTokenUpdateTest, PacksMtp4AndZerosTail) { + const torch::Device device("npu:0"); + const auto i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + const auto i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto base_token = torch::tensor({42}, i32); + std::vector draft_tokens = { + torch::tensor({1}, i64), + torch::tensor({2}, i64), + torch::tensor({3}, i64), + torch::tensor({4}, i64), + }; + auto persistent_tokens = torch::full({8}, -1, i32); + + spec_verify_token_update(base_token, draft_tokens, persistent_tokens); + + EXPECT_TRUE( + torch::equal(persistent_tokens.cpu(), + torch::tensor({42, 1, 2, 3, 4, 0, 0, 0}, torch::kInt32))); +} + +} // namespace +} // namespace xllm::kernel::npu::tilelang diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index 7db963de46..0d7f7986d1 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -929,6 +929,74 @@ TEST(AclGraphPersistentParamTest, SpecVerifyMetadataUsesTokenCapacity) { speculative_config.enable_atb_spec_kernel(original_enable_atb_spec_kernel); } +TEST(AclGraphPersistentParamTest, SpecVerifyGraphUpdateSupportsRuntimeWidth) { + constexpr int32_t kSpecWidth = 4; + ModelArgs args; + args.model_type("deepseek_v4"); + args.dtype("float32"); + args.hidden_size(8); + args.max_position_embeddings(32); + + runtime::Options options; + options.block_size(4); + options.max_seqs_per_batch(4); + options.max_tokens_per_batch(16); + options.num_decoding_tokens(kSpecWidth); + options.enable_speculative_decode(true); + options.is_draft_engine(false); + + const torch::Device device("npu:0"); + const torch::TensorOptions int_options = + torch::dtype(torch::kInt).device(device); + ::xllm::npu::GraphPersistentParam persistent_param( + args, + device, + options, + /*need_update_attn_mask=*/false, + /*is_hybrid_linear_attention=*/true); + + ModelInputParams params; + params.is_spec_verify = true; + params.meta.batch_forward_type = BatchForwardType::CHUNKED_PREFILL; + params.meta.num_sequences = 1; + params.meta.q_max_seq_len = kSpecWidth; + params.graph.spec_verify_source_addresses_stable = true; + params.graph.input_tokens_override = + torch::tensor({11, 12, 13, 14}, int_options); + params.graph.expanded_kv_seq_lens = + torch::tensor({17, 18, 19, 20}, int_options); + params.graph.expanded_block_tables = + torch::arange(8, int_options).view({kSpecWidth, 2}); + params.attention.device.q_seq_lens = torch::tensor({kSpecWidth}, int_options); + params.attention.device.kv_seq_lens = torch::tensor({20}, int_options); + params.attention.device.new_cache_slots = + torch::tensor({21, 22, 23, 24}, int_options); + params.attention.device.block_tables = torch::tensor({{31, 32}}, int_options); + params.attention.device.q_cu_seq_lens = + torch::tensor({0, kSpecWidth}, int_options); + params.embedding.linear_state_indices = torch::tensor({3}, int_options); + params.num_accepted_tokens = torch::tensor({2}, int_options); + const torch::Tensor positions = + torch::tensor({101, 102, 103, 104}, int_options); + + const std::vector sources = + persistent_param.capture_spec_verify_input_update( + params.graph.input_tokens_override, + positions, + params, + /*padded_num_tokens=*/kSpecWidth); + + EXPECT_EQ(sources.size(), 11); + EXPECT_TRUE(torch::equal(persistent_param.persistent_tokens(kSpecWidth).cpu(), + params.graph.input_tokens_override.cpu())); + EXPECT_TRUE( + torch::equal(persistent_param.persistent_positions(kSpecWidth).cpu(), + positions.cpu())); + EXPECT_TRUE(torch::equal( + persistent_param.persistent_new_cache_slots(kSpecWidth).cpu(), + params.attention.device.new_cache_slots.cpu())); +} + TEST(AclGraphExecutorHybridTest, KvCacheSupportsLinearOnlyLayers) { auto conv_cache = torch::zeros({4, 32, 3}, torch::dtype(torch::kFloat32)); auto ssm_cache = torch::zeros({4, 8, 64, 64}, torch::dtype(torch::kFloat32)); diff --git a/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_attention_tiling_update.py b/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_attention_tiling_update.py new file mode 100644 index 0000000000..c938d9b44f --- /dev/null +++ b/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_attention_tiling_update.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 + +# Copyright 2026 The xLLM Authors. All Rights Reserved. +# +# 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://www.apache.org/licenses/LICENSE-2.0 +# +# 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. +import tilelang +import tilelang.language as T + +from .utils import DEFAULT_ASCEND_PASS_CONFIGS +from ....common.spec import DispatchField, TilelangKernel, register_kernel + +SUPPORTED_SPEC_WIDTHS = (4, 5, 6) + + +def build_spec_verify_attention_tiling_update_kernel( + spec_width: int, block_size: int +): + if spec_width not in SUPPORTED_SPEC_WIDTHS: + raise ValueError(f"unsupported MTP tiling width: {spec_width}") + + @T.prim_func + def spec_verify_attention_tiling_update( + src_kv_seq_lens: T.Tensor((spec_width,), "int32"), + tiling_data: T.Tensor((262144,), "int32"), + ): + with T.Kernel(1, is_npu=True): + # Keep the reduction explicit. The Ascend scalar lowering drops + # assignments nested under a T.serial conditional here. + max_kv_4 = T.max( + T.max(src_kv_seq_lens[0], src_kv_seq_lens[1]), + T.max(src_kv_seq_lens[2], src_kv_seq_lens[3]), + ) + # Write each specialized reduction inside its branch. TileLang's + # parser scopes branch-local scalar assignments, so assigning + # max_kv here and consuming it after the branch is not valid. + if spec_width == 4: + max_kv = max_kv_4 + tiling_data[22] = max_kv + if max_kv % block_size == 0: + tiling_data[23] = max_kv + else: + tiling_data[23] = T.ceildiv(max_kv, block_size) * block_size + elif spec_width == 5: + max_kv = T.max(max_kv_4, src_kv_seq_lens[4]) + tiling_data[22] = max_kv + if max_kv % block_size == 0: + tiling_data[23] = max_kv + else: + tiling_data[23] = T.ceildiv(max_kv, block_size) * block_size + else: + max_kv = T.max( + T.max(max_kv_4, src_kv_seq_lens[4]), + src_kv_seq_lens[5], + ) + tiling_data[22] = max_kv + if max_kv % block_size == 0: + tiling_data[23] = max_kv + else: + tiling_data[23] = T.ceildiv(max_kv, block_size) * block_size + for i in T.serial(spec_width): + kv_len = src_kv_seq_lens[i] + tiling_data[44 + i * 17 + 1] = kv_len + + return spec_verify_attention_tiling_update + + +@register_kernel +class SpecVerifyAttentionTilingUpdateKernel(TilelangKernel): + DISPATCH_SCHEMA = [ + DispatchField("spec_width", "int32"), + DispatchField("block_size", "int32"), + ] + SPECIALIZATIONS = [ + { + "variant_key": f"w{spec_width}_bs128", + "spec_width": spec_width, + "block_size": 128, + } + for spec_width in SUPPORTED_SPEC_WIDTHS + ] + + @staticmethod + def generate_source(spec_width: int, block_size: int) -> str: + tilelang.disable_cache() + kernel = build_spec_verify_attention_tiling_update_kernel( + spec_width, block_size + ) + with tilelang.tvm.transform.PassContext( + opt_level=3, config=DEFAULT_ASCEND_PASS_CONFIGS + ): + lowered = tilelang.engine.lower(kernel) + return lowered.kernel_source diff --git a/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_metadata_update.py b/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_metadata_update.py new file mode 100644 index 0000000000..2802bffe78 --- /dev/null +++ b/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_metadata_update.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 + +# Copyright 2026 The xLLM Authors. All Rights Reserved. +# +# 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://www.apache.org/licenses/LICENSE-2.0 +# +# 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. +import tilelang +import tilelang.language as T + +from .utils import DEFAULT_ASCEND_PASS_CONFIGS +from ....common.spec import DispatchField, TilelangKernel, register_kernel + +SUPPORTED_SPEC_WIDTHS = (4, 5, 6) +PADDED_WIDTH = 8 + + +def build_spec_verify_metadata_update_kernel( + spec_width: int, + block_table_width: int, +): + if spec_width not in SUPPORTED_SPEC_WIDTHS: + raise ValueError(f"unsupported speculative verify width: {spec_width}") + + @T.prim_func + def spec_verify_metadata_update( + positions: T.Tensor((spec_width,), "int32"), + kv_seq_lens: T.Tensor((1,), "int32"), + new_cache_slots: T.Tensor((spec_width,), "int32"), + block_table: T.Tensor((block_table_width,), "int32"), + linear_state_index: T.Tensor((1,), "int32"), + num_accepted: T.Tensor((1,), "int32"), + persistent_position_0: T.Tensor((PADDED_WIDTH,), "int32"), + persistent_position_1: T.Tensor((PADDED_WIDTH,), "int32"), + persistent_position_2: T.Tensor((PADDED_WIDTH,), "int32"), + persistent_q_seq_lens: T.Tensor((1,), "int32"), + persistent_kv_seq_lens: T.Tensor((1,), "int32"), + persistent_new_cache_slots: T.Tensor((PADDED_WIDTH,), "int32"), + persistent_block_table: T.Tensor((block_table_width,), "int32"), + persistent_linear_state_index: T.Tensor((1,), "int32"), + persistent_num_accepted: T.Tensor((1,), "int32"), + persistent_q_cu_seq_lens: T.Tensor((2,), "int32"), + persistent_expanded_kv_seq_lens: T.Tensor( + (PADDED_WIDTH,), "int32" + ), + persistent_expanded_block_table_0: T.Tensor( + (block_table_width,), "int32" + ), + persistent_expanded_block_table_1: T.Tensor( + (block_table_width,), "int32" + ), + persistent_expanded_block_table_2: T.Tensor( + (block_table_width,), "int32" + ), + persistent_expanded_block_table_3: T.Tensor( + (block_table_width,), "int32" + ), + persistent_expanded_block_table_4: T.Tensor( + (block_table_width,), "int32" + ), + persistent_expanded_block_table_5: T.Tensor( + (block_table_width,), "int32" + ), + ): + with T.Kernel(1, is_npu=True): + with T.Scope("V"): + positions_ub = T.alloc_ub((PADDED_WIDTH,), "int32") + slots_ub = T.alloc_ub((PADDED_WIDTH,), "int32") + block_table_ub = T.alloc_ub( + (block_table_width,), "int32" + ) + q_seq_lens_ub = T.alloc_ub((1,), "int32") + kv_seq_lens_ub = T.alloc_ub((1,), "int32") + linear_state_ub = T.alloc_ub((1,), "int32") + num_accepted_ub = T.alloc_ub((1,), "int32") + q_cu_seq_lens_ub = T.alloc_ub((2,), "int32") + expanded_kv_seq_lens_ub = T.alloc_ub( + (PADDED_WIDTH,), "int32" + ) + + T.copy(positions[0], positions_ub[0:spec_width]) + T.copy(new_cache_slots[0], slots_ub[0:spec_width]) + T.copy(block_table[0], block_table_ub) + T.copy(kv_seq_lens[0], kv_seq_lens_ub) + T.copy(linear_state_index[0], linear_state_ub) + T.copy(num_accepted[0], num_accepted_ub) + + for i in range(spec_width, PADDED_WIDTH): + positions_ub[i] = 0 + slots_ub[i] = 0 + + q_seq_lens_ub[0] = spec_width + q_cu_seq_lens_ub[0] = 0 + q_cu_seq_lens_ub[1] = spec_width + for i in T.serial(spec_width): + expanded_kv_seq_lens_ub[i] = ( + kv_seq_lens_ub[0] - spec_width + i + 1 + ) + for i in range(spec_width, PADDED_WIDTH): + expanded_kv_seq_lens_ub[i] = 0 + + T.copy(positions_ub, persistent_position_0[0]) + T.copy(positions_ub, persistent_position_1[0]) + T.copy(positions_ub, persistent_position_2[0]) + T.copy(q_seq_lens_ub, persistent_q_seq_lens[0]) + T.copy(kv_seq_lens_ub, persistent_kv_seq_lens[0]) + T.copy(slots_ub, persistent_new_cache_slots[0]) + T.copy(block_table_ub, persistent_block_table[0]) + T.copy(linear_state_ub, persistent_linear_state_index[0]) + T.copy(num_accepted_ub, persistent_num_accepted[0]) + T.copy(q_cu_seq_lens_ub, persistent_q_cu_seq_lens[0]) + T.copy( + expanded_kv_seq_lens_ub, + persistent_expanded_kv_seq_lens[0], + ) + T.copy( + block_table_ub, + persistent_expanded_block_table_0[0], + ) + T.copy( + block_table_ub, + persistent_expanded_block_table_1[0], + ) + T.copy( + block_table_ub, + persistent_expanded_block_table_2[0], + ) + T.copy( + block_table_ub, + persistent_expanded_block_table_3[0], + ) + T.copy( + block_table_ub, + persistent_expanded_block_table_4[0], + ) + T.copy( + block_table_ub, + persistent_expanded_block_table_5[0], + ) + + return spec_verify_metadata_update + + +@register_kernel +class SpecVerifyMetadataUpdateKernel(TilelangKernel): + DISPATCH_SCHEMA = [ + DispatchField("spec_width", "int32"), + DispatchField("block_table_width", "int32"), + ] + SPECIALIZATIONS = [ + { + "variant_key": f"w{spec_width}_bt64", + "spec_width": spec_width, + "block_table_width": 64, + } + for spec_width in SUPPORTED_SPEC_WIDTHS + ] + + @staticmethod + def generate_source(spec_width: int, block_table_width: int) -> str: + tilelang.disable_cache() + kernel = build_spec_verify_metadata_update_kernel( + spec_width=spec_width, + block_table_width=block_table_width, + ) + with tilelang.tvm.transform.PassContext( + opt_level=3, config=DEFAULT_ASCEND_PASS_CONFIGS + ): + lowered = tilelang.engine.lower(kernel) + return lowered.kernel_source diff --git a/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_token_update.py b/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_token_update.py new file mode 100644 index 0000000000..0e6689364c --- /dev/null +++ b/xllm/compiler/tilelang/targets/ascend/kernels/spec_verify_token_update.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 + +# Copyright 2026 The xLLM Authors. All Rights Reserved. +# +# 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://www.apache.org/licenses/LICENSE-2.0 +# +# 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. + +import tilelang +import tilelang.language as T + +from .utils import DEFAULT_ASCEND_PASS_CONFIGS +from ....common.spec import DispatchField, TilelangKernel, register_kernel + +SUPPORTED_SPEC_WIDTHS = (4, 5, 6) +PADDED_WIDTH = 8 + + +def build_spec_verify_token_update_kernel(spec_width: int): + if spec_width not in SUPPORTED_SPEC_WIDTHS: + raise ValueError(f"unsupported speculative verify width: {spec_width}") + + @T.prim_func + def spec_verify_token_update( + base_token: T.Tensor((1,), "int32"), + draft_token_0: T.Tensor((1,), "int64"), + draft_token_1: T.Tensor((1,), "int64"), + draft_token_2: T.Tensor((1,), "int64"), + draft_token_3: T.Tensor((1,), "int64"), + draft_token_4: T.Tensor((1,), "int64"), + persistent_tokens: T.Tensor((PADDED_WIDTH,), "int32"), + ): + with T.Kernel(1, is_npu=True): + with T.Scope("V"): + packed_ub = T.alloc_ub((PADDED_WIDTH,), "int32") + packed_ub[0] = base_token[0] + packed_ub[1] = T.Cast("int32", draft_token_0[0]) + packed_ub[2] = T.Cast("int32", draft_token_1[0]) + packed_ub[3] = T.Cast("int32", draft_token_2[0]) + if spec_width >= 5: + packed_ub[4] = T.Cast("int32", draft_token_3[0]) + else: + packed_ub[4] = 0 + if spec_width >= 6: + packed_ub[5] = T.Cast("int32", draft_token_4[0]) + else: + packed_ub[5] = 0 + packed_ub[6] = 0 + packed_ub[7] = 0 + T.copy(packed_ub, persistent_tokens[0]) + + return spec_verify_token_update + + +@register_kernel +class SpecVerifyTokenUpdateKernel(TilelangKernel): + DISPATCH_SCHEMA = [DispatchField("spec_width", "int32")] + SPECIALIZATIONS = [ + {"variant_key": f"w{spec_width}", "spec_width": spec_width} + for spec_width in SUPPORTED_SPEC_WIDTHS + ] + + @staticmethod + def generate_source(spec_width: int) -> str: + tilelang.disable_cache() + kernel = build_spec_verify_token_update_kernel(spec_width) + with tilelang.tvm.transform.PassContext( + opt_level=3, config=DEFAULT_ASCEND_PASS_CONFIGS + ): + lowered = tilelang.engine.lower(kernel) + return lowered.kernel_source diff --git a/xllm/core/framework/model/causal_lm.h b/xllm/core/framework/model/causal_lm.h index 8ac6fb26b7..97518b12ee 100644 --- a/xllm/core/framework/model/causal_lm.h +++ b/xllm/core/framework/model/causal_lm.h @@ -42,6 +42,7 @@ limitations under the License. #include "model_input_params.h" #include "model_output.h" #include "model_traits.h" +#include "speculative_verify_capabilities.h" namespace xllm { @@ -69,6 +70,11 @@ class CausalLM : public torch::nn::Module { virtual bool is_hybrid_linear_attention() { return false; } + virtual SpeculativeVerifyCapabilities speculative_verify_capabilities() + const { + return {}; + } + virtual std::unique_ptr create_graph_forward_metadata_state() { return nullptr; @@ -210,6 +216,15 @@ class CausalLMImpl : public CausalLM { } } + SpeculativeVerifyCapabilities speculative_verify_capabilities() + const override { + if constexpr (detail::has_speculative_verify_capabilities::value) { + return model_->speculative_verify_capabilities(); + } else { + return CausalLM::speculative_verify_capabilities(); + } + } + std::unique_ptr create_graph_forward_metadata_state() override { if constexpr (detail::has_create_graph_forward_metadata_state< diff --git a/xllm/core/framework/model/model_input_params.h b/xllm/core/framework/model/model_input_params.h index 1ed4263e2a..2705dce588 100644 --- a/xllm/core/framework/model/model_input_params.h +++ b/xllm/core/framework/model/model_input_params.h @@ -410,6 +410,12 @@ struct AttentionDeviceInput { }; struct AttentionInput { + struct PackedIntInput { + const std::vector* values = nullptr; + torch::Tensor* host_view = nullptr; + torch::Tensor* device_view = nullptr; + }; + AttentionHostInput host; AttentionDeviceInput device; torch::Tensor attention_host_buffer; @@ -430,11 +436,14 @@ struct AttentionInput { return out; } - bool rebuild_device_buffer(const torch::Device& target_device) { + bool rebuild_device_buffer( + const torch::Device& target_device, + const std::vector& extra_int_inputs = {}) { struct Entry { const void* source = nullptr; std::vector sizes; torch::ScalarType dtype = torch::kUInt8; + torch::Tensor* host_target = nullptr; torch::Tensor* target = nullptr; uint64_t offset = 0; uint64_t bytes = 0; @@ -456,15 +465,17 @@ struct AttentionInput { std::vector sizes, torch::ScalarType dtype, uint64_t bytes, + torch::Tensor* host_target, torch::Tensor* target) { if (source == nullptr || bytes == 0) { return; } - entries.push_back( - Entry{source, std::move(sizes), dtype, target, 0, bytes, 0}); + entries.push_back(Entry{ + source, std::move(sizes), dtype, host_target, target, 0, bytes, 0}); }; auto add_int_vector = [&add_raw](const std::vector& values, + torch::Tensor* host_target, torch::Tensor* target) { if (values.empty()) { return; @@ -473,6 +484,7 @@ struct AttentionInput { {static_cast(values.size())}, torch::kInt, static_cast(values.size() * sizeof(int32_t)), + host_target, target); }; @@ -495,20 +507,26 @@ struct AttentionInput { entries.push_back(Entry{source.data_ptr(), source.sizes().vec(), source.scalar_type(), + nullptr, target, 0, bytes, 0}); }; - add_int_vector(host.q_seq_lens, &device.q_seq_lens); - add_int_vector(host.kv_seq_lens, &device.kv_seq_lens); - add_int_vector(host.q_cu_seq_lens, &device.q_cu_seq_lens); - add_int_vector(host.new_cache_slots, &device.new_cache_slots); + for (const PackedIntInput& extra : extra_int_inputs) { + CHECK(extra.values != nullptr); + add_int_vector(*extra.values, extra.host_view, extra.device_view); + } + add_int_vector(host.q_seq_lens, nullptr, &device.q_seq_lens); + add_int_vector(host.kv_seq_lens, nullptr, &device.kv_seq_lens); + add_int_vector(host.q_cu_seq_lens, nullptr, &device.q_cu_seq_lens); + add_int_vector(host.new_cache_slots, nullptr, &device.new_cache_slots); add_cpu_tensor(host.block_tables, &device.block_tables); - add_int_vector(host.kv_cache_tokens_nums, &device.kv_cache_tokens_nums); - add_int_vector(host.ring_cur_seqlen, &device.ring_cur_seqlen); - add_int_vector(host.ring_cache_seqlen, &device.ring_cache_seqlen); + add_int_vector( + host.kv_cache_tokens_nums, nullptr, &device.kv_cache_tokens_nums); + add_int_vector(host.ring_cur_seqlen, nullptr, &device.ring_cur_seqlen); + add_int_vector(host.ring_cache_seqlen, nullptr, &device.ring_cache_seqlen); add_cpu_tensor(device.paged_kv_indptr, &device.paged_kv_indptr); add_cpu_tensor(device.paged_kv_indices, &device.paged_kv_indices); @@ -563,6 +581,13 @@ struct AttentionInput { const char* device_base = static_cast(attention_device_buffer.data_ptr()); for (const auto& entry : entries) { + if (entry.host_target != nullptr) { + void* host_ptr = host_base + entry.offset; + *entry.host_target = torch::from_blob( + host_ptr, + entry.sizes, + torch::TensorOptions().dtype(entry.dtype).device(torch::kCPU)); + } if (entry.target == nullptr) { continue; } @@ -895,6 +920,19 @@ struct GraphInput { std::shared_ptr acl_graph_task_update_context; #endif torch::Tensor input_tokens_override; + // Device token sources produced by a speculative proposer. The graph input + // updater may fuse them into stable target-verify storage when a matching + // backend specialization exists. Their addresses are not part of the replay + // signature unless spec_verify_source_addresses_stable is true. + std::vector spec_verify_draft_token_sources; + // All dynamic target-verify source tensors have stable addresses across + // replay generations. This contract is algorithm-independent; individual + // backends may still select model- or shape-specific fused implementations. + bool spec_verify_source_addresses_stable = false; + // The static causal-conv task-ready event has already been queued behind + // the final draft's device event. Replay may consume it without issuing a + // host-side task update on the final-draft-to-target critical path. + bool spec_verify_static_graph_tasks_prepared = false; GraphInput to(const torch::Device& device) const { GraphInput out; @@ -914,6 +952,16 @@ struct GraphInput { #endif out.input_tokens_override = safe_to(input_tokens_override, device, /*non_blocking=*/true); + out.spec_verify_draft_token_sources.reserve( + spec_verify_draft_token_sources.size()); + for (const auto& token : spec_verify_draft_token_sources) { + out.spec_verify_draft_token_sources.push_back( + safe_to(token, device, /*non_blocking=*/true)); + } + out.spec_verify_source_addresses_stable = + spec_verify_source_addresses_stable; + out.spec_verify_static_graph_tasks_prepared = + spec_verify_static_graph_tasks_prepared; return out; } }; diff --git a/xllm/core/framework/model/model_traits.h b/xllm/core/framework/model/model_traits.h index 13d3b2d154..74b7859da0 100644 --- a/xllm/core/framework/model/model_traits.h +++ b/xllm/core/framework/model/model_traits.h @@ -25,6 +25,7 @@ limitations under the License. namespace xllm { struct ModelInputParams; struct ModelGraphMetadataState; +struct SpeculativeVerifyCapabilities; class KVCache; namespace layer { @@ -138,6 +139,15 @@ struct has_is_hybrid_linear_attention< std::void_t()->is_hybrid_linear_attention())>> : std::true_type {}; +template +struct has_speculative_verify_capabilities : std::false_type {}; + +template +struct has_speculative_verify_capabilities< + T, + std::void_t()->speculative_verify_capabilities())>> + : std::true_type {}; + template struct has_create_graph_forward_metadata_state : std::false_type {}; diff --git a/xllm/core/framework/model/speculative_verify_capabilities.h b/xllm/core/framework/model/speculative_verify_capabilities.h new file mode 100644 index 0000000000..09fadfd4d2 --- /dev/null +++ b/xllm/core/framework/model/speculative_verify_capabilities.h @@ -0,0 +1,29 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +namespace xllm { + +// Describes target-model requirements at the speculative verification seam. +// A model that enables in_graph_input_update guarantees that its validation +// input follows the expanded-attention and single linear-state layout consumed +// by the NPU graph updater. Unsupported models keep the conservative defaults. +struct SpeculativeVerifyCapabilities { + bool requires_causal_chunked_prefill = false; + bool supports_in_graph_input_update = false; +}; + +} // namespace xllm diff --git a/xllm/core/kernels/npu/tilelang/CMakeLists.txt b/xllm/core/kernels/npu/tilelang/CMakeLists.txt index 3309b0737d..289815da17 100644 --- a/xllm/core/kernels/npu/tilelang/CMakeLists.txt +++ b/xllm/core/kernels/npu/tilelang/CMakeLists.txt @@ -127,6 +127,21 @@ tilelang_register_runtime_kernel( WRAPPER_SRCS rope_wrapper.cpp ) +tilelang_register_runtime_kernel( + NAME spec_verify_attention_tiling_update + WRAPPER_SRCS spec_verify_attention_tiling_update_wrapper.cpp +) + +tilelang_register_runtime_kernel( + NAME spec_verify_token_update + WRAPPER_SRCS spec_verify_token_update_wrapper.cpp +) + +tilelang_register_runtime_kernel( + NAME spec_verify_metadata_update + WRAPPER_SRCS spec_verify_metadata_update_wrapper.cpp +) + tilelang_register_runtime_kernel( NAME fused_gdn_gating WRAPPER_SRCS fused_gdn_gating_wrapper.cpp diff --git a/xllm/core/kernels/npu/tilelang/spec_verify_attention_tiling_update_wrapper.cpp b/xllm/core/kernels/npu/tilelang/spec_verify_attention_tiling_update_wrapper.cpp new file mode 100644 index 0000000000..3dc73cce8a --- /dev/null +++ b/xllm/core/kernels/npu/tilelang/spec_verify_attention_tiling_update_wrapper.cpp @@ -0,0 +1,65 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include "acl/acl.h" +#include "core/kernels/npu/tilelang/dispatch_registry.h" +#include "core/kernels/npu/tilelang/tilelang_ops_api.h" + +#ifndef XLLM_TL_SPEC_VERIFY_ATTENTION_TILING_UPDATE_REGISTRY_INC +#error "XLLM_TL_SPEC_VERIFY_ATTENTION_TILING_UPDATE_REGISTRY_INC is not defined" +#endif + +namespace xllm::kernel::npu::tilelang { +namespace { +#include XLLM_TL_SPEC_VERIFY_ATTENTION_TILING_UPDATE_REGISTRY_INC +} + +void spec_verify_attention_tiling_update(const torch::Tensor& src_kv_seq_lens, + torch::Tensor& tiling_data, + int64_t block_size) { + CHECK_EQ(src_kv_seq_lens.device().type(), c10::DeviceType::PrivateUse1); + CHECK_EQ(tiling_data.device().type(), c10::DeviceType::PrivateUse1); + CHECK_EQ(src_kv_seq_lens.scalar_type(), torch::kInt32); + CHECK_EQ(tiling_data.scalar_type(), torch::kInt32); + const int64_t spec_width = src_kv_seq_lens.numel(); + CHECK_GE(spec_width, 4); + CHECK_LE(spec_width, 6); + CHECK_GE(tiling_data.numel(), 44 + spec_width * 17); + CHECK(src_kv_seq_lens.is_contiguous()); + CHECK(tiling_data.is_contiguous()); + CHECK_EQ(block_size, 128); + const auto specialization = + make_spec_verify_attention_tiling_update_specialization( + SpecVerifyAttentionTilingUpdateSpecWidth{ + static_cast(spec_width)}, + SpecVerifyAttentionTilingUpdateBlockSize{128}); + const auto* entry = + find_spec_verify_attention_tiling_update_kernel_entry(specialization); + CHECK(entry != nullptr) + << available_spec_verify_attention_tiling_update_variant_keys(); + aclrtStream stream = + c10_npu::getCurrentNPUStream(src_kv_seq_lens.device().index()).stream(); + entry->fn( + reinterpret_cast(const_cast(src_kv_seq_lens.data_ptr())), + reinterpret_cast(tiling_data.data_ptr()), + stream); +} + +} // namespace xllm::kernel::npu::tilelang diff --git a/xllm/core/kernels/npu/tilelang/spec_verify_metadata_update_wrapper.cpp b/xllm/core/kernels/npu/tilelang/spec_verify_metadata_update_wrapper.cpp new file mode 100644 index 0000000000..25a3b0fa56 --- /dev/null +++ b/xllm/core/kernels/npu/tilelang/spec_verify_metadata_update_wrapper.cpp @@ -0,0 +1,169 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +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 "dispatch_registry.h" +#include "tilelang_ops_api.h" + +#ifndef XLLM_TL_SPEC_VERIFY_METADATA_UPDATE_REGISTRY_INC +#error "XLLM_TL_SPEC_VERIFY_METADATA_UPDATE_REGISTRY_INC is not defined" +#endif + +namespace xllm::kernel::npu::tilelang { +namespace { + +#include XLLM_TL_SPEC_VERIFY_METADATA_UPDATE_REGISTRY_INC + +void check_npu_tensor(const torch::Tensor& tensor, + torch::ScalarType dtype, + int64_t min_numel, + const char* name) { + CHECK(tensor.defined() && + tensor.device().type() == c10::DeviceType::PrivateUse1) + << name << " must be an NPU tensor"; + CHECK_EQ(tensor.scalar_type(), dtype) << name << " dtype mismatch"; + CHECK_GE(tensor.numel(), min_numel) << name << " is too small"; + CHECK(tensor.is_contiguous()) << name << " must be contiguous"; +} + +uint8_t* ptr(const torch::Tensor& tensor) { + return reinterpret_cast(const_cast(tensor.data_ptr())); +} + +} // namespace + +bool has_spec_verify_metadata_update_specialization(int64_t spec_width, + int64_t block_table_width) { + const auto specialization = make_spec_verify_metadata_update_specialization( + SpecVerifyMetadataUpdateSpecWidth{static_cast(spec_width)}, + SpecVerifyMetadataUpdateBlockTableWidth{ + static_cast(block_table_width)}); + return find_spec_verify_metadata_update_kernel_entry(specialization) != + nullptr; +} + +void spec_verify_metadata_update( + const torch::Tensor& positions, + const torch::Tensor& kv_seq_lens, + const torch::Tensor& new_cache_slots, + const torch::Tensor& block_table, + const torch::Tensor& linear_state_index, + const torch::Tensor& num_accepted, + const std::vector& persistent_position_rows, + torch::Tensor& persistent_q_seq_lens, + torch::Tensor& persistent_kv_seq_lens, + torch::Tensor& persistent_new_cache_slots, + torch::Tensor& persistent_block_table, + torch::Tensor& persistent_linear_state_index, + torch::Tensor& persistent_num_accepted, + torch::Tensor& persistent_q_cu_seq_lens, + torch::Tensor& persistent_expanded_kv_seq_lens, + const std::vector& persistent_expanded_block_rows) { + const int64_t spec_width = positions.numel(); + constexpr int64_t kPaddedWidth = 8; + CHECK(has_spec_verify_metadata_update_specialization(spec_width, + block_table.numel())) + << "speculative verify metadata update has no compiled variant for width " + << spec_width << " and block table width " << block_table.numel(); + check_npu_tensor(positions, torch::kInt32, spec_width, "positions"); + check_npu_tensor(kv_seq_lens, torch::kInt32, 1, "kv_seq_lens"); + check_npu_tensor( + new_cache_slots, torch::kInt32, spec_width, "new_cache_slots"); + check_npu_tensor(block_table, torch::kInt32, 1, "block_table"); + check_npu_tensor(linear_state_index, torch::kInt32, 1, "linear_state_index"); + check_npu_tensor(num_accepted, torch::kInt32, 1, "num_accepted"); + + CHECK_EQ(persistent_position_rows.size(), 3); + CHECK_EQ(persistent_expanded_block_rows.size(), 6) + << "fixed TileLang ABI requires six expanded block-table pointers"; + for (const auto& row : persistent_position_rows) { + check_npu_tensor( + row, torch::kInt32, kPaddedWidth, "persistent_position_row"); + } + check_npu_tensor( + persistent_q_seq_lens, torch::kInt32, 1, "persistent_q_seq_lens"); + check_npu_tensor( + persistent_kv_seq_lens, torch::kInt32, 1, "persistent_kv_seq_lens"); + check_npu_tensor(persistent_new_cache_slots, + torch::kInt32, + kPaddedWidth, + "persistent_new_cache_slots"); + check_npu_tensor(persistent_block_table, + torch::kInt32, + block_table.numel(), + "persistent_block_table"); + check_npu_tensor(persistent_linear_state_index, + torch::kInt32, + 1, + "persistent_linear_state_index"); + check_npu_tensor( + persistent_num_accepted, torch::kInt32, 1, "persistent_num_accepted"); + check_npu_tensor( + persistent_q_cu_seq_lens, torch::kInt32, 2, "persistent_q_cu_seq_lens"); + check_npu_tensor(persistent_expanded_kv_seq_lens, + torch::kInt32, + kPaddedWidth, + "persistent_expanded_kv_seq_lens"); + for (const auto& row : persistent_expanded_block_rows) { + check_npu_tensor(row, + torch::kInt32, + block_table.numel(), + "persistent_expanded_block_row"); + } + + const auto specialization = make_spec_verify_metadata_update_specialization( + SpecVerifyMetadataUpdateSpecWidth{static_cast(spec_width)}, + SpecVerifyMetadataUpdateBlockTableWidth{ + static_cast(block_table.numel())}); + const auto* entry = + find_spec_verify_metadata_update_kernel_entry(specialization); + CHECK(entry != nullptr) + << "speculative verify metadata update has no compiled variant: " + << available_spec_verify_metadata_update_variant_keys(); + + aclrtStream stream = + c10_npu::getCurrentNPUStream(positions.device().index()).stream(); + entry->fn(ptr(positions), + ptr(kv_seq_lens), + ptr(new_cache_slots), + ptr(block_table), + ptr(linear_state_index), + ptr(num_accepted), + ptr(persistent_position_rows[0]), + ptr(persistent_position_rows[1]), + ptr(persistent_position_rows[2]), + ptr(persistent_q_seq_lens), + ptr(persistent_kv_seq_lens), + ptr(persistent_new_cache_slots), + ptr(persistent_block_table), + ptr(persistent_linear_state_index), + ptr(persistent_num_accepted), + ptr(persistent_q_cu_seq_lens), + ptr(persistent_expanded_kv_seq_lens), + ptr(persistent_expanded_block_rows[0]), + ptr(persistent_expanded_block_rows[1]), + ptr(persistent_expanded_block_rows[2]), + ptr(persistent_expanded_block_rows[3]), + ptr(persistent_expanded_block_rows[4]), + ptr(persistent_expanded_block_rows[5]), + stream); +} + +} // namespace xllm::kernel::npu::tilelang diff --git a/xllm/core/kernels/npu/tilelang/spec_verify_token_update_wrapper.cpp b/xllm/core/kernels/npu/tilelang/spec_verify_token_update_wrapper.cpp new file mode 100644 index 0000000000..57156b2e51 --- /dev/null +++ b/xllm/core/kernels/npu/tilelang/spec_verify_token_update_wrapper.cpp @@ -0,0 +1,98 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +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 + +#include "dispatch_registry.h" +#include "tilelang_ops_api.h" + +#ifndef XLLM_TL_SPEC_VERIFY_TOKEN_UPDATE_REGISTRY_INC +#error "XLLM_TL_SPEC_VERIFY_TOKEN_UPDATE_REGISTRY_INC is not defined" +#endif + +namespace xllm::kernel::npu::tilelang { +namespace { + +#include XLLM_TL_SPEC_VERIFY_TOKEN_UPDATE_REGISTRY_INC + +void check_token(const torch::Tensor& token, torch::ScalarType dtype) { + CHECK(token.defined() && + token.device().type() == c10::DeviceType::PrivateUse1) + << "speculative verify token update requires NPU tensors"; + CHECK_EQ(token.scalar_type(), dtype); + CHECK_EQ(token.numel(), 1); + CHECK(token.is_contiguous()); +} + +} // namespace + +bool has_spec_verify_token_update_specialization(int64_t spec_width) { + const auto specialization = make_spec_verify_token_update_specialization( + SpecVerifyTokenUpdateSpecWidth{static_cast(spec_width)}); + return find_spec_verify_token_update_kernel_entry(specialization) != nullptr; +} + +void spec_verify_token_update(const torch::Tensor& base_token, + const std::vector& draft_tokens, + torch::Tensor& persistent_tokens) { + const int64_t spec_width = static_cast(draft_tokens.size()) + 1; + check_token(base_token, torch::kInt32); + CHECK(has_spec_verify_token_update_specialization(spec_width)) + << "speculative verify token update has no compiled variant for width " + << spec_width << ": " + << available_spec_verify_token_update_variant_keys(); + for (const auto& token : draft_tokens) { + check_token(token, torch::kInt64); + } + CHECK(persistent_tokens.defined() && + persistent_tokens.device().type() == c10::DeviceType::PrivateUse1); + CHECK_EQ(persistent_tokens.scalar_type(), torch::kInt32); + CHECK_GE(persistent_tokens.numel(), 8); + CHECK(persistent_tokens.is_contiguous()); + + const auto specialization = make_spec_verify_token_update_specialization( + SpecVerifyTokenUpdateSpecWidth{static_cast(spec_width)}); + const auto* entry = + find_spec_verify_token_update_kernel_entry(specialization); + CHECK(entry != nullptr) + << "speculative verify token update has no compiled variant: " + << available_spec_verify_token_update_variant_keys(); + + aclrtStream stream = + c10_npu::getCurrentNPUStream(base_token.device().index()).stream(); + std::array draft_ptrs; + draft_ptrs.fill( + reinterpret_cast(const_cast(base_token.data_ptr()))); + for (size_t index = 0; index < draft_tokens.size(); ++index) { + draft_ptrs[index] = reinterpret_cast( + const_cast(draft_tokens[index].data_ptr())); + } + entry->fn( + reinterpret_cast(const_cast(base_token.data_ptr())), + draft_ptrs[0], + draft_ptrs[1], + draft_ptrs[2], + draft_ptrs[3], + draft_ptrs[4], + reinterpret_cast(persistent_tokens.data_ptr()), + stream); +} + +} // namespace xllm::kernel::npu::tilelang diff --git a/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h b/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h index 99a0dc9f8d..06d84e4066 100644 --- a/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h +++ b/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h @@ -24,6 +24,37 @@ limitations under the License. namespace xllm::kernel::npu::tilelang { +void spec_verify_token_update(const torch::Tensor& base_token, + const std::vector& draft_tokens, + torch::Tensor& persistent_tokens); + +bool has_spec_verify_token_update_specialization(int64_t spec_width); + +bool has_spec_verify_metadata_update_specialization(int64_t spec_width, + int64_t block_table_width); + +void spec_verify_metadata_update( + const torch::Tensor& positions, + const torch::Tensor& kv_seq_lens, + const torch::Tensor& new_cache_slots, + const torch::Tensor& block_table, + const torch::Tensor& linear_state_index, + const torch::Tensor& num_accepted, + const std::vector& persistent_position_rows, + torch::Tensor& persistent_q_seq_lens, + torch::Tensor& persistent_kv_seq_lens, + torch::Tensor& persistent_new_cache_slots, + torch::Tensor& persistent_block_table, + torch::Tensor& persistent_linear_state_index, + torch::Tensor& persistent_num_accepted, + torch::Tensor& persistent_q_cu_seq_lens, + torch::Tensor& persistent_expanded_kv_seq_lens, + const std::vector& persistent_expanded_block_rows); + +void spec_verify_attention_tiling_update(const torch::Tensor& src_kv_seq_lens, + torch::Tensor& tiling_data, + int64_t block_size); + // Public TileLang kernel APIs exported to the xLLM NPU runtime. // // Apply TileLang RoPE kernel in-place on a single input tensor. diff --git a/xllm/core/runtime/acl_graph_executor_impl.cpp b/xllm/core/runtime/acl_graph_executor_impl.cpp index 3a61d746a3..f1bdf9e179 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.cpp +++ b/xllm/core/runtime/acl_graph_executor_impl.cpp @@ -45,6 +45,58 @@ namespace xllm::npu { namespace { constexpr uint64_t kSpecVerifyGraphKeyMask = 1ull << 63; constexpr uint64_t kSpecVerifyQMaxSeqLenShift = 32; +constexpr uint64_t kSpecVerifyBucketMask = (1ull << 16) - 1; +constexpr uint64_t kSpecVerifyFieldMask = (1ull << 16) - 1; +constexpr uint64_t kSpecVerifyExpandedBlockMask = (1ull << 15) - 1; +constexpr uint64_t kStaticGraphTaskHashSeed = 0x6a09e667f3bcc909ull; + +bool uses_static_mtp_graph_task_variant(const ModelInputParams& params, + uint32_t bucket_num_tokens) { + return params.is_spec_verify && + params.meta.batch_forward_type.is_chunked_prefill() && + params.graph.use_expanded_decode_for_spec_verify_attention && + params.graph.spec_verify_source_addresses_stable && + bucket_num_tokens >= 4 && bucket_num_tokens <= 6 && + params.meta.num_sequences == 1 && + params.embedding.linear_state_ids.size() == 1 && + params.num_accepted_tokens_host.size() == 1; +} + +uint64_t mix_graph_key(uint64_t hash, uint64_t value) { + value += 0x9e3779b97f4a7c15ull; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ull; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebull; + value ^= value >> 31; + return hash ^ (value + 0x9e3779b97f4a7c15ull + (hash << 6) + (hash >> 2)); +} + +uint64_t static_mtp_graph_task_key(uint64_t base_key, + const ModelInputParams& params) { + uint64_t hash = mix_graph_key(kStaticGraphTaskHashSeed, base_key); + hash = mix_graph_key( + hash, static_cast(params.embedding.linear_state_ids.front())); + hash = mix_graph_key( + hash, static_cast(params.num_accepted_tokens_host.front())); + for (const int64_t value : params.parallel.query_start_loc) { + hash = mix_graph_key(hash, static_cast(value)); + } + // Retain the spec-verify namespace bit. A signature comparison on replay + // guards correctness even in the astronomically unlikely event of a hash + // collision. + return hash | kSpecVerifyGraphKeyMask; +} + +uint64_t static_mtp_graph_task_key(uint64_t base_key, + int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width) { + uint64_t hash = mix_graph_key(kStaticGraphTaskHashSeed, base_key); + hash = mix_graph_key(hash, static_cast(linear_state_id)); + hash = mix_graph_key(hash, static_cast(num_accepted_tokens)); + hash = mix_graph_key(hash, 0); + hash = mix_graph_key(hash, static_cast(spec_width)); + return hash | kSpecVerifyGraphKeyMask; +} std::pair find_attention_plan_kv_cache( const std::vector& kv_caches) { @@ -59,6 +111,61 @@ std::pair find_attention_plan_kv_cache( return {torch::Tensor(), torch::Tensor()}; } +std::vector spec_verify_input_sources( + const torch::Tensor& tokens, + const torch::Tensor& positions, + const ModelInputParams& params, + bool fused_metadata_update) { + const torch::Tensor& graph_tokens = + params.graph.input_tokens_override.defined() + ? params.graph.input_tokens_override + : tokens; + if (fused_metadata_update) { + return {positions, + params.attention.device.kv_seq_lens, + params.attention.device.new_cache_slots, + params.attention.device.block_tables, + params.embedding.linear_state_indices, + params.num_accepted_tokens}; + } + return {graph_tokens, + positions, + params.attention.device.q_seq_lens, + params.attention.device.kv_seq_lens, + params.attention.device.new_cache_slots, + params.attention.device.block_tables, + params.embedding.linear_state_indices, + params.num_accepted_tokens, + params.attention.device.q_cu_seq_lens, + params.graph.expanded_kv_seq_lens, + params.graph.expanded_block_tables}; +} + +bool same_tensor_sources(const std::vector& captured, + const std::vector& current) { + if (captured.size() != current.size()) { + LOG(ERROR) << "speculative verify input source count changed: captured=" + << captured.size() << ", current=" << current.size(); + return false; + } + for (size_t i = 0; i < captured.size(); ++i) { + if (!captured[i].defined() || !current[i].defined() || + captured[i].data_ptr() != current[i].data_ptr() || + captured[i].sizes() != current[i].sizes() || + captured[i].strides() != current[i].strides()) { + LOG(ERROR) << "speculative verify input update source changed: index=" + << i << ", captured_ptr=" << captured[i].data_ptr() + << ", current_ptr=" << current[i].data_ptr() + << ", captured_sizes=" << captured[i].sizes() + << ", current_sizes=" << current[i].sizes() + << ", captured_strides=" << captured[i].strides() + << ", current_strides=" << current[i].strides(); + return false; + } + } + return true; +} + } // namespace bool AclGraph::capture(CausalLM* model, @@ -93,13 +200,20 @@ bool AclGraph::capture(CausalLM* model, static_cast(tokens.size(/*dim=*/0)); CHECK_GE(num_tokens_, actual_num_tokens) << "num_tokens_ >= actual_num_tokens"; + const bool fused_spec_verify_token_update = + persistent_param_.supports_fused_spec_verify_token_update(params); auto graph_params = persistent_param_.update(tokens, k_cache, v_cache, positions, params, num_tokens_, - /*return_capture_params=*/true); + /*return_capture_params=*/true, + /*skip_token_update=*/ + fused_spec_verify_token_update); + if (fused_spec_verify_token_update) { + persistent_param_.run_fused_spec_verify_token_update(params); + } // Use the returned ModelInputParams for graph capture CHECK(graph_params.has_value()) @@ -115,7 +229,8 @@ bool AclGraph::capture(CausalLM* model, graph_task_context_->begin_capture(); graph_params->graph.acl_graph_task_update_context = graph_task_context_; } - + const bool capture_static_graph_tasks = + uses_static_mtp_graph_task_variant(graph_params.value(), num_tokens_); // Synchronize stream to ensure all data is copied to graph persistent buffers aclrtSynchronizeStream(stream); @@ -148,6 +263,18 @@ bool AclGraph::capture(CausalLM* model, // other threads to execute synchronous operations graph_.capture_begin( {0, 0}, aclmdlRICaptureMode::ACL_MODEL_RI_CAPTURE_MODE_THREAD_LOCAL); + has_internal_spec_verify_input_update_ = + model->is_hybrid_linear_attention() && + params.graph.spec_verify_source_addresses_stable && + params.is_spec_verify && + params.meta.batch_forward_type.is_chunked_prefill() && + actual_num_tokens == bucket_num_tokens && + params.meta.q_max_seq_len == actual_num_tokens; + if (has_internal_spec_verify_input_update_) { + internal_spec_verify_input_sources_ = + persistent_param_.capture_spec_verify_input_update( + tokens, positions, params, bucket_num_tokens); + } // Execute forward pass - NPUGraph mempool manages temporary tensors auto forward_result = model->forward({persistent_param_.persistent_tokens(num_tokens_)}, @@ -176,14 +303,17 @@ bool AclGraph::capture(CausalLM* model, aclrtSynchronizeStream(stream); graph_.replay(); update_graph_tasks(graph_params.value()); + if (capture_static_graph_tasks) { + capture_static_graph_task_signature(graph_params.value()); + } make_current_stream_wait_for_graph(stream); return true; } -void AclGraph::update_graph_tasks(const ModelInputParams& params) { +bool AclGraph::update_graph_tasks(const ModelInputParams& params) { if (graph_task_context_ == nullptr || graph_task_context_->causal_conv1d_tasks.empty()) { - return; + return false; } const std::vector empty_host_args; @@ -235,6 +365,38 @@ void AclGraph::update_graph_tasks(const ModelInputParams& params) { task.event->record(update_stream); } } + return true; +} + +void AclGraph::signal_static_graph_tasks( + const c10_npu::NPUStream& signal_stream) { + CHECK(graph_task_context_ != nullptr); + for (auto& task : graph_task_context_->causal_conv1d_tasks) { + CHECK(task.event != nullptr) + << "static graph-task replay requires a captured ready event"; + task.event->record(signal_stream); + } +} + +bool AclGraph::static_graph_task_signature_matches( + const ModelInputParams& params) const { + if (!has_static_graph_task_signature_) { + return false; + } + return params.parallel.query_start_loc == static_query_start_loc_ && + params.embedding.linear_state_ids == static_linear_state_ids_ && + params.num_accepted_tokens_host == static_num_accepted_tokens_; +} + +void AclGraph::capture_static_graph_task_signature( + const ModelInputParams& params) { + static_query_start_loc_ = params.parallel.query_start_loc; + static_linear_state_ids_ = params.embedding.linear_state_ids; + static_num_accepted_tokens_ = params.num_accepted_tokens_host; + has_static_graph_task_signature_ = true; + LOG(INFO) << "Captured static MTP graph-task signature: linear_state_id=" + << static_linear_state_ids_.front() + << ", accepted_tokens=" << static_num_accepted_tokens_.front(); } AclGraph::~AclGraph() { @@ -247,6 +409,10 @@ AclGraph::~AclGraph() { aclrtDestroyEvent(replay_done_event_); replay_done_event_ = nullptr; } + if (replay_input_ready_event_ != nullptr) { + aclrtDestroyEvent(replay_input_ready_event_); + replay_input_ready_event_ = nullptr; + } } void AclGraph::initialize_capture_stream(c10::DeviceIndex device_index) { @@ -257,6 +423,9 @@ void AclGraph::initialize_capture_stream(c10::DeviceIndex device_index) { capture_stream_ = c10_npu::getStreamFromPool(true, device_index); update_stream_ = c10_npu::getStreamFromPool(true, device_index); device_index_ = device_index; + CHECK_EQ(aclrtCreateEventWithFlag(&replay_input_ready_event_, ACL_EVENT_SYNC), + ACL_SUCCESS) + << "Failed to create ACL graph replay input-ready event"; CHECK_EQ(aclrtCreateEventWithFlag(&replay_done_event_, ACL_EVENT_SYNC), ACL_SUCCESS) << "Failed to create ACL graph replay completion event"; @@ -265,6 +434,21 @@ void AclGraph::initialize_capture_stream(c10::DeviceIndex device_index) { << ", device_index: " << device_index; } +void AclGraph::make_graph_wait_for_current_stream(aclrtStream current_stream) { + CHECK_NE(graph_stream_, nullptr) << "graph_stream is not initialized"; + CHECK_NE(replay_input_ready_event_, nullptr) + << "replay_input_ready_event is not initialized"; + if (current_stream == graph_stream_) { + return; + } + CHECK_EQ(aclrtRecordEvent(replay_input_ready_event_, current_stream), + ACL_SUCCESS) + << "aclrtRecordEvent(replay_input_ready_event) failed"; + CHECK_EQ(aclrtStreamWaitEvent(graph_stream_, replay_input_ready_event_), + ACL_SUCCESS) + << "aclrtStreamWaitEvent(graph_stream, replay_input_ready_event) failed"; +} + void AclGraph::make_current_stream_wait_for_graph(aclrtStream current_stream) { CHECK_NE(graph_stream_, nullptr) << "graph_stream is not initialized"; CHECK_NE(replay_done_event_, nullptr) @@ -320,7 +504,24 @@ ModelOutput AclGraph::replay(CausalLM* model, replay_inputs_prepared && params.graph.input_tokens_override.defined() && !needs_graph_metadata; std::optional graph_params; - if (can_use_prepared_inputs) { + if (has_internal_spec_verify_input_update_) { + const auto current_sources = spec_verify_input_sources( + tokens, + positions, + params, + persistent_param_.supports_fused_spec_verify_metadata_update(params)); + CHECK(same_tensor_sources(internal_spec_verify_input_sources_, + current_sources)) + << "speculative verify graph input update source address changed " + "across replay"; + if (persistent_param_.supports_fused_spec_verify_token_update(params)) { + persistent_param_.run_fused_spec_verify_token_update(params); + } + // The leading captured graph nodes copy all dynamic device fields into + // persistent graph buffers. Host-only task parameters remain current and + // are consumed by update_graph_tasks() below. + graph_params = params; + } else if (can_use_prepared_inputs) { persistent_param_.update_tokens( tokens, params, actual_num_tokens, num_tokens_); } else { @@ -346,11 +547,35 @@ ModelOutput AclGraph::replay(CausalLM* model, // Get current NPU stream from libtorch NPU API aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + if (has_internal_spec_verify_input_update_) { + // The graph's leading input update reads sources finalized by the current + // compute stream (including the last draft token). Express that dependency + // on device; a host synchronize here would recreate the bubble we remove. + make_graph_wait_for_current_stream(stream); + } + const bool use_static_graph_tasks = + graph_params.has_value() && + static_graph_task_signature_matches(graph_params.value()); + const bool static_graph_tasks_prepared = + params.graph.spec_verify_static_graph_tasks_prepared; + CHECK(!static_graph_tasks_prepared || use_static_graph_tasks) + << "prepared static graph tasks do not match the replay signature"; + if (use_static_graph_tasks && !static_graph_tasks_prepared) { + // Cold/fallback path: the final-draft pre-submit could not find this graph + // variant. Ring its ready event immediately before replay; steady MTP5 + // cycles use the compute-stream pre-submit path instead. + CHECK(update_stream_.has_value()); + signal_static_graph_tasks(update_stream_.value()); + } graph_.replay(); if (model->is_hybrid_linear_attention()) { CHECK(graph_params.has_value()) << "update() should return ModelInputParams for graph task update"; - update_graph_tasks(graph_params.value()); + if (use_static_graph_tasks) { + // This graph variant's task-ready event was recorded before replay. + } else { + update_graph_tasks(graph_params.value()); + } } make_current_stream_wait_for_graph(stream); @@ -381,6 +606,36 @@ void AclGraph::prepare_replay_inputs(const torch::Tensor& tokens, replay_inputs_prepared_.store(true, std::memory_order_release); } +bool AclGraph::prepare_static_graph_tasks( + const ModelInputParams& params, + const c10_npu::NPUStream& signal_stream) { + if (!static_graph_task_signature_matches(params)) { + return false; + } + // The final draft graph has already been submitted to signal_stream. Queue + // each captured task-ready event behind it, so the target graph observes a + // fresh generation without a cross-stream credit/wait or host synchronize. + signal_static_graph_tasks(signal_stream); + return true; +} + +bool AclGraph::prepare_static_mtp_graph_tasks( + int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + const c10_npu::NPUStream& signal_stream) { + if (!has_static_graph_task_signature_ || + static_query_start_loc_ != std::vector({0, spec_width}) || + static_linear_state_ids_ != + std::vector({static_cast(linear_state_id)}) || + static_num_accepted_tokens_ != + std::vector({num_accepted_tokens})) { + return false; + } + signal_static_graph_tasks(signal_stream); + return true; +} + AclGraphExecutorImpl::AclGraphExecutorImpl(CausalLM* model, const ModelArgs& args, const torch::Device& device, @@ -686,6 +941,73 @@ void AclGraphExecutorImpl::prepare_graph_input(const torch::Tensor& tokens, graph->prepare_replay_inputs(tokens, positions, kv_caches, params); } +bool AclGraphExecutorImpl::prepare_static_graph_tasks( + const torch::Tensor& tokens, + const ModelInputParams& params, + const Stream& signal_stream) { + if (!model_->is_hybrid_linear_attention() || graph_slot_count_ != 1 || + !params.is_spec_verify || + !params.meta.batch_forward_type.is_chunked_prefill() || + !params.graph.use_expanded_decode_for_spec_verify_attention) { + return false; + } + uint32_t graph_num_tokens = tokens.size(0); + if (params.parallel.dp_global_token_nums.size() > 1) { + graph_num_tokens = util::max(params.parallel.dp_global_token_nums); + } + if (graph_num_tokens == 0) { + return false; + } + const uint32_t bucket_num_tokens = get_bucket_num_tokens(graph_num_tokens); + const uint64_t graph_key = get_graph_key(bucket_num_tokens, params); + AclGraph* graph = nullptr; + { + std::lock_guard lock(graph_slots_mutex_); + auto& graphs = graph_slots_[0].graphs; + auto it = graphs.find(graph_key); + if (it == graphs.end()) { + return false; + } + graph = it->second.get(); + } + return graph->prepare_static_graph_tasks(params, *signal_stream.get_stream()); +} + +bool AclGraphExecutorImpl::prepare_static_mtp_graph_tasks( + int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + int64_t block_table_width, + const Stream& signal_stream) { + if (!model_->is_hybrid_linear_attention() || graph_slot_count_ != 1 || + spec_width < 4 || spec_width > 6 || block_table_width < 1 || + block_table_width > kSpecVerifyExpandedBlockMask) { + return false; + } + const uint64_t bucket_num_tokens = static_cast(spec_width); + const uint64_t q_max_seq_len = static_cast(spec_width); + const uint64_t width = static_cast(block_table_width); + const uint64_t base_key = kSpecVerifyGraphKeyMask | (width << 48) | + (width << 32) | (q_max_seq_len << 16) | + bucket_num_tokens; + const uint64_t graph_key = static_mtp_graph_task_key( + base_key, linear_state_id, num_accepted_tokens, spec_width); + AclGraph* graph = nullptr; + { + std::lock_guard lock(graph_slots_mutex_); + auto& graphs = graph_slots_[0].graphs; + auto it = graphs.find(graph_key); + if (it == graphs.end()) { + return false; + } + graph = it->second.get(); + } + return graph->prepare_static_mtp_graph_tasks(linear_state_id, + num_accepted_tokens, + spec_width, + *signal_stream.get_stream()); +} + void AclGraph::print_graph_tensors() const { VLOG(kGraphExecutorLogVerboseLevel) << "graph persistent_tokens_: " << persistent_param_.persistent_tokens(); @@ -734,6 +1056,30 @@ uint64_t AclGraphExecutorImpl::get_graph_key( params.meta.batch_forward_type.is_chunked_prefill()) { const uint64_t q_max_seq_len = static_cast(std::max(params.meta.q_max_seq_len, 1)); + if (params.graph.spec_verify_source_addresses_stable) { + CHECK(params.attention.device.block_tables.defined()); + CHECK(params.graph.expanded_block_tables.defined()); + const uint64_t block_table_width = + static_cast(params.attention.device.block_tables.size(1)); + const uint64_t expanded_block_table_width = + static_cast(params.graph.expanded_block_tables.size(1)); + CHECK_LE(bucket_num_tokens, kSpecVerifyBucketMask); + CHECK_LE(q_max_seq_len, kSpecVerifyFieldMask); + CHECK_LE(block_table_width, kSpecVerifyFieldMask); + CHECK_LE(expanded_block_table_width, kSpecVerifyExpandedBlockMask); + // Captured D2D copies encode tensor shapes. Specialize the MTP target + // graph by both block-table widths so a synthetic warmup graph cannot be + // replayed with a real request's wider table view. Source storage may be + // shared; its view shape and stride are still part of graph identity. + const uint64_t base_key = + kSpecVerifyGraphKeyMask | (expanded_block_table_width << 48) | + (block_table_width << 32) | (q_max_seq_len << 16) | + static_cast(bucket_num_tokens); + if (uses_static_mtp_graph_task_variant(params, bucket_num_tokens)) { + return static_mtp_graph_task_key(base_key, params); + } + return base_key; + } return static_cast(bucket_num_tokens) | kSpecVerifyGraphKeyMask | (q_max_seq_len << kSpecVerifyQMaxSeqLenShift); } diff --git a/xllm/core/runtime/acl_graph_executor_impl.h b/xllm/core/runtime/acl_graph_executor_impl.h index 969fbe8548..be28d103c0 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.h +++ b/xllm/core/runtime/acl_graph_executor_impl.h @@ -85,6 +85,13 @@ class AclGraph { std::vector& kv_cache, const ModelInputParams& params); + bool prepare_static_graph_tasks(const ModelInputParams& params, + const c10_npu::NPUStream& signal_stream); + bool prepare_static_mtp_graph_tasks(int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + const c10_npu::NPUStream& signal_stream); + // Get the hidden states from the last capture torch::Tensor get_hidden_states(uint32_t actual_num_tokens = 0) const { return persistent_param_.hidden_states(actual_num_tokens); @@ -96,12 +103,17 @@ class AclGraph { // Initialize capture stream if not already initialized void initialize_capture_stream(c10::DeviceIndex device_index); + void make_graph_wait_for_current_stream(aclrtStream current_stream); void make_current_stream_wait_for_graph(aclrtStream current_stream); void prepare_model_graph_metadata(CausalLM* model, const torch::Tensor& positions, ModelInputParams& params); - void update_graph_tasks(const ModelInputParams& params); + bool update_graph_tasks(const ModelInputParams& params); + void signal_static_graph_tasks(const c10_npu::NPUStream& signal_stream); + bool static_graph_task_signature_matches( + const ModelInputParams& params) const; + void capture_static_graph_task_signature(const ModelInputParams& params); // NPUGraph with mempool for managing temporary tensors during forward pass c10_npu::NPUGraph graph_; @@ -115,11 +127,18 @@ class AclGraph { // Fallback non-default stream for capture when callers are on default stream. std::optional capture_stream_; aclrtStream graph_stream_ = nullptr; + aclrtEvent replay_input_ready_event_ = nullptr; aclrtEvent replay_done_event_ = nullptr; c10::DeviceIndex device_index_; std::shared_ptr graph_task_context_; std::optional update_stream_; std::atomic replay_inputs_prepared_{false}; + bool has_static_graph_task_signature_ = false; + std::vector static_query_start_loc_; + std::vector static_linear_state_ids_; + std::vector static_num_accepted_tokens_; + bool has_internal_spec_verify_input_update_ = false; + std::vector internal_spec_verify_input_sources_; }; // Executor implementation using ACL graph optimization @@ -146,6 +165,15 @@ class AclGraphExecutorImpl : public ExecutorImpl { std::vector& kv_caches, const ModelInputParams& params) override; + bool prepare_static_graph_tasks(const torch::Tensor& tokens, + const ModelInputParams& params, + const Stream& signal_stream) override; + bool prepare_static_mtp_graph_tasks(int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + int64_t block_table_width, + const Stream& signal_stream) override; + [[nodiscard]] int32_t graph_slot_count_for_test() const { return graph_slot_count_; } diff --git a/xllm/core/runtime/acl_graph_persistent_param.cpp b/xllm/core/runtime/acl_graph_persistent_param.cpp index f84b3a7fe5..e8d331245b 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.cpp +++ b/xllm/core/runtime/acl_graph_persistent_param.cpp @@ -30,6 +30,8 @@ limitations under the License. #include "core/common/constants.h" #include "core/common/global_flags.h" #include "core/framework/config/speculative_config.h" +#include "core/kernels/npu/tilelang/tilelang_ops_api.h" +#include "core/util/env_var.h" #include "core/util/utils.h" // ATB includes @@ -132,6 +134,25 @@ int64_t infer_actual_batch_size(const ModelInputParams& params) { return 0; } +bool can_skip_unused_expanded_verify_mask(const ModelInputParams& params, + bool is_hybrid_linear_attention) { + return util::get_bool_env("XLLM_NPU_MTP_SKIP_UNUSED_VERIFY_MASK", true) && + params.is_spec_verify && + params.meta.batch_forward_type.is_chunked_prefill() && + is_hybrid_linear_attention && + params.graph.use_expanded_decode_for_spec_verify_attention; +} + +int64_t get_spec_verify_width(const ModelInputParams& params) { + CHECK(params.graph.input_tokens_override.defined()) + << "speculative verify graph update requires token inputs"; + CHECK_EQ(params.graph.input_tokens_override.dim(), 1) + << "speculative verify token inputs must be flat"; + const int64_t spec_width = params.graph.input_tokens_override.numel(); + CHECK_GT(spec_width, 0) << "speculative verify width must be positive"; + return spec_width; +} + } // namespace // GraphPersistentParam implementation @@ -551,10 +572,29 @@ GraphPersistentParam::update_expanded_spec_decode_attention( } } - torch::Tensor expanded_kv_tensor = - torch::tensor(expanded_kv_seq_lens_vec, torch::kInt).to(device_); - expanded_kv_seq_lens_.slice(/*dim=*/0, /*start=*/0, /*end=*/padded_num_tokens) - .copy_(expanded_kv_tensor, /*non_blocking=*/true); + if (util::get_bool_env("XLLM_NPU_MTP_EXPANDED_KV_DIRECT_COPY", true)) { + // The worker has already packed the real rows into a device tensor. Copy + // it directly into the stable graph buffer instead of rebuilding the same + // data through torch::tensor(host).to(device), which introduces a + // synchronous H2D allocation/copy on every target replay. + expanded_kv_seq_lens_ + .slice(/*dim=*/0, /*start=*/0, /*end=*/actual_num_tokens) + .copy_(input_params.graph.expanded_kv_seq_lens, + /*non_blocking=*/true); + if (padded_num_tokens > actual_num_tokens) { + expanded_kv_seq_lens_ + .slice(/*dim=*/0, + /*start=*/actual_num_tokens, + /*end=*/padded_num_tokens) + .fill_(1); + } + } else { + torch::Tensor expanded_kv_tensor = + torch::tensor(expanded_kv_seq_lens_vec, torch::kInt).to(device_); + expanded_kv_seq_lens_ + .slice(/*dim=*/0, /*start=*/0, /*end=*/padded_num_tokens) + .copy_(expanded_kv_tensor, /*non_blocking=*/true); + } const int64_t block_table_len = input_params.graph.expanded_block_tables.size(1); @@ -594,6 +634,213 @@ void GraphPersistentParam::update_tokens(const torch::Tensor& tokens, } } +bool GraphPersistentParam::supports_fused_spec_verify_metadata_update( + const ModelInputParams& params) const { + if (!use_mrope_ || !params.graph.spec_verify_source_addresses_stable || + !params.graph.input_tokens_override.defined() || !params.is_spec_verify || + !params.meta.batch_forward_type.is_chunked_prefill()) { + return false; + } + + const int64_t spec_width = get_spec_verify_width(params); + if (params.graph.spec_verify_draft_token_sources.size() + 1 != + static_cast(spec_width)) { + return false; + } + const auto& attention = params.attention.device; + if (!attention.kv_seq_lens.defined() || + !attention.new_cache_slots.defined() || + !attention.block_tables.defined() || attention.block_tables.dim() != 2 || + attention.kv_seq_lens.numel() < 1 || + attention.new_cache_slots.numel() < spec_width || + attention.block_tables.size(0) < 1 || + !params.embedding.linear_state_indices.defined() || + params.embedding.linear_state_indices.numel() < 1 || + !params.num_accepted_tokens.defined() || + params.num_accepted_tokens.numel() < 1) { + return false; + } + + return kernel::npu::tilelang::has_spec_verify_metadata_update_specialization( + spec_width, attention.block_tables.size(1)); +} + +bool GraphPersistentParam::supports_fused_spec_verify_token_update( + const ModelInputParams& params) const { + if (!params.graph.spec_verify_source_addresses_stable || + !params.graph.input_tokens_override.defined() || !params.is_spec_verify || + !params.meta.batch_forward_type.is_chunked_prefill()) { + return false; + } + const int64_t spec_width = get_spec_verify_width(params); + return params.graph.spec_verify_draft_token_sources.size() + 1 == + static_cast(spec_width) && + kernel::npu::tilelang::has_spec_verify_token_update_specialization( + spec_width); +} + +void GraphPersistentParam::run_fused_spec_verify_token_update( + const ModelInputParams& params) { + CHECK(supports_fused_spec_verify_token_update(params)); + CHECK(params.graph.input_tokens_override.defined()); + kernel::npu::tilelang::spec_verify_token_update( + params.graph.input_tokens_override.narrow(0, 0, 1), + params.graph.spec_verify_draft_token_sources, + persistent_tokens_); +} + +std::vector +GraphPersistentParam::capture_spec_verify_input_update( + const torch::Tensor& tokens, + const torch::Tensor& positions, + const ModelInputParams& params, + uint32_t padded_num_tokens) { + CHECK(params.graph.spec_verify_source_addresses_stable) + << "speculative verify graph update requires stable source tensors"; + CHECK(params.is_spec_verify && + params.meta.batch_forward_type.is_chunked_prefill()) + << "graph input update only supports spec-verify chunked prefill"; + const torch::Tensor& graph_tokens = + params.graph.input_tokens_override.defined() + ? params.graph.input_tokens_override + : tokens; + const int64_t spec_width = graph_tokens.numel(); + CHECK_EQ(padded_num_tokens, spec_width); + CHECK_EQ(positions.numel(), spec_width); + CHECK_EQ(params.meta.q_max_seq_len, spec_width); + + const auto& attention = params.attention.device; + CHECK(attention.q_seq_lens.defined()); + CHECK(attention.kv_seq_lens.defined()); + CHECK(attention.new_cache_slots.defined()); + CHECK(attention.block_tables.defined()); + CHECK(attention.q_cu_seq_lens.defined()); + CHECK(params.embedding.linear_state_indices.defined()); + CHECK(params.num_accepted_tokens.defined()); + CHECK(params.graph.expanded_kv_seq_lens.defined()); + CHECK(params.graph.expanded_block_tables.defined()); + CHECK_GE(attention.q_seq_lens.numel(), 1); + CHECK_GE(attention.kv_seq_lens.numel(), 1); + CHECK_GE(attention.new_cache_slots.numel(), spec_width); + CHECK_GE(attention.block_tables.size(0), 1); + CHECK_GE(attention.q_cu_seq_lens.numel(), 2); + CHECK_GE(params.embedding.linear_state_indices.numel(), 1); + CHECK_GE(params.num_accepted_tokens.numel(), 1); + CHECK_GE(params.graph.expanded_kv_seq_lens.numel(), spec_width); + CHECK_GE(params.graph.expanded_block_tables.size(0), spec_width); + + const int64_t block_table_len = attention.block_tables.size(1); + const int64_t expanded_block_table_len = + params.graph.expanded_block_tables.size(1); + CHECK_LE(block_table_len, persistent_block_tables_.size(1)); + CHECK_LE(expanded_block_table_len, persistent_expanded_block_tables_.size(1)); + + const bool fused_metadata_update = + supports_fused_spec_verify_metadata_update(params); + const bool fused_token_update = + supports_fused_spec_verify_token_update(params); + LOG(INFO) << "Captured speculative verify input update: spec_width=" + << spec_width << ", fused_token=" << fused_token_update + << ", fused_metadata=" << fused_metadata_update; + + if (fused_metadata_update) { + std::vector position_rows; + position_rows.reserve(3); + for (int64_t row = 0; row < 3; ++row) { + position_rows.emplace_back(persistent_positions_.select(0, row)); + } + std::vector expanded_block_rows; + constexpr int64_t kMaxFusedSpecWidth = 6; + CHECK_GE(persistent_expanded_block_tables_.size(0), kMaxFusedSpecWidth); + expanded_block_rows.reserve(kMaxFusedSpecWidth); + for (int64_t row = 0; row < kMaxFusedSpecWidth; ++row) { + expanded_block_rows.emplace_back( + persistent_expanded_block_tables_.select(0, row).narrow( + 0, 0, block_table_len)); + } + auto persistent_q_seq_lens = q_seq_lens_.narrow(0, 0, 1); + auto persistent_kv_seq_lens = kv_seq_lens_.narrow(0, 0, 1); + torch::Tensor persistent_new_cache_slots = persistent_new_cache_slots_; + auto persistent_block_table = + persistent_block_tables_.select(0, 0).narrow(0, 0, block_table_len); + auto persistent_linear_state_index = + persistent_linear_state_indices_.narrow(0, 0, 1); + auto persistent_num_accepted = + persistent_num_accepted_tokens_.narrow(0, 0, 1); + auto persistent_q_cu_seq_lens = q_cu_seq_lens_.narrow(0, 0, 2); + torch::Tensor persistent_expanded_kv_seq_lens = expanded_kv_seq_lens_; + + kernel::npu::tilelang::spec_verify_metadata_update( + positions.narrow(0, 0, spec_width), + attention.kv_seq_lens.narrow(0, 0, 1), + attention.new_cache_slots.narrow(0, 0, spec_width), + attention.block_tables.select(0, 0).narrow(0, 0, block_table_len), + params.embedding.linear_state_indices.narrow(0, 0, 1), + params.num_accepted_tokens.narrow(0, 0, 1), + position_rows, + persistent_q_seq_lens, + persistent_kv_seq_lens, + persistent_new_cache_slots, + persistent_block_table, + persistent_linear_state_index, + persistent_num_accepted, + persistent_q_cu_seq_lens, + persistent_expanded_kv_seq_lens, + expanded_block_rows); + + return {positions, + attention.kv_seq_lens, + attention.new_cache_slots, + attention.block_tables, + params.embedding.linear_state_indices, + params.num_accepted_tokens}; + } + + if (!fused_token_update) { + persistent_tokens_.narrow(0, 0, spec_width).copy_(graph_tokens, true); + } + if (use_mrope_) { + // Hybrid mRoPE models store graph positions as [3, max_tokens]. The MTP + // worker + // supplies the compact [num_tokens] position vector, matching the normal + // persistent update path where copy_ broadcasts it across the mRoPE rows. + persistent_positions_.narrow(1, 0, spec_width).copy_(positions, true); + } else { + persistent_positions_.narrow(0, 0, spec_width).copy_(positions, true); + } + q_seq_lens_.narrow(0, 0, 1).copy_(attention.q_seq_lens.narrow(0, 0, 1), true); + kv_seq_lens_.narrow(0, 0, 1).copy_(attention.kv_seq_lens.narrow(0, 0, 1), + true); + persistent_new_cache_slots_.narrow(0, 0, spec_width) + .copy_(attention.new_cache_slots.narrow(0, 0, spec_width), true); + persistent_block_tables_.narrow(0, 0, 1) + .narrow(1, 0, block_table_len) + .copy_(attention.block_tables.narrow(0, 0, 1), true); + persistent_linear_state_indices_.narrow(0, 0, 1).copy_( + params.embedding.linear_state_indices.narrow(0, 0, 1), true); + persistent_num_accepted_tokens_.narrow(0, 0, 1).copy_( + params.num_accepted_tokens.narrow(0, 0, 1), true); + q_cu_seq_lens_.narrow(0, 0, 2).copy_(attention.q_cu_seq_lens.narrow(0, 0, 2), + true); + expanded_kv_seq_lens_.narrow(0, 0, spec_width) + .copy_(params.graph.expanded_kv_seq_lens.narrow(0, 0, spec_width), true); + persistent_expanded_block_tables_.narrow(0, 0, spec_width) + .narrow(1, 0, expanded_block_table_len) + .copy_(params.graph.expanded_block_tables, true); + + return {graph_tokens, + positions, + attention.q_seq_lens, + attention.kv_seq_lens, + attention.new_cache_slots, + attention.block_tables, + params.embedding.linear_state_indices, + params.num_accepted_tokens, + attention.q_cu_seq_lens, + params.graph.expanded_kv_seq_lens, + params.graph.expanded_block_tables}; +} + std::optional GraphPersistentParam::update( const torch::Tensor& tokens, const torch::Tensor& k_cache, @@ -608,11 +855,11 @@ std::optional GraphPersistentParam::update( const bool is_decode = params.meta.batch_forward_type.is_decode(); const bool is_chunked_prefill = params.meta.batch_forward_type.is_chunked_prefill(); - const bool is_qwen3_5_spec_verify_chunked_prefill = + const bool is_hybrid_spec_verify_chunked_prefill = params.is_spec_verify && is_chunked_prefill && is_hybrid_linear_attention_; - CHECK(is_decode || is_qwen3_5_spec_verify_chunked_prefill) - << "ACL graph persistent param only supports decode or Qwen3.5 " + CHECK(is_decode || is_hybrid_spec_verify_chunked_prefill) + << "ACL graph persistent param only supports decode or hybrid " "spec-verify chunked prefill"; const int64_t decode_tokens = is_decode ? std::max(options_.num_decoding_tokens(), 1) : 1; @@ -850,15 +1097,16 @@ std::optional GraphPersistentParam::update( ? params.attention.device.q_cu_seq_lens.size(0) : 0; if (has_q_cu && q_cu_size > 0) { - const bool use_qwen3_5_query_start_loc = is_hybrid_linear_attention_; + const bool use_hybrid_query_start_loc = is_hybrid_linear_attention_; const bool input_has_leading_zero = - params.is_spec_verify && use_qwen3_5_query_start_loc; + params.is_spec_verify && use_hybrid_query_start_loc; const int64_t required_q_cu_seq_lens = actual_seq_len_rows + (input_has_leading_zero ? 1 : 0); CHECK_GE(params.attention.device.q_cu_seq_lens.numel(), required_q_cu_seq_lens) - << "q_cu_seq_lens does not have enough entries for ACL graph execution"; - if (use_qwen3_5_query_start_loc && !input_has_leading_zero) { + << "q_cu_seq_lens does not have enough entries for ACL graph " + "execution"; + if (use_hybrid_query_start_loc && !input_has_leading_zero) { q_cu_seq_lens_.slice(/*dim=*/0, /*start=*/0, /*end=*/1).zero_(); q_cu_seq_lens_ .slice(/*dim=*/0, /*start=*/1, /*end=*/actual_seq_len_rows + 1) @@ -885,9 +1133,9 @@ std::optional GraphPersistentParam::update( padded_q_cu_seq_lens.emplace_back(offset); } const int64_t padding_start = - actual_seq_len_rows + (use_qwen3_5_query_start_loc ? 1 : 0); + actual_seq_len_rows + (use_hybrid_query_start_loc ? 1 : 0); const int64_t padding_end = - padded_batch_size + (use_qwen3_5_query_start_loc ? 1 : 0); + padded_batch_size + (use_hybrid_query_start_loc ? 1 : 0); q_cu_seq_lens_ .slice(/*dim=*/0, /*start=*/padding_start, @@ -897,8 +1145,14 @@ std::optional GraphPersistentParam::update( } } - // Update attention mask only if needed - if (need_update_attn_mask_) { + // Expanded Qwen hybrid spec verification is represented as chunked prefill + // to the scheduler, but AttentionImpl routes it through decoder_forward(). + // That path consumes expanded KV lengths, block tables and tiling data, not + // the dense graph attention mask. Avoid rebuilding the unused mask on every + // target replay while preserving all other decode/prefill paths. + const bool skip_unused_expanded_verify_mask = + can_skip_unused_expanded_verify_mask(params, is_hybrid_linear_attention_); + if (need_update_attn_mask_ && !skip_unused_expanded_verify_mask) { update_attention_mask(params); } @@ -924,9 +1178,9 @@ std::optional GraphPersistentParam::update( } const bool use_expanded_spec_decode_attention = params.graph.use_expanded_decode_for_spec_verify_attention; - if (is_qwen3_5_spec_verify_chunked_prefill) { + if (is_hybrid_spec_verify_chunked_prefill) { CHECK(use_expanded_spec_decode_attention) - << "Qwen3.5 spec-verify ACL graph requires MTP worker expanded " + << "hybrid spec-verify ACL graph requires expanded " "decode attention input"; } std::vector expanded_kv_seq_lens_vec; @@ -1027,9 +1281,12 @@ std::optional GraphPersistentParam::update( static_cast(padded_batch_size), 0); } - // Only set attn_mask if need_update_attn_mask_ is true - if (need_update_attn_mask_) { + // Keep the capture contract aligned with replay: the expanded decoder does + // not consume graph.attn_mask, so do not capture a stale persistent mask. + if (need_update_attn_mask_ && !skip_unused_expanded_verify_mask) { params_for_capture->graph.attn_mask = persistent_mask(padded_num_tokens); + } else if (skip_unused_expanded_verify_mask) { + params_for_capture->graph.attn_mask = torch::Tensor(); } if (uses_paged_attention_tiling()) { params_for_capture->graph.tiling_data = tiling_data(); @@ -1060,11 +1317,11 @@ std::optional GraphPersistentParam::update( expanded_kv_seq_lens_vec; } if (params.attention.device.q_cu_seq_lens.defined()) { - const bool use_qwen3_5_query_start_loc = is_hybrid_linear_attention_; + const bool use_hybrid_query_start_loc = is_hybrid_linear_attention_; params_for_capture->attention.device.q_cu_seq_lens = q_cu_seq_lens_.slice( /*dim=*/0, /*start=*/0, - /*end=*/padded_batch_size + (use_qwen3_5_query_start_loc ? 1 : 0)); + /*end=*/padded_batch_size + (use_hybrid_query_start_loc ? 1 : 0)); } // Replace dp/cp ep padding with slices of persistent buffers so that @@ -1100,14 +1357,22 @@ std::optional GraphPersistentParam::update( if (params.num_accepted_tokens.defined() && params.num_accepted_tokens.numel() > 0) { - torch::Tensor nat_host = params.num_accepted_tokens.to(torch::kCPU) - .to(torch::kLong) - .contiguous(); - const int64_t copy_size = - std::min(actual_batch_size, nat_host.numel()); - const int64_t* data = nat_host.data_ptr(); - params_for_capture->num_accepted_tokens_host.assign(data, - data + copy_size); + if (!params.num_accepted_tokens_host.empty()) { + const int64_t copy_size = std::min( + actual_batch_size, params.num_accepted_tokens_host.size()); + params_for_capture->num_accepted_tokens_host.assign( + params.num_accepted_tokens_host.begin(), + params.num_accepted_tokens_host.begin() + copy_size); + } else { + torch::Tensor nat_host = params.num_accepted_tokens.to(torch::kCPU) + .to(torch::kLong) + .contiguous(); + const int64_t copy_size = + std::min(actual_batch_size, nat_host.numel()); + const int64_t* data = nat_host.data_ptr(); + params_for_capture->num_accepted_tokens_host.assign(data, + data + copy_size); + } params_for_capture->num_accepted_tokens_host.resize( static_cast(padded_batch_size), 1); } @@ -1430,6 +1695,24 @@ void GraphPersistentParam::plan_paged_attention_tiling( } custom_variantPack.outTensors.push_back(atb_query); + constexpr size_t kPagedAttentionTilingHeaderWords = 44; + constexpr size_t kPagedAttentionTilingBatchWords = 17; + const int64_t spec_width = input_params.attention.device.kv_seq_lens.numel(); + const size_t required_tiling_words = + kPagedAttentionTilingHeaderWords + + kPagedAttentionTilingBatchWords * static_cast(spec_width); + const bool has_compatible_tiling_layout = + spec_width >= 4 && spec_width <= 6 && + paged_attention_tiling_template_.size() >= required_tiling_words && + paged_attention_tiling_template_.front() == + static_cast(spec_width); + const bool use_graph_tiling_patch = + input_params.graph.use_expanded_decode_for_spec_verify_attention && + has_compatible_tiling_layout; + if (use_graph_tiling_patch) { + return; + } + uint64_t custom_workspace_size = 0; atb::Status status = custom_pa_op_for_plan_->Setup( custom_variantPack, custom_workspace_size, context_for_plan_); @@ -1445,6 +1728,15 @@ void GraphPersistentParam::plan_paged_attention_tiling( CHECK_GT(tiling_buffer_info.tilingBufferSize, 0) << "Tiling buffer size is zero"; + if (input_params.graph.use_expanded_decode_for_spec_verify_attention) { + const size_t word_count = + static_cast(tiling_buffer_info.tilingBufferSize) / + sizeof(uint32_t); + const auto* words = + reinterpret_cast(tiling_buffer_info.tilingBuffer); + paged_attention_tiling_template_.assign(words, words + word_count); + } + if (VLOG_IS_ON(kGraphExecutorLogVerboseLevel)) { parse_pa_host_tiling_buffer(tiling_buffer_info.tilingBuffer, tiling_buffer_info.tilingBufferSize); diff --git a/xllm/core/runtime/acl_graph_persistent_param.h b/xllm/core/runtime/acl_graph_persistent_param.h index b63c8790ed..b63673b630 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.h +++ b/xllm/core/runtime/acl_graph_persistent_param.h @@ -69,6 +69,22 @@ class GraphPersistentParam final { uint32_t actual_num_tokens, uint32_t padded_num_tokens); + // Record device-side updates from speculative-verify sources into persistent + // graph inputs while ACL graph capture is active. The returned tensors define + // the stable-address replay signature. TileLang fusion is an optional + // specialization hidden behind this interface. + std::vector capture_spec_verify_input_update( + const torch::Tensor& tokens, + const torch::Tensor& positions, + const ModelInputParams& params, + uint32_t padded_num_tokens); + + bool supports_fused_spec_verify_metadata_update( + const ModelInputParams& params) const; + bool supports_fused_spec_verify_token_update( + const ModelInputParams& params) const; + void run_fused_spec_verify_token_update(const ModelInputParams& params); + // Getter methods for persistent tensors torch::Tensor persistent_tokens(uint32_t actual_tokens = 0) const { if (actual_tokens > 0) { @@ -250,6 +266,7 @@ class GraphPersistentParam final { // Persistent paged attention tiling tensor on device torch::Tensor tiling_data_; + std::vector paged_attention_tiling_template_; // Cached attention parameters int32_t num_head_; diff --git a/xllm/core/runtime/executor.cpp b/xllm/core/runtime/executor.cpp index e2e529bcda..5fea20bc36 100644 --- a/xllm/core/runtime/executor.cpp +++ b/xllm/core/runtime/executor.cpp @@ -58,4 +58,22 @@ void Executor::prepare_graph_input(const torch::Tensor& tokens, impl_->prepare_graph_input(tokens, positions, kv_caches, params); } +bool Executor::prepare_static_graph_tasks(const torch::Tensor& tokens, + const ModelInputParams& params, + const Stream& signal_stream) { + return impl_->prepare_static_graph_tasks(tokens, params, signal_stream); +} + +bool Executor::prepare_static_mtp_graph_tasks(int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + int64_t block_table_width, + const Stream& signal_stream) { + return impl_->prepare_static_mtp_graph_tasks(linear_state_id, + num_accepted_tokens, + spec_width, + block_table_width, + signal_stream); +} + } // namespace xllm diff --git a/xllm/core/runtime/executor.h b/xllm/core/runtime/executor.h index ddc6548fd5..d575814b26 100644 --- a/xllm/core/runtime/executor.h +++ b/xllm/core/runtime/executor.h @@ -51,6 +51,16 @@ class Executor final { std::vector& kv_caches, const ModelInputParams& params); + bool prepare_static_graph_tasks(const torch::Tensor& tokens, + const ModelInputParams& params, + const Stream& signal_stream); + + bool prepare_static_mtp_graph_tasks(int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + int64_t block_table_width, + const Stream& signal_stream); + private: std::unique_ptr impl_; }; diff --git a/xllm/core/runtime/executor_impl.h b/xllm/core/runtime/executor_impl.h index 3d946f3743..206c70fa39 100644 --- a/xllm/core/runtime/executor_impl.h +++ b/xllm/core/runtime/executor_impl.h @@ -27,6 +27,8 @@ limitations under the License. #include "framework/model/model_input_params.h" #include "framework/model/model_output.h" #include "options.h" +#include "platform/stream.h" +#include "platform/stream_event.h" namespace xllm { @@ -48,6 +50,20 @@ class ExecutorImpl { const torch::Tensor& /*positions*/, std::vector& /*kv_caches*/, const ModelInputParams& /*params*/) {} + + virtual bool prepare_static_graph_tasks(const torch::Tensor& /*tokens*/, + const ModelInputParams& /*params*/, + const Stream& /*signal_stream*/) { + return false; + } + + virtual bool prepare_static_mtp_graph_tasks(int64_t /*linear_state_id*/, + int64_t /*num_accepted_tokens*/, + int64_t /*spec_width*/, + int64_t /*block_table_width*/, + const Stream& /*signal_stream*/) { + return false; + } }; } // namespace xllm diff --git a/xllm/core/runtime/forward_params.h b/xllm/core/runtime/forward_params.h index 76fe08e374..f94eb8c3d7 100644 --- a/xllm/core/runtime/forward_params.h +++ b/xllm/core/runtime/forward_params.h @@ -527,6 +527,7 @@ struct ForwardInput { inputs.step_decode = step_decode; inputs.skip_sampling_for_logits_only = skip_sampling_for_logits_only; inputs.cp_partitioned = cp_partitioned; + inputs.metadata_ready_event = metadata_ready_event; } void set_host_views(ForwardInput& inputs) const { diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index 24d3df340d..39ec496afd 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -131,6 +131,33 @@ bool LLMWorkerImpl::init_model(ModelContext& context) { return true; } +#if defined(USE_NPU) +bool LLMWorkerImpl::prepare_static_graph_tasks(const ForwardInput& input, + const Stream& signal_stream) { + if (model_executor_ == nullptr || !input.token_ids.defined()) { + return false; + } + return model_executor_->prepare_static_graph_tasks( + input.token_ids, input.input_params, signal_stream); +} + +bool LLMWorkerImpl::prepare_static_mtp_graph_tasks( + int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + int64_t block_table_width, + const Stream& signal_stream) { + if (model_executor_ == nullptr) { + return false; + } + return model_executor_->prepare_static_mtp_graph_tasks(linear_state_id, + num_accepted_tokens, + spec_width, + block_table_width, + signal_stream); +} +#endif + std::optional LLMWorkerImpl::step_no_sync( const ForwardInput& input) { ForwardInput input_on_device; diff --git a/xllm/core/runtime/llm_worker_impl.h b/xllm/core/runtime/llm_worker_impl.h index bb9757e012..e4afb318ef 100644 --- a/xllm/core/runtime/llm_worker_impl.h +++ b/xllm/core/runtime/llm_worker_impl.h @@ -70,6 +70,14 @@ class LLMWorkerImpl : public WorkerImpl { public: #if defined(USE_NPU) + bool prepare_static_graph_tasks(const ForwardInput& input, + const Stream& signal_stream); + bool prepare_static_mtp_graph_tasks(int64_t linear_state_id, + int64_t num_accepted_tokens, + int64_t spec_width, + int64_t block_table_width, + const Stream& signal_stream); + layer::NpuLmHead get_npu_lm_head() { return model_->get_npu_lm_head(); }; void set_npu_lm_head(layer::NpuLmHead& head) { @@ -97,6 +105,10 @@ class LLMWorkerImpl : public WorkerImpl { model_->set_word_embedding(embedding); }; + SpeculativeVerifyCapabilities speculative_verify_capabilities() const { + return model_->speculative_verify_capabilities(); + } + // DFlash-specific delegate: eagerly project target hidden into the draft's // per-layer KV cache. Runs outside the executor because the pass has no // attention and its shape doesn't match the decode graph. See CausalLM. diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index 4854fdf7b3..b58960fd55 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -32,6 +32,7 @@ limitations under the License. #include "core/framework/kv_cache/kv_cache_estimation.h" #include "core/framework/multimodal/mm_data.h" #include "spec_input_builder.h" +#include "util/env_var.h" #include "util/pretty_print.h" #include "util/slice.h" #include "util/timer.h" @@ -42,6 +43,10 @@ constexpr uint64_t MBUF_SIZE = 128 * 1024 * 1024; namespace { +bool enable_spec_verify_graph_update() { + return util::get_bool_env("XLLM_NPU_SPEC_VERIFY_GRAPH_UPDATE", true); +} + ProcessGroup* spec_broadcast_group(const ParallelArgs& parallel_args) { return parallel_args.tp_group_ != nullptr ? parallel_args.tp_group_ : parallel_args.process_group_; @@ -182,12 +187,12 @@ void clear_expanded_spec_verify_graph_input(ModelInputParams& input_params) { input_params.graph.expanded_kv_seq_lens_vec.clear(); } -void build_expanded_spec_verify_graph_input(ModelInputParams& input_params, - const torch::Device& device) { +bool build_expanded_spec_verify_graph_host_input( + ModelInputParams& input_params) { clear_expanded_spec_verify_graph_input(input_params); if (!input_params.is_spec_verify || !input_params.meta.batch_forward_type.is_chunked_prefill()) { - return; + return false; } const std::vector& q_seq_lens = @@ -195,19 +200,13 @@ void build_expanded_spec_verify_graph_input(ModelInputParams& input_params, const std::vector& kv_seq_lens = input_params.attention.host.kv_seq_lens; if (q_seq_lens.empty() || kv_seq_lens.empty()) { - return; + return false; } CHECK_EQ(q_seq_lens.size(), kv_seq_lens.size()) << "spec verify q/kv seq lens must both be sequence-scoped"; - CHECK(input_params.attention.device.block_tables.defined()) - << "spec verify block tables must be rebuilt before graph input"; - const int64_t batch_size = static_cast(q_seq_lens.size()); - CHECK_GE(input_params.attention.device.block_tables.size(0), batch_size) - << "spec verify block table rows are fewer than sequences"; std::vector expanded_kv_seq_lens; - std::vector expanded_block_rows; for (int64_t seq_idx = 0; seq_idx < batch_size; ++seq_idx) { const int32_t q_len = q_seq_lens[static_cast(seq_idx)]; const int32_t kv_len = kv_seq_lens[static_cast(seq_idx)]; @@ -215,27 +214,75 @@ void build_expanded_spec_verify_graph_input(ModelInputParams& input_params, CHECK_GE(kv_len, q_len) << "kv_len must include validate query tokens"; for (int32_t token_idx = 0; token_idx < q_len; ++token_idx) { expanded_kv_seq_lens.emplace_back(kv_len - q_len + token_idx + 1); - expanded_block_rows.emplace_back( - input_params.attention.device.block_tables.select(/*dim=*/0, - seq_idx)); } } if (expanded_kv_seq_lens.empty()) { - return; + return false; } - torch::Tensor expanded_kv_seq_lens_host = - torch::tensor(expanded_kv_seq_lens, - torch::TensorOptions() - .dtype(torch::kInt) - .device(torch::kCPU) - .pinned_memory(true)); input_params.graph.use_expanded_decode_for_spec_verify_attention = true; - input_params.graph.expanded_kv_seq_lens = - expanded_kv_seq_lens_host.to(device, /*non_blocking=*/true); - input_params.graph.expanded_block_tables = - torch::stack(expanded_block_rows, 0); input_params.graph.expanded_kv_seq_lens_vec = std::move(expanded_kv_seq_lens); + return true; +} + +void bind_expanded_spec_verify_graph_input(ModelInputParams& input_params, + const torch::Device& device, + bool kv_lens_already_bound) { + if (!input_params.graph.use_expanded_decode_for_spec_verify_attention) { + return; + } + CHECK(input_params.attention.device.block_tables.defined()) + << "spec verify block tables must be rebuilt before graph input"; + const auto& q_seq_lens = input_params.attention.host.q_seq_lens; + CHECK_GE(input_params.attention.device.block_tables.size(0), + static_cast(q_seq_lens.size())) + << "spec verify block table rows are fewer than sequences"; + std::vector expanded_block_rows; + for (int64_t seq_idx = 0; seq_idx < static_cast(q_seq_lens.size()); + ++seq_idx) { + for (int32_t token_idx = 0; + token_idx < q_seq_lens[static_cast(seq_idx)]; + ++token_idx) { + expanded_block_rows.emplace_back( + input_params.attention.device.block_tables.select(/*dim=*/0, + seq_idx)); + } + } + + if (!kv_lens_already_bound) { + torch::Tensor expanded_kv_seq_lens_host = + torch::tensor(input_params.graph.expanded_kv_seq_lens_vec, + torch::TensorOptions() + .dtype(torch::kInt) + .device(torch::kCPU) + .pinned_memory(true)); + input_params.graph.expanded_kv_seq_lens = + expanded_kv_seq_lens_host.to(device, /*non_blocking=*/true); + } + + // The single-request expanded rows all alias the same block-table row. + // Preserve that alias as a zero-stride view instead of launching a + // graph-external stack copy; multi-request shapes retain the generic path. + bool all_same_row = true; + const void* first_ptr = expanded_block_rows.front().data_ptr(); + for (const auto& row : expanded_block_rows) { + all_same_row = all_same_row && row.data_ptr() == first_ptr; + } + if (all_same_row) { + input_params.graph.expanded_block_tables = + expanded_block_rows.front().unsqueeze(0).expand( + {static_cast(expanded_block_rows.size()), + expanded_block_rows.front().size(0)}); + } else { + input_params.graph.expanded_block_tables = + torch::stack(expanded_block_rows, 0); + } +} + +void build_expanded_spec_verify_graph_input(ModelInputParams& input_params, + const torch::Device& device) { + build_expanded_spec_verify_graph_host_input(input_params); + bind_expanded_spec_verify_graph_input(input_params, device, false); } #endif @@ -561,20 +608,10 @@ KVCacheShape MTPDraftKVCacheShape(const KVCacheShape& target_shape, return KVCacheShape(draft_capacity, draft_model_args, /*world_size=*/1); } -bool is_qwen3_5_target_model_type(const std::string& model_type) { - return model_type == "qwen3_5" || model_type == "qwen3_5_moe" || - model_type == "qwen3_5_text" || model_type == "qwen3_5_moe_text" || - model_type.rfind("qwen3_5_", 0) == 0; -} - bool is_qwen3_5_draft_model_type(const std::string& model_type) { return model_type == "qwen3_5_mtp" || model_type == "qwen3_5_moe_mtp"; } -bool is_mimo_target_model_type(const std::string& model_type) { - return model_type == "mimo"; -} - } // namespace MTPWorkerImpl::MTPWorkerImpl(const ParallelArgs& parallel_args, @@ -620,6 +657,7 @@ bool MTPWorkerImpl::init_model(const std::string& model_weights_path, if (impl_ != nullptr && impl_->get_status() == WorkerImpl::Status::LOADED) { context_ = impl_->context_; + target_spec_verify_capabilities_ = impl_->speculative_verify_capabilities(); } if (draft_impl_ != nullptr && @@ -653,12 +691,6 @@ bool MTPWorkerImpl::init_model(const std::string& model_weights_path, #endif } } -#if defined(USE_NPU) - if (result && use_qwen3_5_spec_verify_path()) { - CHECK_EQ(::xllm::KernelConfig::get_instance().npu_kernel_backend(), "TORCH") - << "Qwen3.5 MTP only supports NPU Torch backend"; - } -#endif return result; } @@ -719,28 +751,17 @@ int64_t MTPWorkerImpl::get_embedding_placeholder_size() { return static_cast(embedding_size_); } -bool MTPWorkerImpl::use_qwen3_5_spec_verify_path() const { - return impl_ != nullptr && - impl_->get_status() != WorkerImpl::Status::UNINITIALIZED && - is_qwen3_5_target_model_type( - impl_->context_.get_model_args().model_type()); +SpeculativeVerifyCapabilities MTPWorkerImpl::speculative_verify_capabilities() + const { + return target_spec_verify_capabilities_; } -// MiMo MTP validation requires chunked-prefill mode (same as Qwen3.5) to -// avoid the read-before-write race in FlashInfer batch-decode: validation -// token 1 (at position p+1) must attend to the KV written by token 0 (at -// position p) within the same batch call, but all-reads-then-all-writes -// ordering means it reads garbage. Chunked prefill executes tokens causally -// so token 1 always sees token 0's committed KV. -bool MTPWorkerImpl::use_mimo_spec_verify_path() const { - return impl_ != nullptr && - impl_->get_status() != WorkerImpl::Status::UNINITIALIZED && - is_mimo_target_model_type( - impl_->context_.get_model_args().model_type()); +bool MTPWorkerImpl::supports_spec_verify_graph_input_update() const { + return speculative_verify_capabilities().supports_in_graph_input_update; } bool MTPWorkerImpl::use_chunked_prefill_spec_verify_path() const { - return use_qwen3_5_spec_verify_path() || use_mimo_spec_verify_path(); + return speculative_verify_capabilities().requires_causal_chunked_prefill; } bool MTPWorkerImpl::allocate_kv_cache(const KVCacheShape& kv_cache_shape) { @@ -1160,6 +1181,10 @@ std::optional MTPWorkerImpl::step_decode( draft_impl_->context_.get_model_args(), enable_schedule_overlap()); torch::Tensor mtp_topk_indices; for (int32_t draft_idx = 0; draft_idx < num_speculative_tokens; ++draft_idx) { + const bool is_final_draft = draft_idx == num_speculative_tokens - 1; + const bool static_graph_tasks_prepared = + is_final_draft && + prepare_static_mtp_graph_tasks_before_final_draft(input); if (reuse_mtp_topk_indices) { current_draft_input.input_params.dsa_topk_indices = mtp_topk_indices; } @@ -1171,8 +1196,9 @@ std::optional MTPWorkerImpl::step_decode( draft_prepared[draft_idx]); // Overlap next-step input preparation with async draft forward. - if (draft_idx == num_speculative_tokens - 1) { - prepare_validate_inputs(input, validate_input); + if (is_final_draft) { + prepare_validate_inputs( + input, validate_input, static_graph_tasks_prepared); } else { prepare_draft_inputs(input, next_step_input, draft_idx + 1); } @@ -1245,18 +1271,49 @@ void MTPWorkerImpl::fill_validate_input_from_draft_outputs( validate_input.token_ids.view({num_sequences, num_val_tokens}); validate_input.device_tensors_ready = false; + auto& fused_draft_tokens = + validate_input.input_params.graph.spec_verify_draft_token_sources; + fused_draft_tokens.clear(); +#if defined(USE_NPU) + const bool use_fused_verify_token_update = + validate_input.input_params.graph.spec_verify_source_addresses_stable && + validate_input.input_params.graph.input_tokens_override.defined() && + supports_spec_verify_graph_input_update() && num_sequences == 1 && + num_val_tokens >= 4 && num_val_tokens <= 6; + if (use_fused_verify_token_update) { + fused_draft_tokens.reserve(draft_outputs.size()); + for (const ForwardOutput& draft_output : draft_outputs) { + const torch::Tensor& next_tokens = draft_output.sample_output.next_tokens; + CHECK(next_tokens.defined() && next_tokens.numel() == 1 && + next_tokens.scalar_type() == torch::kInt64 && + next_tokens.device() == validate_input.token_ids.device() && + next_tokens.is_contiguous()) + << "fused speculative verify input update requires one contiguous " + "int64 token " + "per draft step"; + fused_draft_tokens.emplace_back(next_tokens.flatten()); + } + validate_input.device_tensors_ready = true; + return; + } +#endif + std::vector draft_token_columns; + draft_token_columns.reserve(num_speculative_tokens); for (int32_t i = 0; i < num_speculative_tokens; ++i) { - const auto& draft_output = draft_outputs[i]; - const torch::Tensor& next_tokens = draft_output.sample_output.next_tokens; + const torch::Tensor& next_tokens = + draft_outputs[i].sample_output.next_tokens; CHECK(next_tokens.defined()) << "draft next_tokens must be defined for validate token fill"; torch::Tensor draft_tokens = safe_to(next_tokens.flatten(), token_options, /*non_blocking=*/true); CHECK_EQ(draft_tokens.numel(), num_sequences) << "draft token count must match validate sequence count"; - validate_token_rows.select(/*dim=*/1, /*index=*/i + 1) - .copy_(draft_tokens, /*non_blocking=*/true); + draft_token_columns.emplace_back(std::move(draft_tokens)); } + torch::Tensor packed_draft_tokens = + torch::stack(draft_token_columns, /*dim=*/1); + validate_token_rows.slice(/*dim=*/1, /*start=*/1, /*end=*/num_val_tokens) + .copy_(packed_draft_tokens, /*non_blocking=*/true); validate_input.device_tensors_ready = true; record_metadata_ready_event(compute_stream, validate_input); } @@ -1267,9 +1324,9 @@ std::optional MTPWorkerImpl::run_validate( ForwardInput& validate_input) { // run the target model to get the verification scores Timer timer; - ForwardInput target_prepared; fill_validate_input_from_draft_outputs( draft_outputs, validate_input, *compute_stream_); + ForwardInput target_prepared; ForwardOutput target_output = run_llm_no_sync_impl(*impl_, validate_input, *prepare_stream_, @@ -1461,7 +1518,8 @@ void MTPWorkerImpl::update_decode_step_input( } void MTPWorkerImpl::prepare_validate_inputs(const ForwardInput& input, - ForwardInput& validate_input) { + ForwardInput& validate_input, + bool static_graph_tasks_prepared) { c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); validate_input = input; clear_ready_events(validate_input); @@ -1476,6 +1534,13 @@ void MTPWorkerImpl::prepare_validate_inputs(const ForwardInput& input, const int32_t num_val_tokens = num_speculative_tokens + 1; const int32_t total_num_val_tokens = num_sequences * num_val_tokens; const int32_t block_size = options_.block_size(); +#if defined(USE_NPU) + const bool use_graph_internal_verify_update = + enable_spec_verify_graph_update() && + supports_spec_verify_graph_input_update() && num_sequences == 1; +#else + const bool use_graph_internal_verify_update = false; +#endif specBuilder::DecodeRowContext row_ctx = specBuilder::make_decode_row_context(input); Slice token_ids = { @@ -1537,11 +1602,51 @@ void MTPWorkerImpl::prepare_validate_inputs(const ForwardInput& input, CHECK_EQ(buf.out_positions.size(), buf.out_token_ids.size()) << "validate positions/tokens mismatch"; - specBuilder::set_token_position_tensors(validate_input, - buf.out_token_ids, - buf.out_positions, - token_options, - position_options); + if (use_graph_internal_verify_update) { +#if defined(USE_NPU) + const auto host_options = torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCPU) + .pinned_memory(true); + if (!spec_verify_tokens_host_.defined() || + spec_verify_tokens_host_.numel() != total_num_val_tokens) { + spec_verify_tokens_host_ = + torch::empty({total_num_val_tokens}, host_options); + spec_verify_tokens_device_ = + torch::empty({total_num_val_tokens}, token_options); + spec_verify_positions_host_ = + torch::empty({total_num_val_tokens}, host_options); + spec_verify_positions_device_ = + torch::empty({total_num_val_tokens}, position_options); + } + std::memcpy(spec_verify_tokens_host_.data_ptr(), + buf.out_token_ids.data(), + buf.out_token_ids.size() * sizeof(int32_t)); + std::memcpy(spec_verify_positions_host_.data_ptr(), + buf.out_positions.data(), + buf.out_positions.size() * sizeof(int32_t)); + spec_verify_tokens_device_.copy_(spec_verify_tokens_host_, + /*non_blocking=*/true); + spec_verify_positions_device_.copy_(spec_verify_positions_host_, + /*non_blocking=*/true); + validate_input.token_ids_host = spec_verify_tokens_host_; + validate_input.positions_host = spec_verify_positions_host_; + validate_input.token_ids = spec_verify_tokens_device_; + validate_input.positions = spec_verify_positions_device_; + // Expose stable base-token storage to the fused graph input update. + // Without this binding the fused gate falls back to five casts + stack on + // every final-draft boundary even though token_ids already has a stable + // address suitable for ACL graph replay. + input_params.graph.input_tokens_override = spec_verify_tokens_device_; + validate_input.device_tensors_ready = true; +#endif + } else { + specBuilder::set_token_position_tensors(validate_input, + buf.out_token_ids, + buf.out_positions, + token_options, + position_options); + } if (!use_atb_spec_kernel) { input_params.meta.num_sequences = total_num_val_tokens; input_params.meta.batch_forward_type = BatchForwardType::DECODE; @@ -1567,13 +1672,61 @@ void MTPWorkerImpl::prepare_validate_inputs(const ForwardInput& input, /*update_block_tables=*/true); } + auto& validate_sampling_params = validate_input.sampling_params; +#if defined(USE_NPU) + // update_sampling_params() uses repeat_interleave on device tensors. For a + // single greedy sequence that work only expands [0] and [false] into fixed + // validation controls, yet it lands on the final-draft -> target dependency + // chain. Reuse stable controls after warmup and retain the generic builder + // for sampling, penalties, multi-sequence batches and other backends. + const bool use_stable_greedy_validate_sampling = + use_graph_internal_verify_update && + validate_sampling_params.all_greedy_sample && + validate_sampling_params.selected_token_idxes.defined() && + validate_sampling_params.selected_token_idxes.numel() == 1 && + validate_sampling_params.sample_idxes.defined() && + validate_sampling_params.sample_idxes.numel() == 1 && + validate_sampling_params.do_sample.defined() && + validate_sampling_params.do_sample.numel() == 1 && + !validate_sampling_params.frequency_penalties.defined() && + !validate_sampling_params.presence_penalties.defined() && + !validate_sampling_params.repetition_penalties.defined() && + !validate_sampling_params.temperatures.defined() && + !validate_sampling_params.top_p.defined() && + !validate_sampling_params.top_k.defined() && + !validate_sampling_params.unique_token_ids.defined() && + !validate_sampling_params.unique_token_counts.defined() && + !validate_sampling_params.unique_token_ids_lens.defined(); + if (use_stable_greedy_validate_sampling) { + if (!mtp_validate_greedy_indices_.defined() || + mtp_validate_greedy_indices_.numel() != total_num_val_tokens) { + mtp_validate_greedy_indices_ = torch::arange( + total_num_val_tokens, + torch::TensorOptions().dtype(torch::kInt).device(device_)); + mtp_validate_greedy_do_sample_ = torch::zeros( + {total_num_val_tokens}, + torch::TensorOptions().dtype(torch::kBool).device(device_)); + } + validate_sampling_params.selected_token_idxes = + mtp_validate_greedy_indices_; + validate_sampling_params.sample_idxes = mtp_validate_greedy_indices_; + validate_sampling_params.do_sample = mtp_validate_greedy_do_sample_; + validate_sampling_params.all_random_sample = false; + validate_sampling_params.all_greedy_sample = true; + } else { + update_sampling_params( + validate_sampling_params, num_val_tokens, total_num_val_tokens); + } +#else update_sampling_params( - validate_input.sampling_params, num_val_tokens, total_num_val_tokens); + validate_sampling_params, num_val_tokens, total_num_val_tokens); +#endif for (int32_t& token_num : input_params.parallel.dp_global_token_nums) { token_num *= num_val_tokens; } + std::vector accepted_prefix_lengths; if (use_chunked_prefill_spec_verify_path()) { input_params.embedding.input_embedding = torch::Tensor(); input_params.is_spec_verify = true; @@ -1586,28 +1739,159 @@ void MTPWorkerImpl::prepare_validate_inputs(const ForwardInput& input, } input_params.attention.host.q_cu_seq_lens = std::move(q_cu_seq_lens_vec); } - std::vector accepted_prefix_lengths(num_sequences, 1); + accepted_prefix_lengths.assign(num_sequences, 1); if (embedding_cache_ != nullptr && !input.input_params.embedding.embedding_ids.empty()) { accepted_prefix_lengths = embedding_cache_->read_accepted_prefix_lengths( input.input_params.embedding.embedding_ids); } - input_params.num_accepted_tokens = - torch::tensor(accepted_prefix_lengths, token_options); input_params.num_accepted_tokens_host.assign( accepted_prefix_lengths.begin(), accepted_prefix_lengths.end()); + if (!use_graph_internal_verify_update) { + input_params.num_accepted_tokens = + torch::tensor(accepted_prefix_lengths, token_options); + } } - input_params.attention.rebuild_device_buffer(device_); #if defined(USE_NPU) - if (use_qwen3_5_spec_verify_path()) { - build_expanded_spec_verify_graph_input(input_params, device_); + if (use_graph_internal_verify_update) { + build_expanded_spec_verify_graph_host_input(input_params); + + auto& attention = input_params.attention; + CHECK(attention.host.block_tables.defined()); + CHECK_EQ(attention.host.block_tables.dim(), 2); + CHECK_EQ(attention.host.block_tables.size(0), 1); + CHECK_EQ(attention.host.block_tables.scalar_type(), torch::kInt32); + const int64_t active_block_table_width = + attention.host.block_tables.size(1); + int64_t verify_block_table_width = 64; + while (verify_block_table_width < active_block_table_width) { + verify_block_table_width *= 2; + } + if (active_block_table_width != verify_block_table_width) { + torch::Tensor padded_block_tables = torch::zeros( + {1, verify_block_table_width}, + attention.host.block_tables.options().device(torch::kCPU)); + padded_block_tables.narrow(1, 0, active_block_table_width) + .copy_(attention.host.block_tables); + attention.host.block_tables = std::move(padded_block_tables); + } + + if (spec_verify_attention_host_buffer_.defined()) { + attention.attention_host_buffer = spec_verify_attention_host_buffer_; + attention.attention_device_buffer = spec_verify_attention_device_buffer_; + attention.attention_buffer_capacity = + spec_verify_attention_buffer_capacity_; + } + // The stable buffers are owned by MTPWorkerImpl; this temporary owner must + // remain unique so rebuild_device_buffer reuses rather than detaches them. + attention.attention_buffer_owner = std::make_shared(0); + + const int64_t expanded_block_rows = static_cast( + input_params.graph.expanded_kv_seq_lens_vec.size()); + CHECK_EQ(expanded_block_rows, num_val_tokens); + const torch::Tensor dense_block_source = + attention.host.block_tables.contiguous(); + std::vector expanded_block_tables_dense( + static_cast(expanded_block_rows * verify_block_table_width)); + const int32_t* block_row = dense_block_source.data_ptr(); + for (int64_t row = 0; row < expanded_block_rows; ++row) { + std::memcpy( + expanded_block_tables_dense.data() + row * verify_block_table_width, + block_row, + static_cast(verify_block_table_width) * sizeof(int32_t)); + } + torch::Tensor expanded_block_tables_flat; + std::vector extra_int_inputs; + if (!input_params.embedding.linear_state_ids.empty()) { + extra_int_inputs.push_back( + {&input_params.embedding.linear_state_ids, + nullptr, + &input_params.embedding.linear_state_indices}); + } + if (!input_params.num_accepted_tokens_host.empty()) { + extra_int_inputs.push_back({&accepted_prefix_lengths, + nullptr, + &input_params.num_accepted_tokens}); + } + if (!input_params.graph.expanded_kv_seq_lens_vec.empty()) { + extra_int_inputs.push_back({&input_params.graph.expanded_kv_seq_lens_vec, + nullptr, + &input_params.graph.expanded_kv_seq_lens}); + } + extra_int_inputs.push_back( + {&expanded_block_tables_dense, nullptr, &expanded_block_tables_flat}); + attention.rebuild_device_buffer(device_, extra_int_inputs); + spec_verify_attention_host_buffer_ = attention.attention_host_buffer; + spec_verify_attention_device_buffer_ = attention.attention_device_buffer; + spec_verify_attention_buffer_capacity_ = + attention.attention_buffer_capacity; + input_params.graph.expanded_block_tables = expanded_block_tables_flat.view( + {expanded_block_rows, verify_block_table_width}); + input_params.graph.spec_verify_source_addresses_stable = true; + } else { + input_params.attention.rebuild_device_buffer(device_); + if (supports_spec_verify_graph_input_update()) { + build_expanded_spec_verify_graph_input(input_params, device_); + } } +#else + input_params.attention.rebuild_device_buffer(device_); #endif validate_input.device_tensors_ready = true; +#if defined(USE_NPU) + input_params.graph.spec_verify_static_graph_tasks_prepared = false; + if (use_graph_internal_verify_update && static_graph_tasks_prepared) { + input_params.graph.spec_verify_static_graph_tasks_prepared = true; + } else if (use_graph_internal_verify_update && + impl_->prepare_static_graph_tasks(validate_input, + *compute_stream_)) { + input_params.graph.spec_verify_static_graph_tasks_prepared = true; + } +#endif finish_metadata_prepare(*prepare_stream_, validate_input); } +bool MTPWorkerImpl::prepare_static_mtp_graph_tasks_before_final_draft( + const ForwardInput& input) { +#if defined(USE_NPU) + if (!enable_spec_verify_graph_update() || + !supports_spec_verify_graph_input_update() || + options_.num_speculative_tokens() < 3 || + options_.num_speculative_tokens() > 5 || + input.input_params.meta.num_sequences != 1 || + input.input_params.embedding.linear_state_ids.size() != 1 || + embedding_cache_ == nullptr || + input.input_params.embedding.embedding_ids.empty()) { + return false; + } + const auto& block_tables = input.input_params.attention.host.block_tables; + if (!block_tables.defined() || block_tables.dim() != 2 || + block_tables.size(0) != 1) { + return false; + } + const std::vector accepted_prefix_lengths = + embedding_cache_->read_accepted_prefix_lengths( + input.input_params.embedding.embedding_ids); + if (accepted_prefix_lengths.size() != 1) { + return false; + } + int64_t verify_block_table_width = 64; + while (verify_block_table_width < block_tables.size(1)) { + verify_block_table_width *= 2; + } + return impl_->prepare_static_mtp_graph_tasks( + input.input_params.embedding.linear_state_ids.front(), + accepted_prefix_lengths.front(), + options_.num_speculative_tokens() + 1, + verify_block_table_width, + *compute_stream_); +#else + (void)input; + return false; +#endif +} + void MTPWorkerImpl::prepare_draft_extend_inputs( const ForwardInput& base_input, const std::vector& last_states, @@ -1730,11 +2014,9 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( CHECK_EQ(expanded_embeddings.size(), buf.out_positions.size()) << "draft extend embeddings/positions mismatch"; - specBuilder::set_token_position_tensors(extend_input, - buf.out_token_ids, - buf.out_positions, - token_options, - position_options); + CHECK(token_options.dtype_opt() == position_options.dtype_opt()); + CHECK(token_options.device_opt() == position_options.device_opt()); + extend_input.device_tensors_ready = false; if (use_chunked_prefill) { input_params.meta.num_sequences = num_sequences; input_params.meta.batch_forward_type = BatchForwardType::CHUNKED_PREFILL; @@ -1766,7 +2048,7 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( std::move(buf.out_kv_seq_lens), /*update_block_tables=*/true); } - if (use_qwen3_5_spec_verify_path()) { + if (supports_spec_verify_graph_input_update()) { input_params.attention.host.q_cu_seq_lens.clear(); input_params.attention.host.q_cu_seq_lens.reserve( input_params.meta.num_sequences + 1); @@ -1777,7 +2059,16 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( input_params.get_q_seq_len(i)); } } - input_params.attention.rebuild_device_buffer(device_); + const std::vector extra_int_inputs = { + {&buf.out_token_ids, + &extend_input.token_ids_host, + &extend_input.token_ids}, + {&buf.out_positions, + &extend_input.positions_host, + &extend_input.positions}, + }; + input_params.attention.rebuild_device_buffer(device_, extra_int_inputs); + extend_input.device_tensors_ready = true; // Establish cross-stream ordering before stacking the draft-extend // embeddings. The stack runs on prepare_stream_, but the embedding rows it @@ -1785,7 +2076,10 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( // prepare_stream_ wait on compute_stream_ via an event. prepare_stream_->wait_stream(*compute_stream_); - input_params.embedding.input_embedding = torch::stack(expanded_embeddings); + input_params.embedding.input_embedding = + expanded_embeddings.size() == 1 + ? expanded_embeddings.front().unsqueeze(/*dim=*/0) + : torch::stack(expanded_embeddings); if (!input_params.parallel.dp_global_token_nums.empty()) { if (use_chunked_prefill) { @@ -1808,10 +2102,26 @@ 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 (num_sequences == 1 && + (selected_row_idx[0] == 0 || selected_row_idx[0] == 1)) { + torch::Tensor& persistent_idx = selected_row_idx[0] == 0 + ? draft_selected_row_zero_ + : draft_selected_row_one_; + if (!persistent_idx.defined() || + persistent_idx.scalar_type() != idx_options.dtype_opt().value() || + persistent_idx.device() != idx_options.device_opt().value()) { + persistent_idx = + safe_to(specBuilder::make_cpu_int_tensor({selected_row_idx[0]}), + idx_options, + /*non_blocking=*/true); + } + params.selected_token_idxes = persistent_idx; + } 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()) { std::vector sample_idxes_vec; sample_idxes_vec.reserve(num_sequences); @@ -1869,7 +2179,7 @@ void MTPWorkerImpl::prepare_draft_inputs(const ForwardInput& input, std::move(input_params.attention.host.q_cu_seq_lens), buf.meta.kv_max_seq_len, std::move(buf.out_kv_seq_lens)); - if (use_qwen3_5_spec_verify_path()) { + if (supports_spec_verify_graph_input_update()) { input_params.attention.host.q_cu_seq_lens.clear(); input_params.attention.host.q_cu_seq_lens.reserve( input_params.meta.num_sequences + 1); diff --git a/xllm/core/runtime/mtp_worker_impl.h b/xllm/core/runtime/mtp_worker_impl.h index e961019cee..a6b5cb12d2 100644 --- a/xllm/core/runtime/mtp_worker_impl.h +++ b/xllm/core/runtime/mtp_worker_impl.h @@ -15,6 +15,10 @@ limitations under the License. #pragma once +#include +#include +#include + #include "framework/kv_cache/embedding_cache.h" #include "framework/kv_cache_transfer/kv_cache_transfer.h" #if defined(USE_NPU) @@ -103,15 +107,18 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { // prepare inputs for draft model at Prefill phase. void prepare_prefill_inputs(const ForwardInput& inputs, ForwardInput& prefill_inputs); - bool use_qwen3_5_spec_verify_path() const; - bool use_mimo_spec_verify_path() const; + SpeculativeVerifyCapabilities speculative_verify_capabilities() const; + bool supports_spec_verify_graph_input_update() const; // Returns true when validation must use chunked-prefill to avoid the // FlashInfer batch-decode read-before-write race on the bonus token. bool use_chunked_prefill_spec_verify_path() const; // Prepare target validate input from cached target context. void prepare_validate_inputs(const ForwardInput& inputs, - ForwardInput& validate_inputs); + ForwardInput& validate_inputs, + bool static_graph_tasks_prepared = false); + bool prepare_static_mtp_graph_tasks_before_final_draft( + const ForwardInput& input); // prepare inputs for draft model at Decode phase. void prepare_draft_inputs(const ForwardInput& inputs, @@ -141,6 +148,34 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { // If false, selected-only cache values are restored to dense [B, S, V]. bool enable_opt_validate_probs_ = false; + // Cached once when the target model is loaded. Decode-path decisions must + // not repeatedly traverse the model's virtual capability interface. + SpeculativeVerifyCapabilities target_spec_verify_capabilities_; + + // Immutable single-request draft row selectors. The steady workload + // alternates between the legal one-row (index 0) and two-row (index 1) + // layouts; keeping both indices on device avoids rebuilding/H2D each cycle. + torch::Tensor draft_selected_row_zero_; + torch::Tensor draft_selected_row_one_; +#if defined(USE_NPU) + // Stable-address sources consumed by the target ACL graph's leading input + // update. The existing H2D preparation overlaps with the final draft, so no + // extra graph-external D2D launch is introduced. + torch::Tensor spec_verify_tokens_host_; + torch::Tensor spec_verify_tokens_device_; + torch::Tensor spec_verify_positions_host_; + torch::Tensor spec_verify_positions_device_; + torch::Tensor spec_verify_attention_host_buffer_; + torch::Tensor spec_verify_attention_device_buffer_; + uint64_t spec_verify_attention_buffer_capacity_ = 0; + + // Stable validate-sampling controls for the common single-sequence greedy + // path. Their values depend on speculative width, not tensor-parallel + // topology, and are rebuilt only when that width changes. + torch::Tensor mtp_validate_greedy_indices_; + torch::Tensor mtp_validate_greedy_do_sample_; +#endif + #if defined(USE_NPU) || defined(USE_MLU) std::shared_ptr kv_cache_transfer_; #endif diff --git a/xllm/models/llm/mimo.h b/xllm/models/llm/mimo.h index ab835e1db1..51325a9838 100644 --- a/xllm/models/llm/mimo.h +++ b/xllm/models/llm/mimo.h @@ -15,6 +15,7 @@ limitations under the License. #pragma once +#include "core/framework/model/speculative_verify_capabilities.h" #include "core/layers/qwen2_decoder_layer.h" #include "llm_model_base.h" @@ -81,6 +82,14 @@ class MiMoForCausalLMImpl final : public LlmForCausalLMImplBase { public: explicit MiMoForCausalLMImpl(const ModelContext& context) : LlmForCausalLMImplBase(context) {} + + SpeculativeVerifyCapabilities speculative_verify_capabilities() const { + // MiMo also needs causal chunked prefill for correct bonus-token ordering, + // but it does not expose the hybrid linear-state layout consumed by the + // current NPU graph input updater. + return {/*requires_causal_chunked_prefill=*/true, + /*supports_in_graph_input_update=*/false}; + } }; TORCH_MODULE(MiMoForCausalLM); diff --git a/xllm/models/llm/qwen3_next_hybrid_base.h b/xllm/models/llm/qwen3_next_hybrid_base.h index be8b0a5e9a..8c5c78898b 100644 --- a/xllm/models/llm/qwen3_next_hybrid_base.h +++ b/xllm/models/llm/qwen3_next_hybrid_base.h @@ -25,6 +25,7 @@ limitations under the License. #include "core/framework/kv_cache/kv_cache.h" #include "core/framework/model/model_input_params.h" #include "core/framework/model/model_output.h" +#include "core/framework/model/speculative_verify_capabilities.h" #include "core/framework/model_context.h" #include "core/framework/model_loader.h" #include "core/layers/common/attention_mask.h" @@ -33,6 +34,9 @@ limitations under the License. #include "core/layers/common/qwen3_next_rms_norm.h" #include "core/layers/common/word_embedding.h" #include "core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h" +#if defined(USE_NPU) +#include "core/kernels/npu/tilelang/tilelang_ops_api.h" +#endif namespace xllm { @@ -97,6 +101,18 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { input_params, model_args_.enable_mla(), build_attention_mask(input_params)); +#if defined(USE_NPU) + if (attn_metadata.use_expanded_decode_for_spec_verify_attention && + attn_metadata.expanded_kv_seq_lens.defined() && + attn_metadata.expanded_kv_seq_lens.numel() >= 4 && + attn_metadata.expanded_kv_seq_lens.numel() <= 6 && + attn_metadata.expanded_paged_attention_tiling_data.defined()) { + kernel::npu::tilelang::spec_verify_attention_tiling_update( + attn_metadata.expanded_kv_seq_lens, + attn_metadata.expanded_paged_attention_tiling_data, + /*block_size=*/128); + } +#endif torch::Tensor h; if (input_params.embedding.input_embedding.defined()) { h = input_params.embedding.input_embedding; @@ -309,6 +325,14 @@ class Qwen3HybridForCausalLMImplBase : public torch::nn::Module { bool is_hybrid_linear_attention() { return true; } + SpeculativeVerifyCapabilities speculative_verify_capabilities() const { + // Hybrid linear-attention verification uses causal chunked prefill and + // exposes the expanded attention plus linear-state layout required by the + // NPU graph input updater. + return {/*requires_causal_chunked_prefill=*/true, + /*supports_in_graph_input_update=*/true}; + } + layer::LmHead get_lm_head() { return lm_head_; } void set_lm_head(layer::LmHead& head) { lm_head_ = head; }