From 20b1aee3237edd15edc206df7cb6f2b88def0148 Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Wed, 5 Aug 2026 09:33:53 +0800 Subject: [PATCH 01/22] feat: add decode context parallel (DCP) for Qwen3.5 GQA dense. Shard standard attention KV cache along sequence within a TP group during decode, reusing TP cards without expanding world size. GDN layers and prefill are untouched, aligning functionally with vllm-ascend DCP. - DCP-0a/0b: independent decode_context_parallel_size flag, startup GQA-topology validation, and TP-internal KV-replica subgroup (dcp_rank = tp_rank % dcp). - DCP-1c: owner-mask cache-slot remap keeping original physical slots (parallel_state::remap_dcp_cache_slots, integer floor_divide) + local block table selection by original allocator id. - DCP-2: decode via FIA with softmax_lse, zero-shard normalization, all-gather partials, fp32 online-softmax merge. - First-version startup compat gates (dcp_compat.h): fail-closed on chunked prefill, prefix cache, schedule overlap, P/D, speculative, and unvalidated MoE. Wording is "does not yet support", not "incompatible". - Tests: cp_group_ranks (incl. floor_divide regression), dcp_compat (12), fia_decode_lse probe, dcp_attention. Validated: Qwen3.5-2B tp=4/dcp=2 dense, per-token aligned with dcp=1 across short/boundary/L513/multi-sequence cases. MoE and HBM savings are follow-ups. Note: committed with --no-verify; pre-commit clang-format hook could not run (container virtualenv broken / physical host lacks pre-commit). All staged C/C++ verified clean via clang-format --dry-run --Werror manually. --- tests/core/common/options_test.cpp | 3 + tests/core/distributed_runtime/CMakeLists.txt | 10 + .../distributed_runtime/dcp_compat_test.cpp | 141 ++++++ .../spawn_worker_protocol_test.cpp | 4 + .../framework/config/config_json_test.cpp | 20 +- .../parallel_state/cp_group_ranks_test.cpp | 246 ++++++++++ tests/core/kernels/npu/CMakeLists.txt | 20 + .../kernels/npu/fia_decode_lse_probe_test.cpp | 456 ++++++++++++++++++ tests/core/layers/npu_torch/CMakeLists.txt | 26 + .../layers/npu_torch/dcp_attention_test.cpp | 152 ++++++ xllm/core/common/global_flags.h | 2 + xllm/core/common/options.cpp | 1 + xllm/core/common/options.h | 2 + xllm/core/distributed_runtime/dcp_compat.h | 74 +++ xllm/core/distributed_runtime/master.cpp | 174 ++++++- .../spawn_worker_protocol.h | 3 +- .../spawn_worker_server.cpp | 6 +- .../spawn_worker_server/spawn_worker_server.h | 3 +- .../spawn_worker_server_process.cpp | 20 +- .../distributed_runtime/worker_server.cpp | 5 + .../core/framework/config/parallel_config.cpp | 10 + xllm/core/framework/config/parallel_config.h | 3 + .../collective_communicator.cpp | 47 ++ .../parallel_state/collective_communicator.h | 1 + .../framework/parallel_state/parallel_args.h | 21 + .../parallel_state/parallel_state.cpp | 95 ++++ .../framework/parallel_state/parallel_state.h | 30 ++ xllm/core/layers/npu_torch/attention.cpp | 214 +++++++- xllm/core/layers/npu_torch/attention.h | 17 +- .../layers/npu_torch/qwen3_next_attention.cpp | 5 +- xllm/core/runtime/forward_params.h | 1 + xllm/core/runtime/options.h | 2 + xllm/core/runtime/worker_impl.cpp | 48 ++ xllm/core/runtime/worker_impl.h | 1 + xllm/xllm.cpp | 2 + 35 files changed, 1845 insertions(+), 20 deletions(-) create mode 100644 tests/core/distributed_runtime/dcp_compat_test.cpp create mode 100644 tests/core/kernels/npu/fia_decode_lse_probe_test.cpp create mode 100644 tests/core/layers/npu_torch/dcp_attention_test.cpp create mode 100644 xllm/core/distributed_runtime/dcp_compat.h diff --git a/tests/core/common/options_test.cpp b/tests/core/common/options_test.cpp index 3aa2ef6705..702241eaa2 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..6bfec89fec --- /dev/null +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -0,0 +1,141 @@ +/* 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_prefix_cache(false) + .enable_schedule_overlap(false) + .enable_disagg_pd(false) + .instance_role(InstanceRole::DEFAULT) + .num_speculative_tokens(0); + 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, AllowsSupportedFirstVersionFeatureFlags) { + const Options options = dcp_options_with_supported_feature_flags(); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, RejectsDefaultChunkedPrefillFirst) { + Options options; + options.decode_context_parallel_size(2); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_chunked_prefill=false"); +} + +TEST(DcpCompatTest, RejectsPrefixCache) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_prefix_cache(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_prefix_cache=false"); +} + +TEST(DcpCompatTest, RejectsScheduleOverlap) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_schedule_overlap(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_schedule_overlap=false"); +} + +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"); +} + +TEST(DcpCompatTest, AllowsDenseQwen35ModelType) { + EXPECT_FALSE( + validate_dcp_first_version_model_type("qwen3_5_text").has_value()); +} + +TEST(DcpCompatTest, RejectsUnvalidatedQwen35MoeModelType) { + expect_error_contains( + validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); +} + +} // 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 0ad450919e..0614a25c43 100644 --- a/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp +++ b/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp @@ -88,5 +88,9 @@ TEST(SpawnWorkerProtocolTest, PreservesExplicitEmptyDtype) { EXPECT_TRUE(indexer_cache_dtype->empty()); } +TEST(SpawnWorkerProtocolTest, AppendsDecodeContextParallelSizeAtTail) { + 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 a279b1643e..14adc1c49c 100644 --- a/tests/core/framework/config/config_json_test.cpp +++ b/tests/core/framework/config/config_json_test.cpp @@ -84,11 +84,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 { @@ -274,18 +280,22 @@ 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"); const std::string removed_flag = std::string("enable_") + "prefill_sp"; EXPECT_FALSE( 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 062d78137a..700e718127 100644 --- a/tests/core/framework/parallel_state/cp_group_ranks_test.cpp +++ b/tests/core/framework/parallel_state/cp_group_ranks_test.cpp @@ -36,6 +36,14 @@ 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; +} + TEST(ComputeCpGroupRanks, CpSizeTwoTpFourDpOne) { const int32_t world_size = 8; const int32_t dp_size = 1; @@ -139,6 +147,244 @@ 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(ComputeDcpCacheSlot, PreservesOwnerPhysicalSlots) { + const int32_t block_size = 4; + const int32_t dcp_size = 2; + const int32_t interleave_size = block_size; + + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/151, + /*position=*/0, + block_size, + dcp_size, + /*dcp_rank=*/0, + interleave_size), + 151); + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/23, + /*position=*/4, + block_size, + dcp_size, + /*dcp_rank=*/0, + interleave_size), + -1); + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/23, + /*position=*/4, + block_size, + dcp_size, + /*dcp_rank=*/1, + interleave_size), + 23); + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/359, + /*position=*/8, + block_size, + dcp_size, + /*dcp_rank=*/0, + interleave_size), + 359); +} + +TEST(ComputeDcpCacheSlot, RejectsSubBlockInterleave) { + const int32_t block_size = 4; + const int32_t dcp_size = 2; + EXPECT_DEATH(compute_dcp_cache_slot(/*logical_slot=*/0, + /*position=*/0, + block_size, + dcp_size, + /*dcp_rank=*/0, + /*interleave_size=*/1), + ""); +} + +TEST(ComputeDcpCacheSlot, PreservesNegativeSlots) { + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/-1, + /*position=*/0, + /*block_size=*/4, + /*dcp_size=*/2, + /*dcp_rank=*/0, + /*interleave_size=*/4), + -1); +} + +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 = + compute_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); + } + } +} + +// 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 f09f050864..997ea0befb 100644 --- a/tests/core/kernels/npu/CMakeLists.txt +++ b/tests/core/kernels/npu/CMakeLists.txt @@ -51,4 +51,24 @@ 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() + add_subdirectory(tilelang) 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..273cf42cfb --- /dev/null +++ b/tests/core/kernels/npu/fia_decode_lse_probe_test.cpp @@ -0,0 +1,456 @@ +/* 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 "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).view({-1}); + const torch::Tensor actual_cpu = actual.cpu().to(torch::kFloat32).view({-1}); + CHECK_EQ(expected_cpu.numel(), actual_cpu.numel()); + return (expected_cpu - actual_cpu).abs().max().item(); +} + +// 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"; +} + +// 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 09162f4f2e..6b1eb78eb3 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/tests/core/layers/npu_torch/CMakeLists.txt @@ -57,3 +57,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..3596111bb6 --- /dev/null +++ b/tests/core/layers/npu_torch/dcp_attention_test.cpp @@ -0,0 +1,152 @@ +/* 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 "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); +} + +} // namespace +} // namespace xllm::layer::test diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 2b3265e2da..4b18dfaa19 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -107,6 +107,8 @@ DECLARE_int32(ep_size); DECLARE_int32(cp_size); +DECLARE_int32(decode_context_parallel_size); + DECLARE_int64(tp_size); DECLARE_int64(sp_size); diff --git a/xllm/core/common/options.cpp b/xllm/core/common/options.cpp index 885245baa2..57fbe6124e 100644 --- a/xllm/core/common/options.cpp +++ b/xllm/core/common/options.cpp @@ -59,6 +59,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 c6765999bb..c73ccaf5a1 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -140,6 +140,8 @@ class Options { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, decode_context_parallel_size) = 1; + 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..6a8501cd5e --- /dev/null +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -0,0 +1,74 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://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" + +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 (options.enable_chunked_prefill()) { + return "decode_context_parallel_size first version does not yet support " + "chunked prefill; set --enable_chunked_prefill=false or set " + "--decode_context_parallel_size=1"; + } + if (options.enable_prefix_cache()) { + return "decode_context_parallel_size first version does not yet support " + "prefix cache; set --enable_prefix_cache=false or set " + "--decode_context_parallel_size=1"; + } + if (options.enable_schedule_overlap()) { + return "decode_context_parallel_size first version does not yet support " + "schedule overlap; set --enable_schedule_overlap=false 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; +} + +inline std::optional validate_dcp_first_version_model_type( + const std::string& model_type) { + if (model_type == "qwen3_5_moe_text") { + return "decode_context_parallel_size first version does not yet support " + "Qwen3.5 MoE; use dense Qwen3.5 or set " + "--decode_context_parallel_size=1"; + } + return std::nullopt; +} + +} // namespace xllm diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 9a1b36de3b..b58042d624 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -34,6 +34,7 @@ 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" @@ -54,6 +55,7 @@ limitations under the License. #include "rec_engine.h" #include "rec_master.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" @@ -69,6 +71,154 @@ 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 (options.enable_graph()) { + return "decode_context_parallel_size does not support graph capture yet; " + "disable graph or set decode_context_parallel_size=1"; + } + 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 (std::optional dcp_model_error = + validate_dcp_first_version_model_type(model_config->model_type)) { + return dcp_model_error; + } + 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; +} + std::optional validate_model_cp(const Options& options, EngineType engine_type, const std::string& model_type, @@ -297,10 +447,19 @@ 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 cp_model_type; - if (options_.cp_size() > 1 && Platform::uses_model_cp_sharding()) { - cp_model_type = util::get_model_type(model_path, options_.backend()); +#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()); } + const std::optional dcp_error = + validate_model_dcp(options_, type, dcp_model_config, global_world_size); + CHECK(!dcp_error.has_value()) << dcp_error.value(); + const std::string cp_model_type = + dcp_model_config.has_value() ? dcp_model_config->model_type : ""; const std::optional cp_error = validate_model_cp(options_, type, cp_model_type, global_world_size); CHECK(!cp_error.has_value()) << cp_error.value(); @@ -308,10 +467,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() @@ -351,7 +514,6 @@ Master::Master(const Options& options, EngineType type) if (options.eplb_update_threshold().has_value()) { eplb_config.eplb_update_threshold(options.eplb_update_threshold().value()); } - resolve_npu_kernel_backend_for_options(&options_); #endif ParallelConfig::get_instance().enable_multi_stream_parallel( options.enable_multi_stream_parallel() && (options.nnodes() > 1)); @@ -393,6 +555,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()) .npu_kernel_backend(options_.npu_kernel_backend()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .enable_offline_inference(options_.enable_offline_inference()) @@ -471,6 +634,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()) @@ -529,6 +693,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()) @@ -595,6 +760,7 @@ Master::Master(const Options& options, EngineType type) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .cp_size(options_.cp_size()) + .decode_context_parallel_size(options_.decode_context_parallel_size()) .max_seqs_per_batch(options_.max_seqs_per_batch()) .beam_width(options_.beam_width()) .max_tokens_per_batch(options_.max_tokens_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 b51c5d84b4..b50dec0565 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,10 +21,11 @@ limitations under the License. namespace xllm::spawn_worker_protocol { -inline constexpr int32_t kArgumentCount = 36; +inline constexpr int32_t kArgumentCount = 37; inline constexpr int32_t kMinimumArgumentCount = 34; inline constexpr int32_t kIndexerCacheDtypeArgumentIndex = 34; inline constexpr int32_t kEnableMtpDraftBodyTp1ArgumentIndex = 35; +inline constexpr int32_t kDecodeContextParallelSizeArgumentIndex = 36; inline constexpr char kDefaultIndexerCacheDtype[] = "auto"; inline std::optional parse_indexer_cache_dtype( 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 5d621351fb..a0243373a9 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 @@ -91,7 +91,8 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, int32_t cp_size, int32_t ep_size, const InstanceRole& instance_role, - bool enable_mtp_draft_body_tp1) { + bool enable_mtp_draft_body_tp1, + 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); @@ -116,6 +117,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) @@ -139,6 +141,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) @@ -183,6 +186,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 3b4935f2a2..8bcccd429d 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 @@ -62,7 +62,8 @@ class SpawnWorkerServer final { int32_t cp_size, int32_t ep_size, const InstanceRole& instance_role, - bool enable_mtp_draft_body_tp1); + bool enable_mtp_draft_body_tp1, + 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 a2cdbc2c18..20d5a9a3a5 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 @@ -60,6 +60,7 @@ limitations under the License. // @instance_role // @indexer_cache_dtype // @enable_mtp_draft_body_tp1 +// @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); @@ -117,13 +118,21 @@ int main(int argc, char* argv[]) { static_cast( atoi(argv[xllm::spawn_worker_protocol:: kEnableMtpDraftBodyTp1ArgumentIndex])) > 0; + 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 || + cp_size < 1 || ep_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 + << ", ep_size=" << ep_size << ", decode_context_parallel_size=" + << decode_context_parallel_size << ", instance_role=" << instance_role_str; return 1; } @@ -160,7 +169,9 @@ int main(int argc, char* argv[]) { << ", dp_size = " << dp_size << ", tp_size = " << tp_size << ", sp_size = " << sp_size << ", cfg_size = " << cfg_size << ", indexer_cache_dtype = " << indexer_cache_dtype - << ", enable_mtp_draft_body_tp1 = " << enable_mtp_draft_body_tp1 << "\n"; + << ", enable_mtp_draft_body_tp1 = " << enable_mtp_draft_body_tp1 + << ", decode_context_parallel_size = " << decode_context_parallel_size + << "\n"; xllm::SpawnWorkerServer worker(master_node_addr, local_rank, @@ -196,7 +207,8 @@ int main(int argc, char* argv[]) { cp_size, ep_size, instance_role, - enable_mtp_draft_body_tp1); + enable_mtp_draft_body_tp1, + 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 42598464c5..3e461efd48 100644 --- a/xllm/core/distributed_runtime/worker_server.cpp +++ b/xllm/core/distributed_runtime/worker_server.cpp @@ -271,6 +271,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(); @@ -372,6 +376,7 @@ void WorkerServer::create_spawn_server(int32_t local_rank, instance_role_ptr, indexer_cache_dtype_ptr, enable_mtp_draft_body_tp1_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 d29e1e2786..24e5ef3c92 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -26,6 +26,12 @@ 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_int32(kv_split_size, 1, "KV-cache split width. 0 falls back to cp_size (legacy); 1 means " @@ -75,6 +81,7 @@ 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(kv_split_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(tp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(sp_size); @@ -91,6 +98,7 @@ 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(tp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(cfg_size); @@ -108,6 +116,8 @@ 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, tp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, sp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 22159a1c2a..08717a3b11 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/xllm/core/framework/config/parallel_config.h @@ -44,6 +44,7 @@ class ParallelConfig final { {"dp_size", "ep_size", "cp_size", + "decode_context_parallel_size", "tp_size", "sp_size", "cfg_size", @@ -62,6 +63,8 @@ class ParallelConfig final { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, decode_context_parallel_size) = 1; + // 0 means follow cp_size (legacy KV-split width). PROPERTY(int32_t, kv_split_size) = 1; diff --git a/xllm/core/framework/parallel_state/collective_communicator.cpp b/xllm/core/framework/parallel_state/collective_communicator.cpp index 511beead40..a6a869229e 100644 --- a/xllm/core/framework/parallel_state/collective_communicator.cpp +++ b/xllm/core/framework/parallel_state/collective_communicator.cpp @@ -227,6 +227,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()); return; } @@ -284,11 +286,15 @@ 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()); #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()); #endif } @@ -417,6 +423,47 @@ void CollectiveCommunicator::create_process_groups( parallel_args_->cp_group_ = tp_group_.get(); port += dp_size + single_rank_group_port_gap + single_rank_group_count; + 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) { port_offset = global_rank % tp_size + 1; dp_local_process_group_ = create_process_group(global_rank, diff --git a/xllm/core/framework/parallel_state/collective_communicator.h b/xllm/core/framework/parallel_state/collective_communicator.h index 08a0415c6d..980bf5467f 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 69cdf97221..01543b9a80 100644 --- a/xllm/core/framework/parallel_state/parallel_args.h +++ b/xllm/core/framework/parallel_state/parallel_args.h @@ -148,6 +148,8 @@ struct ParallelArgs { // cp size PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, dcp_size) = 1; + // Derived: CP rank of the current process within its DP group. // rank layout: dp_rank * (cp_size * tp_size) + cp_rank * tp_size + tp_rank [[nodiscard]] int32_t cp_rank() const noexcept { @@ -165,6 +167,24 @@ struct ParallelArgs { return kv_split_size_ > 0 ? kv_split_size_ : cp_size_; } + [[nodiscard]] int32_t dcp_size_effective() const noexcept { + return dcp_size_ > 0 ? dcp_size_ : 1; + } + + [[nodiscard]] int32_t dcp_rank() const noexcept { + 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 { const int32_t kv = kv_split_size_effective(); if (kv <= 1) { @@ -211,6 +231,7 @@ struct ParallelArgs { ProcessGroup* single_rank_group_ = nullptr; // CP ProcessGroup for prefill AllGather (NPU standalone; MLU aliases TP). ProcessGroup* cp_group_ = nullptr; + ProcessGroup* dcp_group_ = nullptr; ProcessGroup* moe_ep_group_ = nullptr; // Dedicated group for EPLB weight migration. It has the same rank set as // moe_ep_group_ but isolates migration P2P from forward collectives. diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index 313ff262b6..5eabdc195a 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/xllm/core/framework/parallel_state/parallel_state.cpp @@ -275,6 +275,101 @@ 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; +} + +int64_t compute_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; +} + +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 b40ddfa34d..1860cf76e4 100644 --- a/xllm/core/framework/parallel_state/parallel_state.h +++ b/xllm/core/framework/parallel_state/parallel_state.h @@ -74,6 +74,36 @@ 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); + +// Remap a logical KV cache slot to this DCP rank's local physical slot. Returns +// -1 when the token is owned by a different DCP rank. +int64_t compute_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); + +torch::Tensor select_dcp_local_block_table(const torch::Tensor& block_table, + int32_t dcp_size, + int32_t dcp_rank); + +// Batched tensor form of compute_dcp_cache_slot for the production remap path. +// 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 +#include +#include +#include + +#include "framework/parallel_state/parallel_state.h" #include "kernels/npu/npu_ops_api.h" #include "kernels/ops_api.h" +namespace { + +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 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()); + } + } +} + +torch::Tensor merge_dcp_partials(const torch::Tensor& all_partial_out, + const torch::Tensor& all_partial_lse) { + CHECK_EQ(all_partial_out.scalar_type(), torch::kFloat32); + CHECK_EQ(all_partial_lse.scalar_type(), torch::kFloat32); + 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 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(torch::isfinite(all_partial_lse), + torch::exp(all_partial_lse - safe_max_lse), + torch::zeros_like(all_partial_lse)); + 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 * all_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 + namespace xllm { namespace layer { @@ -25,12 +132,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; } @@ -134,11 +250,105 @@ 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.use_expanded_decode_for_spec_verify_attention) + << "DCP-2 does not support speculative decode attention."; + CHECK(!attn_metadata.paged_attention_tiling_data.defined()) + << "DCP-2 does not support graph-captured 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 = compute_dcp_local_kv_seq_lens( + global_kv_seq_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_); + 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; + 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, + 0, + "TND", + true); + torch::Tensor partial_out = std::get<0>(fia_result).to(torch::kFloat32); + torch::Tensor partial_lse = std::get<1>(fia_result).to(torch::kFloat32); + 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 = + 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::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_}); diff --git a/xllm/core/layers/npu_torch/attention.h b/xllm/core/layers/npu_torch/attention.h index bd81823023..9b5f9c30c8 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,20 @@ 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); + 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/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index c1dec2e90e..21edba4438 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/runtime/forward_params.h b/xllm/core/runtime/forward_params.h index a2f5e9de7f..5297000e61 100644 --- a/xllm/core/runtime/forward_params.h +++ b/xllm/core/runtime/forward_params.h @@ -406,6 +406,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 d361faec68..df28ba6e06 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -122,6 +122,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/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index c190f2d78a..7ad502a485 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -73,6 +73,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" @@ -750,6 +751,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, @@ -879,6 +909,24 @@ 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_decode() || + batch_forward_type.is_empty()) + << "DCP-1c supports only normal full prefill and decode cache " + "writes; chunked and mixed batches require DCP-2 layout " + "support."; + CHECK(!processed_input.input_params.is_spec_verify) + << "DCP-1c does not support speculative verification cache writes."; + CHECK(!processed_input.input_params.enable_graph) + << "DCP-1c does not support graph-captured 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 60f1f002db..8d161a7b39 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -119,6 +119,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; #endif // False on MTP composite: only leaf workers run NpuCpPlan::prepare. diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index d4fa0295c4..712deff676 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -179,6 +179,8 @@ 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()) .ep_size(parallel_config.ep_size()) .tp_size(static_cast(parallel_config.tp_size())) .sp_size(static_cast(parallel_config.sp_size())) From 45275cbf8757c5838c5f1dbb2038e54a343249fd Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Wed, 5 Aug 2026 09:33:53 +0800 Subject: [PATCH 02/22] feat: add decode context parallel (DCP) for Qwen3.5 GQA dense. Shard standard attention KV cache along sequence within a TP group during decode, reusing TP cards without expanding world size. GDN layers and prefill are untouched, aligning functionally with vllm-ascend DCP. - DCP-0a/0b: independent decode_context_parallel_size flag, startup GQA-topology validation, and TP-internal KV-replica subgroup (dcp_rank = tp_rank % dcp). - DCP-1c: owner-mask cache-slot remap keeping original physical slots (parallel_state::remap_dcp_cache_slots, integer floor_divide) + local block table selection by original allocator id. - DCP-2: decode via FIA with softmax_lse, zero-shard normalization, all-gather partials, fp32 online-softmax merge. - First-version startup compat gates (dcp_compat.h): fail-closed on chunked prefill, prefix cache, schedule overlap, P/D, speculative, and unvalidated MoE. Wording is "does not yet support", not "incompatible". - Tests: cp_group_ranks (incl. floor_divide regression), dcp_compat (12), fia_decode_lse probe, dcp_attention. Validated: Qwen3.5-2B tp=4/dcp=2 dense, per-token aligned with dcp=1 across short/boundary/L513/multi-sequence cases. MoE and HBM savings are follow-ups. Note: committed with --no-verify; pre-commit clang-format hook could not run (container virtualenv broken / physical host lacks pre-commit). All staged C/C++ verified clean via clang-format --dry-run --Werror manually. --- tests/core/common/options_test.cpp | 3 + tests/core/distributed_runtime/CMakeLists.txt | 10 + .../distributed_runtime/dcp_compat_test.cpp | 141 ++++++ .../spawn_worker_protocol_test.cpp | 4 + .../framework/config/config_json_test.cpp | 20 +- .../parallel_state/cp_group_ranks_test.cpp | 246 ++++++++++ tests/core/kernels/npu/CMakeLists.txt | 20 + .../kernels/npu/fia_decode_lse_probe_test.cpp | 456 ++++++++++++++++++ tests/core/layers/npu_torch/CMakeLists.txt | 26 + .../layers/npu_torch/dcp_attention_test.cpp | 152 ++++++ xllm/core/common/global_flags.h | 2 + xllm/core/common/options.cpp | 1 + xllm/core/common/options.h | 2 + xllm/core/distributed_runtime/dcp_compat.h | 74 +++ xllm/core/distributed_runtime/master.cpp | 174 ++++++- .../spawn_worker_protocol.h | 3 +- .../spawn_worker_server.cpp | 6 +- .../spawn_worker_server/spawn_worker_server.h | 3 +- .../spawn_worker_server_process.cpp | 20 +- .../distributed_runtime/worker_server.cpp | 5 + .../core/framework/config/parallel_config.cpp | 10 + xllm/core/framework/config/parallel_config.h | 3 + .../collective_communicator.cpp | 47 ++ .../parallel_state/collective_communicator.h | 1 + .../framework/parallel_state/parallel_args.h | 21 + .../parallel_state/parallel_state.cpp | 95 ++++ .../framework/parallel_state/parallel_state.h | 30 ++ xllm/core/layers/npu_torch/attention.cpp | 214 +++++++- xllm/core/layers/npu_torch/attention.h | 17 +- .../layers/npu_torch/qwen3_next_attention.cpp | 5 +- xllm/core/runtime/forward_params.h | 1 + xllm/core/runtime/options.h | 2 + xllm/core/runtime/worker_impl.cpp | 48 ++ xllm/core/runtime/worker_impl.h | 1 + xllm/xllm.cpp | 2 + 35 files changed, 1845 insertions(+), 20 deletions(-) create mode 100644 tests/core/distributed_runtime/dcp_compat_test.cpp create mode 100644 tests/core/kernels/npu/fia_decode_lse_probe_test.cpp create mode 100644 tests/core/layers/npu_torch/dcp_attention_test.cpp create mode 100644 xllm/core/distributed_runtime/dcp_compat.h diff --git a/tests/core/common/options_test.cpp b/tests/core/common/options_test.cpp index 3aa2ef6705..702241eaa2 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..6bfec89fec --- /dev/null +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -0,0 +1,141 @@ +/* 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_prefix_cache(false) + .enable_schedule_overlap(false) + .enable_disagg_pd(false) + .instance_role(InstanceRole::DEFAULT) + .num_speculative_tokens(0); + 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, AllowsSupportedFirstVersionFeatureFlags) { + const Options options = dcp_options_with_supported_feature_flags(); + + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); +} + +TEST(DcpCompatTest, RejectsDefaultChunkedPrefillFirst) { + Options options; + options.decode_context_parallel_size(2); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_chunked_prefill=false"); +} + +TEST(DcpCompatTest, RejectsPrefixCache) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_prefix_cache(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_prefix_cache=false"); +} + +TEST(DcpCompatTest, RejectsScheduleOverlap) { + Options options = dcp_options_with_supported_feature_flags(); + options.enable_schedule_overlap(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_schedule_overlap=false"); +} + +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"); +} + +TEST(DcpCompatTest, AllowsDenseQwen35ModelType) { + EXPECT_FALSE( + validate_dcp_first_version_model_type("qwen3_5_text").has_value()); +} + +TEST(DcpCompatTest, RejectsUnvalidatedQwen35MoeModelType) { + expect_error_contains( + validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); +} + +} // 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 0ad450919e..0614a25c43 100644 --- a/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp +++ b/tests/core/distributed_runtime/spawn_worker_protocol_test.cpp @@ -88,5 +88,9 @@ TEST(SpawnWorkerProtocolTest, PreservesExplicitEmptyDtype) { EXPECT_TRUE(indexer_cache_dtype->empty()); } +TEST(SpawnWorkerProtocolTest, AppendsDecodeContextParallelSizeAtTail) { + 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 a279b1643e..14adc1c49c 100644 --- a/tests/core/framework/config/config_json_test.cpp +++ b/tests/core/framework/config/config_json_test.cpp @@ -84,11 +84,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 { @@ -274,18 +280,22 @@ 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"); const std::string removed_flag = std::string("enable_") + "prefill_sp"; EXPECT_FALSE( 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 062d78137a..700e718127 100644 --- a/tests/core/framework/parallel_state/cp_group_ranks_test.cpp +++ b/tests/core/framework/parallel_state/cp_group_ranks_test.cpp @@ -36,6 +36,14 @@ 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; +} + TEST(ComputeCpGroupRanks, CpSizeTwoTpFourDpOne) { const int32_t world_size = 8; const int32_t dp_size = 1; @@ -139,6 +147,244 @@ 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(ComputeDcpCacheSlot, PreservesOwnerPhysicalSlots) { + const int32_t block_size = 4; + const int32_t dcp_size = 2; + const int32_t interleave_size = block_size; + + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/151, + /*position=*/0, + block_size, + dcp_size, + /*dcp_rank=*/0, + interleave_size), + 151); + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/23, + /*position=*/4, + block_size, + dcp_size, + /*dcp_rank=*/0, + interleave_size), + -1); + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/23, + /*position=*/4, + block_size, + dcp_size, + /*dcp_rank=*/1, + interleave_size), + 23); + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/359, + /*position=*/8, + block_size, + dcp_size, + /*dcp_rank=*/0, + interleave_size), + 359); +} + +TEST(ComputeDcpCacheSlot, RejectsSubBlockInterleave) { + const int32_t block_size = 4; + const int32_t dcp_size = 2; + EXPECT_DEATH(compute_dcp_cache_slot(/*logical_slot=*/0, + /*position=*/0, + block_size, + dcp_size, + /*dcp_rank=*/0, + /*interleave_size=*/1), + ""); +} + +TEST(ComputeDcpCacheSlot, PreservesNegativeSlots) { + EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/-1, + /*position=*/0, + /*block_size=*/4, + /*dcp_size=*/2, + /*dcp_rank=*/0, + /*interleave_size=*/4), + -1); +} + +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 = + compute_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); + } + } +} + +// 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 f09f050864..997ea0befb 100644 --- a/tests/core/kernels/npu/CMakeLists.txt +++ b/tests/core/kernels/npu/CMakeLists.txt @@ -51,4 +51,24 @@ 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() + add_subdirectory(tilelang) 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..273cf42cfb --- /dev/null +++ b/tests/core/kernels/npu/fia_decode_lse_probe_test.cpp @@ -0,0 +1,456 @@ +/* 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 "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).view({-1}); + const torch::Tensor actual_cpu = actual.cpu().to(torch::kFloat32).view({-1}); + CHECK_EQ(expected_cpu.numel(), actual_cpu.numel()); + return (expected_cpu - actual_cpu).abs().max().item(); +} + +// 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"; +} + +// 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 09162f4f2e..6b1eb78eb3 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/tests/core/layers/npu_torch/CMakeLists.txt @@ -57,3 +57,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..3596111bb6 --- /dev/null +++ b/tests/core/layers/npu_torch/dcp_attention_test.cpp @@ -0,0 +1,152 @@ +/* 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 "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); +} + +} // namespace +} // namespace xllm::layer::test diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 2b3265e2da..4b18dfaa19 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -107,6 +107,8 @@ DECLARE_int32(ep_size); DECLARE_int32(cp_size); +DECLARE_int32(decode_context_parallel_size); + DECLARE_int64(tp_size); DECLARE_int64(sp_size); diff --git a/xllm/core/common/options.cpp b/xllm/core/common/options.cpp index 885245baa2..57fbe6124e 100644 --- a/xllm/core/common/options.cpp +++ b/xllm/core/common/options.cpp @@ -59,6 +59,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 c6765999bb..c73ccaf5a1 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -140,6 +140,8 @@ class Options { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, decode_context_parallel_size) = 1; + 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..6a8501cd5e --- /dev/null +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -0,0 +1,74 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://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" + +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 (options.enable_chunked_prefill()) { + return "decode_context_parallel_size first version does not yet support " + "chunked prefill; set --enable_chunked_prefill=false or set " + "--decode_context_parallel_size=1"; + } + if (options.enable_prefix_cache()) { + return "decode_context_parallel_size first version does not yet support " + "prefix cache; set --enable_prefix_cache=false or set " + "--decode_context_parallel_size=1"; + } + if (options.enable_schedule_overlap()) { + return "decode_context_parallel_size first version does not yet support " + "schedule overlap; set --enable_schedule_overlap=false 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; +} + +inline std::optional validate_dcp_first_version_model_type( + const std::string& model_type) { + if (model_type == "qwen3_5_moe_text") { + return "decode_context_parallel_size first version does not yet support " + "Qwen3.5 MoE; use dense Qwen3.5 or set " + "--decode_context_parallel_size=1"; + } + return std::nullopt; +} + +} // namespace xllm diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 9a1b36de3b..b58042d624 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -34,6 +34,7 @@ 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" @@ -54,6 +55,7 @@ limitations under the License. #include "rec_engine.h" #include "rec_master.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" @@ -69,6 +71,154 @@ 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 (options.enable_graph()) { + return "decode_context_parallel_size does not support graph capture yet; " + "disable graph or set decode_context_parallel_size=1"; + } + 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 (std::optional dcp_model_error = + validate_dcp_first_version_model_type(model_config->model_type)) { + return dcp_model_error; + } + 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; +} + std::optional validate_model_cp(const Options& options, EngineType engine_type, const std::string& model_type, @@ -297,10 +447,19 @@ 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 cp_model_type; - if (options_.cp_size() > 1 && Platform::uses_model_cp_sharding()) { - cp_model_type = util::get_model_type(model_path, options_.backend()); +#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()); } + const std::optional dcp_error = + validate_model_dcp(options_, type, dcp_model_config, global_world_size); + CHECK(!dcp_error.has_value()) << dcp_error.value(); + const std::string cp_model_type = + dcp_model_config.has_value() ? dcp_model_config->model_type : ""; const std::optional cp_error = validate_model_cp(options_, type, cp_model_type, global_world_size); CHECK(!cp_error.has_value()) << cp_error.value(); @@ -308,10 +467,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() @@ -351,7 +514,6 @@ Master::Master(const Options& options, EngineType type) if (options.eplb_update_threshold().has_value()) { eplb_config.eplb_update_threshold(options.eplb_update_threshold().value()); } - resolve_npu_kernel_backend_for_options(&options_); #endif ParallelConfig::get_instance().enable_multi_stream_parallel( options.enable_multi_stream_parallel() && (options.nnodes() > 1)); @@ -393,6 +555,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()) .npu_kernel_backend(options_.npu_kernel_backend()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .enable_offline_inference(options_.enable_offline_inference()) @@ -471,6 +634,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()) @@ -529,6 +693,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()) @@ -595,6 +760,7 @@ Master::Master(const Options& options, EngineType type) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .cp_size(options_.cp_size()) + .decode_context_parallel_size(options_.decode_context_parallel_size()) .max_seqs_per_batch(options_.max_seqs_per_batch()) .beam_width(options_.beam_width()) .max_tokens_per_batch(options_.max_tokens_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 b51c5d84b4..b50dec0565 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,10 +21,11 @@ limitations under the License. namespace xllm::spawn_worker_protocol { -inline constexpr int32_t kArgumentCount = 36; +inline constexpr int32_t kArgumentCount = 37; inline constexpr int32_t kMinimumArgumentCount = 34; inline constexpr int32_t kIndexerCacheDtypeArgumentIndex = 34; inline constexpr int32_t kEnableMtpDraftBodyTp1ArgumentIndex = 35; +inline constexpr int32_t kDecodeContextParallelSizeArgumentIndex = 36; inline constexpr char kDefaultIndexerCacheDtype[] = "auto"; inline std::optional parse_indexer_cache_dtype( 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 5d621351fb..a0243373a9 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 @@ -91,7 +91,8 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, int32_t cp_size, int32_t ep_size, const InstanceRole& instance_role, - bool enable_mtp_draft_body_tp1) { + bool enable_mtp_draft_body_tp1, + 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); @@ -116,6 +117,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) @@ -139,6 +141,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) @@ -183,6 +186,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 3b4935f2a2..8bcccd429d 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 @@ -62,7 +62,8 @@ class SpawnWorkerServer final { int32_t cp_size, int32_t ep_size, const InstanceRole& instance_role, - bool enable_mtp_draft_body_tp1); + bool enable_mtp_draft_body_tp1, + 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 a2cdbc2c18..20d5a9a3a5 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 @@ -60,6 +60,7 @@ limitations under the License. // @instance_role // @indexer_cache_dtype // @enable_mtp_draft_body_tp1 +// @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); @@ -117,13 +118,21 @@ int main(int argc, char* argv[]) { static_cast( atoi(argv[xllm::spawn_worker_protocol:: kEnableMtpDraftBodyTp1ArgumentIndex])) > 0; + 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 || + cp_size < 1 || ep_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 + << ", ep_size=" << ep_size << ", decode_context_parallel_size=" + << decode_context_parallel_size << ", instance_role=" << instance_role_str; return 1; } @@ -160,7 +169,9 @@ int main(int argc, char* argv[]) { << ", dp_size = " << dp_size << ", tp_size = " << tp_size << ", sp_size = " << sp_size << ", cfg_size = " << cfg_size << ", indexer_cache_dtype = " << indexer_cache_dtype - << ", enable_mtp_draft_body_tp1 = " << enable_mtp_draft_body_tp1 << "\n"; + << ", enable_mtp_draft_body_tp1 = " << enable_mtp_draft_body_tp1 + << ", decode_context_parallel_size = " << decode_context_parallel_size + << "\n"; xllm::SpawnWorkerServer worker(master_node_addr, local_rank, @@ -196,7 +207,8 @@ int main(int argc, char* argv[]) { cp_size, ep_size, instance_role, - enable_mtp_draft_body_tp1); + enable_mtp_draft_body_tp1, + 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 42598464c5..3e461efd48 100644 --- a/xllm/core/distributed_runtime/worker_server.cpp +++ b/xllm/core/distributed_runtime/worker_server.cpp @@ -271,6 +271,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(); @@ -372,6 +376,7 @@ void WorkerServer::create_spawn_server(int32_t local_rank, instance_role_ptr, indexer_cache_dtype_ptr, enable_mtp_draft_body_tp1_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 d29e1e2786..24e5ef3c92 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -26,6 +26,12 @@ 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_int32(kv_split_size, 1, "KV-cache split width. 0 falls back to cp_size (legacy); 1 means " @@ -75,6 +81,7 @@ 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(kv_split_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(tp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(sp_size); @@ -91,6 +98,7 @@ 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(tp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(cfg_size); @@ -108,6 +116,8 @@ 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, tp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, sp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 22159a1c2a..08717a3b11 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/xllm/core/framework/config/parallel_config.h @@ -44,6 +44,7 @@ class ParallelConfig final { {"dp_size", "ep_size", "cp_size", + "decode_context_parallel_size", "tp_size", "sp_size", "cfg_size", @@ -62,6 +63,8 @@ class ParallelConfig final { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, decode_context_parallel_size) = 1; + // 0 means follow cp_size (legacy KV-split width). PROPERTY(int32_t, kv_split_size) = 1; diff --git a/xllm/core/framework/parallel_state/collective_communicator.cpp b/xllm/core/framework/parallel_state/collective_communicator.cpp index 511beead40..a6a869229e 100644 --- a/xllm/core/framework/parallel_state/collective_communicator.cpp +++ b/xllm/core/framework/parallel_state/collective_communicator.cpp @@ -227,6 +227,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()); return; } @@ -284,11 +286,15 @@ 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()); #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()); #endif } @@ -417,6 +423,47 @@ void CollectiveCommunicator::create_process_groups( parallel_args_->cp_group_ = tp_group_.get(); port += dp_size + single_rank_group_port_gap + single_rank_group_count; + 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) { port_offset = global_rank % tp_size + 1; dp_local_process_group_ = create_process_group(global_rank, diff --git a/xllm/core/framework/parallel_state/collective_communicator.h b/xllm/core/framework/parallel_state/collective_communicator.h index 08a0415c6d..980bf5467f 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 69cdf97221..01543b9a80 100644 --- a/xllm/core/framework/parallel_state/parallel_args.h +++ b/xllm/core/framework/parallel_state/parallel_args.h @@ -148,6 +148,8 @@ struct ParallelArgs { // cp size PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, dcp_size) = 1; + // Derived: CP rank of the current process within its DP group. // rank layout: dp_rank * (cp_size * tp_size) + cp_rank * tp_size + tp_rank [[nodiscard]] int32_t cp_rank() const noexcept { @@ -165,6 +167,24 @@ struct ParallelArgs { return kv_split_size_ > 0 ? kv_split_size_ : cp_size_; } + [[nodiscard]] int32_t dcp_size_effective() const noexcept { + return dcp_size_ > 0 ? dcp_size_ : 1; + } + + [[nodiscard]] int32_t dcp_rank() const noexcept { + 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 { const int32_t kv = kv_split_size_effective(); if (kv <= 1) { @@ -211,6 +231,7 @@ struct ParallelArgs { ProcessGroup* single_rank_group_ = nullptr; // CP ProcessGroup for prefill AllGather (NPU standalone; MLU aliases TP). ProcessGroup* cp_group_ = nullptr; + ProcessGroup* dcp_group_ = nullptr; ProcessGroup* moe_ep_group_ = nullptr; // Dedicated group for EPLB weight migration. It has the same rank set as // moe_ep_group_ but isolates migration P2P from forward collectives. diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index 313ff262b6..5eabdc195a 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/xllm/core/framework/parallel_state/parallel_state.cpp @@ -275,6 +275,101 @@ 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; +} + +int64_t compute_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; +} + +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 b40ddfa34d..1860cf76e4 100644 --- a/xllm/core/framework/parallel_state/parallel_state.h +++ b/xllm/core/framework/parallel_state/parallel_state.h @@ -74,6 +74,36 @@ 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); + +// Remap a logical KV cache slot to this DCP rank's local physical slot. Returns +// -1 when the token is owned by a different DCP rank. +int64_t compute_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); + +torch::Tensor select_dcp_local_block_table(const torch::Tensor& block_table, + int32_t dcp_size, + int32_t dcp_rank); + +// Batched tensor form of compute_dcp_cache_slot for the production remap path. +// 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 +#include +#include +#include + +#include "framework/parallel_state/parallel_state.h" #include "kernels/npu/npu_ops_api.h" #include "kernels/ops_api.h" +namespace { + +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 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()); + } + } +} + +torch::Tensor merge_dcp_partials(const torch::Tensor& all_partial_out, + const torch::Tensor& all_partial_lse) { + CHECK_EQ(all_partial_out.scalar_type(), torch::kFloat32); + CHECK_EQ(all_partial_lse.scalar_type(), torch::kFloat32); + 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 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(torch::isfinite(all_partial_lse), + torch::exp(all_partial_lse - safe_max_lse), + torch::zeros_like(all_partial_lse)); + 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 * all_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 + namespace xllm { namespace layer { @@ -25,12 +132,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; } @@ -134,11 +250,105 @@ 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.use_expanded_decode_for_spec_verify_attention) + << "DCP-2 does not support speculative decode attention."; + CHECK(!attn_metadata.paged_attention_tiling_data.defined()) + << "DCP-2 does not support graph-captured 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 = compute_dcp_local_kv_seq_lens( + global_kv_seq_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_); + 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; + 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, + 0, + "TND", + true); + torch::Tensor partial_out = std::get<0>(fia_result).to(torch::kFloat32); + torch::Tensor partial_lse = std::get<1>(fia_result).to(torch::kFloat32); + 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 = + 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::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_}); diff --git a/xllm/core/layers/npu_torch/attention.h b/xllm/core/layers/npu_torch/attention.h index bd81823023..9b5f9c30c8 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,20 @@ 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); + 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/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index c1dec2e90e..21edba4438 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/runtime/forward_params.h b/xllm/core/runtime/forward_params.h index a2f5e9de7f..5297000e61 100644 --- a/xllm/core/runtime/forward_params.h +++ b/xllm/core/runtime/forward_params.h @@ -406,6 +406,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 d361faec68..df28ba6e06 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -122,6 +122,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/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index c190f2d78a..7ad502a485 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -73,6 +73,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" @@ -750,6 +751,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, @@ -879,6 +909,24 @@ 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_decode() || + batch_forward_type.is_empty()) + << "DCP-1c supports only normal full prefill and decode cache " + "writes; chunked and mixed batches require DCP-2 layout " + "support."; + CHECK(!processed_input.input_params.is_spec_verify) + << "DCP-1c does not support speculative verification cache writes."; + CHECK(!processed_input.input_params.enable_graph) + << "DCP-1c does not support graph-captured 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 60f1f002db..8d161a7b39 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -119,6 +119,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; #endif // False on MTP composite: only leaf workers run NpuCpPlan::prepare. diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index d4fa0295c4..712deff676 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -179,6 +179,8 @@ 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()) .ep_size(parallel_config.ep_size()) .tp_size(static_cast(parallel_config.tp_size())) .sp_size(static_cast(parallel_config.sp_size())) From 44cfe3d17a7b06de2ae3a28b2b3e86481f6775fc Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Wed, 5 Aug 2026 13:16:10 +0800 Subject: [PATCH 03/22] refactor: extract DCP attention math into testable helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move compute_dcp_local_kv_seq_lens and merge_dcp_partials out of attention.cpp into layers/npu_torch/dcp_attention_utils.{h,cpp} (namespace xllm::layer::detail). These are the production DCP-2 local-length and online-softmax merge implementations — attention.cpp now calls detail::* rather than defining them inline. The extraction makes them linkable from a CPU-only float64 unit test (dcp_attention_utils_test) that verifies sharded-KV merge equals full attention at 1e-10, invalid/all-invalid LSE shards, and dcp=4 tail allocation — without pulling in torch_npu / libpython. Note: committed with --no-verify; pre-commit clang-format hook cannot run (container virtualenv broken). All 4 C/C++ files verified clean via clang-format --dry-run --Werror manually. --- cmake/cc_test.cmake | 15 +- tests/core/layers/npu_torch/CMakeLists.txt | 17 +++ .../npu_torch/dcp_attention_utils_test.cpp | 130 ++++++++++++++++++ xllm/core/layers/npu_torch/CMakeLists.txt | 13 ++ xllm/core/layers/npu_torch/attention.cpp | 63 +-------- .../layers/npu_torch/dcp_attention_utils.cpp | 84 +++++++++++ .../layers/npu_torch/dcp_attention_utils.h | 34 +++++ 7 files changed, 296 insertions(+), 60 deletions(-) create mode 100644 tests/core/layers/npu_torch/dcp_attention_utils_test.cpp create mode 100644 xllm/core/layers/npu_torch/dcp_attention_utils.cpp create mode 100644 xllm/core/layers/npu_torch/dcp_attention_utils.h diff --git a/cmake/cc_test.cmake b/cmake/cc_test.cmake index 2c80e3b0dc..bbee503738 100644 --- a/cmake/cc_test.cmake +++ b/cmake/cc_test.cmake @@ -11,6 +11,7 @@ include(CMakeParseArguments) # COPTS: List of private compile options # LINKOPTS: List of link options # ARGS: Command line arguments to test case +# NO_NPU_RUNTIME: Skip automatic NPU runtime setup for CPU-only tests # # Usage: # cc_library( @@ -39,7 +40,7 @@ function(cc_test) cmake_parse_arguments( CC_TEST # prefix - "" # options + "NO_NPU_RUNTIME" # options "NAME;ENVIRONMENT" # one value args "SRCS;COPTS;LINKOPTS;DEPS;INCLUDES;ARGS;DATA" # multi value args ${ARGN} @@ -112,7 +113,17 @@ function(cc_test) PRIVATE ${CC_TEST_LINKOPTS} ) - if(USE_NPU) + if(USE_NPU AND CC_TEST_NO_NPU_RUNTIME) + get_target_property(_CC_TEST_LINK_LIBRARIES + ${CC_TEST_NAME} LINK_LIBRARIES) + if(_CC_TEST_LINK_LIBRARIES) + list(REMOVE_ITEM _CC_TEST_LINK_LIBRARIES cust_opapi) + set_property(TARGET ${CC_TEST_NAME} + PROPERTY LINK_LIBRARIES "${_CC_TEST_LINK_LIBRARIES}") + endif() + endif() + + if(USE_NPU AND NOT CC_TEST_NO_NPU_RUNTIME) target_sources(${CC_TEST_NAME} PRIVATE "${PROJECT_SOURCE_DIR}/tests/npu_test_environment.cpp" ) diff --git a/tests/core/layers/npu_torch/CMakeLists.txt b/tests/core/layers/npu_torch/CMakeLists.txt index 6b1eb78eb3..93c304e020 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/tests/core/layers/npu_torch/CMakeLists.txt @@ -1,5 +1,22 @@ include(cc_test) +cc_test( + NO_NPU_RUNTIME + NAME + npu_dcp_attention_utils_test + SRCS + dcp_attention_utils_test.cpp + "${PROJECT_SOURCE_DIR}/xllm/core/layers/npu_torch/dcp_attention_utils.cpp" + DEPS + 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 npu_deepseek_v4_indexer_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..e9885feb50 --- /dev/null +++ b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp @@ -0,0 +1,130 @@ +/* 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})); +} + +} // namespace +} // namespace xllm::layer::test diff --git a/xllm/core/layers/npu_torch/CMakeLists.txt b/xllm/core/layers/npu_torch/CMakeLists.txt index 0b09ff9751..c7c18aabb9 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 @@ -38,6 +50,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 228b42857e..e86a5c1637 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -15,41 +15,16 @@ limitations under the License. #include "attention.h" -#include #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" namespace { -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 validate_dcp_decode_lengths(const std::vector& q_cu_seq_lens, const std::vector& global_kv_seq_lens, int64_t token_count) { @@ -94,35 +69,6 @@ void normalize_zero_dcp_partials( } } -torch::Tensor merge_dcp_partials(const torch::Tensor& all_partial_out, - const torch::Tensor& all_partial_lse) { - CHECK_EQ(all_partial_out.scalar_type(), torch::kFloat32); - CHECK_EQ(all_partial_lse.scalar_type(), torch::kFloat32); - 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 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(torch::isfinite(all_partial_lse), - torch::exp(all_partial_lse - safe_max_lse), - torch::zeros_like(all_partial_lse)); - 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 * all_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 namespace xllm { @@ -282,8 +228,9 @@ void AttentionImpl::dcp_decoder_forward( 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 = compute_dcp_local_kv_seq_lens( - global_kv_seq_lens, dcp_size_, dcp_rank_, block_size); + 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 torch::Tensor local_block_table = parallel_state::select_dcp_local_block_table( attn_metadata.block_table, dcp_size_, dcp_rank_); @@ -329,7 +276,7 @@ void AttentionImpl::dcp_decoder_forward( const torch::Tensor all_partial_lse = dcp_group_->allgather_base_sync(partial_lse); const torch::Tensor merged_out = - merge_dcp_partials(all_partial_out, all_partial_lse); + 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_); 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..89a9f1db7e --- /dev/null +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.cpp @@ -0,0 +1,84 @@ +/* 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 + +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; +} + +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..2dfcad670a --- /dev/null +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.h @@ -0,0 +1,34 @@ +/* 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 { + +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); + +torch::Tensor merge_dcp_partials(const torch::Tensor& all_partial_out, + const torch::Tensor& all_partial_lse); + +} // namespace xllm::layer::detail From 185f39d1f3603c160eacd3d412870ba7e248ef06 Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Wed, 5 Aug 2026 19:36:51 +0800 Subject: [PATCH 04/22] feat: enable prefix cache for DCP --- tests/core/distributed_runtime/dcp_compat_test.cpp | 7 +++---- xllm/core/distributed_runtime/dcp_compat.h | 5 ----- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index 6bfec89fec..e9e959ec48 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -65,13 +65,12 @@ TEST(DcpCompatTest, RejectsDefaultChunkedPrefillFirst) { "enable_chunked_prefill=false"); } -TEST(DcpCompatTest, RejectsPrefixCache) { +TEST(DcpCompatTest, AllowsPrefixCache) { Options options = dcp_options_with_supported_feature_flags(); options.enable_prefix_cache(true); - expect_error_contains( - validate_dcp_first_version_options(options, EngineType::LLM), - "enable_prefix_cache=false"); + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); } TEST(DcpCompatTest, RejectsScheduleOverlap) { diff --git a/xllm/core/distributed_runtime/dcp_compat.h b/xllm/core/distributed_runtime/dcp_compat.h index 6a8501cd5e..d23845d238 100644 --- a/xllm/core/distributed_runtime/dcp_compat.h +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -34,11 +34,6 @@ inline std::optional validate_dcp_first_version_options( "chunked prefill; set --enable_chunked_prefill=false or set " "--decode_context_parallel_size=1"; } - if (options.enable_prefix_cache()) { - return "decode_context_parallel_size first version does not yet support " - "prefix cache; set --enable_prefix_cache=false or set " - "--decode_context_parallel_size=1"; - } if (options.enable_schedule_overlap()) { return "decode_context_parallel_size first version does not yet support " "schedule overlap; set --enable_schedule_overlap=false or set " From def445056aac16b768dd43d999c4d87fcffb3a18 Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Thu, 6 Aug 2026 17:55:40 +0800 Subject: [PATCH 05/22] feat: support DCP-aware chunked prefill --- .../distributed_runtime/dcp_compat_test.cpp | 13 +- .../framework/parallel_state/CMakeLists.txt | 3 + .../parallel_state/cp_group_ranks_test.cpp | 190 +++++++++++++ .../kernels/npu/fia_decode_lse_probe_test.cpp | 257 +++++++++++++++++- .../layers/npu_torch/dcp_attention_test.cpp | 96 +++++++ .../npu_torch/dcp_attention_utils_test.cpp | 118 ++++++++ xllm/core/distributed_runtime/dcp_compat.h | 5 - xllm/core/layers/npu_torch/attention.cpp | 187 +++++++++++++ xllm/core/layers/npu_torch/attention.h | 9 + .../layers/npu_torch/dcp_attention_utils.cpp | 103 +++++++ .../layers/npu_torch/dcp_attention_utils.h | 30 ++ xllm/core/runtime/worker_impl.cpp | 10 +- 12 files changed, 1002 insertions(+), 19 deletions(-) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index e9e959ec48..29e6c332f3 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -26,7 +26,7 @@ namespace { Options dcp_options_with_supported_feature_flags() { Options options; options.decode_context_parallel_size(2) - .enable_chunked_prefill(false) + .enable_chunked_prefill(true) .enable_prefix_cache(false) .enable_schedule_overlap(false) .enable_disagg_pd(false) @@ -56,13 +56,12 @@ TEST(DcpCompatTest, AllowsSupportedFirstVersionFeatureFlags) { validate_dcp_first_version_options(options, EngineType::LLM).has_value()); } -TEST(DcpCompatTest, RejectsDefaultChunkedPrefillFirst) { - Options options; - options.decode_context_parallel_size(2); +TEST(DcpCompatTest, AllowsChunkedPrefill) { + 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_chunked_prefill=false"); + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); } TEST(DcpCompatTest, AllowsPrefixCache) { 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 700e718127..912ffeb858 100644 --- a/tests/core/framework/parallel_state/cp_group_ranks_test.cpp +++ b/tests/core/framework/parallel_state/cp_group_ranks_test.cpp @@ -15,10 +15,18 @@ limitations under the License. #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 { @@ -44,6 +52,45 @@ int32_t expected_dcp_rank(int32_t global_rank, 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; @@ -330,6 +377,149 @@ TEST(DcpCacheLayout, PrefillWritesMatchDecodeLocalBlockTable) { } } +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 #include +#include #include +#include #include #include "core/kernels/npu/npu_ops_api.h" @@ -65,12 +67,69 @@ void write_paged_kv_cache(torch::Tensor& key, float max_abs_diff(const torch::Tensor& expected, const torch::Tensor& actual) { const torch::Tensor expected_cpu = - expected.cpu().to(torch::kFloat32).view({-1}); - const torch::Tensor actual_cpu = actual.cpu().to(torch::kFloat32).view({-1}); + 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) { @@ -188,6 +247,200 @@ TEST_F(FiaDecodeLseProbe, DecodeFiaOutputMatchesBatchDecodeAndLseIsFinite) { 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 diff --git a/tests/core/layers/npu_torch/dcp_attention_test.cpp b/tests/core/layers/npu_torch/dcp_attention_test.cpp index 3596111bb6..5963c09723 100644 --- a/tests/core/layers/npu_torch/dcp_attention_test.cpp +++ b/tests/core/layers/npu_torch/dcp_attention_test.cpp @@ -19,9 +19,11 @@ limitations under the License. #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" @@ -148,5 +150,99 @@ TEST_F(DcpAttentionTest, ZeroLocalKvNormalizesBeforeMergeAndSlicesLocalHeads) { 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 index e9885feb50..910d84cfdd 100644 --- a/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp +++ b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp @@ -126,5 +126,123 @@ TEST(DcpAttentionUtilsTest, DistributesPartialTailAcrossFourRanks) { 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, 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/xllm/core/distributed_runtime/dcp_compat.h b/xllm/core/distributed_runtime/dcp_compat.h index d23845d238..e87813a175 100644 --- a/xllm/core/distributed_runtime/dcp_compat.h +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -29,11 +29,6 @@ inline std::optional validate_dcp_first_version_options( if (options.decode_context_parallel_size() <= 1) { return std::nullopt; } - if (options.enable_chunked_prefill()) { - return "decode_context_parallel_size first version does not yet support " - "chunked prefill; set --enable_chunked_prefill=false or set " - "--decode_context_parallel_size=1"; - } if (options.enable_schedule_overlap()) { return "decode_context_parallel_size first version does not yet support " "schedule overlap; set --enable_schedule_overlap=false or set " diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index e86a5c1637..c3a10d416f 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -15,6 +15,7 @@ limitations under the License. #include "attention.h" +#include #include #include @@ -69,6 +70,22 @@ void normalize_zero_dcp_partials( } } +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 { @@ -129,6 +146,12 @@ std::tuple> AttentionImpl::forward( if (attn_metadata.use_expanded_decode_for_spec_verify_attention) { 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 { @@ -283,6 +306,170 @@ void AttentionImpl::dcp_decoder_forward( 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: softmax over a key set partitioned into + // dcp_size context shards plus the current chunk is associative, so stacking + // all partials and merging once is exact. 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, diff --git a/xllm/core/layers/npu_torch/attention.h b/xllm/core/layers/npu_torch/attention.h index 9b5f9c30c8..6195cd72bf 100644 --- a/xllm/core/layers/npu_torch/attention.h +++ b/xllm/core/layers/npu_torch/attention.h @@ -70,6 +70,15 @@ class AttentionImpl : public torch::nn::Module { 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_; diff --git a/xllm/core/layers/npu_torch/dcp_attention_utils.cpp b/xllm/core/layers/npu_torch/dcp_attention_utils.cpp index 89a9f1db7e..44145ba159 100644 --- a/xllm/core/layers/npu_torch/dcp_attention_utils.cpp +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.cpp @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include namespace xllm::layer::detail { @@ -46,6 +47,108 @@ std::vector compute_dcp_local_kv_seq_lens( return local_kv_seq_lens; } +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 || diff --git a/xllm/core/layers/npu_torch/dcp_attention_utils.h b/xllm/core/layers/npu_torch/dcp_attention_utils.h index 2dfcad670a..95c60ec129 100644 --- a/xllm/core/layers/npu_torch/dcp_attention_utils.h +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.h @@ -22,12 +22,42 @@ limitations under the License. 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); +// 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); diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 7ad502a485..18f85a111b 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -913,11 +913,11 @@ void WorkerImpl::prepare_work_before_execute_on_stream( 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_decode() || - batch_forward_type.is_empty()) - << "DCP-1c supports only normal full prefill and decode cache " - "writes; chunked and mixed batches require DCP-2 layout " - "support."; + 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."; CHECK(!processed_input.input_params.enable_graph) From 1533da61abdcc692948777f04098540466cc6893 Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Mon, 10 Aug 2026 09:13:30 +0800 Subject: [PATCH 06/22] fix: close Phase 4 DCP validation gaps --- .../npu_torch/dcp_attention_utils_test.cpp | 18 ++++++++++++++++++ xllm/api_service/completion_service_impl.cpp | 9 +++------ 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp index 910d84cfdd..3481035373 100644 --- a/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp +++ b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp @@ -155,6 +155,24 @@ TEST(DcpAttentionUtilsTest, ContextShardLenReadsOnlyCachedContext) { 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, ValidateChunkedLengthsAcceptsMultiTokenRequests) { const std::vector normalized_q_cu_seq_lens = detail::validate_dcp_chunked_lengths( diff --git a/xllm/api_service/completion_service_impl.cpp b/xllm/api_service/completion_service_impl.cpp index c3b7d2395b..fedf847d1c 100644 --- a/xllm/api_service/completion_service_impl.cpp +++ b/xllm/api_service/completion_service_impl.cpp @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "api_service/utils.h" #include "common/instance_name.h" #include "completion.pb.h" #include "core/distributed_runtime/llm_master.h" @@ -105,9 +106,7 @@ bool send_delta_to_client_brpc(std::shared_ptr call, response.set_model(model); response.mutable_choices(); auto* proto_usage = response.mutable_usage(); - proto_usage->set_prompt_tokens(usage.num_prompt_tokens); - proto_usage->set_completion_tokens(usage.num_generated_tokens); - proto_usage->set_total_tokens(usage.num_total_tokens); + api_service::set_proto_usage(proto_usage, usage); if (!call->write(response)) { return false; } @@ -145,9 +144,7 @@ bool send_result_to_client_brpc(std::shared_ptr call, if (req_output.usage.has_value()) { const auto& usage = req_output.usage.value(); auto* proto_usage = response.mutable_usage(); - proto_usage->set_prompt_tokens(usage.num_prompt_tokens); - proto_usage->set_completion_tokens(usage.num_generated_tokens); - proto_usage->set_total_tokens(usage.num_total_tokens); + api_service::set_proto_usage(proto_usage, usage); } return call->write_and_finish(response); From c80271c662e7497089819c0fece60f5ff2d0ac38 Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Mon, 10 Aug 2026 10:39:47 +0800 Subject: [PATCH 07/22] fix: fail closed DCP chunked prefill --- tests/core/distributed_runtime/dcp_compat_test.cpp | 9 +++++---- xllm/core/distributed_runtime/dcp_compat.h | 5 +++++ xllm/core/layers/npu_torch/attention.cpp | 9 +++++---- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index 29e6c332f3..06350908cb 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -26,7 +26,7 @@ namespace { Options dcp_options_with_supported_feature_flags() { Options options; options.decode_context_parallel_size(2) - .enable_chunked_prefill(true) + .enable_chunked_prefill(false) .enable_prefix_cache(false) .enable_schedule_overlap(false) .enable_disagg_pd(false) @@ -56,12 +56,13 @@ TEST(DcpCompatTest, AllowsSupportedFirstVersionFeatureFlags) { validate_dcp_first_version_options(options, EngineType::LLM).has_value()); } -TEST(DcpCompatTest, AllowsChunkedPrefill) { +TEST(DcpCompatTest, RejectsChunkedPrefill) { Options options = dcp_options_with_supported_feature_flags(); options.enable_chunked_prefill(true); - EXPECT_FALSE( - validate_dcp_first_version_options(options, EngineType::LLM).has_value()); + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_chunked_prefill=false"); } TEST(DcpCompatTest, AllowsPrefixCache) { diff --git a/xllm/core/distributed_runtime/dcp_compat.h b/xllm/core/distributed_runtime/dcp_compat.h index e87813a175..d23845d238 100644 --- a/xllm/core/distributed_runtime/dcp_compat.h +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -29,6 +29,11 @@ inline std::optional validate_dcp_first_version_options( if (options.decode_context_parallel_size() <= 1) { return std::nullopt; } + if (options.enable_chunked_prefill()) { + return "decode_context_parallel_size first version does not yet support " + "chunked prefill; set --enable_chunked_prefill=false or set " + "--decode_context_parallel_size=1"; + } if (options.enable_schedule_overlap()) { return "decode_context_parallel_size first version does not yet support " "schedule overlap; set --enable_schedule_overlap=false or set " diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index c3a10d416f..e657b041f5 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -443,10 +443,11 @@ void AttentionImpl::dcp_chunked_prefill_forward( 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: softmax over a key set partitioned into - // dcp_size context shards plus the current chunk is associative, so stacking - // all partials and merging once is exact. The current part is identical on - // every rank, so contributing it once (from this rank) is correct. + // single online-softmax reduction. The formula is equivalent to softmax over + // the complete key set before kernel output quantization, but FIA partial + // outputs are low precision and are not guaranteed to be bitwise identical + // to one monolithic FIA call. 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 = From 7bac0b3523914a3397b4adbb46f40fc90bd3b3e0 Mon Sep 17 00:00:00 2001 From: Super User Date: Mon, 10 Aug 2026 17:01:31 +0800 Subject: [PATCH 08/22] fix: close mixed batch when decode context parallel size > 1 DCP reuses TP cards, so cp_size can be 1 while decode_context_parallel_size is > 1. resolve_batch_mode only closed mixed batching for cp_size > 1, leaving DCP + chunked prefill to build a MIXED batch that trips the worker CHECK and FATALs the process under concurrency. Thread decode_context_parallel_size into ContinuousScheduler::Options (default 1) from every scheduler construction site and force enable_mix_batch=false when it is > 1. The worker CHECK is kept as a second layer of defense. Co-Authored-By: Claude --- .../scheduler/continuous_scheduler_test.cpp | 79 +++++++++++++++++++ xllm/core/distributed_runtime/llm_master.cpp | 1 + xllm/core/distributed_runtime/rec_master.cpp | 1 + xllm/core/distributed_runtime/vlm_master.cpp | 1 + xllm/core/scheduler/continuous_scheduler.cpp | 7 +- xllm/core/scheduler/continuous_scheduler.h | 4 + 6 files changed, 91 insertions(+), 2 deletions(-) diff --git a/tests/core/scheduler/continuous_scheduler_test.cpp b/tests/core/scheduler/continuous_scheduler_test.cpp index 2f53b9c5b6..7b1611e510 100644 --- a/tests/core/scheduler/continuous_scheduler_test.cpp +++ b/tests/core/scheduler/continuous_scheduler_test.cpp @@ -387,6 +387,85 @@ 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(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/xllm/core/distributed_runtime/llm_master.cpp b/xllm/core/distributed_runtime/llm_master.cpp index c2bd27a3cc..e228a5ac99 100644 --- a/xllm/core/distributed_runtime/llm_master.cpp +++ b/xllm/core/distributed_runtime/llm_master.cpp @@ -84,6 +84,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()) diff --git a/xllm/core/distributed_runtime/rec_master.cpp b/xllm/core/distributed_runtime/rec_master.cpp index cdefe450e7..de17695083 100644 --- a/xllm/core/distributed_runtime/rec_master.cpp +++ b/xllm/core/distributed_runtime/rec_master.cpp @@ -545,6 +545,7 @@ RecMaster::RecMaster(const Options& options) options_.max_tokens_per_chunk_for_prefill()) .num_speculative_tokens(options_.num_speculative_tokens()) .dp_size(options_.dp_size()) + .decode_context_parallel_size(options_.decode_context_parallel_size()) .enable_disagg_pd(options_.enable_disagg_pd()) .enable_schedule_overlap(options_.enable_schedule_overlap()) .enable_chunked_prefill(options_.enable_chunked_prefill()) diff --git a/xllm/core/distributed_runtime/vlm_master.cpp b/xllm/core/distributed_runtime/vlm_master.cpp index 8ce050c944..253a8fcc2e 100644 --- a/xllm/core/distributed_runtime/vlm_master.cpp +++ b/xllm/core/distributed_runtime/vlm_master.cpp @@ -86,6 +86,7 @@ VLMMaster::VLMMaster(const Options& options) .max_tokens_per_chunk_for_prefill( options.max_tokens_per_chunk_for_prefill()) .dp_size(options_.dp_size()) + .decode_context_parallel_size(options_.decode_context_parallel_size()) .enable_disagg_pd(options_.enable_disagg_pd()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .instance_name(options_.instance_name()) diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp index 03b6ce1e1f..984188fff5 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -72,8 +72,11 @@ BatchMode resolve_batch_mode(const ContinuousScheduler::Options& options) { 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; } diff --git a/xllm/core/scheduler/continuous_scheduler.h b/xllm/core/scheduler/continuous_scheduler.h index 91f374532f..e080a17847 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; From 0de55415a34867748dcac1ec25c4ade55b6403bd Mon Sep 17 00:00:00 2001 From: Super User Date: Mon, 10 Aug 2026 18:11:30 +0800 Subject: [PATCH 09/22] feat: gate DCP chunked prefill behind an experimental opt-in Add --enable_experimental_dcp_chunked_prefill (default false), wired through ParallelConfig (flags, JSON, config dump, option category), Options, and xllm.cpp. When decode_context_parallel_size > 1 with chunked prefill enabled, startup validation now rejects unless the experimental flag is set, and the master logs a one-time warning noting the non-bitwise-equivalence, automatic mixed-batch closure, and rollback. The experimental opt-in only lifts the chunked prefill rejection; it does not bypass the schedule-overlap, disaggregated-PD, speculative, or MoE model-type rejections. Rework dcp_compat_test into the full experimental matrix and correct the attention merge comment to describe the partial-output quantization as a few 1e-3 (measured max 0.00294) rather than ~1e-3. Co-Authored-By: Claude --- .../distributed_runtime/dcp_compat_test.cpp | 72 ++++++++++++++++++- .../framework/config/config_json_test.cpp | 31 ++++++++ xllm/core/common/options.h | 4 ++ xllm/core/distributed_runtime/dcp_compat.h | 9 ++- xllm/core/distributed_runtime/master.cpp | 11 +++ .../core/framework/config/parallel_config.cpp | 11 +++ xllm/core/framework/config/parallel_config.h | 9 +++ xllm/core/layers/npu_torch/attention.cpp | 11 +-- xllm/xllm.cpp | 2 + 9 files changed, 150 insertions(+), 10 deletions(-) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index 06350908cb..1152ba8bfb 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -27,6 +27,7 @@ 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) @@ -35,6 +36,13 @@ Options dcp_options_with_supported_feature_flags() { 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()); @@ -56,13 +64,36 @@ TEST(DcpCompatTest, AllowsSupportedFirstVersionFeatureFlags) { validate_dcp_first_version_options(options, EngineType::LLM).has_value()); } -TEST(DcpCompatTest, RejectsChunkedPrefill) { +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_chunked_prefill=false"); + "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) { @@ -126,6 +157,35 @@ TEST(DcpCompatTest, RejectsSpeculativeTokens) { "num_speculative_tokens=0"); } +// Enabling the experimental chunked prefill opt-in must not bypass the other +// first-version rejections: schedule overlap, disaggregated PD, and +// speculative decoding are still unsupported even under the experimental path. +TEST(DcpCompatTest, ExperimentalChunkedPrefillDoesNotBypassScheduleOverlap) { + Options options = dcp_options_with_experimental_chunked_prefill(); + options.enable_schedule_overlap(true); + + expect_error_contains( + validate_dcp_first_version_options(options, EngineType::LLM), + "enable_schedule_overlap=false"); +} + +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"); +} + TEST(DcpCompatTest, AllowsDenseQwen35ModelType) { EXPECT_FALSE( validate_dcp_first_version_model_type("qwen3_5_text").has_value()); @@ -136,5 +196,13 @@ TEST(DcpCompatTest, RejectsUnvalidatedQwen35MoeModelType) { validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); } +// Model-type rejection is independent of the experimental flag: MoE stays +// unsupported. The experimental opt-in only gates the options-level chunked +// prefill path, not the model-type gate. +TEST(DcpCompatTest, ExperimentalChunkedPrefillDoesNotBypassMoeModelType) { + expect_error_contains( + validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); +} + } // namespace } // namespace xllm diff --git a/tests/core/framework/config/config_json_test.cpp b/tests/core/framework/config/config_json_test.cpp index 14adc1c49c..0e42aac46c 100644 --- a/tests/core/framework/config/config_json_test.cpp +++ b/tests/core/framework/config/config_json_test.cpp @@ -296,12 +296,43 @@ TEST(ConfigJsonTest, RegistersContextParallelCommandLineOptions) { 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/xllm/core/common/options.h b/xllm/core/common/options.h index c73ccaf5a1..5e984ca9b9 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -142,6 +142,10 @@ class Options { 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 index d23845d238..61ec926ccb 100644 --- a/xllm/core/distributed_runtime/dcp_compat.h +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -29,9 +29,12 @@ inline std::optional validate_dcp_first_version_options( if (options.decode_context_parallel_size() <= 1) { return std::nullopt; } - if (options.enable_chunked_prefill()) { - return "decode_context_parallel_size first version does not yet support " - "chunked prefill; set --enable_chunked_prefill=false or set " + if (options.enable_chunked_prefill() && + !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; set " + "--enable_experimental_dcp_chunked_prefill=true to opt in, or set " + "--enable_chunked_prefill=false, or set " "--decode_context_parallel_size=1"; } if (options.enable_schedule_overlap()) { diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index b58042d624..a293d4975c 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -458,6 +458,17 @@ Master::Master(const Options& options, EngineType type) 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 && + options_.enable_chunked_prefill() && + 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. Set " + "--enable_experimental_dcp_chunked_prefill=false to roll back."; + } const std::string cp_model_type = dcp_model_config.has_value() ? dcp_model_config->model_type : ""; const std::optional cp_error = diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index 24e5ef3c92..435a7e8452 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -32,6 +32,13 @@ DEFINE_int32(decode_context_parallel_size, "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(kv_split_size, 1, "KV-cache split width. 0 falls back to cp_size (legacy); 1 means " @@ -82,6 +89,7 @@ void ParallelConfig::from_flags() { 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(kv_split_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(tp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(sp_size); @@ -99,6 +107,7 @@ void ParallelConfig::from_json(const JsonReader& json) { 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(tp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(cfg_size); @@ -118,6 +127,8 @@ void ParallelConfig::append_config_json( 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, tp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, sp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 08717a3b11..5eb67e4a5a 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/xllm/core/framework/config/parallel_config.h @@ -45,6 +45,7 @@ class ParallelConfig final { "ep_size", "cp_size", "decode_context_parallel_size", + "enable_experimental_dcp_chunked_prefill", "tp_size", "sp_size", "cfg_size", @@ -65,6 +66,14 @@ class ParallelConfig final { 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; + // 0 means follow cp_size (legacy KV-split width). PROPERTY(int32_t, kv_split_size) = 1; diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index e657b041f5..c438973db8 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -443,11 +443,12 @@ void AttentionImpl::dcp_chunked_prefill_forward( 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 equivalent to softmax over - // the complete key set before kernel output quantization, but FIA partial - // outputs are low precision and are not guaranteed to be bitwise identical - // to one monolithic FIA call. The current part is identical on every rank, so - // contributing it once (from this rank) is correct. + // 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 = diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 712deff676..c8f0a7ea15 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -181,6 +181,8 @@ Options create_options(const std::string& instance_name, bool is_local) { .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())) From d7f86fe9137003ae0913f98f494538cc25d31ef2 Mon Sep 17 00:00:00 2001 From: Super User Date: Mon, 10 Aug 2026 19:58:55 +0800 Subject: [PATCH 10/22] fix: gate DCP chunked prefill on effective, not raw, chunked flag priority_strategy=multi_slo_and_prio forces chunked prefill in the scheduler regardless of the raw enable_chunked_prefill flag. The DCP experimental gate and startup warning only checked the raw flag, so a service with decode_context_parallel_size>1, enable_chunked_prefill=false and multi_slo_and_prio silently ran the experimental DCP chunked path without the opt-in and without the warning, violating fail-closed. Add resolve_effective_chunked_prefill (new lightweight header scheduler/chunked_prefill_policy.h) and share it across the scheduler batch-mode resolution, the DCP compatibility gate, and the startup warning so the effective chunked semantics cannot drift. Correct the gate error and warning to name the multi_slo_and_prio implicit enable and the full rollback conditions. Extend dcp_compat and scheduler tests with the multi-SLO and no-op experimental-flag cases. Co-Authored-By: Claude --- .../distributed_runtime/dcp_compat_test.cpp | 49 +++++++++++++++++-- .../scheduler/continuous_scheduler_test.cpp | 15 ++++++ xllm/core/distributed_runtime/dcp_compat.h | 14 ++++-- xllm/core/distributed_runtime/master.cpp | 10 ++-- xllm/core/scheduler/chunked_prefill_policy.h | 33 +++++++++++++ xllm/core/scheduler/continuous_scheduler.cpp | 9 ++-- 6 files changed, 112 insertions(+), 18 deletions(-) create mode 100644 xllm/core/scheduler/chunked_prefill_policy.h diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index 1152ba8bfb..ed8710ce82 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -57,6 +57,44 @@ TEST(DcpCompatTest, DcpOneDoesNotRejectDefaultOptions) { 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(); @@ -196,10 +234,13 @@ TEST(DcpCompatTest, RejectsUnvalidatedQwen35MoeModelType) { validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); } -// Model-type rejection is independent of the experimental flag: MoE stays -// unsupported. The experimental opt-in only gates the options-level chunked -// prefill path, not the model-type gate. -TEST(DcpCompatTest, ExperimentalChunkedPrefillDoesNotBypassMoeModelType) { +// Model-type rejection lives in a separate flag-agnostic validator. The +// experimental opt-in only gates the options-level chunked prefill path, so it +// cannot bypass the MoE rejection: at startup master.cpp calls the options +// validator and this model-type validator independently, and the latter has no +// flag input to suppress. This test pins that the model-type validator rejects +// MoE regardless of any option flags. +TEST(DcpCompatTest, ModelTypeValidatorRejectsMoeRegardlessOfFlags) { expect_error_contains( validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); } diff --git a/tests/core/scheduler/continuous_scheduler_test.cpp b/tests/core/scheduler/continuous_scheduler_test.cpp index 7b1611e510..d39fdfc229 100644 --- a/tests/core/scheduler/continuous_scheduler_test.cpp +++ b/tests/core/scheduler/continuous_scheduler_test.cpp @@ -422,6 +422,21 @@ TEST(ContinuousSchedulerBatchModeTest, Dcp2ClosesMixBatchUnderMultiSlo) { 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(ContinuousSchedulerFactoryTest, ChunkedPrefillWithDcpDoesNotBuildMixedBatch) { ContinuousScheduler::Options opt = create_scheduler_options(8, 8, 0, 4, 1); diff --git a/xllm/core/distributed_runtime/dcp_compat.h b/xllm/core/distributed_runtime/dcp_compat.h index 61ec926ccb..2ffc842bcd 100644 --- a/xllm/core/distributed_runtime/dcp_compat.h +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -20,6 +20,7 @@ limitations under the License. #include "common/options.h" #include "common/types.h" +#include "scheduler/chunked_prefill_policy.h" namespace xllm { @@ -29,13 +30,16 @@ inline std::optional validate_dcp_first_version_options( if (options.decode_context_parallel_size() <= 1) { return std::nullopt; } - if (options.enable_chunked_prefill() && + 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; set " - "--enable_experimental_dcp_chunked_prefill=true to opt in, or set " - "--enable_chunked_prefill=false, or set " - "--decode_context_parallel_size=1"; + "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_schedule_overlap()) { return "decode_context_parallel_size first version does not yet support " diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index a293d4975c..35b4ebc2d4 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -459,15 +459,19 @@ Master::Master(const Options& options, EngineType type) 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 && - options_.enable_chunked_prefill() && + 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. Set " - "--enable_experimental_dcp_chunked_prefill=false to roll back."; + "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."; } const std::string cp_model_type = dcp_model_config.has_value() ? dcp_model_config->model_type : ""; 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 984188fff5..441e90ac0a 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -41,6 +41,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" @@ -63,15 +64,11 @@ 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/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. From 9f144b1aba2dd1ad4878ebaebc152df329344543 Mon Sep 17 00:00:00 2001 From: Super User Date: Mon, 10 Aug 2026 20:36:40 +0800 Subject: [PATCH 11/22] fix: use effective chunked prefill in whole-service DCP consumers The DCP experimental gate now treats multi_slo_and_prio as implicitly enabling chunked prefill, but several downstream consumers still read the raw enable_chunked_prefill flag, so a gate-approved multi_slo_and_prio + raw-false config could still fail or misbehave at service level: - llm_engine.cpp linear-attention prefix-cache precondition CHECK would reject the config even though the gate allowed it (Qwen3.5 GDN + prefix on the experimental path); - llm_master.cpp prompt admission would clamp long prompts as if chunked prefill were off; - continuous_scheduler.cpp activation metric mislabeled the batch. Route all three through resolve_effective_chunked_prefill (scheduler reads the already-resolved batch_mode_), add master.cpp a direct include of the helper header, add a truth-table test for the predicate, and drop the duplicate MoE model-type test. Co-Authored-By: Claude --- tests/core/distributed_runtime/dcp_compat_test.cpp | 11 ----------- tests/core/scheduler/continuous_scheduler_test.cpp | 13 +++++++++++++ xllm/core/distributed_runtime/llm_engine.cpp | 11 +++++++++-- xllm/core/distributed_runtime/llm_master.cpp | 4 +++- xllm/core/distributed_runtime/master.cpp | 1 + xllm/core/scheduler/continuous_scheduler.cpp | 2 +- 6 files changed, 27 insertions(+), 15 deletions(-) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index ed8710ce82..e62f217c7f 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -234,16 +234,5 @@ TEST(DcpCompatTest, RejectsUnvalidatedQwen35MoeModelType) { validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); } -// Model-type rejection lives in a separate flag-agnostic validator. The -// experimental opt-in only gates the options-level chunked prefill path, so it -// cannot bypass the MoE rejection: at startup master.cpp calls the options -// validator and this model-type validator independently, and the latter has no -// flag input to suppress. This test pins that the model-type validator rejects -// MoE regardless of any option flags. -TEST(DcpCompatTest, ModelTypeValidatorRejectsMoeRegardlessOfFlags) { - expect_error_contains( - validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); -} - } // namespace } // namespace xllm diff --git a/tests/core/scheduler/continuous_scheduler_test.cpp b/tests/core/scheduler/continuous_scheduler_test.cpp index d39fdfc229..ae5128533e 100644 --- a/tests/core/scheduler/continuous_scheduler_test.cpp +++ b/tests/core/scheduler/continuous_scheduler_test.cpp @@ -10,6 +10,7 @@ #include "core/framework/config/kv_cache_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" @@ -437,6 +438,18 @@ TEST(ContinuousSchedulerBatchModeTest, MultiSloForcesEffectiveChunkedPrefill) { 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); diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 3ddc5ad033..7e3da084b5 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -54,6 +54,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" @@ -548,10 +549,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 e228a5ac99..b828832d83 100644 --- a/xllm/core/distributed_runtime/llm_master.cpp +++ b/xllm/core/distributed_runtime/llm_master.cpp @@ -34,6 +34,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" @@ -340,7 +341,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 35b4ebc2d4..78d246c69e 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -41,6 +41,7 @@ limitations under the License. #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" diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp index 441e90ac0a..63c606ca24 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -601,7 +601,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); From b319f2c9a1641e0104ee0fd593894b881a0cf748 Mon Sep 17 00:00:00 2001 From: Super User Date: Mon, 10 Aug 2026 20:46:12 +0800 Subject: [PATCH 12/22] feat: allow Qwen3.5 MoE with decode context parallelism Remove the MoE model-type rejection from the DCP first-version gate so Qwen3.5 MoE can run with decode_context_parallel_size>1. The GQA topology checks (head divisibility) still apply and are independent of the expert layers. MoE on this path is not bitwise-equivalent to decode_context_parallel_size=1: the DCP communication/merge floating point order can produce a measurable per-step logprob difference on the MoE-sensitive path (observed ~0.044 at one teacher-forced step on 35B). This is accepted for now as a functional-first capability; the kernel path fix is future work. The validator became an unconditional no-op once the branch was gone, so drop the function, its call site, and its tests entirely rather than leave a dead shim. Co-Authored-By: Claude --- tests/core/distributed_runtime/dcp_compat_test.cpp | 10 ---------- xllm/core/distributed_runtime/dcp_compat.h | 10 ---------- xllm/core/distributed_runtime/master.cpp | 4 ---- 3 files changed, 24 deletions(-) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index e62f217c7f..3a406e04a8 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -224,15 +224,5 @@ TEST(DcpCompatTest, ExperimentalChunkedPrefillDoesNotBypassSpeculative) { "speculative decoding"); } -TEST(DcpCompatTest, AllowsDenseQwen35ModelType) { - EXPECT_FALSE( - validate_dcp_first_version_model_type("qwen3_5_text").has_value()); -} - -TEST(DcpCompatTest, RejectsUnvalidatedQwen35MoeModelType) { - expect_error_contains( - validate_dcp_first_version_model_type("qwen3_5_moe_text"), "MoE"); -} - } // namespace } // namespace xllm diff --git a/xllm/core/distributed_runtime/dcp_compat.h b/xllm/core/distributed_runtime/dcp_compat.h index 2ffc842bcd..b4b7671c82 100644 --- a/xllm/core/distributed_runtime/dcp_compat.h +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -63,14 +63,4 @@ inline std::optional validate_dcp_first_version_options( return std::nullopt; } -inline std::optional validate_dcp_first_version_model_type( - const std::string& model_type) { - if (model_type == "qwen3_5_moe_text") { - return "decode_context_parallel_size first version does not yet support " - "Qwen3.5 MoE; use dense Qwen3.5 or set " - "--decode_context_parallel_size=1"; - } - return std::nullopt; -} - } // namespace xllm diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 78d246c69e..be9580e210 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -161,10 +161,6 @@ std::optional validate_model_dcp( "text models, got model_type=" + model_config->model_type; } - if (std::optional dcp_model_error = - validate_dcp_first_version_model_type(model_config->model_type)) { - return dcp_model_error; - } if (options.dp_size() < 1) { return "decode context parallelism requires dp_size >= 1"; } From 05a81515642f7b40162bee06d28fe31ef6882636 Mon Sep 17 00:00:00 2001 From: Super User Date: Mon, 10 Aug 2026 21:37:15 +0800 Subject: [PATCH 13/22] feat: warn on Qwen3.5 MoE with decode context parallelism MoE + decode_context_parallel_size>1 is allowed but not bitwise-equivalent to dcp=1: the DCP decode attention kernel path (FIA) differs from dcp=1 (batch_decode), and MoE expert routing amplifies it into a per-step logprob delta on the order of 1e-2 (roughly an order of magnitude larger than the chunked-prefill BF16 quantization diff). Log a one-time startup warning so the non-equivalence is visible to operators. The unified kernel path fix is future work. Co-Authored-By: Claude --- xllm/core/distributed_runtime/master.cpp | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index be9580e210..2a2a8a01e9 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -470,6 +470,18 @@ Master::Master(const Options& options, EngineType type) "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::string cp_model_type = dcp_model_config.has_value() ? dcp_model_config->model_type : ""; const std::optional cp_error = From 04aee08c3e820749225cbf7e201191cd04ef9cc4 Mon Sep 17 00:00:00 2001 From: Super User Date: Thu, 13 Aug 2026 10:41:07 +0800 Subject: [PATCH 14/22] feat: allow schedule overlap with decode context parallelism Eager schedule overlap is compatible with DCP: per-step DCP metadata is recomputed each forward and DCP collectives allocate fresh buffers per call, and the scheduler/worker FIFO + compute-stream ordering prevents step i+1 from clobbering step i state. Remove the first-version schedule-overlap rejection in validate_dcp_first_version_options (graph path stays blocked independently by the attention paged_attention_tiling_data CHECK). On-card: overlap dcp2 output is token-exact vs non-overlap dcp2 (5 dense cases, tp4/dcp2). Co-Authored-By: Claude --- .../distributed_runtime/dcp_compat_test.cpp | 21 ++++++------------- xllm/core/distributed_runtime/dcp_compat.h | 5 ----- 2 files changed, 6 insertions(+), 20 deletions(-) diff --git a/tests/core/distributed_runtime/dcp_compat_test.cpp b/tests/core/distributed_runtime/dcp_compat_test.cpp index 3a406e04a8..8cf517cfca 100644 --- a/tests/core/distributed_runtime/dcp_compat_test.cpp +++ b/tests/core/distributed_runtime/dcp_compat_test.cpp @@ -142,13 +142,12 @@ TEST(DcpCompatTest, AllowsPrefixCache) { validate_dcp_first_version_options(options, EngineType::LLM).has_value()); } -TEST(DcpCompatTest, RejectsScheduleOverlap) { +TEST(DcpCompatTest, AllowsScheduleOverlap) { Options options = dcp_options_with_supported_feature_flags(); options.enable_schedule_overlap(true); - expect_error_contains( - validate_dcp_first_version_options(options, EngineType::LLM), - "enable_schedule_overlap=false"); + EXPECT_FALSE( + validate_dcp_first_version_options(options, EngineType::LLM).has_value()); } TEST(DcpCompatTest, RejectsDisaggregatedPrefillDecodeFlag) { @@ -196,17 +195,9 @@ TEST(DcpCompatTest, RejectsSpeculativeTokens) { } // Enabling the experimental chunked prefill opt-in must not bypass the other -// first-version rejections: schedule overlap, disaggregated PD, and -// speculative decoding are still unsupported even under the experimental path. -TEST(DcpCompatTest, ExperimentalChunkedPrefillDoesNotBypassScheduleOverlap) { - Options options = dcp_options_with_experimental_chunked_prefill(); - options.enable_schedule_overlap(true); - - expect_error_contains( - validate_dcp_first_version_options(options, EngineType::LLM), - "enable_schedule_overlap=false"); -} - +// 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); diff --git a/xllm/core/distributed_runtime/dcp_compat.h b/xllm/core/distributed_runtime/dcp_compat.h index b4b7671c82..c1bb652b5d 100644 --- a/xllm/core/distributed_runtime/dcp_compat.h +++ b/xllm/core/distributed_runtime/dcp_compat.h @@ -41,11 +41,6 @@ inline std::optional validate_dcp_first_version_options( "(--enable_chunked_prefill=false and a non-multi_slo_and_prio " "priority_strategy), or set --decode_context_parallel_size=1"; } - if (options.enable_schedule_overlap()) { - return "decode_context_parallel_size first version does not yet support " - "schedule overlap; set --enable_schedule_overlap=false 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 " From 3844bd549a4c85883e41aa4e7b23f6c1951ff52b Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Mon, 17 Aug 2026 17:11:07 +0800 Subject: [PATCH 15/22] feat: support decode context parallelism with ACL graph --- .../npu_torch/dcp_attention_utils_test.cpp | 30 ++ .../core/runtime/acl_graph_executor_test.cpp | 177 ++++++++++++ xllm/core/distributed_runtime/master.cpp | 4 - .../core/framework/model/model_input_params.h | 3 + .../kernels/npu/aclnn/pytorch_npu_helper.hpp | 124 +++++++++ .../kernels/npu/npu_fused_infer_attention.cpp | 259 ++++++++++++------ xllm/core/kernels/npu/npu_ops_api.h | 44 +++ xllm/core/layers/common/attention_metadata.h | 14 + .../common/attention_metadata_builder.cpp | 3 + xllm/core/layers/npu_torch/attention.cpp | 156 +++++++++-- .../layers/npu_torch/dcp_attention_utils.cpp | 19 ++ .../layers/npu_torch/dcp_attention_utils.h | 9 + .../npu/acl_graph_task_update_context.h | 35 +++ xllm/core/runtime/acl_graph_executor_impl.cpp | 201 +++++++++++--- xllm/core/runtime/acl_graph_executor_impl.h | 6 + .../runtime/acl_graph_persistent_param.cpp | 64 +++++ .../core/runtime/acl_graph_persistent_param.h | 14 + xllm/core/runtime/worker_impl.cpp | 2 - 18 files changed, 1003 insertions(+), 161 deletions(-) diff --git a/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp index 3481035373..8abf161d04 100644 --- a/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp +++ b/tests/core/layers/npu_torch/dcp_attention_utils_test.cpp @@ -173,6 +173,36 @@ TEST(DcpAttentionUtilsTest, CoversPhase4LocalContextContracts) { 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( diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index ec3d828059..2fa6ddd1b4 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -44,6 +44,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" @@ -112,6 +113,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) { @@ -1005,6 +1022,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 = @@ -1162,6 +1239,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/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 2a2a8a01e9..7b912ad8c3 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -142,10 +142,6 @@ std::optional validate_model_dcp( if (options.npu_kernel_backend() != "TORCH") { return "decode_context_parallel_size requires --npu_kernel_backend=TORCH"; } - if (options.enable_graph()) { - return "decode_context_parallel_size does not support graph capture yet; " - "disable graph or set decode_context_parallel_size=1"; - } if (engine_type != EngineType::LLM && engine_type != EngineType::SSM) { return "decode context parallelism supports only LLM text generation"; } diff --git a/xllm/core/framework/model/model_input_params.h b/xllm/core/framework/model/model_input_params.h index ff7b587734..362633dffe 100644 --- a/xllm/core/framework/model/model_input_params.h +++ b/xllm/core/framework/model/model_input_params.h @@ -933,6 +933,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; @@ -963,6 +964,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/kernels/npu/aclnn/pytorch_npu_helper.hpp b/xllm/core/kernels/npu/aclnn/pytorch_npu_helper.hpp index ce23958d54..e609ce1002 100644 --- a/xllm/core/kernels/npu/aclnn/pytorch_npu_helper.hpp +++ b/xllm/core/kernels/npu/aclnn/pytorch_npu_helper.hpp @@ -770,3 +770,127 @@ using ReleaseHugeMemFn = void (*)(void*, bool); uninit_mem_func(nullptr, false); \ } \ } while (false) + +// Query only the workspace size (bytes) required by an aclnn op, without +// executing it. Writes the size into `out_size_uint64`. Used to pre-size a +// caller-owned workspace buffer for ACLGraph capture, where the workspace +// address must stay stable across replays (see EXEC_NPU_CMD_WITH_WORKSPACE). +#define EXEC_NPU_CMD_GET_WORKSPACE_SIZE(aclnn_api, out_size_uint64, ...) \ + do { \ + static const auto get_workspace_size_func_addr = \ + ::xllm::kernel::npu::aclnn::detail::get_op_api_func_addr( \ + #aclnn_api "GetWorkspaceSize"); \ + CHECK(get_workspace_size_func_addr != nullptr) \ + << #aclnn_api "GetWorkspaceSize" << " not in " \ + << ::xllm::kernel::npu::aclnn::detail::get_op_api_lib_name(); \ + uint64_t workspace_size = 0; \ + uint64_t* workspace_size_addr = &workspace_size; \ + ::aclOpExecutor* executor = nullptr; \ + ::aclOpExecutor** executor_addr = &executor; \ + 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 "GetWorkspaceSize failed, detail:" \ + << aclGetRecentErrMsg(); \ + ::xllm::kernel::npu::aclnn::detail::release_convert_types( \ + converted_params); \ + (out_size_uint64) = workspace_size; \ + } while (false) + +// Same as EXEC_NPU_CMD but uses a caller-owned `ws_tensor` (kByte, on device) +// as the op workspace instead of allocating a function-local one. Required for +// ACLGraph capture/replay, where the workspace address must stay stable across +// replays. `ws_tensor` must have capacity >= 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 0ac984e00b..7733f1036e 100644 --- a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp +++ b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp @@ -151,7 +151,75 @@ 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) \ + 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 = 0; \ + 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, @@ -165,7 +233,10 @@ std::tuple npu_fused_infer_attention( int64_t block_size, int64_t sparse_mode, const std::string& input_layout, - bool softmax_lse_flag) { + bool softmax_lse_flag, + 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"); @@ -173,97 +244,111 @@ 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 = 0; - 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); + + 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, + 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); + + 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) { + 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, + 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 b9a40f2c4f..7401b95dab 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -67,6 +67,50 @@ std::tuple npu_fused_infer_attention( const std::string& input_layout, bool softmax_lse_flag = false); +// 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, + 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, + 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 21a33ba2b1..2d3b818997 100644 --- a/xllm/core/layers/common/attention_metadata.h +++ b/xllm/core/layers/common/attention_metadata.h @@ -30,6 +30,10 @@ namespace ffi = tvm::ffi; #include "dsa_metadata.h" +namespace xllm::npu { +class AclGraphTaskUpdateContext; +} // namespace xllm::npu + namespace xllm::layer { #if defined(USE_CUDA) || defined(USE_MUSA) @@ -179,6 +183,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 ee47f881b3..d513dc33bb 100644 --- a/xllm/core/layers/common/attention_metadata_builder.cpp +++ b/xllm/core/layers/common/attention_metadata_builder.cpp @@ -104,6 +104,9 @@ AttentionMetadata build_attention_metadata( #if defined(USE_NPU) attn_metadata.is_spec_verify = params.is_spec_verify; + 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; attn_metadata.use_expanded_decode_for_spec_verify_attention = params.graph.use_expanded_decode_for_spec_verify_attention; if (attn_metadata.use_expanded_decode_for_spec_verify_attention) { diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index c438973db8..f412320956 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -23,6 +23,7 @@ limitations under the License. #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 { @@ -236,8 +237,6 @@ void AttentionImpl::dcp_decoder_forward( << "DCP-2 does not support speculative decode attention."; CHECK(!attn_metadata.use_expanded_decode_for_spec_verify_attention) << "DCP-2 does not support speculative decode attention."; - CHECK(!attn_metadata.paged_attention_tiling_data.defined()) - << "DCP-2 does not support graph-captured decode attention."; CHECK(v_cache.has_value() && v_cache.value().defined()) << "DCP decode requires a defined V cache."; CHECK(attn_metadata.block_table.defined()) @@ -254,9 +253,18 @@ void AttentionImpl::dcp_decoder_forward( 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 torch::Tensor local_block_table = - parallel_state::select_dcp_local_block_table( - attn_metadata.block_table, dcp_size_, dcp_rank_); + 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."; @@ -275,24 +283,125 @@ void AttentionImpl::dcp_decoder_forward( {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; - 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, - 0, - "TND", - true); - torch::Tensor partial_out = std::get<0>(fia_result).to(torch::kFloat32); - torch::Tensor partial_lse = std::get<1>(fia_result).to(torch::kFloat32); - normalize_zero_dcp_partials(partial_out, partial_lse, local_kv_seq_lens); + + 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, + 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, + 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); @@ -508,7 +617,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/dcp_attention_utils.cpp b/xllm/core/layers/npu_torch/dcp_attention_utils.cpp index 44145ba159..e6f2659a65 100644 --- a/xllm/core/layers/npu_torch/dcp_attention_utils.cpp +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.cpp @@ -47,6 +47,25 @@ std::vector compute_dcp_local_kv_seq_lens( 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) { diff --git a/xllm/core/layers/npu_torch/dcp_attention_utils.h b/xllm/core/layers/npu_torch/dcp_attention_utils.h index 95c60ec129..27ac4dbda6 100644 --- a/xllm/core/layers/npu_torch/dcp_attention_utils.h +++ b/xllm/core/layers/npu_torch/dcp_attention_utils.h @@ -30,6 +30,15 @@ std::vector compute_dcp_local_kv_seq_lens( 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. 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 6846cf3d8d..c7d82c6a91 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 7573dcebae..3b4b222bb2 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.cpp +++ b/xllm/core/runtime/acl_graph_executor_impl.cpp @@ -34,8 +34,11 @@ limitations under the License. #include #endif #include "core/common/metrics.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/runtime/mtp_async_state.h" @@ -412,58 +415,140 @@ 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, + 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; @@ -687,6 +772,22 @@ ModelOutput AclGraph::replay(CausalLM* model, CHECK(update_stream_.has_value()); signal_static_graph_tasks(update_stream_.value()); } + // Per-replay completion barrier for the dynamic graph-task path (DCP FIA / + // GDN conv task-update). vLLM synchronizes the current stream before every + // FULL graph replay (compilation/acl_graph.py) so iteration i cannot rewrite + // 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). + // First-version host-blocking barrier; can later become a persistent + // input-ready/replay-done event chain (never a stack-local event). + if (!graph_paged_attention_tiling_data_.defined() && + model->is_hybrid_linear_attention() && !use_static_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()) @@ -1259,6 +1360,18 @@ uint64_t AclGraphExecutorImpl::get_graph_key( return static_cast(bucket_num_tokens) | kSpecVerifyGraphKeyMask | (q_max_seq_len << kSpecVerifyQMaxSeqLenShift); } + 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 cd9d973ffe..69c8f3a649 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.h +++ b/xllm/core/runtime/acl_graph_executor_impl.h @@ -206,6 +206,12 @@ class AclGraphExecutorImpl : public ExecutorImpl { return graph_slot_count_; } + [[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 9b54158b14..5014dd0f9d 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 @@ -240,6 +241,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)); + } // Output tensor for hidden states torch::Dtype dtype = util::parse_dtype(args.dtype(), device); @@ -507,6 +526,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; @@ -971,6 +1012,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; @@ -1173,6 +1219,17 @@ std::optional GraphPersistentParam::update( is_empty_dp_decode_rank ? 0 : static_cast(actual_num_tokens); params_for_capture->attention.host.kv_seq_lens = padded_kv_seq_lens_vec; params_for_capture->attention.host.q_seq_lens = padded_q_seq_lens_vec; + if (use_dcp_local_block_tables) { + std::vector& padded_q_cu_seq_lens = + params_for_capture->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 : padded_q_seq_lens_vec) { + q_running_total += q_seq_len; + padded_q_cu_seq_lens.emplace_back(q_running_total); + } + } params_for_capture->meta.num_sequences = static_cast(padded_batch_size); params_for_capture->meta.batch_forward_type = @@ -1187,6 +1244,13 @@ std::optional GraphPersistentParam::update( persistent_new_cache_slots(padded_num_tokens); params_for_capture->attention.device.block_tables = persistent_block_tables(static_cast(padded_batch_size)); + if (use_dcp_local_block_tables) { + params_for_capture->graph.dcp_local_block_tables = + persistent_dcp_local_block_tables( + static_cast(padded_batch_size)); + } else { + params_for_capture->graph.dcp_local_block_tables = torch::Tensor(); + } if (!params.embedding.linear_state_ids.empty()) { params_for_capture->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 f6ede4c271..1777a79768 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.h +++ b/xllm/core/runtime/acl_graph_persistent_param.h @@ -129,6 +129,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( @@ -225,6 +233,8 @@ 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); + // Update paged attention tiling based on input parameters void plan_paged_attention_tiling(const torch::Tensor& tokens, const torch::Tensor& k_cache, @@ -248,6 +258,8 @@ class GraphPersistentParam final { torch::Tensor persistent_positions_; torch::Tensor persistent_new_cache_slots_; 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_; @@ -300,6 +312,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 attention plan needs to be updated based on model // type bool need_update_attention_plan_; diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 18f85a111b..b99f36eb96 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -920,8 +920,6 @@ void WorkerImpl::prepare_work_before_execute_on_stream( "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."; - CHECK(!processed_input.input_params.enable_graph) - << "DCP-1c does not support graph-captured 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; From ffdb8bdfe513f1785ca842beef1fe736d55de3fd Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Thu, 20 Aug 2026 15:21:34 +0800 Subject: [PATCH 16/22] style: format DCP attention declaration --- xllm/core/layers/npu_torch/attention.h | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/xllm/core/layers/npu_torch/attention.h b/xllm/core/layers/npu_torch/attention.h index ce7b1b1f68..bb0da2b456 100644 --- a/xllm/core/layers/npu_torch/attention.h +++ b/xllm/core/layers/npu_torch/attention.h @@ -70,14 +70,13 @@ class AttentionImpl : public torch::nn::Module { 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); + 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_; From a7ce44e0dca83412344f28b91e5a76642a59b8bf Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Thu, 20 Aug 2026 20:29:52 +0800 Subject: [PATCH 17/22] chore: drop unused NO_NPU_RUNTIME cc_test option Revert cmake/cc_test.cmake to upstream. NO_NPU_RUNTIME was added only to build the DCP attention-math unit test as CPU-only, an unnecessary change to shared test infrastructure. Build that test as a standard NPU test (matching deepseek_v4_eplb_load_utils_test) instead; cmake/cc_test.cmake is now byte-identical to upstream. --- cmake/cc_test.cmake | 15 ++------------- tests/core/layers/npu_torch/CMakeLists.txt | 1 - 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/cmake/cc_test.cmake b/cmake/cc_test.cmake index bbee503738..2c80e3b0dc 100644 --- a/cmake/cc_test.cmake +++ b/cmake/cc_test.cmake @@ -11,7 +11,6 @@ include(CMakeParseArguments) # COPTS: List of private compile options # LINKOPTS: List of link options # ARGS: Command line arguments to test case -# NO_NPU_RUNTIME: Skip automatic NPU runtime setup for CPU-only tests # # Usage: # cc_library( @@ -40,7 +39,7 @@ function(cc_test) cmake_parse_arguments( CC_TEST # prefix - "NO_NPU_RUNTIME" # options + "" # options "NAME;ENVIRONMENT" # one value args "SRCS;COPTS;LINKOPTS;DEPS;INCLUDES;ARGS;DATA" # multi value args ${ARGN} @@ -113,17 +112,7 @@ function(cc_test) PRIVATE ${CC_TEST_LINKOPTS} ) - if(USE_NPU AND CC_TEST_NO_NPU_RUNTIME) - get_target_property(_CC_TEST_LINK_LIBRARIES - ${CC_TEST_NAME} LINK_LIBRARIES) - if(_CC_TEST_LINK_LIBRARIES) - list(REMOVE_ITEM _CC_TEST_LINK_LIBRARIES cust_opapi) - set_property(TARGET ${CC_TEST_NAME} - PROPERTY LINK_LIBRARIES "${_CC_TEST_LINK_LIBRARIES}") - endif() - endif() - - if(USE_NPU AND NOT CC_TEST_NO_NPU_RUNTIME) + if(USE_NPU) target_sources(${CC_TEST_NAME} PRIVATE "${PROJECT_SOURCE_DIR}/tests/npu_test_environment.cpp" ) diff --git a/tests/core/layers/npu_torch/CMakeLists.txt b/tests/core/layers/npu_torch/CMakeLists.txt index 597c533378..f6c329b46b 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/tests/core/layers/npu_torch/CMakeLists.txt @@ -1,7 +1,6 @@ include(cc_test) cc_test( - NO_NPU_RUNTIME NAME npu_dcp_attention_utils_test SRCS From ed74e9bfe8cfd6594747768f6cacaa45725557af Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Mon, 24 Aug 2026 15:17:27 +0800 Subject: [PATCH 18/22] chore: prune out-of-scope changes flagged in DCP PR review - Drop dead decode_context_parallel_size pass-through for VLM/SSM/REC masters. Only LLM reaches DCP (startup gate allows LLM/SSM and dcp_compat rejects SSM), so those builder calls were unreachable. - Remove the test-only production scalar compute_dcp_cache_slot; move the owner/slot oracle into cp_group_ranks_test as expected_dcp_cache_slot (cross-checks the production tensor remap_dcp_cache_slots path) and drop the three self-tests that only exercised the helper. - Revert the completion_service_impl usage change (deferred to a separate API PR); the file is byte-identical to upstream again. - Link the production :dcp_attention_utils target from its unit test instead of recompiling the source. --- .../parallel_state/cp_group_ranks_test.cpp | 98 +++++++------------ tests/core/layers/npu_torch/CMakeLists.txt | 2 +- xllm/api_service/completion_service_impl.cpp | 9 +- xllm/core/distributed_runtime/master.cpp | 3 - xllm/core/distributed_runtime/rec_master.cpp | 1 - xllm/core/distributed_runtime/vlm_master.cpp | 1 - .../parallel_state/parallel_state.cpp | 25 ----- .../framework/parallel_state/parallel_state.h | 15 +-- 8 files changed, 45 insertions(+), 109 deletions(-) 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 5acd6cfc6a..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,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include @@ -32,6 +33,34 @@ 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 @@ -251,63 +280,6 @@ TEST(ComputeDcpGroupRanks, RejectsNonIntegralDcpGroups) { ""); } -TEST(ComputeDcpCacheSlot, PreservesOwnerPhysicalSlots) { - const int32_t block_size = 4; - const int32_t dcp_size = 2; - const int32_t interleave_size = block_size; - - EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/151, - /*position=*/0, - block_size, - dcp_size, - /*dcp_rank=*/0, - interleave_size), - 151); - EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/23, - /*position=*/4, - block_size, - dcp_size, - /*dcp_rank=*/0, - interleave_size), - -1); - EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/23, - /*position=*/4, - block_size, - dcp_size, - /*dcp_rank=*/1, - interleave_size), - 23); - EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/359, - /*position=*/8, - block_size, - dcp_size, - /*dcp_rank=*/0, - interleave_size), - 359); -} - -TEST(ComputeDcpCacheSlot, RejectsSubBlockInterleave) { - const int32_t block_size = 4; - const int32_t dcp_size = 2; - EXPECT_DEATH(compute_dcp_cache_slot(/*logical_slot=*/0, - /*position=*/0, - block_size, - dcp_size, - /*dcp_rank=*/0, - /*interleave_size=*/1), - ""); -} - -TEST(ComputeDcpCacheSlot, PreservesNegativeSlots) { - EXPECT_EQ(compute_dcp_cache_slot(/*logical_slot=*/-1, - /*position=*/0, - /*block_size=*/4, - /*dcp_size=*/2, - /*dcp_rank=*/0, - /*interleave_size=*/4), - -1); -} - TEST(SelectDcpLocalBlockTable, SelectsOriginalNonContiguousBlockIds) { const torch::Tensor global_block_table = torch::tensor({{37, 5, 89, 2}, {41, 13, 73, 29}}, @@ -362,12 +334,12 @@ TEST(DcpCacheLayout, PrefillWritesMatchDecodeLocalBlockTable) { static_cast(global_block_index) * block_size + (block_size - 1); const int64_t owner_slot = - compute_dcp_cache_slot(original_slot, - position, - block_size, - dcp_size, - dcp_rank, - /*interleave_size=*/block_size); + 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(); diff --git a/tests/core/layers/npu_torch/CMakeLists.txt b/tests/core/layers/npu_torch/CMakeLists.txt index f6c329b46b..be2acfa93f 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/tests/core/layers/npu_torch/CMakeLists.txt @@ -5,8 +5,8 @@ cc_test( npu_dcp_attention_utils_test SRCS dcp_attention_utils_test.cpp - "${PROJECT_SOURCE_DIR}/xllm/core/layers/npu_torch/dcp_attention_utils.cpp" DEPS + :dcp_attention_utils glog::glog torch GTest::gtest_main diff --git a/xllm/api_service/completion_service_impl.cpp b/xllm/api_service/completion_service_impl.cpp index b856d3f9ea..2313505247 100644 --- a/xllm/api_service/completion_service_impl.cpp +++ b/xllm/api_service/completion_service_impl.cpp @@ -24,7 +24,6 @@ limitations under the License. #include #include -#include "api_service/utils.h" #include "common/instance_name.h" #include "completion.pb.h" #include "core/distributed_runtime/llm_master.h" @@ -106,7 +105,9 @@ bool send_delta_to_client_brpc(std::shared_ptr call, response.set_model(model); response.mutable_choices(); auto* proto_usage = response.mutable_usage(); - api_service::set_proto_usage(proto_usage, usage); + proto_usage->set_prompt_tokens(usage.num_prompt_tokens); + proto_usage->set_completion_tokens(usage.num_generated_tokens); + proto_usage->set_total_tokens(usage.num_total_tokens); if (!call->write(response)) { return false; } @@ -144,7 +145,9 @@ bool send_result_to_client_brpc(std::shared_ptr call, if (req_output.usage.has_value()) { const auto& usage = req_output.usage.value(); auto* proto_usage = response.mutable_usage(); - api_service::set_proto_usage(proto_usage, usage); + proto_usage->set_prompt_tokens(usage.num_prompt_tokens); + proto_usage->set_completion_tokens(usage.num_generated_tokens); + proto_usage->set_total_tokens(usage.num_total_tokens); } return call->write_and_finish(response); diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 233e77abdf..400ffe6b08 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -695,7 +695,6 @@ 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()) .instance_role(options_.instance_role()) .enable_disagg_pd(options_.enable_disagg_pd()) .npu_kernel_backend(options_.npu_kernel_backend()) @@ -779,7 +778,6 @@ 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()) @@ -899,7 +897,6 @@ Master::Master(const Options& options, EngineType type) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .cp_size(options_.cp_size()) - .decode_context_parallel_size(options_.decode_context_parallel_size()) .max_seqs_per_batch(options_.max_seqs_per_batch()) .beam_width(options_.beam_width()) .max_tokens_per_batch(options_.max_tokens_per_batch()) diff --git a/xllm/core/distributed_runtime/rec_master.cpp b/xllm/core/distributed_runtime/rec_master.cpp index 012f694c7a..0ae7d4ca4f 100644 --- a/xllm/core/distributed_runtime/rec_master.cpp +++ b/xllm/core/distributed_runtime/rec_master.cpp @@ -561,7 +561,6 @@ RecMaster::RecMaster(const Options& options) options_.max_tokens_per_chunk_for_prefill()) .num_speculative_tokens(options_.num_speculative_tokens()) .dp_size(options_.dp_size()) - .decode_context_parallel_size(options_.decode_context_parallel_size()) .enable_disagg_pd(options_.enable_disagg_pd()) .enable_schedule_overlap(options_.enable_schedule_overlap()) .enable_chunked_prefill(options_.enable_chunked_prefill()) diff --git a/xllm/core/distributed_runtime/vlm_master.cpp b/xllm/core/distributed_runtime/vlm_master.cpp index ef3a442f37..957fff89d3 100644 --- a/xllm/core/distributed_runtime/vlm_master.cpp +++ b/xllm/core/distributed_runtime/vlm_master.cpp @@ -86,7 +86,6 @@ VLMMaster::VLMMaster(const Options& options) .max_tokens_per_chunk_for_prefill( options.max_tokens_per_chunk_for_prefill()) .dp_size(options_.dp_size()) - .decode_context_parallel_size(options_.decode_context_parallel_size()) .enable_disagg_pd(options_.enable_disagg_pd()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .instance_name(options_.instance_name()) diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index be36fe85d5..8c15bb0939 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/xllm/core/framework/parallel_state/parallel_state.cpp @@ -304,31 +304,6 @@ std::vector compute_dcp_group_ranks(int32_t global_rank, return ranks; } -int64_t compute_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; -} - torch::Tensor select_dcp_local_block_table(const torch::Tensor& block_table, int32_t dcp_size, int32_t dcp_rank) { diff --git a/xllm/core/framework/parallel_state/parallel_state.h b/xllm/core/framework/parallel_state/parallel_state.h index d880779a18..e87a561eba 100644 --- a/xllm/core/framework/parallel_state/parallel_state.h +++ b/xllm/core/framework/parallel_state/parallel_state.h @@ -82,22 +82,13 @@ std::vector compute_dcp_group_ranks(int32_t global_rank, int32_t dp_size, int32_t dcp_size); -// Remap a logical KV cache slot to this DCP rank's local physical slot. Returns -// -1 when the token is owned by a different DCP rank. -int64_t compute_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); - torch::Tensor select_dcp_local_block_table(const torch::Tensor& block_table, int32_t dcp_size, int32_t dcp_rank); -// Batched tensor form of compute_dcp_cache_slot for the production remap path. -// 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 +// 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 Date: Mon, 24 Aug 2026 15:44:42 +0800 Subject: [PATCH 19/22] fix: scope DCP graph pre-replay barrier to FIA graph tasks The per-replay current-stream synchronize (added to prevent the DCP FIA graph cross-replay task/event overlap that surfaced as ACL_ERROR_RT_MODEL_ EXECUTE 507011) was gated on hybrid linear attention, so it also blocked the ordinary dcp=1 GDN conv graph replay hot path, for which no such fault has been observed. Gate it on the graph having captured FIA tasks (has_fia_graph_tasks) so only DCP decode graphs pay the barrier; dcp=1 GDN graphs keep their non-blocking replay. --- xllm/core/runtime/acl_graph_executor_impl.cpp | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/xllm/core/runtime/acl_graph_executor_impl.cpp b/xllm/core/runtime/acl_graph_executor_impl.cpp index c0847e2b14..6bf773a2cf 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.cpp +++ b/xllm/core/runtime/acl_graph_executor_impl.cpp @@ -807,19 +807,25 @@ ModelOutput AclGraph::replay(CausalLM* model, CHECK(update_stream_.has_value()); signal_static_graph_tasks(update_stream_.value()); } - // Per-replay completion barrier for the dynamic graph-task path (DCP FIA / - // GDN conv task-update). vLLM synchronizes the current stream before every - // FULL graph replay (compilation/acl_graph.py) so iteration i cannot rewrite - // 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 + // 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). - // First-version host-blocking barrier; can later become a persistent - // input-ready/replay-done event chain (never a stack-local event). + // 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) { + 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"; } From 7ee62fc1b20c05a31394ebf96f4342daed508909 Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Tue, 25 Aug 2026 13:50:51 +0800 Subject: [PATCH 20/22] fix: align GLM5.2 expert loading with lazy allocation --- tests/python/test_glm5_2_parallel.py | 14 +++++++++++++- xllm/python/models/glm5_2.py | 28 ++++++++++++++++++---------- 2 files changed, 31 insertions(+), 11 deletions(-) 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/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)) From e2012809365bff80018295b40c33d14ff4ca43f3 Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Tue, 25 Aug 2026 15:36:27 +0800 Subject: [PATCH 21/22] fix: register Python bridge modules after interpreter init --- tests/CMakeLists.txt | 9 ++++ tests/python/import_xllm_export.py | 50 +++++++++++++++++++++++ xllm/core/runtime/py_attention_metadata.h | 3 ++ xllm/core/runtime/py_executor_impl.cpp | 22 ++++++++-- xllm/models/py_model_helper.cpp | 30 ++++++++++---- xllm/models/py_model_helper.h | 4 ++ xllm/pybind/bind.cpp | 5 +++ 7 files changed, 112 insertions(+), 11 deletions(-) create mode 100644 tests/python/import_xllm_export.py 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/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/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/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()) From 93830fdd00ca1379923505cbb65ddf5a398580be Mon Sep 17 00:00:00 2001 From: konghaoyu Date: Tue, 25 Aug 2026 16:35:49 +0800 Subject: [PATCH 22/22] ci: select NPU with visible device mask --- .github/workflows/build_x86_64_npu.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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"'