diff --git a/.github/workflows/build_x86_64_npu.yaml b/.github/workflows/build_x86_64_npu.yaml index c724163450..6a1c8edf9a 100644 --- a/.github/workflows/build_x86_64_npu.yaml +++ b/.github/workflows/build_x86_64_npu.yaml @@ -130,4 +130,4 @@ jobs: timeout-minutes: 60 run: | chmod +x ./cibuild/build_npu.sh - bash cibuild/build_npu.sh 'pip install pre-commit -i https://pypi.tuna.tsinghua.edu.cn/simple; python setup.py bdist_wheel; pip install dist/* --force-reinstall; cd /tmp && python /export/home/actions-runner/_work/xllm/xllm/examples/generate.py --model="/export/home/models/Qwen2-7B-Instruct" --devices="npu:7"' + bash cibuild/build_npu.sh 'pip install pre-commit -i https://pypi.tuna.tsinghua.edu.cn/simple; python setup.py bdist_wheel; pip install dist/* --force-reinstall; cd /tmp && ASCEND_RT_VISIBLE_DEVICES=7 python /export/home/actions-runner/_work/xllm/xllm/examples/generate.py --model="/export/home/models/Qwen2-7B-Instruct"' diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c39642ace7..966c00641e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -16,6 +16,15 @@ if(USE_NPU) ) endif() +add_test( + NAME xllm_export_import_test + COMMAND + "${Python3_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/python/import_xllm_export.py" + "$" +) +add_dependencies(all_tests xllm_export) + add_subdirectory(api_service) add_subdirectory(core) add_subdirectory(function_call) diff --git a/tests/core/common/options_test.cpp b/tests/core/common/options_test.cpp index 967b3b47f1..b67b6ea63e 100644 --- a/tests/core/common/options_test.cpp +++ b/tests/core/common/options_test.cpp @@ -31,6 +31,7 @@ TEST(OptionsTest, ContextParallelDefaultsToOneAcrossPublicApis) { const XLLM_InitLLMOptions cc_options; EXPECT_EQ(options.cp_size(), 1); + EXPECT_EQ(options.decode_context_parallel_size(), 1); EXPECT_EQ(cc_options.cp_size, 1); EXPECT_EQ(XLLM_INIT_LLM_OPTIONS_DEFAULT.cp_size, 1U); EXPECT_EQ(XLLM_C_ABI_VERSION_MAJOR, 1); @@ -40,12 +41,14 @@ TEST(OptionsTest, ContextParallelDefaultsToOneAcrossPublicApis) { TEST(OptionsTest, ContextParallelAcceptsExplicitValuesAcrossPublicApis) { Options options; options.cp_size(4); + options.decode_context_parallel_size(2); XLLM_InitLLMOptions cc_options; cc_options.cp_size = 4; XLLM_InitOptions c_options = XLLM_INIT_LLM_OPTIONS_DEFAULT; c_options.cp_size = 4; EXPECT_EQ(options.cp_size(), 4); + EXPECT_EQ(options.decode_context_parallel_size(), 2); EXPECT_EQ(cc_options.cp_size, 4); EXPECT_EQ(c_options.cp_size, 4U); } diff --git a/tests/core/distributed_runtime/CMakeLists.txt b/tests/core/distributed_runtime/CMakeLists.txt index 57f629527a..932ff0d007 100644 --- a/tests/core/distributed_runtime/CMakeLists.txt +++ b/tests/core/distributed_runtime/CMakeLists.txt @@ -11,3 +11,13 @@ cc_test( DEPS GTest::gtest_main ) + +cc_test( + NAME + dcp_compat_test + SRCS + dcp_compat_test.cpp + DEPS + common + GTest::gtest_main +) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp new file mode 100644 index 0000000000..8cf517cfca --- /dev/null +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -0,0 +1,219 @@ +/* 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://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "core/distributed_runtime/dcp_compat.h" + +#include + +#include +#include + +namespace xllm { +namespace { + +Options dcp_options_with_supported_feature_flags() { + Options options; + options.decode_context_parallel_size(2) + .enable_chunked_prefill(false) + .enable_experimental_dcp_chunked_prefill(false) + .enable_prefix_cache(false) + .enable_schedule_overlap(false) + .enable_disagg_pd(false) + .instance_role(InstanceRole::DEFAULT) + .num_speculative_tokens(0); + return options; +} + +Options dcp_options_with_experimental_chunked_prefill() { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_chunked_prefill(true).enable_experimental_dcp_chunked_prefill( + true); + return options; +} + +void expect_error_contains(const std::optional& error, + const std::string& expected) { + ASSERT_TRUE(error.has_value()); + EXPECT_NE(error->find(expected), std::string::npos) << error.value(); +} + +TEST(DcpCompatTest, DcpOneDoesNotRejectDefaultOptions) { + Options options; + options.decode_context_parallel_size(1); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, DcpOneExperimentalFlagHasNoEffect) { + Options options; + options.decode_context_parallel_size(1) + .enable_chunked_prefill(true) + .enable_experimental_dcp_chunked_prefill(true) + .priority_strategy("multi_slo_and_prio"); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, DcpChunkedPrefillFalseWithExperimentalMatchesBaseline) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_experimental_dcp_chunked_prefill(true); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, DcpMultiSloWithoutExperimentalRejected) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_chunked_prefill(false).priority_strategy("multi_slo_and_prio"); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_experimental_dcp_chunked_prefill=true"); +} + +TEST(DcpCompatTest, DcpMultiSloWithExperimentalAllowed) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_chunked_prefill(false) + .priority_strategy("multi_slo_and_prio") + .enable_experimental_dcp_chunked_prefill(true); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, AllowsSupportedFirstVersionFeatureFlags) { + const Options options = dcp_options_with_supported_feature_flags(); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, DcpOneChunkedPrefillWithoutExperimentalAllowed) { + Options options = dcp_options_with_supported_feature_flags(); + options.decode_context_parallel_size(1).enable_chunked_prefill(true); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, DcpChunkedPrefillWithoutExperimentalRejected) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_chunked_prefill(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_experimental_dcp_chunked_prefill=true"); +} + +TEST(DcpCompatTest, DcpChunkedPrefillWithExperimentalAllowed) { + const Options options = dcp_options_with_experimental_chunked_prefill(); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, DcpChunkedPrefillWithPrefixAndExperimentalAllowed) { + Options options = dcp_options_with_experimental_chunked_prefill(); + options.enable_prefix_cache(true); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, AllowsPrefixCache) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_prefix_cache(true); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, AllowsScheduleOverlap) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_schedule_overlap(true); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, RejectsDisaggregatedPrefillDecodeFlag) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_disagg_pd(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_disagg_pd=false"); +} + +TEST(DcpCompatTest, RejectsDisaggregatedPrefillDecodeRole) { + Options options = dcp_options_with_supported_feature_flags(); + options.instance_role(InstanceRole::DECODE); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "instance_role=DEFAULT"); +} + +TEST(DcpCompatTest, RejectsSpeculativeEngineType) { + const Options options = dcp_options_with_supported_feature_flags(); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::SSM), + "speculative decoding"); +} + +TEST(DcpCompatTest, RejectsDraftModelPath) { + Options options = dcp_options_with_supported_feature_flags(); + options.draft_model_path("/tmp/draft-model"); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "draft_model"); +} + +TEST(DcpCompatTest, RejectsSpeculativeTokens) { + Options options = dcp_options_with_supported_feature_flags(); + options.num_speculative_tokens(1); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "num_speculative_tokens=0"); +} + +// Enabling the experimental chunked prefill opt-in must not bypass the other +// first-version rejections: disaggregated PD and speculative decoding are still +// unsupported even under the experimental path. (Schedule overlap is now +// supported for eager DCP and is asserted by AllowsScheduleOverlap.) +TEST(DcpCompatTest, ExperimentalChunkedPrefillDoesNotBypassDisaggPd) { + Options options = dcp_options_with_experimental_chunked_prefill(); + options.enable_disagg_pd(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_disagg_pd=false"); +} + +TEST(DcpCompatTest, ExperimentalChunkedPrefillDoesNotBypassSpeculative) { + const Options options = dcp_options_with_experimental_chunked_prefill(); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::SSM), + "speculative decoding"); +} + +} // namespace +} // namespace xllm diff --git a/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp b/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp index 3d86be63b7..7dbcc3c0e6 100644 --- a/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp +++ b/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp @@ -88,5 +88,10 @@ TEST(SpawnWorkerProtocolTest, PreservesExplicitEmptyDtype) { EXPECT_TRUE(indexer_cache_dtype->empty()); } +TEST(SpawnWorkerProtocolTest, AppendsDecodeContextParallelSizeAtTail) { + EXPECT_EQ(kDraftSamplingModeArgumentIndex, kArgumentCount - 2); + EXPECT_EQ(kDecodeContextParallelSizeArgumentIndex, kArgumentCount - 1); +} + } // namespace } // namespace xllm::spawn_worker_protocol diff --git a/tests/core/framework/config/config_json_test.cpp b/tests/core/framework/config/config_json_test.cpp index 3aba034951..8a3b9fc95c 100644 --- a/tests/core/framework/config/config_json_test.cpp +++ b/tests/core/framework/config/config_json_test.cpp @@ -96,11 +96,17 @@ class DumpConfigJsonFlagGuard final { class CpSizeFlagGuard final { public: - CpSizeFlagGuard() : old_cp_size_(FLAGS_cp_size) {} - ~CpSizeFlagGuard() { FLAGS_cp_size = old_cp_size_; } + CpSizeFlagGuard() + : old_cp_size_(FLAGS_cp_size), + old_decode_context_parallel_size_(FLAGS_decode_context_parallel_size) {} + ~CpSizeFlagGuard() { + FLAGS_cp_size = old_cp_size_; + FLAGS_decode_context_parallel_size = old_decode_context_parallel_size_; + } private: int32_t old_cp_size_; + int32_t old_decode_context_parallel_size_; }; class ConfigFlagGuard final { @@ -322,24 +328,59 @@ TEST(KVCacheConfigValidationTest, RejectsUnsupportedIndexerCacheDtypes) { TEST(ConfigJsonTest, ParallelConfigReadsContextParallelSize) { CpSizeFlagGuard flag_guard; - const JsonReader json = - config::parse_json_string(R"json({"cp_size": 4})json"); + const JsonReader json = config::parse_json_string( + R"json({"cp_size": 4, "decode_context_parallel_size": 2})json"); ParallelConfig parallel_config; parallel_config.from_json(json); EXPECT_EQ(parallel_config.cp_size(), 4); + EXPECT_EQ(parallel_config.decode_context_parallel_size(), 2); } -TEST(ConfigJsonTest, RegistersOnlyContextParallelCommandLineOption) { +TEST(ConfigJsonTest, RegistersContextParallelCommandLineOptions) { google::CommandLineFlagInfo flag_info; EXPECT_TRUE(google::GetCommandLineFlagInfo("cp_size", &flag_info)); EXPECT_EQ(flag_info.default_value, "1"); + EXPECT_TRUE(google::GetCommandLineFlagInfo("decode_context_parallel_size", + &flag_info)); + EXPECT_EQ(flag_info.default_value, "1"); + EXPECT_TRUE(google::GetCommandLineFlagInfo( + "enable_experimental_dcp_chunked_prefill", &flag_info)); + EXPECT_EQ(flag_info.default_value, "false"); const std::string removed_flag = std::string("enable_") + "prefill_sp"; EXPECT_FALSE( google::GetCommandLineFlagInfo(removed_flag.c_str(), &flag_info)); } +TEST(ConfigJsonTest, ParallelConfigReadsExperimentalDcpChunkedPrefill) { + google::FlagSaver flag_saver; + const JsonReader json = config::parse_json_string( + R"json({"enable_experimental_dcp_chunked_prefill": true})json"); + ParallelConfig parallel_config; + parallel_config.from_json(json); + + EXPECT_TRUE(parallel_config.enable_experimental_dcp_chunked_prefill()); +} + +TEST(ConfigJsonTest, ParallelConfigDefaultsExperimentalDcpChunkedPrefillFalse) { + const ParallelConfig parallel_config; + + EXPECT_FALSE(parallel_config.enable_experimental_dcp_chunked_prefill()); +} + +TEST(ConfigJsonTest, ParallelConfigDumpsExperimentalDcpChunkedPrefill) { + ParallelConfig parallel_config; + parallel_config.enable_experimental_dcp_chunked_prefill(true); + + nlohmann::ordered_json config_json; + parallel_config.append_config_json(config_json); + + ASSERT_TRUE(config_json.contains("enable_experimental_dcp_chunked_prefill")); + EXPECT_TRUE( + config_json.at("enable_experimental_dcp_chunked_prefill").get()); +} + TEST(ConfigJsonTest, LoadJsonFileReadsConfigFixture) { // The fixture sets more keys than ConfigFlagGuard restores, and from_json // writes every resolved value back into its FLAGS_ global. FlagSaver reverts diff --git a/tests/core/framework/parallel_state/CMakeLists.txt b/tests/core/framework/parallel_state/CMakeLists.txt index 82aeba8b62..cc40ddfb9d 100644 --- a/tests/core/framework/parallel_state/CMakeLists.txt +++ b/tests/core/framework/parallel_state/CMakeLists.txt @@ -44,7 +44,10 @@ if(USE_NPU) SRCS cp_group_ranks_test.cpp DEPS + block + dcp_attention_utils parallel_state + torch GTest::gtest_main ) diff --git a/tests/core/framework/parallel_state/cp_group_ranks_test.cpp b/tests/core/framework/parallel_state/cp_group_ranks_test.cpp index e9250dbb64..7cb6a183e6 100644 --- a/tests/core/framework/parallel_state/cp_group_ranks_test.cpp +++ b/tests/core/framework/parallel_state/cp_group_ranks_test.cpp @@ -13,17 +13,54 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include +#include #include #include +#include "core/framework/multimodal/mm_data.h" +#include "core/framework/sampling/sampling_params.h" +#include "framework/block/block_manager_pool.h" #include "framework/parallel_state/parallel_state.h" +#include "framework/request/incremental_decoder.h" +#include "framework/request/sequence.h" +#include "framework/request/stopping_checker.h" +#include "layers/npu_torch/dcp_attention_utils.h" namespace xllm { namespace parallel_state { namespace { +// Expected DCP owner/local-slot mapping, used to cross-check the production +// tensor path (select_dcp_local_block_table / remap_dcp_cache_slots). Owner +// formula mirrors production: floor(position / interleave) % dcp_size. +int64_t expected_dcp_cache_slot(int64_t logical_slot, + int64_t position, + int32_t block_size, + int32_t dcp_size, + int32_t dcp_rank, + int32_t interleave_size) { + if (logical_slot < 0) { + return -1; + } + CHECK_GE(position, 0) << "position must be non-negative."; + CHECK_GT(block_size, 0) << "block_size must be positive."; + CHECK_GT(dcp_size, 1) << "dcp_size must be greater than 1."; + CHECK_GE(dcp_rank, 0) << "dcp_rank must be non-negative."; + CHECK_LT(dcp_rank, dcp_size) << "dcp_rank must be smaller than dcp_size."; + CHECK_GT(interleave_size, 0) << "interleave_size must be positive."; + CHECK_EQ(interleave_size, block_size) + << "DCP local block-table selection requires block interleave."; + + const int64_t owner = (position / block_size) % dcp_size; + if (owner != dcp_rank) { + return -1; + } + return logical_slot; +} + // Re-derive the CP rank of a global rank from the documented layout: // rank = dp_rank * (cp_size * attn_tp_size) + cp_rank * attn_tp_size + // tp_rank @@ -36,6 +73,53 @@ int32_t expected_cp_rank(int32_t global_rank, return (global_rank % (cp_size * attn_tp_size)) / attn_tp_size; } +int32_t expected_dcp_rank(int32_t global_rank, + int32_t world_size, + int32_t dp_size, + int32_t dcp_size) { + const int32_t tp_size = world_size / dp_size; + return (global_rank % tp_size) % dcp_size; +} + +class PrefixTestSequence final { + public: + PrefixTestSequence(size_t index, const std::vector& prompt_tokens) { + sampling_param_.beam_width = 0; + sampling_param_.is_embeddings = false; + + SequenceParams params; + params.seq_capacity = prompt_tokens.size() + 8; + params.echo = false; + params.skip_special_tokens = true; + params.streaming = false; + params.enable_schedule_overlap = false; + params.rec_type = RecType::kNone; + params.bos_token_id = 0; + params.request_id = "dcp_prefix_contract_test"; + params.sampling_param = &sampling_param_; + params.stopping_checker = &stopping_checker_; + + IncrementalDecoder decoder( + /*prompt=*/"prompt", + /*num_prompt_tokens=*/prompt_tokens.size(), + /*echo=*/params.echo, + /*skip_special_tokens=*/params.skip_special_tokens); + sequence_ = std::make_unique(index, + prompt_tokens, + /*input_embedding=*/torch::Tensor(), + /*mm_data=*/MMData(), + decoder, + params); + } + + Sequence* get() { return sequence_.get(); } + + private: + RequestSamplingParam sampling_param_; + StoppingChecker stopping_checker_; + std::unique_ptr sequence_; +}; + TEST(ComputeCpGroupRanks, CpSizeTwoTpFourDpOne) { const int32_t world_size = 8; const int32_t dp_size = 1; @@ -139,6 +223,330 @@ TEST(ComputeCpGroupRanks, RejectsNonIntegralAttnTpSize) { ""); } +TEST(ComputeDcpGroupRanks, DcpSizeTwoTpEightDpOne) { + const int32_t world_size = 8; + const int32_t dp_size = 1; + const int32_t dcp_size = 2; + for (int32_t rank = 0; rank < world_size; ++rank) { + const std::vector ranks = + compute_dcp_group_ranks(rank, world_size, dp_size, dcp_size); + ASSERT_EQ(ranks.size(), dcp_size); + EXPECT_EQ(ranks[expected_dcp_rank(rank, world_size, dp_size, dcp_size)], + rank); + + const int32_t tp_rank = rank % (world_size / dp_size); + const int32_t expected_base = (tp_rank / dcp_size) * dcp_size; + for (int32_t dcp_rank = 0; dcp_rank < dcp_size; ++dcp_rank) { + EXPECT_EQ(ranks[dcp_rank], expected_base + dcp_rank); + } + } +} + +TEST(ComputeDcpGroupRanks, DcpSizeTwoTpFourDpTwo) { + const int32_t world_size = 8; + const int32_t dp_size = 2; + const int32_t dcp_size = 2; + const int32_t tp_size = world_size / dp_size; + for (int32_t rank = 0; rank < world_size; ++rank) { + const std::vector ranks = + compute_dcp_group_ranks(rank, world_size, dp_size, dcp_size); + ASSERT_EQ(ranks.size(), dcp_size); + EXPECT_EQ(ranks[expected_dcp_rank(rank, world_size, dp_size, dcp_size)], + rank); + + const int32_t dp_rank = rank / tp_size; + const int32_t dcp_group_base = ((rank % tp_size) / dcp_size) * dcp_size; + for (int32_t member : ranks) { + EXPECT_EQ(member / tp_size, dp_rank); + EXPECT_GE(member % tp_size, dcp_group_base); + EXPECT_LT(member % tp_size, dcp_group_base + dcp_size); + } + } +} + +TEST(ComputeDcpGroupRanks, DocumentsContinuousGroupCounterexample) { + const std::vector ranks = compute_dcp_group_ranks( + /*global_rank=*/2, /*world_size=*/12, /*dp_size=*/1, /*dcp_size=*/2); + ASSERT_EQ(ranks.size(), 2); + EXPECT_EQ(ranks[0], 2); + EXPECT_EQ(ranks[1], 3); +} + +TEST(ComputeDcpGroupRanks, RejectsNonIntegralDcpGroups) { + EXPECT_DEATH(compute_dcp_group_ranks(/*global_rank=*/0, + /*world_size=*/10, + /*dp_size=*/1, + /*dcp_size=*/4), + ""); +} + +TEST(SelectDcpLocalBlockTable, SelectsOriginalNonContiguousBlockIds) { + const torch::Tensor global_block_table = + torch::tensor({{37, 5, 89, 2}, {41, 13, 73, 29}}, + torch::TensorOptions().dtype(torch::kInt64)); + + const torch::Tensor rank_zero_table = select_dcp_local_block_table( + global_block_table, /*dcp_size=*/2, /*dcp_rank=*/0); + const torch::Tensor rank_one_table = select_dcp_local_block_table( + global_block_table, /*dcp_size=*/2, /*dcp_rank=*/1); + + EXPECT_TRUE( + torch::equal(rank_zero_table, + torch::tensor({{37, 89}, {41, 73}}, + torch::TensorOptions().dtype(torch::kInt64)))); + EXPECT_TRUE( + torch::equal(rank_one_table, + torch::tensor({{5, 2}, {13, 29}}, + torch::TensorOptions().dtype(torch::kInt64)))); +} + +TEST(SelectDcpLocalBlockTable, AllowsRankWithoutBlockColumns) { + const torch::Tensor global_block_table = + torch::tensor({{37}}, torch::TensorOptions().dtype(torch::kInt64)); + + const torch::Tensor local_block_table = select_dcp_local_block_table( + global_block_table, /*dcp_size=*/2, /*dcp_rank=*/1); + + EXPECT_EQ(local_block_table.dim(), 2); + EXPECT_EQ(local_block_table.size(0), 1); + EXPECT_EQ(local_block_table.size(1), 0); +} + +TEST(DcpCacheLayout, PrefillWritesMatchDecodeLocalBlockTable) { + const int32_t block_size = 4; + const int32_t dcp_size = 2; + const std::vector global_block_ids = {37, 5, 89, 2}; + const torch::Tensor global_block_table = torch::tensor( + {{37, 5, 89, 2}}, torch::TensorOptions().dtype(torch::kInt64)); + + for (int32_t dcp_rank = 0; dcp_rank < dcp_size; ++dcp_rank) { + const torch::Tensor local_block_table = + select_dcp_local_block_table(global_block_table, dcp_size, dcp_rank); + for (int32_t local_block_index = 0; + local_block_index < local_block_table.size(1); + ++local_block_index) { + const int32_t global_block_index = + dcp_rank + local_block_index * dcp_size; + const int64_t original_block_id = global_block_ids[global_block_index]; + const int64_t original_slot = + original_block_id * block_size + (block_size - 1); + const int64_t position = + static_cast(global_block_index) * block_size + + (block_size - 1); + const int64_t owner_slot = + expected_dcp_cache_slot(original_slot, + position, + block_size, + dcp_size, + dcp_rank, + /*interleave_size=*/block_size); + const int64_t decode_block_id = + local_block_table.index({0, local_block_index}).item(); + + EXPECT_EQ(owner_slot, original_slot); + EXPECT_EQ(owner_slot / block_size, decode_block_id); + } + } +} + +TEST(DcpPrefixCacheContract, + SharedBlocksSuffixWritesAndLocalLengthsStayOwnerAligned) { + const int32_t block_size = 4; + const int32_t dcp_size = 2; + BlockManagerPool::Options options; + options.num_blocks(16) + .host_num_blocks(0) + .block_size(block_size) + .enable_prefix_cache(true); + BlockManagerPool pool(options, /*dp_size=*/1); + + PrefixTestSequence seed( + /*index=*/0, + /*prompt_tokens=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + ASSERT_TRUE(pool.allocate(seed.get())); + seed.get()->kv_state().set_kv_cache_tokens_num(seed.get()->num_tokens()); + const Slice seed_blocks = seed.get()->kv_state().blocks(BlockType::KV); + ASSERT_EQ(seed_blocks.size(), 3); + std::vector seed_block_ids; + seed_block_ids.reserve(seed_blocks.size()); + for (const Block& block : seed_blocks) { + seed_block_ids.emplace_back(block.id()); + } + pool.cache(seed.get()); + pool.deallocate(seed.get()); + + PrefixTestSequence probe( + /*index=*/1, + /*prompt_tokens=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14}); + ASSERT_TRUE(pool.allocate(probe.get(), probe.get()->num_tokens())); + EXPECT_EQ(probe.get()->kv_state().shared_blocks_num(BlockType::KV), 3); + EXPECT_EQ(probe.get()->kv_state().kv_cache_tokens_num(), 12); + + const Slice probe_blocks = + probe.get()->kv_state().blocks(BlockType::KV); + ASSERT_EQ(probe_blocks.size(), 4); + for (int64_t block_index = 0; + block_index < static_cast(seed_block_ids.size()); + ++block_index) { + EXPECT_EQ(probe_blocks[block_index].id(), seed_block_ids[block_index]); + } + + std::vector global_block_ids; + global_block_ids.reserve(probe_blocks.size()); + for (const Block& block : probe_blocks) { + global_block_ids.emplace_back(block.id()); + } + const torch::Tensor global_block_table = + torch::tensor(global_block_ids, + torch::TensorOptions().dtype(torch::kInt64)) + .view({1, static_cast(global_block_ids.size())}); + + const int32_t suffix_start = + static_cast(probe.get()->kv_state().kv_cache_tokens_num()); + const int32_t suffix_end = static_cast(probe.get()->num_tokens()); + const std::vector suffix_slots = probe.get()->kv_state().cache_slots( + BlockType::KV, suffix_start, suffix_end); + const torch::Tensor suffix_positions = torch::arange( + suffix_start, suffix_end, torch::TensorOptions().dtype(torch::kInt64)); + const torch::Tensor suffix_slot_tensor = + torch::tensor(suffix_slots, torch::TensorOptions().dtype(torch::kInt64)); + + const std::vector full_slots = probe.get()->kv_state().cache_slots( + BlockType::KV, /*pos_start=*/0, suffix_end); + const torch::Tensor full_positions = + torch::arange(suffix_end, torch::TensorOptions().dtype(torch::kInt64)); + const torch::Tensor full_slot_tensor = + torch::tensor(full_slots, torch::TensorOptions().dtype(torch::kInt64)); + const std::vector expected_local_kv_seq_lens = {8, 6}; + + for (int32_t dcp_rank = 0; dcp_rank < dcp_size; ++dcp_rank) { + const torch::Tensor local_block_table = + select_dcp_local_block_table(global_block_table, dcp_size, dcp_rank); + int64_t local_block_index = 0; + for (int64_t global_block_index = dcp_rank; + global_block_index < static_cast(global_block_ids.size()); + global_block_index += dcp_size) { + ASSERT_LT(local_block_index, local_block_table.size(1)); + EXPECT_EQ(local_block_table[0][local_block_index].item(), + global_block_ids[global_block_index]); + ++local_block_index; + } + EXPECT_EQ(local_block_index, local_block_table.size(1)); + + const torch::Tensor remapped_suffix = + remap_dcp_cache_slots(suffix_positions, + suffix_slot_tensor, + /*interleave_size=*/block_size, + dcp_size, + dcp_rank); + for (int64_t suffix_index = 0; suffix_index < remapped_suffix.numel(); + ++suffix_index) { + const int64_t remapped_slot = + remapped_suffix[suffix_index].item(); + if (remapped_slot < 0) { + continue; + } + const int64_t position = suffix_positions[suffix_index].item(); + const int64_t global_block_index = position / block_size; + const int64_t local_block_index = global_block_index / dcp_size; + ASSERT_LT(local_block_index, local_block_table.size(1)); + EXPECT_EQ(remapped_slot / block_size, + local_block_table[0][local_block_index].item()); + } + + const torch::Tensor remapped_full = + remap_dcp_cache_slots(full_positions, + full_slot_tensor, + /*interleave_size=*/block_size, + dcp_size, + dcp_rank); + const std::vector local_kv_seq_lens = + layer::detail::compute_dcp_local_kv_seq_lens( + /*global_kv_seq_lens=*/{suffix_end}, + dcp_size, + dcp_rank, + block_size); + ASSERT_EQ(local_kv_seq_lens.size(), 1); + EXPECT_EQ(local_kv_seq_lens.front(), expected_local_kv_seq_lens[dcp_rank]); + EXPECT_EQ(local_kv_seq_lens.front(), + remapped_full.ge(0).sum().item()); + + if (dcp_rank == 0) { + EXPECT_EQ(remapped_full.slice(/*dim=*/0, /*start=*/4, /*end=*/8) + .ge(0) + .sum() + .item(), + 0); + } else { + EXPECT_EQ(remapped_full.slice(/*dim=*/0, /*start=*/0, /*end=*/4) + .ge(0) + .sum() + .item(), + 0); + EXPECT_EQ(remapped_full.slice(/*dim=*/0, /*start=*/8, /*end=*/12) + .ge(0) + .sum() + .item(), + 0); + } + } +} + +// Regression for the owner float-division bug: a plain `/` on an integer +// position tensor is float true-division, so 0 + // owner 1), owner-1 interior (134,137), and an L513-class 2nd-virtual-cycle + // position (523 -> 523/128=4, owner 0). + const torch::Tensor positions = torch::tensor( + {5, 133, 134, 137, 523}, torch::TensorOptions().dtype(torch::kInt32)); + const torch::Tensor slots = torch::tensor( + {5, 133, 134, 137, 523}, torch::TensorOptions().dtype(torch::kInt32)); + + // rank0 owns positions whose (pos/128)%2==0: 5(->0), 523(->4%2=0). Others -1. + const torch::Tensor r0 = remap_dcp_cache_slots(positions, + slots, + /*interleave_size=*/block_size, + dcp_size, + /*dcp_rank=*/0); + EXPECT_EQ(r0[0].item(), 5); // pos 5: float bug would give -1 + EXPECT_EQ(r0[1].item(), -1); // pos 133: owner 1 + EXPECT_EQ(r0[2].item(), -1); // pos 134: owner 1 + EXPECT_EQ(r0[3].item(), -1); // pos 137: owner 1 + EXPECT_EQ(r0[4].item(), 523); // pos 523: owner 0 (2nd cycle) + + // rank1 owns (pos/128)%2==1: 133,134,137. 5 and 523 -> -1. + const torch::Tensor r1 = remap_dcp_cache_slots(positions, + slots, + /*interleave_size=*/block_size, + dcp_size, + /*dcp_rank=*/1); + EXPECT_EQ(r1[0].item(), -1); + EXPECT_EQ(r1[1].item(), 133); + EXPECT_EQ(r1[2].item(), 134); + EXPECT_EQ(r1[3].item(), 137); + EXPECT_EQ(r1[4].item(), -1); +} + +// Negative slots stay -1 regardless of owner (non-owner or unallocated token). +TEST(RemapDcpCacheSlots, NegativeSlotsStayNegative) { + const torch::Tensor positions = + torch::tensor({5, 133}, torch::TensorOptions().dtype(torch::kInt32)); + const torch::Tensor slots = + torch::tensor({-1, -1}, torch::TensorOptions().dtype(torch::kInt32)); + const torch::Tensor r0 = remap_dcp_cache_slots(positions, + slots, + /*interleave_size=*/128, + /*dcp_size=*/2, + /*dcp_rank=*/0); + EXPECT_EQ(r0[0].item(), -1); + EXPECT_EQ(r0[1].item(), -1); +} + } // namespace } // namespace parallel_state } // namespace xllm diff --git a/tests/core/kernels/npu/CMakeLists.txt b/tests/core/kernels/npu/CMakeLists.txt index 29f9e9a85b..b5cce6e156 100644 --- a/tests/core/kernels/npu/CMakeLists.txt +++ b/tests/core/kernels/npu/CMakeLists.txt @@ -51,6 +51,26 @@ if(EXISTS "$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") "-Wl,-rpath-link,$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") endif() +cc_test( + NAME + fia_decode_lse_probe_test + SRCS + fia_decode_lse_probe_test.cpp + DEPS + ascendcl + nnopbase + torch + torch_npu + kernels + npu_kernels + GTest::gtest_main + glog::glog +) +if(EXISTS "$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") + target_link_options(fia_decode_lse_probe_test PRIVATE + "-Wl,-rpath-link,$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") +endif() + # build_cp_context is CompositeExplicitAutograd pure host index math, so this # runs on CPU with no NPU; it only needs the xllm_ops registrations linked in. cc_test( diff --git a/tests/core/kernels/npu/fia_decode_lse_probe_test.cpp b/tests/core/kernels/npu/fia_decode_lse_probe_test.cpp new file mode 100644 index 0000000000..a980b680f1 --- /dev/null +++ b/tests/core/kernels/npu/fia_decode_lse_probe_test.cpp @@ -0,0 +1,709 @@ +/* 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://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. +==============================================================================*/ + +// DCP-2 probe: can npu_fused_infer_attention emit a CORRECT softmax_lse in the +// decode setting (single query token + paged KV + block_table + GQA)? +// DCP decode-merge needs per-rank LSE, but xLLM decode currently runs +// batch_decode (ATB paged attention) which emits no LSE. The plan is to switch +// decode to FIA with softmax_lse_flag=true. This probe verifies, on real NPU: +// (1) FIA decode output matches batch_decode output (attention numerics OK); +// (2) FIA softmax_lse is finite and order-of-magnitude sane. +// It also verifies the DCP-specific batch metadata shape and reports whether +// FIA accepts a zero local-KV shard before DCP-2 production merge is designed. + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "core/kernels/npu/npu_ops_api.h" + +namespace xllm::kernel::npu { +namespace test { +namespace { + +class FiaDecodeLseProbe : public ::testing::Test { + protected: + static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); } + static void TearDownTestSuite() { torch_npu::finalize_npu(); } + + torch::Device device_ = torch::Device("npu:0"); +}; + +torch::Tensor make_slot_mapping(const std::vector& slots_host, + const torch::Device& device) { + return torch::tensor(slots_host, torch::TensorOptions().dtype(torch::kInt32)) + .to(device); +} + +void write_paged_kv_cache(torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& k_cache, + torch::Tensor& v_cache, + const std::vector& slots_host, + const torch::Device& device) { + const torch::Tensor slot_mapping = make_slot_mapping(slots_host, device); + std::optional value_opt = value; + std::optional v_cache_opt = v_cache; + reshape_paged_cache(key, value_opt, k_cache, v_cache_opt, slot_mapping); +} + +float max_abs_diff(const torch::Tensor& expected, const torch::Tensor& actual) { + const torch::Tensor expected_cpu = + expected.cpu().to(torch::kFloat32).reshape({-1}); + const torch::Tensor actual_cpu = + actual.cpu().to(torch::kFloat32).reshape({-1}); + CHECK_EQ(expected_cpu.numel(), actual_cpu.numel()); + return (expected_cpu - actual_cpu).abs().max().item(); +} + +std::pair reference_attention( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + double scale, + bool causal) { + CHECK_EQ(query.dim(), 3); + CHECK_EQ(key.dim(), 3); + CHECK_EQ(value.sizes(), key.sizes()); + CHECK_EQ(query.size(2), key.size(2)); + CHECK_EQ(query.size(1) % key.size(1), 0); + + const torch::Tensor query_cpu = query.cpu().to(torch::kFloat32); + const torch::Tensor key_cpu = key.cpu().to(torch::kFloat32); + const torch::Tensor value_cpu = value.cpu().to(torch::kFloat32); + const int64_t expansion_factor = query_cpu.size(1) / key_cpu.size(1); + const torch::Tensor expanded_key = + key_cpu.unsqueeze(2) + .expand({key_cpu.size(0), + key_cpu.size(1), + expansion_factor, + key_cpu.size(2)}) + .reshape({key_cpu.size(0), query_cpu.size(1), key_cpu.size(2)}); + const torch::Tensor expanded_value = + value_cpu.unsqueeze(2) + .expand({value_cpu.size(0), + value_cpu.size(1), + expansion_factor, + value_cpu.size(2)}) + .reshape({value_cpu.size(0), query_cpu.size(1), value_cpu.size(2)}); + + torch::Tensor scores = + torch::einsum("qhd,khd->qhk", {query_cpu, expanded_key}) * scale; + if (causal) { + CHECK_EQ(query_cpu.size(0), key_cpu.size(0)); + const torch::Tensor causal_mask = + torch::triu(torch::ones({query_cpu.size(0), key_cpu.size(0)}, + torch::TensorOptions().dtype(torch::kBool)), + 1); + scores = scores.masked_fill(causal_mask.unsqueeze(1), + -std::numeric_limits::infinity()); + } + const torch::Tensor lse = torch::logsumexp(scores, -1, true); + const torch::Tensor output = torch::einsum( + "qhk,khd->qhd", {torch::softmax(scores, -1), expanded_value}); + return {output, lse}; +} + +torch::Tensor make_fia_causal_mask(const torch::Device& device) { + const torch::TensorOptions options = + torch::TensorOptions().device(device).dtype(torch::kFloat32); + return torch::triu(torch::ones({2048, 2048}, options), 1) + .to(torch::kInt8) + .contiguous(); +} + +// One decode step: batch=1, q_len=1, ctx_len tokens already in paged KV cache. +// GQA: num_heads=8, num_kv_heads=2 (num_heads > num_kv_heads). +TEST_F(FiaDecodeLseProbe, DecodeFiaOutputMatchesBatchDecodeAndLseIsFinite) { + // Dims mirror real Qwen3.5-2B attention (k cache shape [nblk,128,2,256]): + // block_size=128, num_kv_heads=2, head_dim=256, num_heads=8 (GQA). + const int64_t ctx_len = 200; // history tokens in KV cache (>1 block) + const int64_t block_size = 128; + const int64_t num_blocks = 8; + const int64_t num_heads = 8; + const int64_t num_kv_heads = 2; + const int64_t head_dim = 256; + const double scale = 1.0 / std::sqrt(static_cast(head_dim)); + + auto opts = torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + + // --- Fill paged KV cache with ctx_len tokens via the real write path. --- + torch::Tensor key = + torch::randn({ctx_len, num_kv_heads, head_dim}, opts) * 0.1; + torch::Tensor value = + torch::randn({ctx_len, num_kv_heads, head_dim}, opts) * 0.1; + torch::Tensor k_cache = + torch::zeros({num_blocks, block_size, num_kv_heads, head_dim}, opts); + torch::Tensor v_cache = torch::zeros_like(k_cache); + + std::vector slots_host; + slots_host.reserve(ctx_len); + for (int64_t t = 0; t < ctx_len; ++t) { + slots_host.push_back(static_cast(t)); // contiguous slots 0..ctx-1 + } + write_paged_kv_cache(key, value, k_cache, v_cache, slots_host, device_); + + // block_table: sequence occupies blocks 0..ceil(ctx/block_size)-1. + const int64_t n_used_blocks = (ctx_len + block_size - 1) / block_size; + std::vector bt_host; + for (int64_t b = 0; b < n_used_blocks; ++b) { + bt_host.push_back(static_cast(b)); + } + torch::Tensor block_table = + torch::tensor(bt_host, torch::TensorOptions().dtype(torch::kInt32)) + .to(device_) + .view({1, n_used_blocks}); + + // --- Decode query: 1 token. --- + torch::Tensor query = torch::randn({1, num_heads, head_dim}, opts) * 0.1; + // context_lens must be a CPU host int32 tensor: ATB PagedAttention marks it + // as hostData (Input(context_lens, /*isHost=*/true)); a device tensor makes + // PagedAttentionOperation setup fail. + torch::Tensor seq_lens = + torch::tensor({static_cast(ctx_len)}, + torch::TensorOptions().dtype(torch::kInt32)); + + // --- (A) golden: existing batch_decode (no LSE). --- + torch::Tensor out_golden = torch::zeros({1, num_heads, head_dim}, opts); + batch_decode(query, + k_cache, + v_cache, + static_cast(scale), + block_table, + seq_lens, + out_golden); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + // KV cache viewed to 3D [num_blocks, block_size, num_kv_heads*head_dim] + // (avoids FIA reading head_dim=256 and rejecting it in TND). Decode is + // non-causal (no mask), so sparse_mode MUST be 0 (FIA: "when attnMask is not + // provided, sparseMode must be 0"). This differs from chunked_prefill which + // passes a causal mask + sparse_mode=3. + torch::Tensor k_view = k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + torch::Tensor v_view = v_cache.view({v_cache.size(0), v_cache.size(1), -1}); + std::vector actual_seq_lengths = {1}; // q tokens per seq + std::vector actual_seq_lengths_kv = {ctx_len}; + std::optional no_mask = std::nullopt; + std::optional bt_opt = block_table; + auto [out_fia, lse_fia] = + npu_fused_infer_attention(query, + k_view, + v_view, + no_mask, + bt_opt, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_kv_heads, + scale, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + // (1) output numerics match batch_decode. + const float max_diff = max_abs_diff(out_golden, out_fia); + const float golden_absmax = + out_golden.cpu().to(torch::kFloat32).abs().max().item(); + LOG(INFO) << "[DCP2-probe] output max|golden-fia|=" << max_diff + << " golden_absmax=" << golden_absmax; + EXPECT_LT(max_diff, 2e-2f) + << "FIA decode output diverges from batch_decode (max_diff=" << max_diff + << ")"; + + // (2) LSE finite + sane. + ASSERT_TRUE(lse_fia.defined() && lse_fia.numel() > 0) + << "FIA returned empty softmax_lse under softmax_lse_flag=true"; + const torch::Tensor lse = lse_fia.cpu().to(torch::kFloat32); + const bool all_finite = torch::isfinite(lse).all().item(); + const float lse_min = lse.min().item(); + const float lse_max = lse.max().item(); + LOG(INFO) << "[DCP2-probe] lse shape=" << lse.sizes() << " min=" << lse_min + << " max=" << lse_max << " finite=" << all_finite; + EXPECT_TRUE(all_finite) << "FIA softmax_lse has nan/inf"; + // LSE = log(sum exp(scores)) over ctx_len keys; must be finite real number. + EXPECT_GT(lse_max, -1e30f) << "LSE unreasonably small"; + EXPECT_LT(lse_max, 1e30f) << "LSE unreasonably large"; +} + +TEST_F(FiaDecodeLseProbe, MultiTokenRawKvCausalOutputAndLseMatchReference) { + const int64_t num_heads = 8; + const int64_t num_kv_heads = 2; + const int64_t head_dim = 128; + const int64_t token_count = 7; + const double scale = 1.0 / std::sqrt(static_cast(head_dim)); + const torch::TensorOptions options = + torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + + const torch::Tensor query = + torch::randn({token_count, num_heads, head_dim}, options) * 0.1; + const torch::Tensor key = + torch::randn({token_count, num_kv_heads, head_dim}, options) * 0.1; + const torch::Tensor value = + torch::randn({token_count, num_kv_heads, head_dim}, options) * 0.1; + const torch::Tensor causal_mask = make_fia_causal_mask(device_); + const std::vector cumulative_query_lengths = {3, 7}; + + const auto [output, lse] = + npu_fused_infer_attention(query, + key, + value, + std::make_optional(causal_mask), + /*block_table=*/std::nullopt, + cumulative_query_lengths, + cumulative_query_lengths, + num_heads, + num_kv_heads, + scale, + /*block_size=*/0, + /*sparse_mode=*/3, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + ASSERT_EQ(output.sizes(), + torch::IntArrayRef({token_count, num_heads, head_dim})); + ASSERT_EQ(lse.sizes(), torch::IntArrayRef({token_count, num_heads, 1})); + std::vector reference_outputs; + std::vector reference_lses; + int64_t sequence_begin = 0; + for (const int64_t sequence_end : cumulative_query_lengths) { + const int64_t sequence_length = sequence_end - sequence_begin; + const auto [sequence_output, sequence_lse] = + reference_attention(query.narrow(0, sequence_begin, sequence_length), + key.narrow(0, sequence_begin, sequence_length), + value.narrow(0, sequence_begin, sequence_length), + scale, + /*causal=*/true); + reference_outputs.emplace_back(sequence_output); + reference_lses.emplace_back(sequence_lse); + sequence_begin = sequence_end; + } + const torch::Tensor reference_output = torch::cat(reference_outputs, 0); + const torch::Tensor reference_lse = torch::cat(reference_lses, 0); + const float output_max_diff = max_abs_diff(reference_output, output); + const float lse_max_diff = max_abs_diff(reference_lse, lse); + LOG(INFO) << "[DCP4-probe][raw-multi-token] output_max_diff=" + << output_max_diff << " lse_max_diff=" << lse_max_diff + << " output_shape=" << output.sizes() + << " lse_shape=" << lse.sizes(); + EXPECT_LT(output_max_diff, 3e-2f); + EXPECT_LT(lse_max_diff, 3e-2f); +} + +TEST_F(FiaDecodeLseProbe, MultiTokenPagedContextTruncatesSharedPartialBlock) { + const int64_t num_heads = 8; + const int64_t num_kv_heads = 1; + const int64_t head_dim = 128; + const int64_t block_size = 128; + const int64_t num_blocks = 4; + const int64_t context_len = 130; + const int64_t token_count = 5; + const double scale = 1.0 / std::sqrt(static_cast(head_dim)); + const torch::TensorOptions options = + torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + + const torch::Tensor query = + torch::randn({token_count, num_heads, head_dim}, options) * 0.1; + torch::Tensor context_key = + torch::randn({context_len, num_kv_heads, head_dim}, options) * 0.1; + torch::Tensor context_value = + torch::randn({context_len, num_kv_heads, head_dim}, options) * 0.1; + torch::Tensor k_cache = + torch::zeros({num_blocks, block_size, num_kv_heads, head_dim}, options); + torch::Tensor v_cache = torch::zeros_like(k_cache); + + std::vector context_slots; + context_slots.reserve(context_len); + for (int64_t token = 0; token < context_len; ++token) { + const int64_t physical_block = token < block_size ? 1 : 3; + const int64_t block_offset = token % block_size; + context_slots.emplace_back( + static_cast(physical_block * block_size + block_offset)); + } + write_paged_kv_cache( + context_key, context_value, k_cache, v_cache, context_slots, device_); + + const torch::Tensor block_table = + torch::tensor(std::vector{0, 2, 1, 3}, + torch::TensorOptions().dtype(torch::kInt32)) + .to(device_) + .view({2, 2}); + const std::vector cumulative_query_lengths = {2, 5}; + const std::vector local_context_lengths = {0, context_len}; + const torch::Tensor k_view = + k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + const torch::Tensor v_view = + v_cache.view({v_cache.size(0), v_cache.size(1), -1}); + const auto [baseline_output, baseline_lse] = + npu_fused_infer_attention(query, + k_view, + v_view, + /*atten_mask=*/std::nullopt, + std::make_optional(block_table), + cumulative_query_lengths, + local_context_lengths, + num_heads, + num_kv_heads, + scale, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + torch::Tensor current_key = + torch::full({3, num_kv_heads, head_dim}, 50.0, options); + torch::Tensor current_value = + torch::full({3, num_kv_heads, head_dim}, 100.0, options); + const std::vector current_slots = { + static_cast(3 * block_size + 2), + static_cast(3 * block_size + 3), + static_cast(3 * block_size + 4)}; + write_paged_kv_cache( + current_key, current_value, k_cache, v_cache, current_slots, device_); + const auto [polluted_output, polluted_lse] = + npu_fused_infer_attention(query, + k_view, + v_view, + /*atten_mask=*/std::nullopt, + std::make_optional(block_table), + cumulative_query_lengths, + local_context_lengths, + num_heads, + num_kv_heads, + scale, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + ASSERT_EQ(polluted_output.sizes(), + torch::IntArrayRef({token_count, num_heads, head_dim})); + ASSERT_EQ(polluted_lse.sizes(), + torch::IntArrayRef({token_count, num_heads, 1})); + const torch::Tensor positive_query = query.narrow(0, 2, 3); + const auto [reference_output, reference_lse] = + reference_attention(positive_query, + context_key, + context_value, + scale, + /*causal=*/false); + const torch::Tensor baseline_positive_output = + baseline_output.narrow(0, 2, 3); + const torch::Tensor baseline_positive_lse = baseline_lse.narrow(0, 2, 3); + const torch::Tensor polluted_positive_output = + polluted_output.narrow(0, 2, 3); + const torch::Tensor polluted_positive_lse = polluted_lse.narrow(0, 2, 3); + const float output_pollution_diff = + max_abs_diff(baseline_positive_output, polluted_positive_output); + const float lse_pollution_diff = + max_abs_diff(baseline_positive_lse, polluted_positive_lse); + const float output_reference_diff = + max_abs_diff(reference_output, polluted_positive_output); + const float lse_reference_diff = + max_abs_diff(reference_lse, polluted_positive_lse); + LOG(INFO) << "[DCP4-probe][paged-partial-block] output_pollution_diff=" + << output_pollution_diff + << " lse_pollution_diff=" << lse_pollution_diff + << " output_reference_diff=" << output_reference_diff + << " lse_reference_diff=" << lse_reference_diff; + EXPECT_LT(output_pollution_diff, 1e-3f) + << "FIA over-read nonzero current KV beyond local_context_len"; + EXPECT_LT(lse_pollution_diff, 1e-3f) + << "FIA LSE included current KV beyond local_context_len"; + EXPECT_LT(output_reference_diff, 3e-2f); + EXPECT_LT(lse_reference_diff, 3e-2f); +} + +// DCP gathers Q heads across two ranks before each rank runs FIA against its +// local KV. This models rank 1 of a dcp_size=2 group: two requests have 72 and +// 128 local KV tokens, while FIA sees R * Hq_local = 2 * 4 Q heads and one +// local KV head. For TND decode, actual_seq_lengths must be cumulative Q ends. +TEST_F(FiaDecodeLseProbe, + BatchDecodeUsesCumulativeQueryLengthsAndGatheredGqaHeads) { + const int64_t dcp_size = 2; + const int64_t local_num_q_heads = 4; + const int64_t num_heads = dcp_size * local_num_q_heads; + const int64_t num_kv_heads = 1; + const int64_t head_dim = 256; + const int64_t block_size = 128; + const int64_t num_blocks = 4; + const int64_t first_local_kv_len = 72; + const int64_t second_local_kv_len = 128; + const double scale = 1.0 / std::sqrt(static_cast(head_dim)); + + const auto opts = + torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + const torch::Tensor first_key = + torch::randn({first_local_kv_len, num_kv_heads, head_dim}, opts) * 0.1; + const torch::Tensor first_value = + torch::randn({first_local_kv_len, num_kv_heads, head_dim}, opts) * 0.1; + const torch::Tensor second_key = + torch::randn({second_local_kv_len, num_kv_heads, head_dim}, opts) * 0.1; + const torch::Tensor second_value = + torch::randn({second_local_kv_len, num_kv_heads, head_dim}, opts) * 0.1; + torch::Tensor key = torch::cat({first_key, second_key}, /*dim=*/0); + torch::Tensor value = torch::cat({first_value, second_value}, /*dim=*/0); + torch::Tensor k_cache = + torch::zeros({num_blocks, block_size, num_kv_heads, head_dim}, opts); + torch::Tensor v_cache = torch::zeros_like(k_cache); + + std::vector slots_host; + slots_host.reserve(first_local_kv_len + second_local_kv_len); + for (int64_t token = 0; token < first_local_kv_len; ++token) { + slots_host.push_back(static_cast(block_size + token)); + } + for (int64_t token = 0; token < second_local_kv_len; ++token) { + slots_host.push_back(static_cast(3 * block_size + token)); + } + write_paged_kv_cache(key, value, k_cache, v_cache, slots_host, device_); + + // Local table keeps original physical block ids selected by DCP-1c. + const torch::Tensor block_table = + torch::tensor(std::vector{1, 3}, + torch::TensorOptions().dtype(torch::kInt32)) + .to(device_) + .view({2, 1}); + const torch::Tensor query = + torch::randn({2, num_heads, head_dim}, opts) * 0.1; + const torch::Tensor local_kv_lens = torch::tensor( + std::vector{static_cast(first_local_kv_len), + static_cast(second_local_kv_len)}, + torch::TensorOptions().dtype(torch::kInt32)); + torch::Tensor out_golden = torch::zeros({2, num_heads, head_dim}, opts); + batch_decode(query, + k_cache, + v_cache, + static_cast(scale), + block_table, + local_kv_lens, + out_golden); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + const torch::Tensor k_view = + k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + const torch::Tensor v_view = + v_cache.view({v_cache.size(0), v_cache.size(1), -1}); + const std::vector actual_seq_lengths = {1, 2}; + const std::vector actual_seq_lengths_kv = {first_local_kv_len, + second_local_kv_len}; + const auto [out_fia, lse_fia] = + npu_fused_infer_attention(query, + k_view, + v_view, + std::nullopt, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_kv_heads, + scale, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + const float output_max_diff = max_abs_diff(out_golden, out_fia); + LOG(INFO) << "[DCP2-probe][batch] output max|golden-fia|=" << output_max_diff + << " q_lengths={1,2} kv_lengths={" << first_local_kv_len << "," + << second_local_kv_len << "}"; + EXPECT_LT(output_max_diff, 2e-2f); + + ASSERT_EQ(lse_fia.dim(), 3); + EXPECT_EQ(lse_fia.size(0), 2); + EXPECT_EQ(lse_fia.size(1), num_heads); + EXPECT_EQ(lse_fia.size(2), 1); + EXPECT_TRUE(torch::isfinite(lse_fia.cpu()).all().item()); +} + +// A short request can have no block owned by this rank while a later request +// in the same decode batch has local KV. Probe a leading zero explicitly: FIA +// must at least accept the metadata and preserve the positive row's result. +TEST_F(FiaDecodeLseProbe, LeadingZeroLocalKvDoesNotCorruptPositiveRow) { + const int64_t num_heads = 8; + const int64_t num_kv_heads = 1; + const int64_t head_dim = 256; + const int64_t block_size = 128; + const int64_t num_blocks = 4; + const int64_t positive_local_kv_len = 72; + const double scale = 1.0 / std::sqrt(static_cast(head_dim)); + + const auto opts = + torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + torch::Tensor key = + torch::randn({positive_local_kv_len, num_kv_heads, head_dim}, opts) * 0.1; + torch::Tensor value = + torch::randn({positive_local_kv_len, num_kv_heads, head_dim}, opts) * 0.1; + torch::Tensor k_cache = + torch::zeros({num_blocks, block_size, num_kv_heads, head_dim}, opts); + torch::Tensor v_cache = torch::zeros_like(k_cache); + + std::vector slots_host; + slots_host.reserve(positive_local_kv_len); + for (int64_t token = 0; token < positive_local_kv_len; ++token) { + slots_host.push_back(static_cast(3 * block_size + token)); + } + write_paged_kv_cache(key, value, k_cache, v_cache, slots_host, device_); + + // The first row has no local KV. Its table entry is intentionally ignored by + // actual_seq_lengths_kv; the second row owns physical block 3. + const torch::Tensor block_table = + torch::tensor(std::vector{1, 3}, + torch::TensorOptions().dtype(torch::kInt32)) + .to(device_) + .view({2, 1}); + const torch::Tensor query = + torch::randn({2, num_heads, head_dim}, opts) * 0.1; + const torch::Tensor positive_query = + query.slice(/*dim=*/0, /*start=*/1, /*end=*/2).contiguous(); + const torch::Tensor positive_block_table = + block_table.slice(/*dim=*/0, /*start=*/1, /*end=*/2).contiguous(); + const torch::Tensor positive_kv_len = torch::tensor( + std::vector{static_cast(positive_local_kv_len)}, + torch::TensorOptions().dtype(torch::kInt32)); + torch::Tensor positive_golden = torch::zeros({1, num_heads, head_dim}, opts); + batch_decode(positive_query, + k_cache, + v_cache, + static_cast(scale), + positive_block_table, + positive_kv_len, + positive_golden); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + const torch::Tensor k_view = + k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + const torch::Tensor v_view = + v_cache.view({v_cache.size(0), v_cache.size(1), -1}); + const std::vector actual_seq_lengths = {1, 2}; + const std::vector actual_seq_lengths_kv = {0, positive_local_kv_len}; + const auto [out_fia, lse_fia] = + npu_fused_infer_attention(query, + k_view, + v_view, + std::nullopt, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_kv_heads, + scale, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + const torch::Tensor positive_fia = + out_fia.slice(/*dim=*/0, /*start=*/1, /*end=*/2).contiguous(); + const float positive_max_diff = max_abs_diff(positive_golden, positive_fia); + EXPECT_LT(positive_max_diff, 2e-2f) + << "leading zero local KV corrupted the following positive row"; + + const torch::Tensor zero_out = + out_fia.slice(/*dim=*/0, /*start=*/0, /*end=*/1) + .cpu() + .to(torch::kFloat32); + const torch::Tensor zero_lse = + lse_fia.slice(/*dim=*/0, /*start=*/0, /*end=*/1) + .cpu() + .to(torch::kFloat32); + LOG(INFO) << "[DCP2-probe][leading-zero] positive max|golden-fia|=" + << positive_max_diff + << " zero_out_absmax=" << zero_out.abs().max().item() + << " zero_lse_min=" << zero_lse.min().item() + << " zero_lse_max=" << zero_lse.max().item() + << " zero_lse_finite=" + << torch::isfinite(zero_lse).all().item(); +} + +// DCP-1c can legitimately select no table column when every request is shorter +// than this rank's first owned block. Keep this separate so a possible FIA +// fatal for [batch, 0] does not hide the positive-batch probe result. +TEST_F(FiaDecodeLseProbe, AllZeroLocalKvWithEmptyBlockTableReportsFiaBehavior) { + const int64_t num_heads = 8; + const int64_t num_kv_heads = 1; + const int64_t head_dim = 256; + const int64_t block_size = 128; + const int64_t num_blocks = 1; + const double scale = 1.0 / std::sqrt(static_cast(head_dim)); + + const auto opts = + torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + const torch::Tensor query = + torch::randn({2, num_heads, head_dim}, opts) * 0.1; + const torch::Tensor k_cache = + torch::zeros({num_blocks, block_size, num_kv_heads, head_dim}, opts); + const torch::Tensor v_cache = torch::zeros_like(k_cache); + const torch::Tensor block_table = torch::empty( + {2, 0}, torch::TensorOptions().device(device_).dtype(torch::kInt32)); + const torch::Tensor k_view = + k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + const torch::Tensor v_view = + v_cache.view({v_cache.size(0), v_cache.size(1), -1}); + const std::vector actual_seq_lengths = {1, 2}; + const std::vector actual_seq_lengths_kv = {0, 0}; + + const auto [out_fia, lse_fia] = + npu_fused_infer_attention(query, + k_view, + v_view, + std::nullopt, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_kv_heads, + scale, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + ASSERT_EQ(out_fia.sizes(), torch::IntArrayRef({2, num_heads, head_dim})); + ASSERT_EQ(lse_fia.sizes(), torch::IntArrayRef({2, num_heads, 1})); + const torch::Tensor out_cpu = out_fia.cpu().to(torch::kFloat32); + const torch::Tensor lse_cpu = lse_fia.cpu().to(torch::kFloat32); + LOG(INFO) << "[DCP2-probe][all-zero-empty-table] out_absmax=" + << out_cpu.abs().max().item() + << " lse_min=" << lse_cpu.min().item() + << " lse_max=" << lse_cpu.max().item() + << " lse_finite=" << torch::isfinite(lse_cpu).all().item(); +} + +} // namespace +} // namespace test +} // namespace xllm::kernel::npu diff --git a/tests/core/layers/npu_torch/CMakeLists.txt b/tests/core/layers/npu_torch/CMakeLists.txt index cf3eca53d5..be2acfa93f 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/tests/core/layers/npu_torch/CMakeLists.txt @@ -1,5 +1,21 @@ include(cc_test) +cc_test( + NAME + npu_dcp_attention_utils_test + SRCS + dcp_attention_utils_test.cpp + DEPS + :dcp_attention_utils + glog::glog + torch + GTest::gtest_main +) +if(EXISTS "$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") + target_link_options(npu_dcp_attention_utils_test PRIVATE + "-Wl,-rpath-link,$ENV{PYTORCH_INSTALL_PATH}/../torch.libs") +endif() + cc_test( NAME deepseek_v4_eplb_load_utils_test @@ -93,3 +109,29 @@ target_link_options(npu_linear_w8a8_dynamic_test PRIVATE "-Wl,--whole-archive" "${CMAKE_BINARY_DIR}/third_party/spdlog/libspdlog.a" "-Wl,--no-whole-archive") + +cc_test( + NAME + npu_dcp_attention_test + SRCS + dcp_attention_test.cpp + DEPS + :npu_torch_layers + :parallel_state + glog::glog + torch + GTest::gtest_main +) + +target_link_libraries(npu_dcp_attention_test + PRIVATE + ascendcl + hccl + c_sec + nnopbase + atb) + +target_link_options(npu_dcp_attention_test PRIVATE + "-Wl,--whole-archive" + "${CMAKE_BINARY_DIR}/third_party/spdlog/libspdlog.a" + "-Wl,--no-whole-archive") diff --git a/tests/core/layers/npu_torch/dcp_attention_test.cpp b/tests/core/layers/npu_torch/dcp_attention_test.cpp new file mode 100644 index 0000000000..5963c09723 --- /dev/null +++ b/tests/core/layers/npu_torch/dcp_attention_test.cpp @@ -0,0 +1,248 @@ +/* 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://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "core/kernels/npu/npu_ops_api.h" +#include "framework/kv_cache/kv_cache.h" +#include "framework/parallel_state/process_group.h" +#include "layers/npu_torch/attention.h" + +namespace xllm::layer::test { +namespace { + +class ScriptedDcpProcessGroup final : public ProcessGroup { + public: + ScriptedDcpProcessGroup(const torch::Device& device, + torch::Tensor peer_partial_out, + torch::Tensor peer_partial_lse) + : ProcessGroup(1, 2, device), + peer_partial_out_(std::move(peer_partial_out)), + peer_partial_lse_(std::move(peer_partial_lse)) {} + + torch::Tensor allgather_base_sync(const torch::Tensor& input) override { + if (call_count_ == 0) { + ++call_count_; + return torch::stack({input, input}, 0); + } + if (call_count_ == 1) { + CHECK_EQ(input.sizes(), peer_partial_out_.sizes()); + normalized_out_before_gather_ = + torch::equal(input, torch::zeros_like(input)); + ++call_count_; + return torch::stack({peer_partial_out_, input}, 0); + } + if (call_count_ == 2) { + CHECK_EQ(input.sizes(), peer_partial_lse_.sizes()); + normalized_lse_before_gather_ = torch::equal( + input, + torch::full_like(input, -std::numeric_limits::infinity())); + ++call_count_; + return torch::stack({peer_partial_lse_, input}, 0); + } + LOG(FATAL) << "Unexpected DCP all-gather call " << call_count_; + return torch::Tensor(); + } + + int32_t call_count() const { return call_count_; } + bool normalized_out_before_gather() const { + return normalized_out_before_gather_; + } + bool normalized_lse_before_gather() const { + return normalized_lse_before_gather_; + } + + private: + torch::Tensor peer_partial_out_; + torch::Tensor peer_partial_lse_; + int32_t call_count_ = 0; + bool normalized_out_before_gather_ = false; + bool normalized_lse_before_gather_ = false; +}; + +class DcpAttentionTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { torch_npu::init_npu("npu:0"); } + static void TearDownTestSuite() { torch_npu::finalize_npu(); } + + torch::Device device_ = torch::Device("npu:0"); +}; + +TEST_F(DcpAttentionTest, ZeroLocalKvNormalizesBeforeMergeAndSlicesLocalHeads) { + const int64_t block_size = 128; + const int64_t head_size = 128; + const int64_t local_num_heads = 4; + const int64_t num_kv_heads = 1; + const int64_t group_num_heads = 8; + const float scale = 1.0f / std::sqrt(static_cast(head_size)); + const torch::TensorOptions bf16_options = + torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + const torch::TensorOptions fp32_options = + torch::TensorOptions().device(device_).dtype(torch::kFloat32); + + torch::Tensor peer_partial_out = + torch::ones({1, group_num_heads, head_size}, fp32_options); + peer_partial_out.slice(1, local_num_heads, group_num_heads).fill_(2.0f); + const torch::Tensor peer_partial_lse = + torch::zeros({1, group_num_heads, 1}, fp32_options); + ScriptedDcpProcessGroup dcp_group( + device_, peer_partial_out, peer_partial_lse); + + torch::Tensor key = torch::zeros({1, num_kv_heads, head_size}, bf16_options); + torch::Tensor value = torch::zeros_like(key); + torch::Tensor query = + torch::randn({1, local_num_heads * head_size}, bf16_options); + const torch::Tensor k_cache = + torch::zeros({1, block_size, num_kv_heads, head_size}, bf16_options); + const torch::Tensor v_cache = torch::zeros_like(k_cache); + KVCache kv_cache(KVCacheTensors{k_cache, v_cache}); + + AttentionMetadata attn_metadata{}; + attn_metadata.slot_mapping = + torch::tensor(std::vector{-1}, + torch::TensorOptions().dtype(torch::kInt32)) + .to(device_); + attn_metadata.block_table = + torch::tensor(std::vector{0}, + torch::TensorOptions().dtype(torch::kInt32)) + .to(device_) + .view({1, 1}); + attn_metadata.q_cu_seq_lens_host_vec = {1}; + attn_metadata.kv_seq_lens_host_vec = {1}; + + AttentionImpl attention( + local_num_heads, head_size, scale, num_kv_heads, -1, 2, 1, &dcp_group); + const auto [output, output_lse] = + attention.forward(attn_metadata, query, key, value, kv_cache); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + EXPECT_FALSE(output_lse.has_value()); + EXPECT_EQ(dcp_group.call_count(), 3); + EXPECT_TRUE(dcp_group.normalized_out_before_gather()); + EXPECT_TRUE(dcp_group.normalized_lse_before_gather()); + const torch::Tensor output_cpu = + output.cpu().to(torch::kFloat32).view({1, local_num_heads, head_size}); + const torch::Tensor expected = + torch::full({1, local_num_heads, head_size}, + 2.0f, + torch::TensorOptions().dtype(torch::kFloat32)); + EXPECT_LT((output_cpu - expected).abs().max().item(), 1e-4f); +} + +TEST_F(DcpAttentionTest, + ChunkedFirstChunkNormalizesLeadingZeroAndSkipsContextCollectives) { + const int64_t token_count = 5; + const int64_t block_size = 128; + const int64_t head_size = 128; + const int64_t local_num_heads = 4; + const int64_t num_kv_heads = 1; + const int64_t group_num_heads = 8; + const float scale = 1.0f / std::sqrt(static_cast(head_size)); + const torch::TensorOptions bf16_options = + torch::TensorOptions().device(device_).dtype(torch::kBFloat16); + const torch::TensorOptions fp32_options = + torch::TensorOptions().device(device_).dtype(torch::kFloat32); + + ScriptedDcpProcessGroup dcp_group(device_, + /*peer_partial_out=*/torch::Tensor(), + /*peer_partial_lse=*/torch::Tensor()); + torch::Tensor query = + torch::randn({token_count, local_num_heads * head_size}, bf16_options) * + 0.1; + torch::Tensor key = + torch::randn({token_count, num_kv_heads, head_size}, bf16_options) * 0.1; + torch::Tensor value = + torch::randn({token_count, num_kv_heads, head_size}, bf16_options) * 0.1; + const torch::Tensor k_cache = + torch::zeros({2, block_size, num_kv_heads, head_size}, bf16_options); + const torch::Tensor v_cache = torch::zeros_like(k_cache); + KVCache kv_cache(KVCacheTensors{k_cache, v_cache}); + + AttentionMetadata attn_metadata{}; + attn_metadata.is_chunked_prefill = true; + attn_metadata.slot_mapping = + torch::full({token_count}, + -1, + torch::TensorOptions().dtype(torch::kInt32).device(device_)); + attn_metadata.block_table = + torch::tensor(std::vector{0, 1}, + torch::TensorOptions().dtype(torch::kInt32)) + .to(device_) + .view({2, 1}); + attn_metadata.fia_attn_mask = + torch::triu(torch::ones({2048, 2048}, fp32_options), 1) + .to(torch::kInt8) + .contiguous(); + attn_metadata.q_cu_seq_lens_host_vec = {0, 2, 5}; + attn_metadata.kv_seq_lens_host_vec = {2, 3}; + + const torch::Tensor local_query = + query.view({token_count, local_num_heads, head_size}); + const torch::Tensor query_group = + torch::cat({local_query, local_query}, /*dim=*/1).contiguous(); + const std::vector normalized_q_cu_seq_lens = {2, 5}; + const auto [expected_group_out, expected_group_lse] = + xllm::kernel::npu::npu_fused_infer_attention( + query_group, + key, + value, + std::make_optional(attn_metadata.fia_attn_mask), + /*block_table=*/std::nullopt, + normalized_q_cu_seq_lens, + normalized_q_cu_seq_lens, + group_num_heads, + num_kv_heads, + scale, + /*block_size=*/0, + /*sparse_mode=*/3, + "TND", + /*softmax_lse_flag=*/true); + ASSERT_EQ(expected_group_lse.sizes(), + torch::IntArrayRef({token_count, group_num_heads, 1})); + const torch::Tensor expected = + expected_group_out.slice(1, local_num_heads, group_num_heads) + .contiguous(); + + AttentionImpl attention( + local_num_heads, head_size, scale, num_kv_heads, -1, 2, 1, &dcp_group); + const auto [output, output_lse] = + attention.forward(attn_metadata, query, key, value, kv_cache); + ASSERT_EQ(aclrtSynchronizeStream(c10_npu::getCurrentNPUStream().stream()), + ACL_SUCCESS); + + EXPECT_FALSE(output_lse.has_value()); + EXPECT_EQ(dcp_group.call_count(), 1) + << "all-zero global context must skip context FIA collectives"; + const torch::Tensor output_3d = + output.view({token_count, local_num_heads, head_size}); + EXPECT_LT( + (output_3d.cpu().to(torch::kFloat32) - expected.cpu().to(torch::kFloat32)) + .abs() + .max() + .item(), + 2e-2f); +} + +} // namespace +} // namespace xllm::layer::test diff --git a/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp new file mode 100644 index 0000000000..8abf161d04 --- /dev/null +++ b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp @@ -0,0 +1,296 @@ +/* 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://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "layers/npu_torch/dcp_attention_utils.h" + +#include +#include +#include + +#include +#include +#include +#include + +namespace xllm::layer::test { +namespace { + +std::pair compute_attention_partial( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value) { + CHECK_EQ(query.dim(), 3); + CHECK_EQ(key.dim(), 3); + CHECK_EQ(value.sizes(), key.sizes()); + CHECK_EQ(query.size(1), key.size(1)); + CHECK_EQ(query.size(2), key.size(2)); + + const double scale = 1.0 / std::sqrt(static_cast(query.size(2))); + const torch::Tensor scores = + torch::einsum("qhd,khd->qhk", {query, key}) * scale; + const torch::Tensor partial_lse = torch::logsumexp(scores, -1, true); + const torch::Tensor partial_out = + torch::einsum("qhk,khd->qhd", {torch::softmax(scores, -1), value}); + return {partial_out, partial_lse}; +} + +TEST(DcpAttentionUtilsTest, ShardedKvMergeMatchesFullAttentionInFloat64) { + torch::manual_seed(7); + const torch::TensorOptions options = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat64); + const torch::Tensor query = torch::randn({5, 3, 8}, options); + const torch::Tensor key = torch::randn({12, 3, 8}, options); + const torch::Tensor value = torch::randn({12, 3, 8}, options); + const auto [reference_out, reference_lse] = + compute_attention_partial(query, key, value); + (void)reference_lse; + + const std::vector key_shards = key.chunk(3, 0); + const std::vector value_shards = value.chunk(3, 0); + ASSERT_EQ(key_shards.size(), value_shards.size()); + std::vector partial_outputs; + std::vector partial_lses; + partial_outputs.reserve(key_shards.size()); + partial_lses.reserve(key_shards.size()); + for (int64_t shard_index = 0; + shard_index < static_cast(key_shards.size()); + ++shard_index) { + const auto [partial_out, partial_lse] = compute_attention_partial( + query, key_shards[shard_index], value_shards[shard_index]); + partial_outputs.emplace_back(partial_out); + partial_lses.emplace_back(partial_lse); + } + + const torch::Tensor merged_out = detail::merge_dcp_partials( + torch::stack(partial_outputs, 0), torch::stack(partial_lses, 0)); + EXPECT_TRUE(torch::allclose( + merged_out, reference_out, /*rtol=*/1e-10, /*atol=*/1e-10)); +} + +TEST(DcpAttentionUtilsTest, InvalidLseShardIsIgnored) { + const torch::TensorOptions options = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat64); + const torch::Tensor valid_out = + torch::tensor({3.0, 7.0}, options).view({1, 2, 1}); + const torch::Tensor invalid_out = + torch::full_like(valid_out, std::numeric_limits::quiet_NaN()); + const torch::Tensor valid_lse = torch::zeros({1, 2, 1}, options); + const torch::Tensor invalid_lse = + torch::full_like(valid_lse, -std::numeric_limits::infinity()); + + const torch::Tensor merged_out = + detail::merge_dcp_partials(torch::stack({invalid_out, valid_out}, 0), + torch::stack({invalid_lse, valid_lse}, 0)); + EXPECT_TRUE(torch::equal(merged_out, valid_out)); +} + +TEST(DcpAttentionUtilsTest, AllInvalidLseShardsProduceZero) { + const torch::TensorOptions options = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat64); + const torch::Tensor partial_out = torch::full( + {4, 2, 3, 5}, std::numeric_limits::quiet_NaN(), options); + const torch::Tensor partial_lse = torch::full( + {4, 2, 3, 1}, -std::numeric_limits::infinity(), options); + + const torch::Tensor merged_out = + detail::merge_dcp_partials(partial_out, partial_lse); + EXPECT_TRUE(torch::equal(merged_out, torch::zeros_like(merged_out))); +} + +TEST(DcpAttentionUtilsTest, DistributesPartialTailAcrossFourRanks) { + std::vector local_kv_seq_lens; + local_kv_seq_lens.reserve(4); + for (int32_t dcp_rank = 0; dcp_rank < 4; ++dcp_rank) { + const std::vector rank_local_kv_seq_lens = + detail::compute_dcp_local_kv_seq_lens( + /*global_kv_seq_lens=*/{257}, + /*dcp_size=*/4, + dcp_rank, + /*block_size=*/128); + ASSERT_EQ(rank_local_kv_seq_lens.size(), 1); + local_kv_seq_lens.emplace_back(rank_local_kv_seq_lens.front()); + } + + EXPECT_EQ(local_kv_seq_lens, (std::vector{128, 128, 1, 0})); +} + +TEST(DcpAttentionUtilsTest, ContextLenIsKvMinusCurrentChunkQuery) { + // Two requests packed into one chunked batch: request 0 has 130 query tokens + // over a 130 KV (no cached context, first chunk); request 1 has 6 query + // tokens over a 262 KV (256 cached context + 6 current chunk). + const std::vector context_lens = detail::compute_dcp_context_lens( + /*q_cu_seq_lens=*/{130, 136}, + /*global_kv_seq_lens=*/{130, 262}); + EXPECT_EQ(context_lens, (std::vector{0, 256})); +} + +TEST(DcpAttentionUtilsTest, ContextShardLenReadsOnlyCachedContext) { + // A request with 256 cached context + a current chunk: the local context + // shard length must be derived from context_len (256), not the full KV, so + // the context part never reads the current chunk's own KV. + std::vector context_shard_lens; + context_shard_lens.reserve(2); + for (int32_t dcp_rank = 0; dcp_rank < 2; ++dcp_rank) { + const std::vector rank_shard = + detail::compute_dcp_local_kv_seq_lens( + /*global_kv_seq_lens=*/{256}, + /*dcp_size=*/2, + dcp_rank, + /*block_size=*/128); + ASSERT_EQ(rank_shard.size(), 1); + context_shard_lens.emplace_back(rank_shard.front()); + } + EXPECT_EQ(context_shard_lens, (std::vector{128, 128})); +} + +TEST(DcpAttentionUtilsTest, CoversPhase4LocalContextContracts) { + const std::vector context_lens = {128, 256, 512}; + std::vector> local_context_lens_by_rank; + local_context_lens_by_rank.reserve(2); + for (int32_t dcp_rank = 0; dcp_rank < 2; ++dcp_rank) { + local_context_lens_by_rank.emplace_back( + detail::compute_dcp_local_kv_seq_lens(context_lens, + /*dcp_size=*/2, + dcp_rank, + /*block_size=*/128)); + } + + ASSERT_EQ(local_context_lens_by_rank.size(), 2); + EXPECT_EQ(local_context_lens_by_rank[0], + (std::vector{128, 128, 256})); + EXPECT_EQ(local_context_lens_by_rank[1], (std::vector{0, 128, 256})); +} + +TEST(DcpAttentionUtilsTest, GraphZeroShardMaskTracksTensorKvLengthChanges) { + const torch::TensorOptions fp32_options = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat32); + torch::Tensor global_kv_seq_lens = torch::tensor( + {128}, torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt32)); + torch::Tensor partial_out = torch::full({1, 2, 3}, 7.0, fp32_options); + torch::Tensor partial_lse = torch::full({1, 2, 1}, 9.0, fp32_options); + + detail::normalize_zero_dcp_partials_for_graph(partial_out, + partial_lse, + global_kv_seq_lens, + /*dcp_rank=*/1, + /*block_size=*/128); + EXPECT_TRUE(torch::equal(partial_out, torch::zeros_like(partial_out))); + EXPECT_TRUE(torch::equal( + partial_lse, + torch::full_like(partial_lse, -std::numeric_limits::infinity()))); + + global_kv_seq_lens.fill_(129); + partial_out.fill_(7.0); + partial_lse.fill_(9.0); + detail::normalize_zero_dcp_partials_for_graph(partial_out, + partial_lse, + global_kv_seq_lens, + /*dcp_rank=*/1, + /*block_size=*/128); + EXPECT_TRUE(torch::equal(partial_out, torch::full_like(partial_out, 7.0))); + EXPECT_TRUE(torch::equal(partial_lse, torch::full_like(partial_lse, 9.0))); +} + +TEST(DcpAttentionUtilsTest, ValidateChunkedLengthsAcceptsMultiTokenRequests) { + const std::vector normalized_q_cu_seq_lens = + detail::validate_dcp_chunked_lengths( + /*q_cu_seq_lens=*/{130, 136}, + /*global_kv_seq_lens=*/{130, 262}, + /*token_count=*/136); + EXPECT_EQ(normalized_q_cu_seq_lens, (std::vector{130, 136})); +} + +TEST(DcpAttentionUtilsTest, NormalizesLeadingZeroBeforeValidation) { + const std::vector normalized_q_cu_seq_lens = + detail::validate_dcp_chunked_lengths( + /*q_cu_seq_lens=*/{0, 130, 136}, + /*global_kv_seq_lens=*/{130, 262}, + /*token_count=*/136); + EXPECT_EQ(normalized_q_cu_seq_lens, (std::vector{130, 136})); + EXPECT_EQ(detail::compute_dcp_context_lens(normalized_q_cu_seq_lens, + /*global_kv_seq_lens=*/{130, 262}), + (std::vector{0, 256})); +} + +TEST(DcpAttentionUtilsTest, ValidateChunkedLengthsRejectsTokenCountMismatch) { + EXPECT_DEATH(detail::validate_dcp_chunked_lengths( + /*q_cu_seq_lens=*/{130, 136}, + /*global_kv_seq_lens=*/{130, 262}, + /*token_count=*/135), + "query tokens"); +} + +TEST(DcpAttentionUtilsTest, ValidateChunkedLengthsRejectsKvShorterThanQuery) { + EXPECT_DEATH(detail::validate_dcp_chunked_lengths( + /*q_cu_seq_lens=*/{10}, + /*global_kv_seq_lens=*/{4}, + /*token_count=*/10), + "cover the current chunk query"); +} + +TEST(DcpAttentionUtilsTest, ValidateChunkedLengthsRejectsZeroQueryRequest) { + EXPECT_DEATH(detail::validate_dcp_chunked_lengths( + /*q_cu_seq_lens=*/{0, 0, 1}, + /*global_kv_seq_lens=*/{0, 1}, + /*token_count=*/1), + "at least one query token"); +} + +TEST(DcpAttentionUtilsTest, ValidateChunkedLengthsRejectsQueryAboveMaskLimit) { + EXPECT_DEATH( + detail::validate_dcp_chunked_lengths( + /*q_cu_seq_lens=*/{detail::kMaxDcpChunkedPrefillQueryLen + 1}, + /*global_kv_seq_lens=*/{detail::kMaxDcpChunkedPrefillQueryLen + 1}, + /*token_count=*/ + detail::kMaxDcpChunkedPrefillQueryLen + 1), + "does not yet support chunked query length above"); +} + +TEST(DcpAttentionUtilsTest, BlockTableRowsMatchRequestsNotQueryTokens) { + const torch::Tensor local_block_table = torch::tensor( + {{37, 89}, {41, 73}}, torch::TensorOptions().dtype(torch::kInt64)); + detail::validate_dcp_chunked_block_table(local_block_table, + /*request_count=*/2); + EXPECT_DEATH(detail::validate_dcp_chunked_block_table(local_block_table, + /*request_count=*/136), + "request count"); +} + +TEST(DcpAttentionUtilsTest, NormalizeChunkedZeroesEmptyContextTokenRange) { + const torch::TensorOptions options = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat32); + // token_count=5: request 0 owns tokens [0,2) with empty context on this rank, + // request 1 owns tokens [2,5) with a non-empty context shard. + torch::Tensor partial_out = torch::ones({5, 2, 3}, options); + torch::Tensor partial_lse = torch::ones({5, 2, 1}, options); + + detail::normalize_zero_dcp_chunked_partials(partial_out, + partial_lse, + /*local_context_lens=*/{0, 64}, + /*q_cu_seq_lens=*/{2, 5}); + + EXPECT_TRUE(torch::equal(partial_out.narrow(0, 0, 2), + torch::zeros({2, 2, 3}, options))); + EXPECT_TRUE(torch::equal(partial_out.narrow(0, 2, 3), + torch::ones({3, 2, 3}, options))); + const torch::Tensor expected_neg_inf_lse = + torch::full({2, 2, 1}, -std::numeric_limits::infinity(), options); + EXPECT_TRUE(torch::equal(partial_lse.narrow(0, 0, 2), expected_neg_inf_lse)); + EXPECT_TRUE(torch::equal(partial_lse.narrow(0, 2, 3), + torch::ones({3, 2, 1}, options))); +} + +} // namespace +} // namespace xllm::layer::test diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index ff7042cda7..0b0a8a082e 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -46,6 +46,7 @@ limitations under the License. #include "core/layers/npu/npu_lm_head_impl.h" #include "core/layers/npu/npu_word_embedding_impl.h" #include "core/layers/npu_torch/tests_utils.h" +#include "core/platform/npu/acl_graph_task_update_context.h" #include "core/runtime/acl_graph_executor_impl.h" #include "core/runtime/acl_graph_persistent_param.h" #include "core/runtime/base_executor_impl.h" @@ -148,6 +149,22 @@ TEST(AclGraphStaticGraphTaskSignatureTest, EXPECT_FALSE(npu::make_static_graph_task_signature(params).has_value()); } +TEST(AclGraphTaskUpdateContextTest, CaptureResetsAndRetainsRecordedTasks) { + npu::AclGraphTaskUpdateContext context; + context.causal_conv1d_tasks.emplace_back(); + context.fia_tasks.emplace_back(); + + context.begin_capture(); + EXPECT_TRUE(context.capturing); + EXPECT_TRUE(context.causal_conv1d_tasks.empty()); + EXPECT_TRUE(context.fia_tasks.empty()); + + context.fia_tasks.emplace_back(); + context.end_capture(); + EXPECT_FALSE(context.capturing); + EXPECT_EQ(context.fia_tasks.size(), 1); +} + namespace { const KVCache& first_full_attention_cache( const std::vector& kv_caches) { @@ -1135,6 +1152,66 @@ TEST_F(AclGraphExecutorTest, GraphDoubleBufferFlagControlsSlotCount) { original_enable_graph_double_buffer); } +TEST_F(AclGraphExecutorTest, SingleSlotGraphKeySeparatesPagedAttentionPlans) { + ExecutionConfig& execution_config = ExecutionConfig::get_instance(); + const bool original_enable_graph_double_buffer = + execution_config.enable_graph_double_buffer(); + options_.block_size(128); + options_.decode_context_parallel_size(1); + + ModelInputParams params; + params.meta.batch_forward_type = BatchForwardType::DECODE; + + execution_config.enable_graph_double_buffer(false); + std::unique_ptr<::xllm::npu::AclGraphExecutorImpl> graph_executor = + std::make_unique<::xllm::npu::AclGraphExecutorImpl>( + model_.get(), model_args_, *device_, options_); + params.meta.kv_max_seq_len = 16; + const uint64_t short_context_key = + graph_executor->graph_key_for_test(/*bucket_num_tokens=*/1, params); + params.meta.kv_max_seq_len = 514; + const uint64_t first_long_context_key = + graph_executor->graph_key_for_test(/*bucket_num_tokens=*/1, params); + params.meta.kv_max_seq_len = 515; + const uint64_t second_long_context_key = + graph_executor->graph_key_for_test(/*bucket_num_tokens=*/1, params); + + EXPECT_NE(short_context_key, first_long_context_key); + EXPECT_EQ(first_long_context_key, second_long_context_key); + execution_config.enable_graph_double_buffer( + original_enable_graph_double_buffer); +} + +TEST_F(AclGraphExecutorTest, DoubleBufferGraphKeySeparatesPagedAttentionPlans) { + ExecutionConfig& execution_config = ExecutionConfig::get_instance(); + const bool original_enable_graph_double_buffer = + execution_config.enable_graph_double_buffer(); + options_.block_size(128); + options_.decode_context_parallel_size(1); + + ModelInputParams params; + params.meta.batch_forward_type = BatchForwardType::DECODE; + + execution_config.enable_graph_double_buffer(true); + std::unique_ptr<::xllm::npu::AclGraphExecutorImpl> graph_executor = + std::make_unique<::xllm::npu::AclGraphExecutorImpl>( + model_.get(), model_args_, *device_, options_); + params.meta.kv_max_seq_len = 16; + const uint64_t short_context_key = + graph_executor->graph_key_for_test(/*bucket_num_tokens=*/1, params); + params.meta.kv_max_seq_len = 514; + const uint64_t first_long_context_key = + graph_executor->graph_key_for_test(/*bucket_num_tokens=*/1, params); + params.meta.kv_max_seq_len = 515; + const uint64_t second_long_context_key = + graph_executor->graph_key_for_test(/*bucket_num_tokens=*/1, params); + + EXPECT_NE(short_context_key, first_long_context_key); + EXPECT_EQ(first_long_context_key, second_long_context_key); + execution_config.enable_graph_double_buffer( + original_enable_graph_double_buffer); +} + TEST(AclGraphPersistentParamTest, SpecVerifyMetadataUsesTokenCapacity) { SpeculativeConfig& speculative_config = SpeculativeConfig::get_instance(); const bool original_enable_atb_spec_kernel = @@ -1292,6 +1369,106 @@ TEST(AclGraphPersistentParamTest, kActiveBlockTableWidth); } +TEST(AclGraphPersistentParamTest, DcpDecodeUsesStableRankLocalBlockTable) { + ModelArgs args; + args.model_type("deepseek_v4"); + args.dtype("float32"); + args.hidden_size(8); + args.max_position_embeddings(16); + + runtime::Options options; + options.block_size(4); + options.max_seqs_per_batch(2); + options.max_tokens_per_batch(2); + options.num_decoding_tokens(1); + options.world_size(2); + options.dp_size(1); + options.decode_context_parallel_size(2); + options.node_rank(1); + + 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); + EXPECT_EQ(persistent_param.persistent_block_tables().size(1), 5); + EXPECT_EQ(persistent_param.persistent_dcp_local_block_tables().size(1), 2); + + ModelInputParams params; + params.meta.batch_forward_type = BatchForwardType::DECODE; + params.meta.num_sequences = 2; + params.attention.host.q_seq_lens = {1, 1}; + params.attention.host.kv_seq_lens = {4, 8}; + params.attention.device.q_seq_lens = torch::ones({2}, int_options); + params.attention.device.kv_seq_lens = torch::tensor({4, 8}, int_options); + params.attention.device.new_cache_slots = torch::zeros({2}, int_options); + params.attention.device.block_tables = + torch::tensor({{10, 11, 12, 13, 14}, {20, 21, 22, 23, 24}}, int_options); + const torch::Tensor tokens = torch::tensor({1, 2}, int_options); + const torch::Tensor positions = torch::tensor({3, 7}, int_options); + + auto first = persistent_param.update(tokens, + torch::Tensor(), + torch::Tensor(), + positions, + params, + /*padded_num_tokens=*/2, + /*return_capture_params=*/true); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(first->graph.dcp_local_block_tables.defined()); + EXPECT_EQ(first->graph.dcp_local_block_tables.data_ptr(), + persistent_param.persistent_dcp_local_block_tables(2).data_ptr()); + EXPECT_TRUE(torch::equal(first->graph.dcp_local_block_tables.cpu(), + torch::tensor({{11, 13}, {21, 23}}, torch::kInt))); + const void* local_table_address = + first->graph.dcp_local_block_tables.data_ptr(); + + params.attention.device.block_tables = + torch::tensor({{30, 31, 32, 33, 34}, {40, 41, 42, 43, 44}}, int_options); + auto second = persistent_param.update(tokens, + torch::Tensor(), + torch::Tensor(), + positions, + params, + /*padded_num_tokens=*/2, + /*return_capture_params=*/true); + ASSERT_TRUE(second.has_value()); + EXPECT_EQ(second->graph.dcp_local_block_tables.data_ptr(), + local_table_address); + EXPECT_TRUE(torch::equal(second->graph.dcp_local_block_tables.cpu(), + torch::tensor({{31, 33}, {41, 43}}, torch::kInt))); + + params.meta.num_sequences = 1; + params.attention.host.q_seq_lens = {1}; + params.attention.host.kv_seq_lens = {4}; + params.attention.host.q_cu_seq_lens = {1}; + params.attention.device.q_seq_lens = torch::ones({1}, int_options); + params.attention.device.kv_seq_lens = torch::tensor({4}, int_options); + params.attention.device.new_cache_slots = torch::zeros({1}, int_options); + params.attention.device.block_tables = + torch::tensor({{50, 51, 52, 53, 54}}, int_options); + const torch::Tensor single_token = torch::tensor({3}, int_options); + const torch::Tensor single_position = torch::tensor({3}, int_options); + + auto padded = persistent_param.update(single_token, + torch::Tensor(), + torch::Tensor(), + single_position, + params, + /*padded_num_tokens=*/2, + /*return_capture_params=*/true); + ASSERT_TRUE(padded.has_value()); + EXPECT_EQ(padded->attention.host.q_seq_lens, (std::vector{1, 1})); + EXPECT_EQ(padded->attention.host.kv_seq_lens, (std::vector{4, 1})); + EXPECT_EQ(padded->attention.host.q_cu_seq_lens, (std::vector{1, 2})); + EXPECT_TRUE(torch::equal(padded->graph.dcp_local_block_tables.cpu(), + torch::tensor({{51, 53}, {0, 0}}, torch::kInt))); +} + TEST(AclGraphPersistentParamTest, AuxHiddenStatesUseGraphTokenCapacity) { SpeculativeConfig& speculative_config = SpeculativeConfig::get_instance(); const bool original_enable_atb_spec_kernel = diff --git a/tests/core/scheduler/continuous_scheduler_test.cpp b/tests/core/scheduler/continuous_scheduler_test.cpp index 28aa4c8a55..cc65b00bbe 100644 --- a/tests/core/scheduler/continuous_scheduler_test.cpp +++ b/tests/core/scheduler/continuous_scheduler_test.cpp @@ -11,6 +11,7 @@ #include "core/framework/config/rec_config.h" #include "core/framework/config/scheduler_config.h" #include "distributed_runtime/engine.h" +#include "scheduler/chunked_prefill_policy.h" #include "scheduler_factory.h" #include "util/utils.h" @@ -529,6 +530,112 @@ TEST(ContinuousSchedulerFactoryTest, opt.max_tokens_per_chunk_for_prefill()); } +TEST(ContinuousSchedulerBatchModeTest, Dcp1KeepsMixBatch) { + ContinuousScheduler::Options opt = + create_scheduler_options(10000, 256, 0, 1024, 1); + opt.enable_chunked_prefill() = true; + ::xllm::SchedulerConfig::get_instance().enable_mix_batch() = true; + opt.decode_context_parallel_size() = 1; + + BatchMode mode = resolve_batch_mode(opt); + + EXPECT_TRUE(mode.enable_mix_batch); +} + +TEST(ContinuousSchedulerBatchModeTest, Dcp2ClosesMixBatch) { + ContinuousScheduler::Options opt = + create_scheduler_options(10000, 256, 0, 1024, 1); + opt.enable_chunked_prefill() = true; + ::xllm::SchedulerConfig::get_instance().enable_mix_batch() = true; + opt.decode_context_parallel_size() = 2; + + BatchMode mode = resolve_batch_mode(opt); + + EXPECT_FALSE(mode.enable_mix_batch); +} + +TEST(ContinuousSchedulerBatchModeTest, Dcp2ClosesMixBatchUnderMultiSlo) { + ContinuousScheduler::Options opt = create_scheduler_options( + 10000, 256, 0, 1024, 1, /*priority_strategy=*/"multi_slo_and_prio"); + ::xllm::SchedulerConfig::get_instance().enable_mix_batch() = true; + opt.decode_context_parallel_size() = 2; + + BatchMode mode = resolve_batch_mode(opt); + + EXPECT_FALSE(mode.enable_mix_batch); +} + +TEST(ContinuousSchedulerBatchModeTest, MultiSloForcesEffectiveChunkedPrefill) { + ContinuousScheduler::Options opt = create_scheduler_options( + 10000, 256, 0, 1024, 1, /*priority_strategy=*/"multi_slo_and_prio"); + ::xllm::SchedulerConfig::get_instance().enable_mix_batch() = true; + opt.enable_chunked_prefill() = false; + opt.decode_context_parallel_size() = 2; + + BatchMode mode = resolve_batch_mode(opt); + + // multi_slo_and_prio drives chunked prefill even when the raw flag is false; + // DCP then closes mix batching. + EXPECT_TRUE(mode.enable_chunked_prefill); + EXPECT_FALSE(mode.enable_mix_batch); +} + +TEST(ChunkedPrefillPolicyTest, EffectiveChunkedTruthTable) { + // Raw flag wins when set, regardless of strategy. + EXPECT_TRUE(resolve_effective_chunked_prefill(true, "fcfs")); + EXPECT_TRUE(resolve_effective_chunked_prefill(true, "multi_slo_and_prio")); + // multi_slo_and_prio implicitly enables chunked prefill. + EXPECT_TRUE(resolve_effective_chunked_prefill(false, "multi_slo_and_prio")); + // Otherwise the raw false value is preserved. + EXPECT_FALSE(resolve_effective_chunked_prefill(false, "fcfs")); + EXPECT_FALSE(resolve_effective_chunked_prefill(false, "priority")); + EXPECT_FALSE(resolve_effective_chunked_prefill(false, "deadline")); +} + +TEST(ContinuousSchedulerFactoryTest, + ChunkedPrefillWithDcpDoesNotBuildMixedBatch) { + ContinuousScheduler::Options opt = create_scheduler_options(8, 8, 0, 4, 1); + opt.enable_chunked_prefill() = true; + opt.decode_context_parallel_size() = 2; // DCP > 1 forces exclusive batch + + auto engine = std::make_unique(32, 32); + auto scheduler = create_continuous_scheduler(engine.get(), opt); + ASSERT_NE(scheduler.get(), nullptr); + + auto requests = generate_request({2, 10}, + {8, 8}, + std::nullopt, + std::nullopt, + std::nullopt, + std::nullopt, + 30000); + for (auto& req : requests) { + scheduler->add_request(req); + } + + auto batches = scheduler->prepare_batch_test(); + ASSERT_EQ(batches.size(), 1); + ASSERT_EQ(batches[0].size(), 2); + const auto& allowed_max_tokens = batches[0].get_allowed_max_tokens(); + ASSERT_EQ(allowed_max_tokens.size(), 2); + + make_request_decode_ready(requests[0]); + set_chunk_kv(requests[1], allowed_max_tokens[1]); + + batches = scheduler->prepare_batch_test(); + ASSERT_EQ(batches.size(), 1); + ASSERT_EQ(batches[0].size(), 1); + + const auto forward_input = + batches[0].prepare_forward_input(1, 0, ModelArgs()); + EXPECT_TRUE( + forward_input.input_params.meta.batch_forward_type.is_chunked_prefill()); + EXPECT_FALSE(forward_input.input_params.meta.batch_forward_type.is_mixed()); + EXPECT_EQ(forward_input.input_params.meta.num_sequences, 1); + EXPECT_EQ(batches[0].get_allowed_max_tokens()[0], + opt.max_tokens_per_chunk_for_prefill()); +} + TEST(SchedulerFactoryTest, DisaggPDChunkedPrefillUsesDisaggPD) { ContinuousScheduler::Options opt = create_scheduler_options(10000, 256, 2, 1024, 1); diff --git a/tests/python/import_xllm_export.py b/tests/python/import_xllm_export.py new file mode 100644 index 0000000000..7688a649af --- /dev/null +++ b/tests/python/import_xllm_export.py @@ -0,0 +1,50 @@ +# 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://github.com/xLLM-AI/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. + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _load_extension(module_path: Path) -> ModuleType: + spec = importlib.util.spec_from_file_location("xllm_export", module_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Failed to create import spec for {module_path}") + + module = importlib.util.module_from_spec(spec) + sys.modules["xllm_export"] = module + try: + spec.loader.exec_module(module) + except Exception: + sys.modules.pop("xllm_export", None) + raise + return module + + +def main() -> None: + if len(sys.argv) != 2: + raise RuntimeError("Expected the xllm_export shared-library path") + + module = _load_extension(Path(sys.argv[1])) + if not hasattr(module, "LLMMaster"): + raise RuntimeError("xllm_export loaded without LLMMaster") + if not hasattr(sys.modules.get("xllm_runtime"), "AttentionMetadataView"): + raise RuntimeError("xllm_runtime was not registered") + if not hasattr(sys.modules.get("xllm_weight_loader"), "StateDict"): + raise RuntimeError("xllm_weight_loader was not registered") + + +if __name__ == "__main__": + main() diff --git a/tests/python/test_glm5_2_parallel.py b/tests/python/test_glm5_2_parallel.py index 70de4e59f4..4ddc9cbffe 100644 --- a/tests/python/test_glm5_2_parallel.py +++ b/tests/python/test_glm5_2_parallel.py @@ -74,6 +74,12 @@ def test_full_world_ep_partitions_glm_experts() -> None: assert moe.local_expert_start == 6 assert moe.local_expert_end == 8 + assert moe.experts_w13.numel() == 0 + assert moe.experts_w2.numel() == 0 + + moe.allocate_experts_w13_for_loading() + moe.allocate_experts_w2_for_loading() + assert moe.experts_w13.shape == (2, 16, 16) assert moe.experts_w2.shape == (2, 16, 8) @@ -170,7 +176,10 @@ def load_w8a8_b(self, prefix: str) -> None: def test_glm_weight_loader_reads_only_local_ep_experts(monkeypatch) -> None: model = Glm52ForCausalLM(_config(ep_rank=2)) model.model.layers[0].self_attn.process_weights_after_loading = MagicMock() - model.model.layers[0].mlp.process_weights_after_loading = MagicMock() + moe = model.model.layers[0].mlp + moe.process_experts_w13_after_loading = MagicMock() + moe.process_experts_w2_after_loading = MagicMock() + moe.process_weights_after_loading = MagicMock() monkeypatch.setattr(glm5_2, "W8A8WeightLoader", _RecordingLoader) model.load_weights([], tp_rank=0, tp_size=2) @@ -183,3 +192,6 @@ def test_glm_weight_loader_reads_only_local_ep_experts(monkeypatch) -> None: assert loader.tp_size == 2 assert loader.tp_rank == 0 assert loader.shared_shards == [("model.layers.0.mlp.shared_experts.", 1, 0)] + moe.process_experts_w13_after_loading.assert_called_once_with() + moe.process_experts_w2_after_loading.assert_called_once_with() + moe.process_weights_after_loading.assert_called_once_with(skip_expert_format=True) diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index d92d7a9298..1817d8c9b1 100644 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -110,6 +110,8 @@ DECLARE_int32(ep_size); DECLARE_int32(cp_size); +DECLARE_int32(decode_context_parallel_size); + DECLARE_int32(layerwise_split_size); DECLARE_int64(tp_size); diff --git a/xllm/core/common/options.cpp b/xllm/core/common/options.cpp index 66ea3e8b5c..e1b7d88129 100644 --- a/xllm/core/common/options.cpp +++ b/xllm/core/common/options.cpp @@ -63,6 +63,7 @@ std::string Options::to_string() const { << ", flashcomm1_min_prefill_tokens: " << flashcomm1_min_prefill_tokens() << ", enable_mmrs_fusion: " << enable_mmrs_fusion() << ", mmrs_comm_mode: " << mmrs_comm_mode() << ", cp_size: " << cp_size() + << ", decode_context_parallel_size: " << decode_context_parallel_size() << ", master_node_addr: " << master_node_addr().value_or("null") << ", instance_role: " << instance_role().to_string() << ", transfer_listen_port: " << transfer_listen_port() diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index eb374a1c0d..9163a54974 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -145,6 +145,12 @@ class Options { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, decode_context_parallel_size) = 1; + + // Opt-in for the experimental DCP chunked prefill path. See ParallelConfig + // for the numerical caveat; not bitwise-equivalent to dcp=1. + PROPERTY(bool, enable_experimental_dcp_chunked_prefill) = false; + PROPERTY(int32_t, ep_size) = 1; PROPERTY(int32_t, tp_size) = 1; diff --git a/xllm/core/distributed_runtime/dcp_compat.h b/xllm/core/distributed_runtime/dcp_compat.h new file mode 100644 index 0000000000..c1bb652b5d --- /dev/null +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -0,0 +1,61 @@ +/* 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://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include "common/options.h" +#include "common/types.h" +#include "scheduler/chunked_prefill_policy.h" + +namespace xllm { + +inline std::optional validate_dcp_first_version_options( + const Options& options, + EngineType engine_type) { + if (options.decode_context_parallel_size() <= 1) { + return std::nullopt; + } + if (resolve_effective_chunked_prefill(options.enable_chunked_prefill(), + options.priority_strategy()) && + !options.enable_experimental_dcp_chunked_prefill()) { + return "decode_context_parallel_size with chunked prefill is experimental " + "and not bitwise-equivalent to decode_context_parallel_size=1 " + "(priority_strategy=multi_slo_and_prio also implicitly enables " + "chunked prefill); set --enable_experimental_dcp_chunked_prefill=" + "true to opt in, or disable chunked prefill " + "(--enable_chunked_prefill=false and a non-multi_slo_and_prio " + "priority_strategy), or set --decode_context_parallel_size=1"; + } + if (options.enable_disagg_pd() || + options.instance_role() != InstanceRole::DEFAULT) { + return "decode_context_parallel_size first version does not yet support " + "disaggregated prefill-decode; set --enable_disagg_pd=false, " + "--instance_role=DEFAULT, or set --decode_context_parallel_size=1"; + } + if (engine_type == EngineType::SSM || + !options.draft_model_path().value_or("").empty() || + options.num_speculative_tokens() > 0) { + return "decode_context_parallel_size first version does not yet support " + "speculative decoding; unset --draft_model, set " + "--num_speculative_tokens=0, or set " + "--decode_context_parallel_size=1"; + } + return std::nullopt; +} + +} // namespace xllm diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index aa167eef61..87e59469a7 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -56,6 +56,7 @@ limitations under the License. #include "runtime/llm_worker_impl.h" #include "runtime/params_utils.h" #include "runtime/worker.h" +#include "scheduler/chunked_prefill_policy.h" #include "server/xllm_server_registry.h" #include "util/env_var.h" #include "util/pretty_print.h" @@ -611,10 +612,16 @@ bool LLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { const bool is_decode = options_.instance_role() == InstanceRole::DECODE; if (options_.enable_prefix_cache() && enable_gdn_attention && !is_decode) { const auto& scheduler_config = ::xllm::SchedulerConfig::get_instance(); - CHECK(scheduler_config.enable_chunked_prefill()) + // multi_slo_and_prio implicitly enables chunked prefill, so accept the + // effective value rather than the raw flag; the linear-state checkpoints + // are saved on the same chunked path either way. + CHECK(resolve_effective_chunked_prefill( + scheduler_config.enable_chunked_prefill(), + options_.priority_strategy())) << "Linear-attention prefix cache requires block-aligned chunked " "prefill to save matching linear states. Please set " - "--enable_chunked_prefill=true in your config."; + "--enable_chunked_prefill=true (or use priority_strategy=" + "multi_slo_and_prio) in your config."; CHECK(scheduler_config.max_tokens_per_chunk_for_prefill() % block_size == 0) << "linear-attention prefix cache saves linear-state checkpoints at " "chunk-end boundaries, so max_tokens_per_chunk_for_prefill (" diff --git a/xllm/core/distributed_runtime/llm_master.cpp b/xllm/core/distributed_runtime/llm_master.cpp index 60cefcae01..be3a21374a 100644 --- a/xllm/core/distributed_runtime/llm_master.cpp +++ b/xllm/core/distributed_runtime/llm_master.cpp @@ -37,6 +37,7 @@ limitations under the License. #include "framework/request/request.h" #include "models/model_registry.h" #include "runtime/xservice_client.h" +#include "scheduler/chunked_prefill_policy.h" #include "scheduler/scheduler_factory.h" #include "server/xllm_server_registry.h" #include "speculative_engine.h" @@ -122,6 +123,7 @@ LLMMaster::LLMMaster(const Options& options) .nnodes(options_.nnodes()) .dp_size(options_.dp_size()) .cp_size(options_.cp_size()) + .decode_context_parallel_size(options_.decode_context_parallel_size()) .enable_disagg_pd(options_.enable_disagg_pd()) .enable_pd_ooc(options_.enable_pd_ooc()) .enable_schedule_overlap(options_.enable_schedule_overlap()) @@ -396,7 +398,8 @@ std::shared_ptr LLMMaster::generate_request( const int32_t max_context_len = model_args_.max_position_embeddings(); int32_t prompt_token_limit = max_context_len; - if (!options_.enable_chunked_prefill()) { + if (!resolve_effective_chunked_prefill(options_.enable_chunked_prefill(), + options_.priority_strategy())) { prompt_token_limit = std::min(prompt_token_limit, options_.max_tokens_per_batch()); } diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index aa2dc8d74d..0e5a0294cc 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -35,12 +35,14 @@ limitations under the License. #include "common/metrics.h" #include "common/types.h" #include "core/common/xllm_build_info.h" +#include "core/distributed_runtime/dcp_compat.h" #include "core/framework/config/eplb_config.h" #include "core/framework/config/kernel_config.h" #include "core/framework/config/kv_cache_config.h" #include "core/framework/config/model_config.h" #include "core/framework/config/parallel_config.h" #include "core/framework/config/speculative_config.h" +#include "core/scheduler/chunked_prefill_policy.h" #include "dit_master.h" #if defined(USE_NPU) #include "framework/parallel_state/npu_rank_table_env.h" @@ -57,6 +59,7 @@ limitations under the License. #include "rec_master.h" #include "runtime/options.h" #include "speculative_engine.h" +#include "util/json_reader.h" #include "util/model_config_utils.h" #include "util/scope_guard.h" #include "util/timer.h" @@ -72,6 +75,146 @@ DECLARE_bool(graceful_quit_on_sighup); namespace xllm { namespace { +struct DcpModelConfig { + std::string model_type; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; +}; + +bool is_qwen3_5_text_model_type(const std::string& model_type) { + return model_type == "qwen3_5_text" || model_type == "qwen3_5_moe_text"; +} + +DcpModelConfig load_dcp_model_config( + const std::filesystem::path& model_path, + const std::optional& backend) { + const std::filesystem::path config_json_path = model_path / "config.json"; + CHECK(std::filesystem::exists(config_json_path)) + << "Please check config.json file in model path: " << model_path; + + JsonReader reader; + CHECK(reader.parse(config_json_path.string())) + << "Failed to parse config.json file in model path: " << model_path; + + DcpModelConfig config; + config.model_type = util::get_model_type(reader, model_path, backend); + config.num_attention_heads = reader.value_or( + std::vector{"text_config.num_attention_heads", + "num_attention_heads"}, + int64_t{0}); + config.num_key_value_heads = reader.value_or( + std::vector{"text_config.num_key_value_heads", + "num_key_value_heads"}, + int64_t{0}); + return config; +} + +std::optional validate_model_dcp( + const Options& options, + EngineType engine_type, + const std::optional& model_config, + int32_t global_world_size) { + const int32_t dcp_size = options.decode_context_parallel_size(); + if (dcp_size < 1) { + return "decode_context_parallel_size must be greater than or equal to 1"; + } + + if (model_config.has_value() && + is_qwen3_5_text_model_type(model_config->model_type) && + options.cp_size() > 1) { + return "Qwen3.5 decode context parallelism uses " + "--decode_context_parallel_size, not --cp_size; keep cp_size=1"; + } + + if (dcp_size == 1) { + return std::nullopt; + } + + if (std::optional dcp_option_error = + validate_dcp_first_version_options(options, engine_type)) { + return dcp_option_error; + } + + if (options.cp_size() != 1) { + return "decode_context_parallel_size cannot be combined with cp_size; " + "keep cp_size=1"; + } + if (!Platform::is_npu()) { + return "decode_context_parallel_size is currently supported only on NPU"; + } + if (options.npu_kernel_backend() != "TORCH") { + return "decode_context_parallel_size requires --npu_kernel_backend=TORCH"; + } + if (engine_type != EngineType::LLM && engine_type != EngineType::SSM) { + return "decode context parallelism supports only LLM text generation"; + } + if (options.task_type() != "generate") { + return "decode context parallelism supports only the generate task"; + } + if (!model_config.has_value()) { + return "decode_context_parallel_size requires model config to validate " + "Qwen3.5 GQA topology"; + } + if (!is_qwen3_5_text_model_type(model_config->model_type)) { + return "decode_context_parallel_size currently supports only Qwen3.5 " + "text models, got model_type=" + + model_config->model_type; + } + if (options.dp_size() < 1) { + return "decode context parallelism requires dp_size >= 1"; + } + if (global_world_size < 1) { + return "decode context parallelism requires world_size >= 1"; + } + if (global_world_size % options.dp_size() != 0) { + return "decode context parallelism requires world_size divisible by " + "dp_size"; + } + + const int64_t tp_size = global_world_size / options.dp_size(); + if (tp_size < 1) { + return "decode context parallelism requires tensor parallel size >= 1"; + } + const int64_t num_attention_heads = model_config->num_attention_heads; + const int64_t num_key_value_heads = model_config->num_key_value_heads; + if (num_attention_heads <= 0 || num_key_value_heads <= 0) { + return "decode context parallelism requires positive num_attention_heads " + "and num_key_value_heads in config.json"; + } + if (num_attention_heads % tp_size != 0) { + return "decode context parallelism requires num_attention_heads divisible " + "by tensor parallel size"; + } + if (num_attention_heads % num_key_value_heads != 0) { + return "decode context parallelism requires num_attention_heads divisible " + "by num_key_value_heads"; + } + if (tp_size <= num_key_value_heads) { + return "decode context parallelism for Qwen3.5 GQA requires tensor " + "parallel size greater than num_key_value_heads"; + } + if (tp_size % num_key_value_heads != 0) { + return "decode context parallelism requires tensor parallel size divisible " + "by num_key_value_heads"; + } + + const int64_t num_kv_head_replicas = tp_size / num_key_value_heads; + const int64_t num_q_heads_per_kv = num_attention_heads / num_key_value_heads; + if (dcp_size > num_kv_head_replicas) { + return "decode_context_parallel_size exceeds the number of replicated KV " + "head ranks in the TP group"; + } + if (num_q_heads_per_kv % dcp_size != 0) { + return "num_attention_heads / num_key_value_heads must be divisible by " + "decode_context_parallel_size"; + } + if (num_kv_head_replicas % dcp_size != 0) { + return "tensor parallel KV-head replica count must be divisible by " + "decode_context_parallel_size"; + } + return std::nullopt; +} + void apply_runtime_kv_cache_options(const Options& source, runtime::Options& destination) { destination.host_blocks_factor(source.host_blocks_factor()) @@ -96,6 +239,9 @@ void validate_layerwise_split_size_startup_config(const Options& options, return; } + CHECK_EQ(options.decode_context_parallel_size(), 1) + << "layerwise_split_size > 1 does not support decode context " + "parallelism."; CHECK_GE(options.dp_size(), 1) << "dp_size must be >= 1."; CHECK_GT(global_world_size, 0) << "world_size must be > 0."; const int32_t dp_cp_size = options.dp_size() * options.cp_size(); @@ -399,11 +545,20 @@ Master::Master(const Options& options, EngineType type) const std::vector devices = {visible_devices[device_idx]}; // World size is the node count (one worker per process). const int32_t global_world_size = options_.nnodes(); - std::string model_type; - if ((options_.cp_size() > 1 && Platform::uses_model_cp_sharding()) || - (ModelConfig::is_python_model_impl( - ModelConfig::get_instance().model_impl()) && - options_.num_speculative_tokens() > 0)) { +#if defined(USE_NPU) + resolve_npu_kernel_backend_for_options(&options_); +#endif + std::optional dcp_model_config; + if (options_.decode_context_parallel_size() > 1 || + (options_.cp_size() > 1 && Platform::uses_model_cp_sharding())) { + dcp_model_config = load_dcp_model_config(model_path, options_.backend()); + } + std::string model_type = + dcp_model_config.has_value() ? dcp_model_config->model_type : ""; + if (model_type.empty() && + ModelConfig::is_python_model_impl( + ModelConfig::get_instance().model_impl()) && + options_.num_speculative_tokens() > 0) { model_type = util::get_model_type(model_path, options_.backend()); } const std::optional speculative_error = @@ -412,6 +567,36 @@ Master::Master(const Options& options, EngineType type) model_type, options_.num_speculative_tokens()); CHECK(!speculative_error.has_value()) << speculative_error.value(); + const std::optional dcp_error = + validate_model_dcp(options_, type, dcp_model_config, global_world_size); + CHECK(!dcp_error.has_value()) << dcp_error.value(); + if (options_.decode_context_parallel_size() > 1 && + resolve_effective_chunked_prefill(options_.enable_chunked_prefill(), + options_.priority_strategy()) && + options_.enable_experimental_dcp_chunked_prefill()) { + LOG(WARNING) + << "Experimental DCP chunked prefill is enabled " + "(--enable_experimental_dcp_chunked_prefill=true). This path is not " + "bitwise-equivalent to decode_context_parallel_size=1 (BF16/FP16 " + "partial-output quantization before the FP32 merge), mixed batching " + "is automatically disabled, and it is unsupported for release. To " + "roll back, set --enable_experimental_dcp_chunked_prefill=false and " + "disable chunked prefill (--enable_chunked_prefill=false and a " + "non-multi_slo_and_prio priority_strategy), or set " + "--decode_context_parallel_size=1."; + } + if (options_.decode_context_parallel_size() > 1 && + dcp_model_config.has_value() && + dcp_model_config->model_type == "qwen3_5_moe_text") { + LOG(WARNING) + << "Qwen3.5 MoE with decode_context_parallel_size>1 is not " + "bitwise-equivalent to decode_context_parallel_size=1. The DCP " + "decode attention uses a different kernel path (FIA) than dcp=1 " + "(batch_decode), and MoE expert routing amplifies that difference " + "into a measurable per-step logprob delta (order 1e-2). This is an " + "accepted known limitation; the unified kernel path fix is future " + "work. Set --decode_context_parallel_size=1 to avoid it."; + } const std::optional cp_error = validate_model_cp(options_, type, model_type, global_world_size); CHECK(!cp_error.has_value()) << cp_error.value(); @@ -419,10 +604,14 @@ Master::Master(const Options& options, EngineType type) print_startup_banner(model_path, options_.backend(), options_.node_rank()); LOG(INFO) << "Master init options: " << options_.to_string(); ParallelConfig::get_instance().cp_size(options_.cp_size()); + ParallelConfig::get_instance().decode_context_parallel_size( + options_.decode_context_parallel_size()); // cp_size <= 1 -> "disabled", otherwise "model" (model-side CP). const char* cp_sharding_stage = options_.cp_size() <= 1 ? "disabled" : "model"; LOG(INFO) << "Resolved CP config: cp_size=" << options_.cp_size() + << ", decode_context_parallel_size=" + << options_.decode_context_parallel_size() << ", world_size=" << global_world_size << ", dp_size=" << options_.dp_size() << ", ep_size=" << options_.ep_size() @@ -463,7 +652,6 @@ Master::Master(const Options& options, EngineType type) eplb_config.eplb_min_peak_load_improvement( options.eplb_min_peak_load_improvement().value()); } - resolve_npu_kernel_backend_for_options(&options_); #endif validate_layerwise_split_size_startup_config( options_, model_type, global_world_size); @@ -651,6 +839,7 @@ Master::Master(const Options& options, EngineType type) .enable_mmrs_fusion(options_.enable_mmrs_fusion()) .mmrs_comm_mode(options_.mmrs_comm_mode()) .cp_size(options_.cp_size()) + .decode_context_parallel_size(options_.decode_context_parallel_size()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .max_tokens_per_batch(options_.max_tokens_per_batch()) .max_seqs_per_batch(options_.max_seqs_per_batch()) diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h index 6ef4e309b3..ee738c771d 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h @@ -21,12 +21,13 @@ limitations under the License. namespace xllm::spawn_worker_protocol { -inline constexpr int32_t kArgumentCount = 38; +inline constexpr int32_t kArgumentCount = 39; inline constexpr int32_t kMinimumArgumentCount = 34; inline constexpr int32_t kIndexerCacheDtypeArgumentIndex = 34; inline constexpr int32_t kEnableMtpDraftBodyTp1ArgumentIndex = 35; inline constexpr int32_t kTextEncoderTpSizeArgumentIndex = 36; inline constexpr int32_t kDraftSamplingModeArgumentIndex = 37; +inline constexpr int32_t kDecodeContextParallelSizeArgumentIndex = 38; inline constexpr char kDefaultIndexerCacheDtype[] = "auto"; inline constexpr char kDefaultDraftSamplingMode[] = "greedy"; diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp index 3ff5f118a8..51527785cf 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp @@ -93,7 +93,8 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, int32_t ep_size, const InstanceRole& instance_role, bool enable_mtp_draft_body_tp1, - const std::string& draft_sampling_mode) { + const std::string& draft_sampling_mode, + int32_t decode_context_parallel_size) { // TODO: pass whole xllm::runtime::Options here from main process. xllm::runtime::Options runner_options; const std::string backend = get_backend_from_worker_type(worker_type); @@ -121,6 +122,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, .dp_size(dp_size) .ep_size(ep_size) .cp_size(cp_size) + .decode_context_parallel_size(decode_context_parallel_size) .tp_size(tp_size) .sp_size(effective_sp_size) .cfg_size(effective_cfg_size) @@ -145,6 +147,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, .dp_size(dp_size) .ep_size(ep_size) .cp_size(cp_size) + .decode_context_parallel_size(decode_context_parallel_size) .tp_size(tp_size) .sp_size(effective_sp_size) .cfg_size(effective_cfg_size) @@ -190,6 +193,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, cp_size, /* process_group = */ nullptr, ep_size); + parallel_args.dcp_size(decode_context_parallel_size); worker_server_ = std::make_unique(local_rank, master_node_addr, done_, diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h index 425aab3ef9..860dc0bcc3 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h @@ -64,7 +64,8 @@ class SpawnWorkerServer final { int32_t ep_size, const InstanceRole& instance_role, bool enable_mtp_draft_body_tp1, - const std::string& draft_sampling_mode); + const std::string& draft_sampling_mode, + int32_t decode_context_parallel_size); ~SpawnWorkerServer(); diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp index 9b78310a7f..67fd81f367 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp @@ -62,6 +62,7 @@ limitations under the License. // @enable_mtp_draft_body_tp1 // @text_encoder_tp_size // @draft_sampling_mode +// @decode_context_parallel_size int main(int argc, char* argv[]) { const std::optional parsed_indexer_cache_dtype = xllm::spawn_worker_protocol::parse_indexer_cache_dtype(argc, argv); @@ -127,13 +128,24 @@ int main(int argc, char* argv[]) { atoi(argv[xllm::spawn_worker_protocol:: kTextEncoderTpSizeArgumentIndex])) : 1; + const int32_t decode_context_parallel_size = + argc > xllm::spawn_worker_protocol:: + kDecodeContextParallelSizeArgumentIndex + ? static_cast( + atoi(argv[xllm::spawn_worker_protocol:: + kDecodeContextParallelSizeArgumentIndex])) + : 1; if (world_size < 1 || global_rank < 0 || global_rank >= world_size || cp_size < 1 || ep_size < 1 || text_encoder_tp_size < 1 || + decode_context_parallel_size < 1 || (instance_role_str != "DEFAULT" && instance_role_str != "PREFILL" && instance_role_str != "DECODE")) { LOG(ERROR) << "Invalid spawn worker topology: global_rank=" << global_rank << ", world_size=" << world_size << ", cp_size=" << cp_size << ", ep_size=" << ep_size + << ", text_encoder_tp_size=" << text_encoder_tp_size + << ", decode_context_parallel_size=" + << decode_context_parallel_size << ", instance_role=" << instance_role_str; return 1; } @@ -172,7 +184,9 @@ int main(int argc, char* argv[]) { << ", text_encoder_tp_size = " << text_encoder_tp_size << ", indexer_cache_dtype = " << indexer_cache_dtype << ", enable_mtp_draft_body_tp1 = " << enable_mtp_draft_body_tp1 - << ", draft_sampling_mode = " << draft_sampling_mode << "\n"; + << ", draft_sampling_mode = " << draft_sampling_mode + << ", decode_context_parallel_size = " << decode_context_parallel_size + << "\n"; xllm::SpawnWorkerServer worker(master_node_addr, local_rank, @@ -210,7 +224,8 @@ int main(int argc, char* argv[]) { ep_size, instance_role, enable_mtp_draft_body_tp1, - draft_sampling_mode); + draft_sampling_mode, + decode_context_parallel_size); worker.run(); diff --git a/xllm/core/distributed_runtime/worker_server.cpp b/xllm/core/distributed_runtime/worker_server.cpp index 36615ee83f..14b26b240a 100644 --- a/xllm/core/distributed_runtime/worker_server.cpp +++ b/xllm/core/distributed_runtime/worker_server.cpp @@ -272,6 +272,10 @@ void WorkerServer::create_spawn_server(int32_t local_rank, const char* is_local_ptr = is_local_str.c_str(); std::string cp_size_str = std::to_string(options.cp_size()); const char* cp_size_ptr = cp_size_str.c_str(); + std::string decode_context_parallel_size_str = + std::to_string(options.decode_context_parallel_size()); + const char* decode_context_parallel_size_ptr = + decode_context_parallel_size_str.c_str(); std::string ep_size_str = std::to_string(parallel_args.ep_size()); const char* ep_size_ptr = ep_size_str.c_str(); std::string instance_role_str = options.instance_role().to_string(); @@ -380,6 +384,7 @@ void WorkerServer::create_spawn_server(int32_t local_rank, enable_mtp_draft_body_tp1_ptr, text_encoder_tp_size_ptr, draft_sampling_mode_ptr, + decode_context_parallel_size_ptr, nullptr}; static_assert(std::size(argv) == spawn_worker_protocol::kArgumentCount + 1); pid_t pid; diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index f5624b04fe..880d6a986b 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -26,6 +26,19 @@ DEFINE_int32(ep_size, 1, "Expert parallel size for MoE model."); DEFINE_int32(cp_size, 1, "Context parallel size for DSA attention."); +DEFINE_int32(decode_context_parallel_size, + 1, + "Decode context parallel size. DCP shards decode attention KV " + "cache along sequence within a TP group and does not expand " + "world size."); + +DEFINE_bool(enable_experimental_dcp_chunked_prefill, + false, + "Opt-in for the experimental DCP chunked prefill path. Required " + "when decode_context_parallel_size > 1 is combined with chunked " + "prefill. The path is not bitwise-equivalent to " + "decode_context_parallel_size=1 and closes mixed batching."); + DEFINE_int32( layerwise_split_size, 1, @@ -88,6 +101,8 @@ void ParallelConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(dp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(ep_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(cp_size); + XLLM_CONFIG_ASSIGN_FROM_FLAG(decode_context_parallel_size); + XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_experimental_dcp_chunked_prefill); XLLM_CONFIG_ASSIGN_FROM_FLAG(layerwise_split_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(kv_split_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(tp_size); @@ -106,6 +121,8 @@ void ParallelConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(dp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(ep_size); XLLM_CONFIG_ASSIGN_FROM_JSON(cp_size); + XLLM_CONFIG_ASSIGN_FROM_JSON(decode_context_parallel_size); + XLLM_CONFIG_ASSIGN_FROM_JSON(enable_experimental_dcp_chunked_prefill); XLLM_CONFIG_ASSIGN_FROM_JSON(layerwise_split_size); XLLM_CONFIG_ASSIGN_FROM_JSON(tp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size); @@ -125,6 +142,10 @@ void ParallelConfig::append_config_json( APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, dp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, ep_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, cp_size); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, decode_context_parallel_size); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, enable_experimental_dcp_chunked_prefill); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, layerwise_split_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, tp_size); diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 45f2492242..687530dff9 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/xllm/core/framework/config/parallel_config.h @@ -44,6 +44,8 @@ class ParallelConfig final { {"dp_size", "ep_size", "cp_size", + "decode_context_parallel_size", + "enable_experimental_dcp_chunked_prefill", "layerwise_split_size", "tp_size", "sp_size", @@ -64,6 +66,16 @@ class ParallelConfig final { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, decode_context_parallel_size) = 1; + + // Opt-in flag for the experimental DCP chunked prefill path. When + // decode_context_parallel_size > 1 and chunked prefill is enabled, the DCP + // attention splits into per-shard partial FIA calls whose partial outputs are + // BF16/FP16-quantized before the FP32 online-softmax merge, so results are + // not bitwise-equivalent to dcp=1. Keep false unless explicitly running the + // experimental path. + PROPERTY(bool, enable_experimental_dcp_chunked_prefill) = false; + PROPERTY(int32_t, layerwise_split_size) = 1; // 0 means follow cp_size (legacy KV-split width). diff --git a/xllm/core/framework/model/model_input_params.h b/xllm/core/framework/model/model_input_params.h index 0427ee8d83..2c54a6773d 100644 --- a/xllm/core/framework/model/model_input_params.h +++ b/xllm/core/framework/model/model_input_params.h @@ -954,6 +954,7 @@ struct GraphInput { torch::Tensor expanded_tiling_data; std::vector expanded_kv_seq_lens_vec; #if defined(USE_NPU) + torch::Tensor dcp_local_block_tables; std::shared_ptr acl_graph_task_update_context; #endif torch::Tensor input_tokens_override; @@ -990,6 +991,8 @@ struct GraphInput { out.expanded_tiling_data = safe_to(expanded_tiling_data, device, true); out.expanded_kv_seq_lens_vec = expanded_kv_seq_lens_vec; #if defined(USE_NPU) + out.dcp_local_block_tables = + safe_to(dcp_local_block_tables, device, /*non_blocking=*/true); out.acl_graph_task_update_context = acl_graph_task_update_context; #endif out.input_tokens_override = diff --git a/xllm/core/framework/parallel_state/collective_communicator.cpp b/xllm/core/framework/parallel_state/collective_communicator.cpp index 805624c837..eb7a425ff9 100644 --- a/xllm/core/framework/parallel_state/collective_communicator.cpp +++ b/xllm/core/framework/parallel_state/collective_communicator.cpp @@ -244,6 +244,8 @@ CollectiveCommunicator::CollectiveCommunicator(int global_rank, global_rank, world_size, dp_size, cp_size, nullptr, ep_size); parallel_args_->kv_split_size( ::xllm::ParallelConfig::get_instance().kv_split_size()); + parallel_args_->dcp_size( + ::xllm::ParallelConfig::get_instance().decode_context_parallel_size()); apply_layerwise_split_config(parallel_args_.get()); return; } @@ -304,12 +306,16 @@ CollectiveCommunicator::CollectiveCommunicator(int global_rank, dispatchAndCombineHcclComm); parallel_args_->kv_split_size( ::xllm::ParallelConfig::get_instance().kv_split_size()); + parallel_args_->dcp_size( + ::xllm::ParallelConfig::get_instance().decode_context_parallel_size()); apply_layerwise_split_config(parallel_args_.get()); #else parallel_args_ = std::make_unique( global_rank, world_size, dp_size, cp_size, nullptr, ep_size); parallel_args_->kv_split_size( ::xllm::ParallelConfig::get_instance().kv_split_size()); + parallel_args_->dcp_size( + ::xllm::ParallelConfig::get_instance().decode_context_parallel_size()); apply_layerwise_split_config(parallel_args_.get()); #endif } @@ -511,6 +517,47 @@ void CollectiveCommunicator::create_process_groups( parallel_args_->cp_group_ = tp_group_.get(); } + const int32_t dcp_size = parallel_args_->dcp_size_effective(); + if (dcp_size > 1) { +#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_DCU) + CHECK_EQ(tp_size % dcp_size, 0) + << "DCP requires tp_size divisible by dcp_size, tp_size=" << tp_size + << ", dcp_size=" << dcp_size; + + const int32_t tp_rank = global_rank % tp_size; + const int32_t dcp_rank = tp_rank % dcp_size; + const int32_t dcp_group_base = (tp_rank / dcp_size) * dcp_size; + const int32_t dp_base = global_rank - tp_rank; + const std::vector dcp_group_ranks = + parallel_state::compute_dcp_group_ranks( + global_rank, world_size, dp_size, dcp_size); + + const int32_t dcp_groups_per_dp = tp_size / dcp_size; + const int32_t dp_rank = dp_base / tp_size; + const int32_t dcp_group_index = + dp_rank * dcp_groups_per_dp + dcp_group_base / dcp_size; + std::string dcp_host = host; +#if defined(USE_NPU) + if (::xllm::KernelConfig::get_instance().npu_kernel_backend() == "TORCH") { + dcp_host = get_rank_table_server_host(dcp_group_ranks.front(), host); + } +#endif + dcp_group_ = create_process_group(global_rank, + dcp_rank, + dcp_group_ranks, + world_size, + dcp_size, + port + dcp_group_index + 1, + dcp_host, + "dcp_group", + device); + parallel_args_->dcp_group_ = dcp_group_.get(); + port += world_size / dcp_size; +#else + CHECK(false) << "DCP process group is not supported on this platform"; +#endif + } + if (dp_size > 1) { // A DP group varies dp_rank while preserving the full local model-shard // index. Under orthogonal CP that index spans cp_rank AND tp_rank, so the diff --git a/xllm/core/framework/parallel_state/collective_communicator.h b/xllm/core/framework/parallel_state/collective_communicator.h index 2744a4914a..014060b5c4 100644 --- a/xllm/core/framework/parallel_state/collective_communicator.h +++ b/xllm/core/framework/parallel_state/collective_communicator.h @@ -43,6 +43,7 @@ class CollectiveCommunicator : public CollectiveCommunicatorBase { std::unique_ptr single_rank_group_; // Owns NPU standalone CP ProcessGroup (empty on MLU). std::unique_ptr cp_group_; + std::unique_ptr dcp_group_; std::unique_ptr moe_tp_group_; std::unique_ptr moe_ep_group_; std::unique_ptr mc2_group_; diff --git a/xllm/core/framework/parallel_state/parallel_args.h b/xllm/core/framework/parallel_state/parallel_args.h index a34ebfe7b4..b2a64ed38d 100644 --- a/xllm/core/framework/parallel_state/parallel_args.h +++ b/xllm/core/framework/parallel_state/parallel_args.h @@ -152,6 +152,8 @@ struct ParallelArgs { // cp size PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, dcp_size) = 1; + PROPERTY(int32_t, layerwise_split_size) = 1; // Derived: CP rank of the current process within its DP group. @@ -171,6 +173,30 @@ struct ParallelArgs { return kv_split_size_ > 0 ? kv_split_size_ : cp_size_; } + [[nodiscard]] int32_t dcp_size_effective() const noexcept { + if (dcp_group_ != nullptr) { + return dcp_group_->world_size(); + } + return dcp_size_ > 0 ? dcp_size_ : 1; + } + + [[nodiscard]] int32_t dcp_rank() const noexcept { + if (dcp_group_ != nullptr) { + return dcp_group_->rank(); + } + if (dcp_size_effective() <= 1) { + return 0; + } + if (dp_size_ <= 0) { + return 0; + } + const int32_t tp_sz = world_size_ / dp_size_; + if (tp_sz <= 0) { + return 0; + } + return (rank_ % tp_sz) % dcp_size_effective(); + } + [[nodiscard]] int32_t kv_split_rank() const noexcept { if (dcp_group_ != nullptr) { return dcp_group_->rank(); diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index 69856eced0..8c15bb0939 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/xllm/core/framework/parallel_state/parallel_state.cpp @@ -275,6 +275,76 @@ std::vector compute_cp_group_ranks(int32_t global_rank, return ranks; } +std::vector compute_dcp_group_ranks(int32_t global_rank, + int32_t world_size, + int32_t dp_size, + int32_t dcp_size) { + CHECK_GT(dcp_size, 1) << "compute_dcp_group_ranks requires dcp_size > 1."; + CHECK_GT(dp_size, 0) << "dp_size must be positive."; + CHECK_GT(world_size, 0) << "world_size must be positive."; + CHECK_EQ(world_size % dp_size, 0) + << "world_size (" << world_size << ") must be divisible by dp_size (" + << dp_size << ") so that tp_size is integral."; + const int32_t tp_size = world_size / dp_size; + CHECK_EQ(tp_size % dcp_size, 0) + << "tp_size (" << tp_size << ") must be divisible by dcp_size (" + << dcp_size << ")."; + CHECK_GE(global_rank, 0); + CHECK_LT(global_rank, world_size); + + const int32_t tp_rank = global_rank % tp_size; + const int32_t dp_base = global_rank - tp_rank; + const int32_t dcp_group_base = (tp_rank / dcp_size) * dcp_size; + + std::vector ranks; + ranks.reserve(dcp_size); + for (int32_t member = 0; member < dcp_size; ++member) { + ranks.emplace_back(dp_base + dcp_group_base + member); + } + return ranks; +} + +torch::Tensor select_dcp_local_block_table(const torch::Tensor& block_table, + int32_t dcp_size, + int32_t dcp_rank) { + CHECK(block_table.defined()) << "block_table must be defined."; + CHECK_EQ(block_table.dim(), 2) << "block_table must be two-dimensional."; + CHECK_GT(dcp_size, 1) << "dcp_size must be greater than 1."; + CHECK_GE(dcp_rank, 0) << "dcp_rank must be non-negative."; + CHECK_LT(dcp_rank, dcp_size) << "dcp_rank must be smaller than dcp_size."; + + const int64_t block_table_width = block_table.size(1); + if (block_table_width <= dcp_rank) { + return block_table.slice(/*dim=*/1, /*start=*/0, /*end=*/0); + } + + const torch::TensorOptions index_options = + torch::TensorOptions().dtype(torch::kLong).device(block_table.device()); + const torch::Tensor local_block_indices = + torch::arange(dcp_rank, block_table_width, dcp_size, index_options); + return block_table.index_select(/*dim=*/1, local_block_indices); +} + +torch::Tensor remap_dcp_cache_slots(const torch::Tensor& positions, + const torch::Tensor& slots, + int32_t interleave_size, + int32_t dcp_size, + int32_t dcp_rank) { + CHECK_GT(interleave_size, 0) << "interleave_size must be positive."; + CHECK_GT(dcp_size, 1) << "dcp_size must be greater than 1."; + CHECK_GE(dcp_rank, 0) << "dcp_rank must be non-negative."; + CHECK_LT(dcp_rank, dcp_size) << "dcp_rank must be smaller than dcp_size."; + CHECK_EQ(positions.numel(), slots.numel()) + << "positions and slots must have the same token count."; + + const torch::Tensor pos = positions.to(torch::kCPU).to(torch::kLong); + const torch::Tensor slot = slots.to(torch::kCPU).to(torch::kLong); + const torch::Tensor owner = + torch::floor_divide(pos, interleave_size) % dcp_size; + const torch::Tensor mask = (owner == dcp_rank) & (slot >= 0); + return torch::where(mask, slot, torch::full_like(slot, -1, slot.options())); +} + torch::Tensor scatter(torch::Tensor input, ProcessGroup* process_group, int dim) { diff --git a/xllm/core/framework/parallel_state/parallel_state.h b/xllm/core/framework/parallel_state/parallel_state.h index b5b84095ec..e87a561eba 100644 --- a/xllm/core/framework/parallel_state/parallel_state.h +++ b/xllm/core/framework/parallel_state/parallel_state.h @@ -76,6 +76,27 @@ std::vector compute_cp_group_ranks(int32_t global_rank, int32_t dp_size, int32_t cp_size); +// Global ranks in this rank's DCP group, ordered by DCP rank. +std::vector compute_dcp_group_ranks(int32_t global_rank, + int32_t world_size, + int32_t dp_size, + int32_t dcp_size); + +torch::Tensor select_dcp_local_block_table(const torch::Tensor& block_table, + int32_t dcp_size, + int32_t dcp_rank); + +// Remap logical KV cache slots to this DCP rank's local physical slots. Keeps a +// slot only when this DCP rank owns the token; others become -1. Owner uses +// integer floor division on the position tensor (a plain `/` on an integer +// tensor is float true-division and mis-owns tokens with +// 0= the op's required workspace size +// (checked fail-closed). +#define EXEC_NPU_CMD_WITH_WORKSPACE(aclnn_api, ws_tensor, ...) \ + do { \ + static const auto get_workspace_size_func_addr = \ + ::xllm::kernel::npu::aclnn::detail::get_op_api_func_addr( \ + #aclnn_api "GetWorkspaceSize"); \ + static const auto op_api_func_addr = \ + ::xllm::kernel::npu::aclnn::detail::get_op_api_func_addr(#aclnn_api); \ + static const auto init_mem_addr = \ + ::xllm::kernel::npu::aclnn::detail::get_op_api_func_addr( \ + "InitHugeMemThreadLocal"); \ + static const auto uninit_mem_addr = \ + ::xllm::kernel::npu::aclnn::detail::get_op_api_func_addr( \ + "UnInitHugeMemThreadLocal"); \ + static const auto release_mem_addr = \ + ::xllm::kernel::npu::aclnn::detail::get_op_api_func_addr( \ + "ReleaseHugeMem"); \ + CHECK(get_workspace_size_func_addr != nullptr && \ + op_api_func_addr != nullptr) \ + << #aclnn_api << " or " << #aclnn_api "GetWorkspaceSize" << " not in " \ + << ::xllm::kernel::npu::aclnn::detail::get_op_api_lib_name(); \ + auto acl_stream = c10_npu::getCurrentNPUStream().stream(false); \ + uint64_t workspace_size = 0; \ + uint64_t* workspace_size_addr = &workspace_size; \ + ::aclOpExecutor* executor = nullptr; \ + ::aclOpExecutor** executor_addr = &executor; \ + ::xllm::kernel::npu::aclnn::detail::InitHugeMemThreadLocalFn \ + init_mem_func = reinterpret_cast< \ + ::xllm::kernel::npu::aclnn::detail::InitHugeMemThreadLocalFn>( \ + init_mem_addr); \ + ::xllm::kernel::npu::aclnn::detail::UnInitHugeMemThreadLocalFn \ + uninit_mem_func = reinterpret_cast< \ + ::xllm::kernel::npu::aclnn::detail::UnInitHugeMemThreadLocalFn>( \ + uninit_mem_addr); \ + if (init_mem_func) { \ + init_mem_func(nullptr, false); \ + } \ + auto converted_params = ::xllm::kernel::npu::aclnn::detail::convert_types( \ + __VA_ARGS__, workspace_size_addr, executor_addr); \ + static auto get_workspace_size_func = \ + ::xllm::kernel::npu::aclnn::detail::convert_to_op_api_func( \ + converted_params, get_workspace_size_func_addr); \ + auto workspace_status = ::xllm::kernel::npu::aclnn::detail::call( \ + get_workspace_size_func, converted_params); \ + CHECK(workspace_status == 0) \ + << "call " #aclnn_api " failed, detail:" << aclGetRecentErrMsg(); \ + void* workspace_addr = nullptr; \ + uint64_t workspace_pass_size = workspace_size; \ + if (workspace_size != 0) { \ + CHECK((ws_tensor).defined() && \ + static_cast((ws_tensor).nbytes()) >= workspace_size) \ + << "caller workspace too small for " #aclnn_api ": have " \ + << ((ws_tensor).defined() ? (ws_tensor).nbytes() : 0) << " need " \ + << workspace_size; \ + workspace_addr = const_cast((ws_tensor).storage().data()); \ + /* Pass the full buffer capacity (sized to the max envelope at capture) \ + * as workspace_size so the captured graph node reserves the max and \ + * every replay (larger KV) fits, since graph_task_update does not \ + * change the captured workspace_size. Mirrors vLLM's max-workspace. */ \ + workspace_pass_size = static_cast((ws_tensor).nbytes()); \ + } \ + auto acl_call = [=]() -> int { \ + using OpApiFunc = \ + int (*)(void*, uint64_t, ::aclOpExecutor*, const aclrtStream); \ + OpApiFunc op_api_func = reinterpret_cast(op_api_func_addr); \ + auto api_ret = op_api_func( \ + workspace_addr, workspace_pass_size, executor, acl_stream); \ + CHECK(api_ret == 0) << "call " #aclnn_api " failed, detail:" \ + << aclGetRecentErrMsg(); \ + ::xllm::kernel::npu::aclnn::detail::release_convert_types( \ + converted_params); \ + ::xllm::kernel::npu::aclnn::detail::ReleaseHugeMemFn release_mem_func = \ + reinterpret_cast< \ + ::xllm::kernel::npu::aclnn::detail::ReleaseHugeMemFn>( \ + release_mem_addr); \ + if (release_mem_func) { \ + release_mem_func(nullptr, false); \ + } \ + return api_ret; \ + }; \ + at_npu::native::OpCommand cmd; \ + cmd.Name(#aclnn_api); \ + cmd.SetCustomHandler(acl_call); \ + cmd.Run(); \ + if (uninit_mem_func) { \ + uninit_mem_func(nullptr, false); \ + } \ + } while (false) diff --git a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp index 097c029721..468814e071 100644 --- a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp +++ b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp @@ -151,7 +151,76 @@ std::optional to_optional_tensor( namespace xllm::kernel::npu { -std::tuple npu_fused_infer_attention( +// Declares the local variables consumed by XLLM_FIA_V3_ARGS. Kept as a macro so +// the execute / workspace-size-query paths share one definition and cannot +// drift. convert_types() binds several of these by non-const reference, so they +// must be named lvalues (not temporaries). `layout_str` owns the layout string +// backing the mutable char* aclnn expects. +#define XLLM_FIA_V3_SETUP(key, \ + value, \ + atten_mask, \ + block_table, \ + actual_seq_lengths, \ + actual_seq_lengths_kv, \ + input_layout, \ + is_causal) \ + std::vector key_tensors_vec{key}; \ + std::vector value_tensors_vec{value}; \ + torch::TensorList key_tensors(key_tensors_vec); \ + torch::TensorList value_tensors(value_tensors_vec); \ + std::optional none_tensor = std::nullopt; \ + std::optional atten_mask_tensor = \ + to_optional_tensor(atten_mask); \ + std::optional block_table_tensor = \ + to_optional_tensor(block_table); \ + torch::IntArrayRef actual_seq_lengths_ref(actual_seq_lengths); \ + torch::IntArrayRef actual_seq_lengths_kv_ref(actual_seq_lengths_kv); \ + std::optional actual_seq_lengths_opt = \ + actual_seq_lengths_ref; \ + std::optional actual_seq_lengths_kv_opt = \ + actual_seq_lengths_kv_ref; \ + std::optional none_int_array = std::nullopt; \ + std::string layout_str = input_layout; \ + char* input_layout_ptr = const_cast(layout_str.c_str()); \ + int64_t pre_tokens = kSwaIntMax; \ + int64_t next_tokens = is_causal ? 0 : kSwaIntMax; \ + int64_t inner_precise = 0; \ + int64_t antiquant_mode = 0; \ + int64_t key_antiquant_mode = 0; \ + int64_t value_antiquant_mode = 0 + +// The full aclnnFusedInferAttentionScoreV3 positional argument list, referring +// to the locals declared by XLLM_FIA_V3_SETUP plus the enclosing function's +// query/num_heads/scale/num_key_value_heads/sparse_mode/block_size/ +// softmax_lse_flag/output/softmax_lse. +#define XLLM_FIA_V3_ARGS \ + query, key_tensors, value_tensors, none_tensor, /* pse_shift */ \ + atten_mask_tensor, actual_seq_lengths_opt, actual_seq_lengths_kv_opt, \ + none_tensor, /* dequant_scale1 */ \ + none_tensor, /* quant_scale1 */ \ + none_tensor, /* dequant_scale2 */ \ + none_tensor, /* quant_scale2 */ \ + none_tensor, /* quant_offset2 */ \ + none_tensor, /* antiquant_scale */ \ + none_tensor, /* antiquant_offset */ \ + block_table_tensor, none_tensor, /* query_padding_size */ \ + none_tensor, /* kv_padding_size */ \ + none_tensor, /* key_antiquant_scale */ \ + none_tensor, /* key_antiquant_offset */ \ + none_tensor, /* value_antiquant_scale */ \ + none_tensor, /* value_antiquant_offset */ \ + none_tensor, /* key_shared_prefix */ \ + none_tensor, /* value_shared_prefix */ \ + none_int_array, /* actual_shared_prefix_len */ \ + none_tensor, /* query_rope */ \ + none_tensor, /* key_rope */ \ + none_tensor, /* key_rope_antiquant_scale */ \ + num_heads, scale, pre_tokens, next_tokens, input_layout_ptr, \ + num_key_value_heads, sparse_mode, inner_precise, block_size, \ + antiquant_mode, softmax_lse_flag, key_antiquant_mode, \ + value_antiquant_mode, output, softmax_lse + +void npu_fused_infer_attention_out( const torch::Tensor& query, const torch::Tensor& key, const torch::Tensor& value, @@ -166,7 +235,10 @@ std::tuple npu_fused_infer_attention( int64_t sparse_mode, const std::string& input_layout, bool softmax_lse_flag, - bool is_causal) { + bool is_causal, + torch::Tensor& output, + torch::Tensor& softmax_lse, + const std::optional& workspace) { check_tensor(query, "query", "npu_fused_infer_attention"); check_tensor(key, "key", "npu_fused_infer_attention"); check_tensor(value, "value", "npu_fused_infer_attention"); @@ -174,97 +246,116 @@ std::tuple npu_fused_infer_attention( CHECK(!actual_seq_lengths.empty()) << "actual_seq_lengths must not be empty"; CHECK(!actual_seq_lengths_kv.empty()) << "actual_seq_lengths_kv must not be empty"; - - torch::Tensor output = infer_attention_output( - query, value, block_table, num_heads, input_layout); - torch::Tensor softmax_lse = - infer_softmax_lse(query, num_heads, input_layout, softmax_lse_flag); + CHECK(output.defined()) << "output must be preallocated for the out variant"; if (is_ascend950() && input_layout == "TND" && !block_table.has_value()) { CHECK(!softmax_lse_flag) << "Ascend950 torch attention fallback does not return softmax_lse"; - output = ascend950_packed_causal_attention(query, - key, - value, - actual_seq_lengths, - actual_seq_lengths_kv, - num_heads, - num_key_value_heads, - scale); - return {output, softmax_lse}; + output.copy_(ascend950_packed_causal_attention(query, + key, + value, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_key_value_heads, + scale)); + return; } - std::vector key_tensors_vec{key}; - std::vector value_tensors_vec{value}; - torch::TensorList key_tensors(key_tensors_vec); - torch::TensorList value_tensors(value_tensors_vec); - - std::optional none_tensor = std::nullopt; - std::optional atten_mask_tensor = - to_optional_tensor(atten_mask); - std::optional block_table_tensor = - to_optional_tensor(block_table); - - torch::IntArrayRef actual_seq_lengths_ref(actual_seq_lengths); - torch::IntArrayRef actual_seq_lengths_kv_ref(actual_seq_lengths_kv); - std::optional actual_seq_lengths_opt = - actual_seq_lengths_ref; - std::optional actual_seq_lengths_kv_opt = - actual_seq_lengths_kv_ref; - std::optional none_int_array = std::nullopt; - - std::string layout = input_layout; - char* input_layout_ptr = const_cast(layout.c_str()); - int64_t pre_tokens = kSwaIntMax; - int64_t next_tokens = is_causal ? 0 : kSwaIntMax; - int64_t inner_precise = 0; - int64_t antiquant_mode = 0; - int64_t key_antiquant_mode = 0; - int64_t value_antiquant_mode = 0; - - EXEC_NPU_CMD(aclnnFusedInferAttentionScoreV3, - query, - key_tensors, - value_tensors, - none_tensor, // pse_shift - atten_mask_tensor, - actual_seq_lengths_opt, - actual_seq_lengths_kv_opt, - none_tensor, // dequant_scale1 - none_tensor, // quant_scale1 - none_tensor, // dequant_scale2 - none_tensor, // quant_scale2 - none_tensor, // quant_offset2 - none_tensor, // antiquant_scale - none_tensor, // antiquant_offset - block_table_tensor, - none_tensor, // query_padding_size - none_tensor, // kv_padding_size - none_tensor, // key_antiquant_scale - none_tensor, // key_antiquant_offset - none_tensor, // value_antiquant_scale - none_tensor, // value_antiquant_offset - none_tensor, // key_shared_prefix - none_tensor, // value_shared_prefix - none_int_array, // actual_shared_prefix_len - none_tensor, // query_rope - none_tensor, // key_rope - none_tensor, // key_rope_antiquant_scale - num_heads, - scale, - pre_tokens, - next_tokens, - input_layout_ptr, - num_key_value_heads, - sparse_mode, - inner_precise, - block_size, - antiquant_mode, - softmax_lse_flag, - key_antiquant_mode, - value_antiquant_mode, - output, - softmax_lse); + XLLM_FIA_V3_SETUP(key, + value, + atten_mask, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + input_layout, + is_causal); + + if (workspace.has_value() && workspace.value().defined()) { + EXEC_NPU_CMD_WITH_WORKSPACE( + aclnnFusedInferAttentionScoreV3, workspace.value(), XLLM_FIA_V3_ARGS); + return; + } + + EXEC_NPU_CMD(aclnnFusedInferAttentionScoreV3, XLLM_FIA_V3_ARGS); +} + +uint64_t npu_fused_infer_attention_workspace_size( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const std::optional& atten_mask, + const std::optional& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size, + int64_t sparse_mode, + const std::string& input_layout, + bool softmax_lse_flag, + bool is_causal, + torch::Tensor& output, + torch::Tensor& softmax_lse) { + CHECK(output.defined()) << "output must be preallocated for the size query"; + CHECK(!actual_seq_lengths.empty()) << "actual_seq_lengths must not be empty"; + CHECK(!actual_seq_lengths_kv.empty()) + << "actual_seq_lengths_kv must not be empty"; + + XLLM_FIA_V3_SETUP(key, + value, + atten_mask, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + input_layout, + is_causal); + + uint64_t workspace_size = 0; + EXEC_NPU_CMD_GET_WORKSPACE_SIZE( + aclnnFusedInferAttentionScoreV3, workspace_size, XLLM_FIA_V3_ARGS); + return workspace_size; +} + +std::tuple npu_fused_infer_attention( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const std::optional& atten_mask, + const std::optional& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size, + int64_t sparse_mode, + const std::string& input_layout, + bool softmax_lse_flag, + bool is_causal) { + torch::Tensor output = infer_attention_output( + query, value, block_table, num_heads, input_layout); + torch::Tensor softmax_lse = + infer_softmax_lse(query, num_heads, input_layout, softmax_lse_flag); + + npu_fused_infer_attention_out(query, + key, + value, + atten_mask, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_key_value_heads, + scale, + block_size, + sparse_mode, + input_layout, + softmax_lse_flag, + is_causal, + output, + softmax_lse); return {output, softmax_lse}; } diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index e757e41f3a..326cb219b8 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -74,6 +74,52 @@ std::tuple npu_fused_infer_attention( bool softmax_lse_flag = false, bool is_causal = true); +// Out-variant: writes into caller-preallocated `output`/`softmax_lse` instead +// of allocating them, so the buffers keep stable device addresses across +// ACLGraph replays (required for graph capture of the DCP attention path). +void npu_fused_infer_attention_out( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const std::optional& atten_mask, + const std::optional& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size, + int64_t sparse_mode, + const std::string& input_layout, + bool softmax_lse_flag, + bool is_causal, + torch::Tensor& output, + torch::Tensor& softmax_lse, + const std::optional& workspace = std::nullopt); + +// Queries the aclnn workspace size (bytes) required by the FIA out-variant for +// the given inputs, without executing the op. Used to pre-size a caller-owned +// workspace buffer for ACLGraph capture, where the workspace address must stay +// stable across replays. +uint64_t npu_fused_infer_attention_workspace_size( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const std::optional& atten_mask, + const std::optional& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size, + int64_t sparse_mode, + const std::string& input_layout, + bool softmax_lse_flag, + bool is_causal, + torch::Tensor& output, + torch::Tensor& softmax_lse); + void batch_chunked_paged_prefill(const torch::Tensor& query, const torch::Tensor& k_cache, const torch::Tensor& v_cache, diff --git a/xllm/core/layers/common/attention_metadata.h b/xllm/core/layers/common/attention_metadata.h index 2b02963dcb..c4c725597d 100644 --- a/xllm/core/layers/common/attention_metadata.h +++ b/xllm/core/layers/common/attention_metadata.h @@ -31,6 +31,10 @@ namespace ffi = tvm::ffi; #include "dsa_metadata.h" #include "layers/common/kv_shard_batch_metadata.h" +namespace xllm::npu { +class AclGraphTaskUpdateContext; +} // namespace xllm::npu + namespace xllm::layer { struct ExpandedDecodeMetadata { @@ -212,6 +216,16 @@ struct AttentionMetadata { // For ACL graph execution - fixed-address device tiling data for // CustomPagedAttention replay. torch::Tensor paged_attention_tiling_data; + // Fixed-address DCP-local block table prepared outside ACL graph capture. + // Graph-captured DCP decode consumes this directly so no temporary + // arange/index_select inputs are retained by the graph. + torch::Tensor dcp_local_block_table; + // ACL-graph task-update context (from ModelInputParams::graph). When non-null + // and capturing, the DCP attention path records its FIA call as a + // FiaGraphTask so the executor can re-inject the per-step DCP-local KV + // lengths on replay. + std::shared_ptr + acl_graph_task_update_context; // Pre-computed attention mask for npu_fused_infer_attention. torch::Tensor fia_attn_mask; // Host vectors for npu_fused_infer_attention (kernel requires host memory). diff --git a/xllm/core/layers/common/attention_metadata_builder.cpp b/xllm/core/layers/common/attention_metadata_builder.cpp index acf653afce..b2291ad2dc 100644 --- a/xllm/core/layers/common/attention_metadata_builder.cpp +++ b/xllm/core/layers/common/attention_metadata_builder.cpp @@ -355,6 +355,9 @@ AttentionMetadata build_attention_metadata( #endif #if defined(USE_NPU) + attn_metadata.dcp_local_block_table = params.graph.dcp_local_block_tables; + attn_metadata.acl_graph_task_update_context = + params.graph.acl_graph_task_update_context; // Determine if we should use ACL graph mode: // - --enable_graph=true // - Must be decode phase or spec-verify chunked prefill diff --git a/xllm/core/layers/npu_torch/CMakeLists.txt b/xllm/core/layers/npu_torch/CMakeLists.txt index 6b6bd28746..18969db59e 100755 --- a/xllm/core/layers/npu_torch/CMakeLists.txt +++ b/xllm/core/layers/npu_torch/CMakeLists.txt @@ -1,5 +1,17 @@ include(cc_library) +cc_library( + NAME + dcp_attention_utils + HDRS + dcp_attention_utils.h + SRCS + dcp_attention_utils.cpp + DEPS + glog::glog + torch +) + cc_library( NAME npu_torch_layers @@ -40,6 +52,7 @@ cc_library( deepseek_v4_indexer.cpp deepseek_v4_gate.cpp DEPS + :dcp_attention_utils :common_layers :npu_layers :xllm_ops diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index f881966661..e357fd6afd 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -15,8 +15,79 @@ limitations under the License. #include "attention.h" +#include +#include +#include + +#include "framework/parallel_state/parallel_state.h" #include "kernels/npu/npu_ops_api.h" #include "kernels/ops_api.h" +#include "layers/npu_torch/dcp_attention_utils.h" +#include "platform/npu/acl_graph_task_update_context.h" + +namespace { + +void validate_dcp_decode_lengths(const std::vector& q_cu_seq_lens, + const std::vector& global_kv_seq_lens, + int64_t token_count) { + CHECK(!q_cu_seq_lens.empty()) + << "DCP decode requires host cumulative query lengths."; + CHECK_EQ(q_cu_seq_lens.size(), global_kv_seq_lens.size()) + << "DCP decode requires one query and KV length per request."; + CHECK_EQ(token_count, static_cast(global_kv_seq_lens.size())) + << "DCP supports only one-token decode requests."; + + int64_t previous_q_end = 0; + for (const int64_t q_end : q_cu_seq_lens) { + CHECK_EQ(q_end - previous_q_end, 1) + << "DCP supports only one-token decode requests."; + previous_q_end = q_end; + } + CHECK_EQ(previous_q_end, token_count) + << "DCP cumulative query lengths do not match query tokens."; +} + +void normalize_zero_dcp_partials( + torch::Tensor& partial_out, + torch::Tensor& partial_lse, + const std::vector& local_kv_seq_lens) { + CHECK_EQ(partial_out.scalar_type(), torch::kFloat32); + CHECK_EQ(partial_lse.scalar_type(), torch::kFloat32); + CHECK_EQ(partial_out.dim(), 3); + CHECK_EQ(partial_lse.dim(), 3); + CHECK_EQ(partial_out.size(0), partial_lse.size(0)); + CHECK_EQ(partial_out.size(1), partial_lse.size(1)); + CHECK_EQ(partial_lse.size(2), 1); + CHECK_EQ(partial_out.size(0), static_cast(local_kv_seq_lens.size())); + + for (int64_t request_index = 0; + request_index < static_cast(local_kv_seq_lens.size()); + ++request_index) { + if (local_kv_seq_lens[request_index] == 0) { + partial_out.select(0, request_index).zero_(); + partial_lse.select(0, request_index) + .fill_(-std::numeric_limits::infinity()); + } + } +} + +void validate_dcp_partial_shape(const torch::Tensor& partial_out, + const torch::Tensor& partial_lse, + int64_t token_count, + int64_t num_heads, + int64_t head_size, + const char* partial_name) { + CHECK_EQ(partial_out.dim(), 3) << partial_name; + CHECK_EQ(partial_out.size(0), token_count) << partial_name; + CHECK_EQ(partial_out.size(1), num_heads) << partial_name; + CHECK_EQ(partial_out.size(2), head_size) << partial_name; + CHECK_EQ(partial_lse.dim(), 3) << partial_name; + CHECK_EQ(partial_lse.size(0), token_count) << partial_name; + CHECK_EQ(partial_lse.size(1), num_heads) << partial_name; + CHECK_EQ(partial_lse.size(2), 1) << partial_name; +} + +} // namespace namespace xllm { namespace layer { @@ -25,12 +96,21 @@ AttentionImpl::AttentionImpl(int64_t num_heads, int64_t head_size, float scale, int64_t num_kv_heads, - int64_t sliding_window) + int64_t sliding_window, + int32_t dcp_size, + int32_t dcp_rank, + ProcessGroup* dcp_group) : num_heads_(num_heads), head_size_(head_size), num_kv_heads_(num_kv_heads), sliding_window_(sliding_window), - scale_(scale) { + scale_(scale), + dcp_size_(dcp_size), + dcp_rank_(dcp_rank), + dcp_group_(dcp_group) { + CHECK_GT(dcp_size_, 0) << "dcp_size must be positive."; + CHECK_GE(dcp_rank_, 0) << "dcp_rank must be non-negative."; + CHECK_LT(dcp_rank_, dcp_size_) << "dcp_rank must be smaller than dcp_size."; if (sliding_window_ > -1) { sliding_window_ = sliding_window_ - 1; } @@ -68,6 +148,12 @@ std::tuple> AttentionImpl::forward( if (attn_metadata.expanded_decode.enabled) { decoder_forward(query, output, k_cache, v_cache, attn_metadata); + } else if (dcp_size_ > 1 && attn_metadata.is_chunked_prefill) { + // Mixed batches also set is_chunked_prefill, but the DCP cache-slot gate in + // WorkerImpl rejects them upstream, so only pure chunked prefill reaches + // here under DCP. + dcp_chunked_prefill_forward( + query, key, value, output, k_cache, v_cache, attn_metadata); } else if (only_prefill) { prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata); } else { @@ -135,11 +221,382 @@ void AttentionImpl::prefill_forward(torch::Tensor& query, } } +void AttentionImpl::dcp_decoder_forward( + torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + CHECK(dcp_group_ != nullptr) << "DCP decode requires a DCP process group."; + CHECK_EQ(dcp_group_->world_size(), dcp_size_) + << "DCP process group size does not match attention DCP size."; + CHECK_EQ(dcp_group_->rank(), dcp_rank_) + << "DCP process group rank does not match attention DCP rank."; + CHECK(!attn_metadata.is_prefill); + CHECK(!attn_metadata.is_chunked_prefill); + CHECK(!attn_metadata.is_spec_verify) + << "DCP-2 does not support speculative decode attention."; + CHECK(!attn_metadata.expanded_decode.enabled) + << "DCP-2 does not support speculative decode attention."; + CHECK(v_cache.has_value() && v_cache.value().defined()) + << "DCP decode requires a defined V cache."; + CHECK(attn_metadata.block_table.defined()) + << "DCP decode requires a paged KV block table."; + + const int64_t token_count = query.size(0); + const std::vector& q_cu_seq_lens = + attn_metadata.q_cu_seq_lens_host_vec; + const std::vector& global_kv_seq_lens = + attn_metadata.kv_seq_lens_host_vec; + validate_dcp_decode_lengths(q_cu_seq_lens, global_kv_seq_lens, token_count); + + const int64_t block_size = k_cache.size(1); + const std::vector local_kv_seq_lens = + detail::compute_dcp_local_kv_seq_lens( + global_kv_seq_lens, dcp_size_, dcp_rank_, block_size); + const std::shared_ptr& graph_context = + attn_metadata.acl_graph_task_update_context; + const bool capturing = graph_context != nullptr && graph_context->capturing; + const torch::Tensor local_block_table = [&]() { + if (capturing) { + CHECK(attn_metadata.dcp_local_block_table.defined()) + << "DCP ACL graph capture requires a persistent local block table"; + return attn_metadata.dcp_local_block_table; + } + return parallel_state::select_dcp_local_block_table( + attn_metadata.block_table, dcp_size_, dcp_rank_); + }(); + CHECK_EQ(local_block_table.size(0), token_count) + << "DCP local block table batch size does not match decode tokens."; + + const torch::Tensor query_group = + parallel_state::gather(query, dcp_group_, 1); + const int64_t group_num_heads = num_heads_ * static_cast(dcp_size_); + CHECK_EQ(query_group.dim(), 3); + CHECK_EQ(query_group.size(0), token_count); + CHECK_EQ(query_group.size(1), group_num_heads); + CHECK_EQ(query_group.size(2), head_size_); + CHECK_EQ(group_num_heads % num_kv_heads_, 0) + << "DCP gathered Q heads must preserve the GQA ratio."; + + const torch::Tensor k = k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + const torch::Tensor v = v_cache.value().view( + {v_cache.value().size(0), v_cache.value().size(1), -1}); + const std::optional no_mask = std::nullopt; + const std::optional local_block_table_opt = local_block_table; + + torch::Tensor fia_out; + torch::Tensor fia_lse; + if (capturing) { + // Graph capture: preallocate stable-address output/LSE, wrap the FIA call + // in a task group, and record a FiaGraphTask so the executor can re-inject + // the per-step DCP-local KV lengths (and local block table) on replay. + fia_out = torch::empty({token_count, group_num_heads, head_size_}, + query_group.options()); + fia_lse = torch::empty({token_count, group_num_heads, 1}, + query_group.options().dtype(torch::kFloat32)); + // aclnn's default workspace is function-local and released after the call, + // leaving a dangling address in the captured graph (hangs on replay once + // the KV length crosses the tiling threshold that needs a non-zero + // workspace). Query the workspace for the largest local-KV envelope the + // graph can replay (full local block-table width) and allocate a stable + // caller-owned buffer reused on every replay. + const int64_t max_local_kv = local_block_table.size(1) * block_size; + const std::vector max_local_kv_seq_lens(local_kv_seq_lens.size(), + max_local_kv); + const uint64_t fia_workspace_bytes = + xllm::kernel::npu::npu_fused_infer_attention_workspace_size( + query_group, + k, + v, + no_mask, + local_block_table_opt, + q_cu_seq_lens, + max_local_kv_seq_lens, + group_num_heads, + num_kv_heads_, + scale_, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true, + /*is_causal=*/true, + fia_out, + fia_lse); + torch::Tensor fia_workspace; + if (fia_workspace_bytes > 0) { + fia_workspace = torch::empty({static_cast(fia_workspace_bytes)}, + query_group.options().dtype(torch::kByte)); + } + const std::optional fia_workspace_opt = + fia_workspace.defined() ? std::optional(fia_workspace) + : std::nullopt; + c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream(); + auto event = std::make_shared(ACL_EVENT_EXTERNAL); + event->block(stream); + event->reset(stream); + c10_npu::graph_task_group_begin(stream); + xllm::kernel::npu::npu_fused_infer_attention_out(query_group, + k, + v, + no_mask, + local_block_table_opt, + q_cu_seq_lens, + local_kv_seq_lens, + group_num_heads, + num_kv_heads_, + scale_, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true, + /*is_causal=*/true, + fia_out, + fia_lse, + fia_workspace_opt); + c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream); + xllm::npu::FiaGraphTask task; + task.output = fia_out; + task.softmax_lse = fia_lse; + task.query = query_group; + task.key = k; + task.value = v; + task.block_table = local_block_table_opt; + task.workspace = fia_workspace; + task.num_heads = group_num_heads; + task.num_key_value_heads = num_kv_heads_; + task.scale = scale_; + task.block_size = block_size; + task.sparse_mode = 0; + task.dcp_size = static_cast(dcp_size_); + task.dcp_rank = static_cast(dcp_rank_); + task.input_layout = "TND"; + task.softmax_lse_flag = true; + task.handle = handle; + task.event = std::move(event); + graph_context->fia_tasks.emplace_back(std::move(task)); + } else { + const auto fia_result = + xllm::kernel::npu::npu_fused_infer_attention(query_group, + k, + v, + no_mask, + local_block_table_opt, + q_cu_seq_lens, + local_kv_seq_lens, + group_num_heads, + num_kv_heads_, + scale_, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + fia_out = std::get<0>(fia_result); + fia_lse = std::get<1>(fia_result); + } + torch::Tensor partial_out = fia_out.to(torch::kFloat32); + torch::Tensor partial_lse = fia_lse.to(torch::kFloat32); + if (capturing) { + detail::normalize_zero_dcp_partials_for_graph(partial_out, + partial_lse, + attn_metadata.kv_seq_lens, + dcp_rank_, + block_size); + } else { + normalize_zero_dcp_partials(partial_out, partial_lse, local_kv_seq_lens); + } + + const torch::Tensor all_partial_out = + dcp_group_->allgather_base_sync(partial_out); + const torch::Tensor all_partial_lse = + dcp_group_->allgather_base_sync(partial_lse); + const torch::Tensor merged_out = + detail::merge_dcp_partials(all_partial_out, all_partial_lse); + const int64_t head_begin = static_cast(dcp_rank_) * num_heads_; + const torch::Tensor local_out = + merged_out.slice(1, head_begin, head_begin + num_heads_); + output.copy_(local_out.to(output.scalar_type())); +} + +void AttentionImpl::dcp_chunked_prefill_forward( + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + CHECK(dcp_group_ != nullptr) << "DCP chunked prefill requires a DCP group."; + CHECK_EQ(dcp_group_->world_size(), dcp_size_) + << "DCP process group size does not match attention DCP size."; + CHECK_EQ(dcp_group_->rank(), dcp_rank_) + << "DCP process group rank does not match attention DCP rank."; + CHECK(!attn_metadata.is_spec_verify) + << "DCP chunked prefill does not support speculative decode attention."; + CHECK(!attn_metadata.paged_attention_tiling_data.defined()) + << "DCP chunked prefill does not support graph-captured attention."; + CHECK(v_cache.has_value() && v_cache.value().defined()) + << "DCP chunked prefill requires a defined V cache."; + CHECK(attn_metadata.block_table.defined()) + << "DCP chunked prefill requires a paged KV block table."; + CHECK(attn_metadata.fia_attn_mask.defined()) + << "DCP chunked prefill requires a causal attention mask."; + + query = query.view({-1, num_heads_, head_size_}); + output = output.view({-1, num_heads_, head_size_}); + key = key.view({-1, num_kv_heads_, head_size_}); + value = value.view({-1, num_kv_heads_, head_size_}); + + const int64_t token_count = query.size(0); + const std::vector& global_kv_seq_lens = + attn_metadata.kv_seq_lens_host_vec; + const std::vector q_cu_seq_lens = + detail::validate_dcp_chunked_lengths(attn_metadata.q_cu_seq_lens_host_vec, + global_kv_seq_lens, + token_count); + + const int64_t block_size = k_cache.size(1); + const int64_t group_num_heads = num_heads_ * static_cast(dcp_size_); + CHECK_EQ(group_num_heads % num_kv_heads_, 0) + << "DCP gathered Q heads must preserve the GQA ratio."; + CHECK_EQ(key.size(0), token_count); + CHECK_EQ(value.size(0), token_count); + + // Gather the query heads across the DCP group so every partial covers the + // full head-group; each rank slices back its own head range after merge. + const torch::Tensor query_group = + parallel_state::gather(query, dcp_group_, 1); + CHECK_EQ(query_group.dim(), 3); + CHECK_EQ(query_group.size(0), token_count); + CHECK_EQ(query_group.size(1), group_num_heads); + CHECK_EQ(query_group.size(2), head_size_); + + // Current/diagonal part: the current chunk attends to its own KV with a + // causal mask. The raw key/value projections are identical on every DCP rank + // (KV heads are replicated within the group), so this part is not sharded. + const auto current_result = xllm::kernel::npu::npu_fused_infer_attention( + query_group, + key, + value, + std::make_optional(attn_metadata.fia_attn_mask), + /*block_table=*/std::nullopt, + q_cu_seq_lens, + /*actual_seq_lengths_kv=*/q_cu_seq_lens, + group_num_heads, + num_kv_heads_, + scale_, + /*block_size=*/0, + /*sparse_mode=*/3, + "TND", + /*softmax_lse_flag=*/true); + torch::Tensor current_out = std::get<0>(current_result).to(torch::kFloat32); + torch::Tensor current_lse = std::get<1>(current_result).to(torch::kFloat32); + validate_dcp_partial_shape(current_out, + current_lse, + token_count, + group_num_heads, + head_size_, + "DCP chunked current partial shape mismatch."); + + // Context part: the current chunk attends to the previously-cached context + // KV, which is DCP-sharded round-robin over blocks. Each rank computes a + // partial over only its local context shard (no mask, full history visible). + const std::vector context_lens = + detail::compute_dcp_context_lens(q_cu_seq_lens, global_kv_seq_lens); + const int64_t head_begin = static_cast(dcp_rank_) * num_heads_; + if (std::all_of(context_lens.begin(), + context_lens.end(), + [](int64_t context_len) { return context_len == 0; })) { + const torch::Tensor local_out = + current_out.slice(1, head_begin, head_begin + num_heads_); + CHECK_EQ(local_out.size(0), token_count); + CHECK_EQ(local_out.size(1), num_heads_); + CHECK_EQ(local_out.size(2), head_size_); + output.copy_(local_out.to(output.scalar_type())); + return; + } + + const std::vector local_context_lens = + detail::compute_dcp_local_kv_seq_lens( + context_lens, dcp_size_, dcp_rank_, block_size); + const torch::Tensor local_block_table = + parallel_state::select_dcp_local_block_table( + attn_metadata.block_table, dcp_size_, dcp_rank_); + detail::validate_dcp_chunked_block_table( + local_block_table, static_cast(q_cu_seq_lens.size())); + + const torch::Tensor k = k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + const torch::Tensor v = v_cache.value().view( + {v_cache.value().size(0), v_cache.value().size(1), -1}); + const auto context_result = xllm::kernel::npu::npu_fused_infer_attention( + query_group, + k, + v, + /*atten_mask=*/std::nullopt, + std::make_optional(local_block_table), + q_cu_seq_lens, + local_context_lens, + group_num_heads, + num_kv_heads_, + scale_, + block_size, + /*sparse_mode=*/0, + "TND", + /*softmax_lse_flag=*/true); + torch::Tensor context_out = std::get<0>(context_result).to(torch::kFloat32); + torch::Tensor context_lse = std::get<1>(context_result).to(torch::kFloat32); + validate_dcp_partial_shape(context_out, + context_lse, + token_count, + group_num_heads, + head_size_, + "DCP chunked context partial shape mismatch."); + detail::normalize_zero_dcp_chunked_partials( + context_out, context_lse, local_context_lens, q_cu_seq_lens); + + // Merge context shards from all ranks with the (replicated) current part in a + // single online-softmax reduction. The formula is mathematically equivalent + // to softmax over the complete key set. FIA emits each partial output in the + // query dtype (BF16/FP16) before this fp32 merge, so results match a single + // monolithic FIA call only up to that partial-output quantization (on the + // order of a few 1e-3), not bitwise. The current part is identical on every + // rank, so contributing it once (from this rank) is correct. + const torch::Tensor all_context_out = + dcp_group_->allgather_base_sync(context_out); + const torch::Tensor all_context_lse = + dcp_group_->allgather_base_sync(context_lse); + const torch::Tensor stacked_out = + torch::cat({all_context_out, current_out.unsqueeze(0)}, 0); + const torch::Tensor stacked_lse = + torch::cat({all_context_lse, current_lse.unsqueeze(0)}, 0); + const torch::Tensor merged_out = + detail::merge_dcp_partials(stacked_out, stacked_lse); + + CHECK_EQ(merged_out.dim(), 3); + CHECK_EQ(merged_out.size(0), token_count); + CHECK_EQ(merged_out.size(1), group_num_heads); + CHECK_EQ(merged_out.size(2), head_size_); + const torch::Tensor local_out = + merged_out.slice(1, head_begin, head_begin + num_heads_); + CHECK_EQ(local_out.size(0), token_count); + CHECK_EQ(local_out.size(1), num_heads_); + CHECK_EQ(local_out.size(2), head_size_); + output.copy_(local_out.to(output.scalar_type())); +} + void AttentionImpl::decoder_forward(torch::Tensor& query, torch::Tensor& output, const torch::Tensor& k_cache, const std::optional& v_cache, const AttentionMetadata& attn_metadata) { + query = query.view({-1, num_heads_, head_size_}); + output = output.view({-1, num_heads_, head_size_}); + + if (dcp_size_ > 1) { + dcp_decoder_forward(query, output, k_cache, v_cache, attn_metadata); + return; + } + query = query.view({-1, 1, num_heads_, head_size_}); output = output.view({-1, 1, num_heads_, head_size_}); @@ -163,7 +620,6 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, if (tiling_data.defined()) { // Use CustomPagedAttention for ACL graph mode to avoid .to(kCPU) operations - xllm::kernel::npu::batch_decode_acl_graph(query, k_cache, v_cache.value_or(torch::Tensor()), diff --git a/xllm/core/layers/npu_torch/attention.h b/xllm/core/layers/npu_torch/attention.h index 6fe09e9c99..bb0da2b456 100644 --- a/xllm/core/layers/npu_torch/attention.h +++ b/xllm/core/layers/npu_torch/attention.h @@ -24,6 +24,9 @@ limitations under the License. #include "layers/common/attention_metadata.h" namespace xllm { + +class ProcessGroup; + namespace layer { class AttentionImpl : public torch::nn::Module { @@ -34,7 +37,10 @@ class AttentionImpl : public torch::nn::Module { int64_t head_size, float scale, int64_t num_kv_heads, - int64_t sliding_window); + int64_t sliding_window, + int32_t dcp_size = 1, + int32_t dcp_rank = 0, + ProcessGroup* dcp_group = nullptr); std::tuple> forward( const AttentionMetadata& attn_metadata, @@ -58,11 +64,28 @@ class AttentionImpl : public torch::nn::Module { const AttentionMetadata& attn_metadata); private: + void dcp_decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + void dcp_chunked_prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + int64_t num_heads_; int64_t head_size_; float scale_; int64_t num_kv_heads_; int64_t sliding_window_; + int32_t dcp_size_ = 1; + int32_t dcp_rank_ = 0; + ProcessGroup* dcp_group_ = nullptr; }; TORCH_MODULE(Attention); diff --git a/xllm/core/layers/npu_torch/dcp_attention_utils.cpp b/xllm/core/layers/npu_torch/dcp_attention_utils.cpp new file mode 100644 index 0000000000..e6f2659a65 --- /dev/null +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.cpp @@ -0,0 +1,206 @@ +/* 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://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "layers/npu_torch/dcp_attention_utils.h" + +#include + +#include +#include +#include + +namespace xllm::layer::detail { + +std::vector compute_dcp_local_kv_seq_lens( + const std::vector& global_kv_seq_lens, + int32_t dcp_size, + int32_t dcp_rank, + int64_t block_size) { + CHECK_GT(dcp_size, 1); + CHECK_GE(dcp_rank, 0); + CHECK_LT(dcp_rank, dcp_size); + CHECK_GT(block_size, 0); + + std::vector local_kv_seq_lens; + local_kv_seq_lens.reserve(global_kv_seq_lens.size()); + for (const int64_t global_kv_seq_len : global_kv_seq_lens) { + CHECK_GE(global_kv_seq_len, 0); + const int64_t base = global_kv_seq_len / block_size / dcp_size * block_size; + const int64_t remainder = global_kv_seq_len - base * dcp_size; + const int64_t rank_offset = static_cast(dcp_rank) * block_size; + const int64_t local_remainder = + std::clamp(remainder - rank_offset, int64_t{0}, block_size); + local_kv_seq_lens.emplace_back(base + local_remainder); + } + return local_kv_seq_lens; +} + +void normalize_zero_dcp_partials_for_graph( + torch::Tensor& partial_out, + torch::Tensor& partial_lse, + const torch::Tensor& global_kv_seq_lens, + int32_t dcp_rank, + int64_t block_size) { + CHECK(global_kv_seq_lens.defined()); + CHECK_EQ(global_kv_seq_lens.dim(), 1); + CHECK_EQ(global_kv_seq_lens.size(0), partial_out.size(0)); + const int64_t first_local_block_offset = + static_cast(dcp_rank) * block_size; + const torch::Tensor zero_local_kv_mask = + global_kv_seq_lens.le(first_local_block_offset) + .view({global_kv_seq_lens.size(0), 1, 1}); + partial_out.masked_fill_(zero_local_kv_mask, 0.0); + partial_lse.masked_fill_(zero_local_kv_mask, + -std::numeric_limits::infinity()); +} + +std::vector compute_dcp_context_lens( + const std::vector& q_cu_seq_lens, + const std::vector& global_kv_seq_lens) { + CHECK(!q_cu_seq_lens.empty()) + << "chunked prefill requires cumulative query lengths"; + CHECK_EQ(q_cu_seq_lens.size(), global_kv_seq_lens.size()) + << "chunked prefill requires one query and KV length per request"; + + std::vector context_lens; + context_lens.reserve(global_kv_seq_lens.size()); + int64_t previous_q_end = 0; + for (size_t request_index = 0; request_index < global_kv_seq_lens.size(); + ++request_index) { + const int64_t q_end = q_cu_seq_lens[request_index]; + const int64_t query_len = q_end - previous_q_end; + const int64_t context_len = global_kv_seq_lens[request_index] - query_len; + CHECK_GE(context_len, 0) + << "chunked prefill context length must be non-negative"; + context_lens.emplace_back(context_len); + previous_q_end = q_end; + } + return context_lens; +} + +std::vector validate_dcp_chunked_lengths( + const std::vector& q_cu_seq_lens, + const std::vector& global_kv_seq_lens, + int64_t token_count) { + CHECK(!q_cu_seq_lens.empty()) + << "DCP chunked prefill requires host cumulative query lengths."; + + const size_t query_begin = q_cu_seq_lens.front() == 0 ? 1 : 0; + CHECK_LT(query_begin, q_cu_seq_lens.size()) + << "DCP chunked prefill requires at least one request."; + const std::vector normalized_q_cu_seq_lens( + q_cu_seq_lens.begin() + query_begin, q_cu_seq_lens.end()); + + CHECK_EQ(normalized_q_cu_seq_lens.size(), global_kv_seq_lens.size()) + << "DCP chunked prefill requires one query and KV length per request."; + + int64_t previous_q_end = 0; + for (size_t request_index = 0; + request_index < normalized_q_cu_seq_lens.size(); + ++request_index) { + const int64_t q_end = normalized_q_cu_seq_lens[request_index]; + const int64_t query_len = q_end - previous_q_end; + CHECK_GE(query_len, 1) + << "DCP chunked prefill requires at least one query token per request."; + CHECK_LE(query_len, kMaxDcpChunkedPrefillQueryLen) + << "DCP chunked prefill does not yet support chunked query length " + "above " + << kMaxDcpChunkedPrefillQueryLen << "."; + CHECK_GE(global_kv_seq_lens[request_index], query_len) + << "DCP chunked prefill KV length must cover the current chunk query."; + previous_q_end = q_end; + } + CHECK_EQ(previous_q_end, token_count) + << "DCP chunked cumulative query lengths do not match query tokens."; + return normalized_q_cu_seq_lens; +} + +void validate_dcp_chunked_block_table(const torch::Tensor& local_block_table, + int64_t request_count) { + CHECK(local_block_table.defined()) + << "DCP chunked local block table must be defined."; + CHECK_EQ(local_block_table.dim(), 2) + << "DCP chunked local block table must be two-dimensional."; + CHECK_EQ(local_block_table.size(0), request_count) + << "DCP local block table batch size does not match request count."; +} + +void normalize_zero_dcp_chunked_partials( + torch::Tensor& partial_out, + torch::Tensor& partial_lse, + const std::vector& local_context_lens, + const std::vector& q_cu_seq_lens) { + CHECK_EQ(partial_out.scalar_type(), torch::kFloat32); + CHECK_EQ(partial_lse.scalar_type(), torch::kFloat32); + CHECK_EQ(partial_out.dim(), 3); + CHECK_EQ(partial_lse.dim(), 3); + CHECK_EQ(partial_out.size(0), partial_lse.size(0)); + CHECK_EQ(partial_out.size(1), partial_lse.size(1)); + CHECK_EQ(partial_lse.size(2), 1); + CHECK_EQ(local_context_lens.size(), q_cu_seq_lens.size()); + + int64_t previous_q_end = 0; + for (size_t request_index = 0; request_index < local_context_lens.size(); + ++request_index) { + const int64_t q_end = q_cu_seq_lens[request_index]; + const int64_t query_len = q_end - previous_q_end; + CHECK_GE(query_len, 1); + CHECK_GE(local_context_lens[request_index], 0); + if (local_context_lens[request_index] == 0) { + partial_out.narrow(0, previous_q_end, query_len).zero_(); + partial_lse.narrow(0, previous_q_end, query_len) + .fill_(-std::numeric_limits::infinity()); + } + previous_q_end = q_end; + } + CHECK_EQ(previous_q_end, partial_out.size(0)); +} + +torch::Tensor merge_dcp_partials(const torch::Tensor& all_partial_out, + const torch::Tensor& all_partial_lse) { + CHECK(all_partial_out.scalar_type() == torch::kFloat32 || + all_partial_out.scalar_type() == torch::kFloat64); + CHECK_EQ(all_partial_lse.scalar_type(), all_partial_out.scalar_type()); + CHECK_EQ(all_partial_out.dim(), 4); + CHECK_EQ(all_partial_lse.dim(), 4); + CHECK_EQ(all_partial_out.size(0), all_partial_lse.size(0)); + CHECK_EQ(all_partial_out.size(1), all_partial_lse.size(1)); + CHECK_EQ(all_partial_out.size(2), all_partial_lse.size(2)); + CHECK_EQ(all_partial_lse.size(3), 1); + + const torch::Tensor finite_lse = torch::isfinite(all_partial_lse); + const torch::Tensor max_lse = std::get<0>(all_partial_lse.max(0)); + const torch::Tensor max_lse_is_finite = torch::isfinite(max_lse); + const torch::Tensor safe_max_lse = + torch::where(max_lse_is_finite, max_lse, torch::zeros_like(max_lse)); + const torch::Tensor weights = + torch::where(finite_lse, + torch::exp(all_partial_lse - safe_max_lse), + torch::zeros_like(all_partial_lse)); + const torch::Tensor safe_partial_out = + torch::where(finite_lse.expand_as(all_partial_out), + all_partial_out, + torch::zeros_like(all_partial_out)); + const torch::Tensor denominator = weights.sum(0); + const torch::Tensor safe_denominator = torch::where( + denominator.gt(0), denominator, torch::ones_like(denominator)); + const torch::Tensor merged_out = + (weights * safe_partial_out).sum(0) / safe_denominator; + return torch::where(max_lse_is_finite.expand_as(merged_out), + merged_out, + torch::zeros_like(merged_out)); +} + +} // namespace xllm::layer::detail diff --git a/xllm/core/layers/npu_torch/dcp_attention_utils.h b/xllm/core/layers/npu_torch/dcp_attention_utils.h new file mode 100644 index 0000000000..27ac4dbda6 --- /dev/null +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.h @@ -0,0 +1,73 @@ +/* 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://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include + +namespace xllm::layer::detail { + +inline constexpr int64_t kMaxDcpChunkedPrefillQueryLen = 2048; + +std::vector compute_dcp_local_kv_seq_lens( + const std::vector& global_kv_seq_lens, + int32_t dcp_size, + int32_t dcp_rank, + int64_t block_size); + +// Graph replay must derive the empty-shard mask from the live device KV +// lengths. A host-side branch would be frozen at capture time. +void normalize_zero_dcp_partials_for_graph( + torch::Tensor& partial_out, + torch::Tensor& partial_lse, + const torch::Tensor& global_kv_seq_lens, + int32_t dcp_rank, + int64_t block_size); + +// Per-request cached-context length for chunked prefill: the KV that precedes +// the current chunk, derived as global_kv_seq_len - current_chunk_query_len. +// q_cu_seq_lens is the cumulative host query length per request. +std::vector compute_dcp_context_lens( + const std::vector& q_cu_seq_lens, + const std::vector& global_kv_seq_lens); + +// Chunked-prefill counterpart of the decode length validator: query tokens per +// request may exceed one, so the partial tensors are indexed by token rather +// than by request. +std::vector validate_dcp_chunked_lengths( + const std::vector& q_cu_seq_lens, + const std::vector& global_kv_seq_lens, + int64_t token_count); + +void validate_dcp_chunked_block_table(const torch::Tensor& local_block_table, + int64_t request_count); + +// Zero the context partial for requests whose local context shard is empty on +// this rank. Unlike the decode variant, the partial's leading dimension is the +// flattened token count, so each empty request zeroes its own token range +// [previous_q_end, q_cu_seq_lens[request]). +void normalize_zero_dcp_chunked_partials( + torch::Tensor& partial_out, + torch::Tensor& partial_lse, + const std::vector& local_context_lens, + const std::vector& q_cu_seq_lens); + +torch::Tensor merge_dcp_partials(const torch::Tensor& all_partial_out, + const torch::Tensor& all_partial_lse); + +} // namespace xllm::layer::detail diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index b2f572de4b..ddab2098c9 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -106,7 +106,10 @@ Qwen3NextAttentionImpl::Qwen3NextAttentionImpl( head_dim_, scaling_, num_kv_heads_, - args.sliding_window())); + args.sliding_window(), + parallel_args.dcp_size_effective(), + parallel_args.dcp_rank(), + parallel_args.dcp_group_)); // 7. Fused split_qkv_rmsnorm_mrope kernel setup rotary_dim_ = static_cast(head_dim_ * args.partial_rotary_factor()); diff --git a/xllm/core/platform/npu/acl_graph_task_update_context.h b/xllm/core/platform/npu/acl_graph_task_update_context.h index eeae5f86b0..f2961d6f01 100644 --- a/xllm/core/platform/npu/acl_graph_task_update_context.h +++ b/xllm/core/platform/npu/acl_graph_task_update_context.h @@ -20,6 +20,7 @@ limitations under the License. #include #include #include +#include #include #if defined(__GNUC__) @@ -60,17 +61,51 @@ struct CausalConv1dGraphTask { std::shared_ptr event; }; +// Graph-capture record for a DCP fused-infer-attention (FIA) call. The tensors +// (query/key/value/output/softmax_lse) keep stable device addresses across +// replays; the per-step varying host scalar actual_seq_lengths_kv (the +// DCP-local KV lengths) is recomputed at replay from the step's global +// kv_seq_lens via compute_dcp_local_kv_seq_lens(dcp_size, dcp_rank, +// block_size), so only these static DCP params are stored here. `workspace` is +// a caller-owned buffer sized for the largest local-KV envelope at capture and +// reused (same address) on every replay, because aclnn's default workspace is +// function-local and would leave a dangling address in the captured graph. See +// update_graph_tasks in acl_graph_executor_impl.cpp. +struct FiaGraphTask { + torch::Tensor output; + torch::Tensor softmax_lse; + torch::Tensor query; + torch::Tensor key; + torch::Tensor value; + std::optional atten_mask; + std::optional block_table; + torch::Tensor workspace; + int64_t num_heads = 0; + int64_t num_key_value_heads = 0; + double scale = 1.0; + int64_t block_size = 0; + int64_t sparse_mode = 0; + int32_t dcp_size = 1; + int32_t dcp_rank = 0; + std::string input_layout; + bool softmax_lse_flag = false; + c10_npu::NPUTaskGroupHandle handle{}; + std::shared_ptr event; +}; + class AclGraphTaskUpdateContext final { public: void begin_capture() { capturing = true; causal_conv1d_tasks.clear(); + fia_tasks.clear(); } void end_capture() { capturing = false; } bool capturing = false; std::vector causal_conv1d_tasks; + std::vector fia_tasks; }; } // namespace xllm::npu diff --git a/xllm/core/runtime/acl_graph_executor_impl.cpp b/xllm/core/runtime/acl_graph_executor_impl.cpp index 9ee655b8e9..6bf773a2cf 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.cpp +++ b/xllm/core/runtime/acl_graph_executor_impl.cpp @@ -35,8 +35,11 @@ limitations under the License. #endif #include "core/common/metrics.h" #include "core/framework/speculative/mtp_async_state.h" +#include "core/kernels/npu/npu_ops_api.h" #include "core/kernels/npu/tilelang/tilelang_ops_api.h" #include "core/kernels/ops_api.h" +#include "core/layers/common/attention_metadata.h" +#include "core/layers/npu_torch/dcp_attention_utils.h" #include "core/platform/device.h" #include "core/platform/npu/acl_graph_task_update_context.h" #include "core/util/utils.h" @@ -456,58 +459,141 @@ bool AclGraph::capture(CausalLM* model, } bool AclGraph::update_graph_tasks(const ModelInputParams& params) { - if (graph_task_context_ == nullptr || - graph_task_context_->causal_conv1d_tasks.empty()) { + if (graph_task_context_ == nullptr) { + return false; + } + const bool has_conv = !graph_task_context_->causal_conv1d_tasks.empty(); + const bool has_fia = !graph_task_context_->fia_tasks.empty(); + if (!has_conv && !has_fia) { return false; } - - const std::vector empty_host_args; - CHECK(!params.parallel.query_start_loc.empty()) - << "causal_conv1d graph update requires padded query_start_loc"; - CHECK(!params.embedding.linear_state_ids.empty()) - << "causal_conv1d graph update requires padded cache indices"; - - std::vector linear_state_indices_host( - params.embedding.linear_state_ids.begin(), - params.embedding.linear_state_ids.end()); c10_npu::NPUStream update_stream = update_stream_.value(); c10_npu::NPUStreamGuard stream_guard(update_stream); - for (auto& task : graph_task_context_->causal_conv1d_tasks) { - CHECK_EQ(params.parallel.query_start_loc.back(), task.x.size(0)) - << "causal_conv1d graph update host args must be padded to the " - "capture x.shape[0]"; - CHECK_EQ(linear_state_indices_host.size() + 1, - params.parallel.query_start_loc.size()) - << "cache_indices must be sequence-scoped"; - - const std::vector& num_accepted_tokens = - task.branch == CausalConv1dGraphBranch::kSpecVerify - ? params.num_accepted_tokens_host - : empty_host_args; - if (task.branch == CausalConv1dGraphBranch::kSpecVerify) { - CHECK_EQ(num_accepted_tokens.size(), linear_state_indices_host.size()) - << "spec causal_conv1d graph update requires accepted-token counts"; + if (has_conv) { + const std::vector empty_host_args; + CHECK(!params.parallel.query_start_loc.empty()) + << "causal_conv1d graph update requires padded query_start_loc"; + CHECK(!params.embedding.linear_state_ids.empty()) + << "causal_conv1d graph update requires padded cache indices"; + + std::vector linear_state_indices_host( + params.embedding.linear_state_ids.begin(), + params.embedding.linear_state_ids.end()); + + for (auto& task : graph_task_context_->causal_conv1d_tasks) { + CHECK_EQ(params.parallel.query_start_loc.back(), task.x.size(0)) + << "causal_conv1d graph update host args must be padded to the " + "capture x.shape[0]"; + CHECK_EQ(linear_state_indices_host.size() + 1, + params.parallel.query_start_loc.size()) + << "cache_indices must be sequence-scoped"; + + const std::vector& num_accepted_tokens = + task.branch == CausalConv1dGraphBranch::kSpecVerify + ? params.num_accepted_tokens_host + : empty_host_args; + if (task.branch == CausalConv1dGraphBranch::kSpecVerify) { + CHECK_EQ(num_accepted_tokens.size(), linear_state_indices_host.size()) + << "spec causal_conv1d graph update requires accepted-token counts"; + } + + c10_npu::graph_task_update_begin(update_stream, task.handle); + xllm::kernel::causal_conv1d_out( + task.output, + task.x, + task.weight, + task.conv_state, + task.bias, + torch::IntArrayRef(params.parallel.query_start_loc), + torch::IntArrayRef(linear_state_indices_host), + torch::IntArrayRef(empty_host_args), + torch::IntArrayRef(num_accepted_tokens), + task.activation_mode, + task.pad_slot_id, + task.run_mode); + c10_npu::graph_task_update_end(update_stream); + if (task.event != nullptr) { + task.event->record(update_stream); + } } + } - c10_npu::graph_task_update_begin(update_stream, task.handle); - xllm::kernel::causal_conv1d_out( - task.output, - task.x, - task.weight, - task.conv_state, - task.bias, - torch::IntArrayRef(params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_host), - torch::IntArrayRef(empty_host_args), - torch::IntArrayRef(num_accepted_tokens), - task.activation_mode, - task.pad_slot_id, - task.run_mode); - c10_npu::graph_task_update_end(update_stream); - if (task.event != nullptr) { - task.event->record(update_stream); + if (has_fia) { + // DCP FIA replay: recompute the per-step DCP-local KV lengths from the + // step's global metadata, then re-inject them via the FIA out variant under + // graph_task_update. The block table keeps the persistent DCP-local address + // captured with the graph. See TODO-vllm-ascend-parity/01. + // + // attn_metadata is NOT threaded to the executor; source the current step's + // global KV/query lengths and stable block table from the persistent graph + // params instead. GraphPersistentParam pads + // host.kv_seq_lens/host.q_seq_lens per step; q_cu_seq_lens is not padded, + // so build it from the padded per-request q_seq_lens (decode: each request + // contributes one token). + const std::vector& host_kv_seq_lens = + params.attention.host.kv_seq_lens; + const std::vector& host_q_seq_lens = + params.attention.host.q_seq_lens; + CHECK(!host_kv_seq_lens.empty() && + host_kv_seq_lens.size() == host_q_seq_lens.size()) + << "DCP FIA graph update requires padded host kv/q seq lens"; + const std::vector global_kv_seq_lens(host_kv_seq_lens.begin(), + host_kv_seq_lens.end()); + std::vector q_cu_seq_lens; + q_cu_seq_lens.reserve(host_q_seq_lens.size()); + int64_t q_running_total = 0; + for (const int32_t q_len : host_q_seq_lens) { + q_running_total += q_len; + q_cu_seq_lens.emplace_back(q_running_total); + } + for (auto& task : graph_task_context_->fia_tasks) { + const std::vector local_kv_seq_lens = + xllm::layer::detail::compute_dcp_local_kv_seq_lens(global_kv_seq_lens, + task.dcp_size, + task.dcp_rank, + task.block_size); + CHECK(task.block_table.has_value() && task.block_table.value().defined()) + << "FiaGraphTask must hold a persistent DCP-local block table"; + const int64_t local_kv_capacity = + task.block_table.value().size(1) * task.block_size; + for (const int64_t local_kv_seq_len : local_kv_seq_lens) { + CHECK_GE(local_kv_seq_len, 0); + CHECK_LE(local_kv_seq_len, local_kv_capacity) + << "DCP local KV length exceeds persistent local block-table " + "capacity"; + } + const std::optional& local_block_table_opt = + task.block_table; + const std::optional workspace_opt = + task.workspace.defined() + ? std::optional(task.workspace) + : std::nullopt; + + c10_npu::graph_task_update_begin(update_stream, task.handle); + xllm::kernel::npu::npu_fused_infer_attention_out(task.query, + task.key, + task.value, + task.atten_mask, + local_block_table_opt, + q_cu_seq_lens, + local_kv_seq_lens, + task.num_heads, + task.num_key_value_heads, + task.scale, + task.block_size, + task.sparse_mode, + task.input_layout, + task.softmax_lse_flag, + /*is_causal=*/true, + task.output, + task.softmax_lse, + workspace_opt); + c10_npu::graph_task_update_end(update_stream); + if (task.event != nullptr) { + task.event->record(update_stream); + } } } return true; @@ -721,6 +807,28 @@ ModelOutput AclGraph::replay(CausalLM* model, CHECK(update_stream_.has_value()); signal_static_graph_tasks(update_stream_.value()); } + // Per-replay completion barrier for the DCP FIA graph-task path. vLLM + // synchronizes the current stream before every FULL graph replay + // (compilation/acl_graph.py) so iteration i cannot rewrite FIA task handles / + // record external events while iteration i-1's graph is still running. + // Synchronizing the current stream here waits for both the previous graph + // replay (queued on this stream by the prior + // make_current_stream_wait_for_graph) and this step's persistent-input copies + // before we launch and re-inject task parameters, preventing cross-replay + // task/event generation overlap (delayed ACL_ERROR_RT_MODEL_EXECUTE 507011). + // Scoped to graphs that captured FIA tasks (i.e. DCP decode) so the ordinary + // dcp=1 GDN conv graph keeps its non-blocking replay hot path, for which no + // such fault has been observed. First-version host-blocking barrier; can + // later become a persistent input-ready/replay-done event chain (never a + // stack-local event). + const bool has_fia_graph_tasks = + graph_task_context_ != nullptr && !graph_task_context_->fia_tasks.empty(); + if (!graph_paged_attention_tiling_data_.defined() && + model->is_hybrid_linear_attention() && !use_static_graph_tasks && + has_fia_graph_tasks) { + CHECK_EQ(aclrtSynchronizeStream(stream), ACL_SUCCESS) + << "pre-replay current-stream synchronize failed"; + } graph_.replay(); if (model->is_hybrid_linear_attention()) { CHECK(graph_params.has_value()) @@ -1383,6 +1491,18 @@ uint64_t AclGraphExecutorImpl::get_graph_key( get_mla_capture_kv_seq_len_bucket(params, options_); return get_mla_graph_key(bucket_num_tokens, capture_kv_seq_len_bucket); } + if (params.meta.batch_forward_type.is_decode() && + options_.decode_context_parallel_size() == 1) { + // Keep CustomPagedAttention static and separate ordinary decode graphs by + // vendor plan geometry. Rebuilding this ATB wrapper through graph task + // update is unsafe because run_atb_cmd creates function-local workspace and + // format-conversion tensors whose addresses do not outlive the update call. + const uint64_t plan_bucket = paged_attention_plan_bucket( + params.meta.kv_max_seq_len, options_.block_size()); + return mix_graph_key(static_cast(bucket_num_tokens), + plan_bucket) & + ~kSpecVerifyGraphKeyMask; + } return static_cast(bucket_num_tokens); } diff --git a/xllm/core/runtime/acl_graph_executor_impl.h b/xllm/core/runtime/acl_graph_executor_impl.h index 50e9a4367c..a32c49ba92 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.h +++ b/xllm/core/runtime/acl_graph_executor_impl.h @@ -214,6 +214,12 @@ class AclGraphExecutorImpl : public ExecutorImpl { size_t get_graph_memory_pool_count(); size_t get_graph_capture_stream_count() const; + [[nodiscard]] uint64_t graph_key_for_test( + uint32_t bucket_num_tokens, + const ModelInputParams& params) const { + return get_graph_key(bucket_num_tokens, params); + } + private: // not own CausalLM* model_; diff --git a/xllm/core/runtime/acl_graph_persistent_param.cpp b/xllm/core/runtime/acl_graph_persistent_param.cpp index 420d8600ec..936d66514c 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.cpp +++ b/xllm/core/runtime/acl_graph_persistent_param.cpp @@ -15,6 +15,7 @@ limitations under the License. #include "core/runtime/acl_graph_persistent_param.h" +#include #include #include #include @@ -278,6 +279,24 @@ GraphPersistentParam::GraphPersistentParam(const ModelArgs& args, persistent_expanded_block_tables_ = torch::zeros({max_graph_tokens, max_block_table_len}, torch::dtype(torch::kInt).device(device)); + dcp_size_ = options.decode_context_parallel_size(); + if (dcp_size_ > 1) { + CHECK_GT(options.world_size(), 0); + CHECK_GT(options.dp_size(), 0); + CHECK_EQ(options.world_size() % options.dp_size(), 0); + const int32_t tp_size = options.world_size() / options.dp_size(); + CHECK_EQ(tp_size % dcp_size_, 0); + CHECK_GE(options.node_rank(), 0); + CHECK_LT(options.node_rank(), options.world_size()); + dcp_rank_ = (options.node_rank() % tp_size) % dcp_size_; + const torch::TensorOptions index_options = + torch::TensorOptions().dtype(torch::kLong).device(device); + persistent_dcp_local_block_indices_ = + torch::arange(dcp_rank_, max_block_table_len, dcp_size_, index_options); + persistent_dcp_local_block_tables_ = torch::zeros( + {metadata_capacity, persistent_dcp_local_block_indices_.numel()}, + torch::dtype(torch::kInt).device(device)); + } torch::Dtype dtype = util::parse_dtype(args.dtype(), device); if (args.dtype() == "float" || args.dtype() == "float32") { @@ -544,6 +563,28 @@ GraphPersistentParam::~GraphPersistentParam() { } } +void GraphPersistentParam::update_dcp_local_block_tables( + int64_t padded_batch_size) { + CHECK_GT(dcp_size_, 1); + CHECK(persistent_dcp_local_block_indices_.defined()); + CHECK(persistent_dcp_local_block_tables_.defined()); + CHECK_GE(padded_batch_size, 0); + CHECK_LE(padded_batch_size, persistent_block_tables_.size(0)); + CHECK_LE(padded_batch_size, persistent_dcp_local_block_tables_.size(0)); + if (padded_batch_size == 0) { + return; + } + + torch::Tensor global_block_tables = persistent_block_tables_.slice( + /*dim=*/0, /*start=*/0, /*end=*/padded_batch_size); + torch::Tensor local_block_tables = persistent_dcp_local_block_tables_.slice( + /*dim=*/0, /*start=*/0, /*end=*/padded_batch_size); + at::index_select_out(local_block_tables, + global_block_tables, + /*dim=*/1, + persistent_dcp_local_block_indices_); +} + void GraphPersistentParam::set_aux_hidden_states(const torch::Tensor& value) { if (!value.defined()) { return; @@ -1061,6 +1102,11 @@ std::optional GraphPersistentParam::update( zero_tensor_tail( persistent_block_tables_, actual_seq_len_rows, padded_batch_size); } + const bool use_dcp_local_block_tables = + dcp_size_ > 1 && is_decode && !params.is_spec_verify; + if (use_dcp_local_block_tables) { + update_dcp_local_block_tables(padded_batch_size); + } // Update persistent embedding from input_embedding if available const auto& embedding = params.embedding.input_embedding; @@ -1326,6 +1372,17 @@ std::optional GraphPersistentParam::update( graph_params->attention.host.kv_seq_lens = padded_kv_seq_lens_vec; graph_params->attention.host.q_seq_lens = padded_q_seq_lens_vec; } + if (use_dcp_local_block_tables) { + std::vector& padded_q_cu_seq_lens = + graph_params->attention.host.q_cu_seq_lens; + padded_q_cu_seq_lens.clear(); + padded_q_cu_seq_lens.reserve(static_cast(padded_batch_size)); + int32_t q_running_total = 0; + for (const int32_t q_seq_len : graph_params->attention.host.q_seq_lens) { + q_running_total += q_seq_len; + padded_q_cu_seq_lens.emplace_back(q_running_total); + } + } graph_params->meta.num_sequences = static_cast(padded_batch_size); graph_params->meta.batch_forward_type = params.meta.batch_forward_type; graph_params->enable_graph = true; @@ -1350,6 +1407,13 @@ std::optional GraphPersistentParam::update( persistent_new_cache_slots(padded_num_tokens); graph_params->attention.device.block_tables = persistent_block_tables(static_cast(padded_batch_size)); + if (use_dcp_local_block_tables) { + graph_params->graph.dcp_local_block_tables = + persistent_dcp_local_block_tables( + static_cast(padded_batch_size)); + } else { + graph_params->graph.dcp_local_block_tables = torch::Tensor(); + } if (!params.embedding.linear_state_ids.empty()) { graph_params->embedding.linear_state_ids = params.embedding.linear_state_ids; diff --git a/xllm/core/runtime/acl_graph_persistent_param.h b/xllm/core/runtime/acl_graph_persistent_param.h index 75bd12e89b..aedf471024 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.h +++ b/xllm/core/runtime/acl_graph_persistent_param.h @@ -142,6 +142,14 @@ class GraphPersistentParam final { } return persistent_block_tables_; } + torch::Tensor persistent_dcp_local_block_tables( + uint32_t actual_batch_size = 0) const { + if (actual_batch_size > 0) { + return persistent_dcp_local_block_tables_.slice( + /*dim=*/0, /*start=*/0, /*end=*/actual_batch_size); + } + return persistent_dcp_local_block_tables_; + } torch::Tensor persistent_mask(uint32_t actual_tokens = 0) const { if (actual_tokens > 0) { return persistent_mask_.slice( @@ -248,6 +256,7 @@ class GraphPersistentParam final { // Update attention mask efficiently from input parameters void update_attention_mask(const ModelInputParams& input_params); + void update_dcp_local_block_tables(int64_t padded_batch_size); void update_eplb_decode_token_mask(const ModelInputParams& input_params, uint32_t padded_num_tokens); @@ -275,6 +284,8 @@ class GraphPersistentParam final { torch::Tensor persistent_new_cache_slots_; torch::Tensor persistent_eplb_decode_token_mask_; torch::Tensor persistent_block_tables_; + torch::Tensor persistent_dcp_local_block_indices_; + torch::Tensor persistent_dcp_local_block_tables_; torch::Tensor persistent_new_cache_slots_default_; torch::Tensor persistent_block_tables_default_; torch::Tensor persistent_expanded_block_tables_; @@ -330,6 +341,8 @@ class GraphPersistentParam final { // Flag indicating whether the model uses hybrid linear attention // (e.g., Qwen3.5/Next with gated delta net layers) bool is_hybrid_linear_attention_; + int32_t dcp_size_ = 1; + int32_t dcp_rank_ = 0; // Flag indicating whether MLA graph capture uses KV length bucketing. bool supports_mla_graph_kv_bucketing_; // Flag indicating whether attention plan needs to be updated based on model diff --git a/xllm/core/runtime/forward_params.h b/xllm/core/runtime/forward_params.h index f20e803ca6..0d61210425 100644 --- a/xllm/core/runtime/forward_params.h +++ b/xllm/core/runtime/forward_params.h @@ -410,6 +410,7 @@ class WorkerType { enum class KvSlotLayout : int8_t { LOGICAL_REAL = 0, // Builder slots; input to prepare_cache_slots. NPU_CP_RECOVERED_PHYSICAL = 1, // Already CP-expanded; skip re-prepare. + NPU_DCP_LOCAL_PHYSICAL = 2, // Already DCP-local; non-owned tokens are -1. }; // Step-level decode metadata for Rec multi-round (device loop). diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 281968d6db..d153f64972 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -128,6 +128,8 @@ struct Options { // Context parallelism size PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, decode_context_parallel_size) = 1; + // tensor parallelism size // Default set as 1 PROPERTY(int32_t, tp_size) = 1; diff --git a/xllm/core/runtime/py_attention_metadata.h b/xllm/core/runtime/py_attention_metadata.h index ef33c2b8aa..a6cd423f2f 100644 --- a/xllm/core/runtime/py_attention_metadata.h +++ b/xllm/core/runtime/py_attention_metadata.h @@ -31,6 +31,9 @@ namespace xllm { struct ModelInputParams; +// Registers the internal runtime module in the active interpreter. The caller +// must hold the GIL. +void __attribute__((visibility("hidden"))) ensure_xllm_runtime_module(); void register_attention_metadata_views(pybind11::module_& module); class PyExpandedDecodeMetadataView final { diff --git a/xllm/core/runtime/py_executor_impl.cpp b/xllm/core/runtime/py_executor_impl.cpp index 8cc46c9c47..af86cb2fa5 100644 --- a/xllm/core/runtime/py_executor_impl.cpp +++ b/xllm/core/runtime/py_executor_impl.cpp @@ -60,12 +60,25 @@ void clear_python_object(py::object& object) { } // namespace -PYBIND11_EMBEDDED_MODULE(xllm_runtime, m) { - register_attention_metadata_views(m); +void ensure_xllm_runtime_module() { + py::module_ sys = py::module_::import("sys"); + py::dict modules = py::reinterpret_borrow(sys.attr("modules")); + const py::str module_name("xllm_runtime"); + if (modules.contains(module_name)) { + return; + } + + PyObject* module_object = PyModule_New("xllm_runtime"); + if (module_object == nullptr) { + throw py::error_already_set(); + } + py::module_ module = py::reinterpret_steal(module_object); + register_attention_metadata_views(module); #if defined(USE_NPU) py::class_>(m, "LayerSynchronizer") + std::shared_ptr>(module, + "LayerSynchronizer") .def("record_event", [](NPULayerSynchronizerImpl& self, int64_t layer_id) { int32_t device_id = static_cast( @@ -73,6 +86,8 @@ PYBIND11_EMBEDDED_MODULE(xllm_runtime, m) { return self.record_event(layer_id, device_id); }); #endif + + modules[module_name] = module; } PyExecutorImpl::PyExecutorImpl(CausalLM* model, @@ -87,6 +102,7 @@ PyExecutorImpl::PyExecutorImpl(CausalLM* model, CHECK(py_causal_lm_ != nullptr) << "PyExecutorImpl requires PyCausalLM"; py::gil_scoped_acquire gil; + ensure_xllm_runtime_module(); py::module_::import("xllm_runtime"); py::module_ executor_module = py::module_::import("xllm.python.model_executor.executor"); diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 798c10c032..e32b46a4ff 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -83,6 +83,7 @@ limitations under the License. #include "framework/model/model_input_params.h" #include "framework/model_loader.h" #include "framework/parallel_state/npu_cp_plan.h" +#include "framework/parallel_state/parallel_state.h" #include "framework/sampling/sampler.h" #include "framework/state_dict/state_dict.h" #include "framework/xtensor/global_xtensor.h" @@ -1189,6 +1190,35 @@ void WorkerImpl::prepare_work_before_execute(const ForwardInput& input, input, processed_input, *prepare_stream_); } +#if defined(USE_NPU) +torch::Tensor WorkerImpl::recompute_dcp_cache_slots( + const ForwardInput& input) const { + const int32_t dcp_size = parallel_args_.dcp_size_effective(); + CHECK_GT(dcp_size, 1) << "recompute_dcp_cache_slots requires dcp_size > 1"; + + const torch::Tensor& old_cache_slots = + input.input_params.attention.device.new_cache_slots; + if (!old_cache_slots.defined() || old_cache_slots.numel() == 0) { + return old_cache_slots; + } + + const torch::Tensor& host_positions = input.host_positions(); + CHECK(host_positions.defined()) + << "DCP cache slot remap requires host positions"; + CHECK_EQ(host_positions.numel(), old_cache_slots.numel()) + << "DCP cache slot remap requires positions and cache slots to have the " + "same token count"; + + const int32_t interleave_size = options_.block_size(); + const int32_t dcp_rank = parallel_args_.dcp_rank(); + + const torch::Tensor remapped = parallel_state::remap_dcp_cache_slots( + host_positions, old_cache_slots, interleave_size, dcp_size, dcp_rank); + return remapped.to(old_cache_slots.scalar_type()) + .to(old_cache_slots.device()); +} +#endif + void WorkerImpl::prepare_work_before_execute_on_stream( const ForwardInput& input, ForwardInput& processed_input, @@ -1315,6 +1345,22 @@ void WorkerImpl::prepare_work_before_execute_on_stream( processed_input.input_params.parallel.cp_plan.prepare( processed_input, npu_cp_plan_runtime_config()); + if (parallel_args_.dcp_size_effective() > 1 && + processed_input.kv_slot_layout == KvSlotLayout::LOGICAL_REAL) { + const BatchForwardType& batch_forward_type = + processed_input.input_params.meta.batch_forward_type; + CHECK(batch_forward_type.is_prefill() || + batch_forward_type.is_chunked_prefill() || + batch_forward_type.is_decode() || batch_forward_type.is_empty()) + << "DCP supports normal full prefill, chunked prefill, and decode " + "cache writes; mixed batches require DCP-2 layout support."; + CHECK(!processed_input.input_params.is_spec_verify) + << "DCP-1c does not support speculative verification cache writes."; + processed_input.input_params.attention.device.new_cache_slots = + recompute_dcp_cache_slots(processed_input); + processed_input.kv_slot_layout = KvSlotLayout::NPU_DCP_LOCAL_PHYSICAL; + } + if (can_prepare_npu_graph_decode_input(input_params)) { model_executor_->prepare_graph_input(processed_input.token_ids, processed_input.positions, diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index 29177fd6d0..4c4332fbc9 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -124,6 +124,7 @@ class WorkerImpl { // Per-worker-static configuration handed to NpuCpPlan::prepare(); built once // and cached. const CpPlanRuntimeConfig& npu_cp_plan_runtime_config() const; + torch::Tensor recompute_dcp_cache_slots(const ForwardInput& input) const; // Builds or reuses draft decode padding on the current stream. MTP calls this // while preparing B/2B metadata so the compute stream only observes hits. bool uses_npu_dp_ep_padding() const; diff --git a/xllm/core/scheduler/chunked_prefill_policy.h b/xllm/core/scheduler/chunked_prefill_policy.h new file mode 100644 index 0000000000..b76aca7e2d --- /dev/null +++ b/xllm/core/scheduler/chunked_prefill_policy.h @@ -0,0 +1,33 @@ +/* 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://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +namespace xllm { + +// The multi-SLO priority strategy always drives the chunked prefill scheduling +// path, so a service can run chunked prefill without setting the raw +// enable_chunked_prefill flag. Both the scheduler batch-mode resolution and the +// DCP startup gate must agree on this effective value; sharing this single +// predicate keeps them from drifting apart. +inline bool resolve_effective_chunked_prefill( + bool raw_chunked_prefill, + const std::string& priority_strategy) { + return raw_chunked_prefill || priority_strategy == "multi_slo_and_prio"; +} + +} // namespace xllm diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp index 2c5cb9e47d..95616f8257 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -42,6 +42,7 @@ limitations under the License. #include "framework/request/priority_comparator.h" #include "framework/request/request.h" #include "framework/request/sequence.h" +#include "scheduler/chunked_prefill_policy.h" #include "scheduler/request_priority_queue.h" #include "scheduler/scheduler_policy.h" #include "util/timer.h" @@ -64,17 +65,16 @@ std::vector> CancelRequestQueue::take_all() { BatchMode resolve_batch_mode(const ContinuousScheduler::Options& options) { BatchMode mode; mode.priority_strategy = options.priority_strategy(); - mode.enable_chunked_prefill = options.enable_chunked_prefill(); + mode.enable_chunked_prefill = resolve_effective_chunked_prefill( + options.enable_chunked_prefill(), options.priority_strategy()); mode.enable_mix_batch = ::xllm::SchedulerConfig::get_instance().enable_mix_batch(); - // multi_slo_and_prio requires chunked prefill. - if (mode.priority_strategy == "multi_slo_and_prio") { - mode.enable_chunked_prefill = true; - } - - // CP/MTP: prefill cannot mix with decode in the same batch. - if (options.cp_size() > 1 || options.num_speculative_tokens() > 0) { + // CP/DCP/MTP: prefill cannot mix with decode in the same batch. + // DCP reuses TP cards, so cp_size may be 1 while decode_context_parallel_size + // is > 1; the worker asserts a non-mixed batch under DCP, so close mix here. + if (options.cp_size() > 1 || options.decode_context_parallel_size() > 1 || + options.num_speculative_tokens() > 0) { mode.enable_mix_batch = false; } @@ -692,7 +692,7 @@ void ContinuousScheduler::update_memory_metrics( std::to_string(dp_rank), static_cast(active_kv_cache_size_in_kilobytes)); - if (::xllm::SchedulerConfig::get_instance().enable_chunked_prefill()) { + if (batch_mode_.enable_chunked_prefill) { MULTI_HISTOGRAM_OBSERVE(decode_active_activation_size_in_kilobytes, std::to_string(dp_rank), active_activation_size_in_kilobytes); diff --git a/xllm/core/scheduler/continuous_scheduler.h b/xllm/core/scheduler/continuous_scheduler.h index a5059d073a..7a673c7d58 100644 --- a/xllm/core/scheduler/continuous_scheduler.h +++ b/xllm/core/scheduler/continuous_scheduler.h @@ -104,6 +104,10 @@ class ContinuousScheduler : public Scheduler { PROPERTY(int32_t, cp_size) = 1; + // decode context parallel size; reuses TP cards, so it can coexist with + // cp_size == 1. When > 1, prefill cannot mix with decode in the same batch. + PROPERTY(int32_t, decode_context_parallel_size) = 1; + // enable disaggregated PD mode. PROPERTY(bool, enable_disagg_pd) = false; diff --git a/xllm/models/py_model_helper.cpp b/xllm/models/py_model_helper.cpp index be35a38320..31ca60db4b 100644 --- a/xllm/models/py_model_helper.cpp +++ b/xllm/models/py_model_helper.cpp @@ -15,7 +15,7 @@ limitations under the License. // Infrastructure for the embedded Python model executor: // - Interpreter lifecycle (ensure_python_interpreter) -// - Weight loading (PyStateDict + PYBIND11_EMBEDDED_MODULE) +// - Weight loading (PyStateDict + xllm_weight_loader module) // - Config serialization (dtype_to_string, PyDictVisitor) #include "models/py_model_helper.h" @@ -56,6 +56,26 @@ void prepend_sys_path(const std::string& dir) { } // namespace +void ensure_xllm_weight_loader_module() { + py::module_ sys = py::module_::import("sys"); + py::dict modules = py::reinterpret_borrow(sys.attr("modules")); + const py::str module_name("xllm_weight_loader"); + if (modules.contains(module_name)) { + return; + } + + PyObject* module_object = PyModule_New("xllm_weight_loader"); + if (module_object == nullptr) { + throw py::error_already_set(); + } + py::module_ module = py::reinterpret_steal(module_object); + py::class_(module, "StateDict") + .def("get_tensor", &PyStateDict::get_tensor, py::arg("name")) + .def("has", &PyStateDict::has, py::arg("name")) + .def("keys", &PyStateDict::keys); + modules[module_name] = module; +} + // --------------------------------------------------------------------------- // dtype_to_string // --------------------------------------------------------------------------- @@ -94,6 +114,7 @@ void ensure_python_interpreter() { { py::gil_scoped_acquire gil; + ensure_xllm_weight_loader_module(); std::string model_path = ModelConfig::get_instance().python_model_path(); if (model_path.empty()) { const char* env = std::getenv("XLLM_PYTHON_MODEL_PATH"); @@ -147,11 +168,4 @@ py::list PyStateDict::keys() const { return result; } -PYBIND11_EMBEDDED_MODULE(xllm_weight_loader, m) { - py::class_(m, "StateDict") - .def("get_tensor", &PyStateDict::get_tensor, py::arg("name")) - .def("has", &PyStateDict::has, py::arg("name")) - .def("keys", &PyStateDict::keys); -} - } // namespace xllm diff --git a/xllm/models/py_model_helper.h b/xllm/models/py_model_helper.h index 42419199fd..c7293a92bb 100644 --- a/xllm/models/py_model_helper.h +++ b/xllm/models/py_model_helper.h @@ -30,6 +30,10 @@ namespace xllm { // Initializes the embedded CPython interpreter (idempotent, process-wide). void ensure_python_interpreter(); +// Registers the internal weight-loader module in the active interpreter. The +// caller must hold the GIL. +void __attribute__((visibility("hidden"))) ensure_xllm_weight_loader_module(); + // Convert torch dtype to the string form used by Python model config. std::string dtype_to_string(const torch::TensorOptions& options); diff --git a/xllm/pybind/bind.cpp b/xllm/pybind/bind.cpp index bdf32a1d39..13e5cd35e1 100644 --- a/xllm/pybind/bind.cpp +++ b/xllm/pybind/bind.cpp @@ -29,13 +29,18 @@ limitations under the License. #include "core/framework/request/request_output.h" #include "core/framework/request/request_params.h" #include "core/framework/request/sample_slot.h" +#include "core/runtime/py_attention_metadata.h" #include "models/model_registry.h" +#include "models/py_model_helper.h" namespace xllm { namespace py = pybind11; using namespace pybind11::literals; PYBIND11_MODULE(xllm_export, m) { + ensure_xllm_runtime_module(); + ensure_xllm_weight_loader_module(); + // 1. export Options py::class_(m, "Options") .def(py::init()) diff --git a/xllm/python/models/glm5_2.py b/xllm/python/models/glm5_2.py index dc7c4bddc2..11233bf493 100644 --- a/xllm/python/models/glm5_2.py +++ b/xllm/python/models/glm5_2.py @@ -646,15 +646,13 @@ def load_weights( self.model.layers[i].mlp.process_weights_after_loading() else: se = p + "mlp.experts." + moe = self.model.layers[i].mlp + moe.allocate_experts_w13_for_loading() w13_param = self.get_parameter(p + "mlp.experts_w13") - w2_param = self.get_parameter(p + "mlp.experts_w2") w13_scale = self.get_buffer(p + "mlp.experts_w13_scale") w13_offset = self.get_buffer(p + "mlp.experts_w13_offset") - w2_scale = self.get_buffer(p + "mlp.experts_w2_scale") - w2_offset = self.get_buffer(p + "mlp.experts_w2_offset") - moe_layer = self.model.layers[i].mlp - expert_start = moe_layer.local_expert_start - expert_end = moe_layer.local_expert_end + expert_start = moe.local_expert_start + expert_end = moe.local_expert_end shard_world = cfg.moe_tp_size if cfg.ep_size > 1 else cfg.tp_size shard_rank = cfg.moe_tp_rank if cfg.ep_size > 1 else cfg.tp_rank for j in range(expert_start, expert_end): @@ -665,9 +663,6 @@ def load_weights( uw = loader.load_tensor(se + f"{j}.up_proj.weight") us = loader.load_tensor(se + f"{j}.up_proj.weight_scale") uo = loader.load_tensor(se + f"{j}.up_proj.weight_offset") - dw = loader.load_tensor(se + f"{j}.down_proj.weight") - ds = loader.load_tensor(se + f"{j}.down_proj.weight_scale") - do = loader.load_tensor(se + f"{j}.down_proj.weight_offset") w13_param.data[local_idx].copy_( torch.cat( [ @@ -695,9 +690,22 @@ def load_weights( dim=0, ).contiguous() ) + + moe.process_experts_w13_after_loading() + moe.allocate_experts_w2_for_loading() + w2_param = self.get_parameter(p + "mlp.experts_w2") + w2_scale = self.get_buffer(p + "mlp.experts_w2_scale") + w2_offset = self.get_buffer(p + "mlp.experts_w2_offset") + for j in range(expert_start, expert_end): + local_idx = j - expert_start + dw = loader.load_tensor(se + f"{j}.down_proj.weight") + ds = loader.load_tensor(se + f"{j}.down_proj.weight_scale") + do = loader.load_tensor(se + f"{j}.down_proj.weight_offset") w2_param.data[local_idx].copy_(loader.shard(dw, 1, shard_world, shard_rank).contiguous()) w2_scale.data[local_idx].copy_(ds.contiguous()) w2_offset.data[local_idx].copy_(do.contiguous()) + + moe.process_experts_w2_after_loading() loader.copy_in(p + "mlp.gate.weight", loader.load_tensor(p + "mlp.gate.weight")) loader.copy_in( p + "mlp.e_score_correction_bias", loader.load_tensor(p + "mlp.gate.e_score_correction_bias") @@ -707,7 +715,7 @@ def load_weights( loader.tp_rank = shard_rank loader.load_w8a8_b(p + "mlp.shared_experts.") loader.tp_size, loader.tp_rank = saved_tp - self.model.layers[i].mlp.process_weights_after_loading() + moe.process_weights_after_loading(skip_expert_format=True) loader.copy_in("model.norm.weight", loader.load_tensor("model.norm.weight")) loader.copy_in("lm_head.weight", loader.shard(loader.load_tensor("lm_head.weight"), dim=0)) diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index df69d84918..092cf098a1 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -193,6 +193,10 @@ Options create_options(const std::string& instance_name, bool is_local) { .node_rank(distributed_config.node_rank()) .dp_size(parallel_config.dp_size()) .cp_size(parallel_config.cp_size()) + .decode_context_parallel_size( + parallel_config.decode_context_parallel_size()) + .enable_experimental_dcp_chunked_prefill( + parallel_config.enable_experimental_dcp_chunked_prefill()) .ep_size(parallel_config.ep_size()) .tp_size(static_cast(parallel_config.tp_size())) .sp_size(static_cast(parallel_config.sp_size()))