From 8c586f5bfe9605995e82b6568e9066f7cc252ff3 Mon Sep 17 00:00:00 2001 From: wanghao Date: Tue, 21 Jul 2026 20:13:43 +0800 Subject: [PATCH] feat: support heterogeneous qwen3.5 pd disaggregation. - transfer target, draft, convolution, and SSM state across heterogeneous TP layouts\n- propagate per-rank cache metadata and synchronize staged prefill pushes\n- parallelize Decode shard pulls with a safe serial fallback and request-level staging protection\n- add topology and push-route coverage plus request-level performance instrumentation --- .../pd_topology_guard_test.cpp | 32 +- .../kv_cache_transfer/push_route_test.cpp | 18 + .../core/distributed_runtime/comm_channel.cpp | 24 + xllm/core/distributed_runtime/comm_channel.h | 8 + xllm/core/distributed_runtime/engine.h | 14 + xllm/core/distributed_runtime/llm_engine.cpp | 121 +++- xllm/core/distributed_runtime/llm_engine.h | 11 + .../distributed_runtime/remote_worker.cpp | 15 + xllm/core/distributed_runtime/remote_worker.h | 8 + .../speculative_engine.cpp | 21 + .../distributed_runtime/speculative_engine.h | 11 + .../distributed_runtime/worker_service.cpp | 28 +- .../kv_cache_transfer/kv_cache_transfer.h | 12 + .../llm_data_dist_transfer.cpp | 639 +++++++++++++++++- .../llm_data_dist_transfer.h | 42 ++ .../kv_cache_transfer/pd_topology_guard.cpp | 25 +- .../kv_cache_transfer/push_route.cpp | 25 + .../framework/kv_cache_transfer/push_route.h | 4 + .../spec_kv_cache_transfer.cpp | 271 +++++++- .../spec_kv_cache_transfer.h | 26 +- xllm/core/runtime/dflash_worker_impl.cpp | 3 +- xllm/core/runtime/mtp_worker_impl.cpp | 4 +- xllm/core/runtime/speculative_worker_impl.h | 15 + xllm/core/runtime/worker.cpp | 15 + xllm/core/runtime/worker.h | 8 + xllm/core/runtime/worker_client.cpp | 16 + xllm/core/runtime/worker_client.h | 8 + xllm/core/runtime/worker_impl.cpp | 49 ++ xllm/core/runtime/worker_impl.h | 8 + .../scheduler/async_response_processor.cpp | 19 + xllm/core/scheduler/disagg_pd_scheduler.cpp | 121 +++- xllm/models/llm/qwen3_5_mtp.h | 7 + xllm/proto/worker.proto | 3 + 33 files changed, 1538 insertions(+), 93 deletions(-) diff --git a/tests/core/framework/kv_cache_transfer/pd_topology_guard_test.cpp b/tests/core/framework/kv_cache_transfer/pd_topology_guard_test.cpp index 8d93c46578..7a108916e5 100644 --- a/tests/core/framework/kv_cache_transfer/pd_topology_guard_test.cpp +++ b/tests/core/framework/kv_cache_transfer/pd_topology_guard_test.cpp @@ -82,14 +82,14 @@ TEST(PdTopologyGuardTest, TryGetPdTopoReturnTopo) { EXPECT_TRUE(reason.empty()); } -TEST(PdTopologyGuardTest, HeteroTopoNeedMla) { - const InstanceInfo local_info = make_info(2, {0, 1, 2, 3}); - const InstanceInfo remote_info = make_info(1, {0, 1, 2, 3}); +TEST(PdTopologyGuardTest, HeteroPrefillTpTwoDecodeTpOneAllowsNonMlaPush) { + const InstanceInfo local_info = make_info(1, {0, 1}); + const InstanceInfo remote_info = make_info(1, {2}); const PdTopoResult result = check_pd_topo(local_info, remote_info, "PUSH", false); - EXPECT_EQ(result.status, PdTopoStatus::DENY_HETERO); - EXPECT_EQ(result.reason, "hetero pd requires enable_mla=true"); + EXPECT_EQ(result.status, PdTopoStatus::ALLOW_HETERO); + EXPECT_TRUE(result.reason.empty()); } TEST(PdTopologyGuardTest, HeteroTopoNeedPushKv) { @@ -112,6 +112,28 @@ TEST(PdTopologyGuardTest, HeteroTopoAllowOnPushMla) { EXPECT_TRUE(result.reason.empty()); } +TEST(PdTopologyGuardTest, NonMlaHeteroTopoRequiresEqualDpSize) { + const InstanceInfo local_info = make_info(2, {0, 1, 2, 3}); + const InstanceInfo remote_info = make_info(1, {4}); + + const PdTopoResult result = + check_pd_topo(local_info, remote_info, "PUSH", false); + EXPECT_EQ(result.status, PdTopoStatus::DENY_HETERO); + EXPECT_EQ(result.reason, "non-mla hetero pd requires equal dp_size"); +} + +TEST(PdTopologyGuardTest, NonMlaHeteroTopoRequiresPrefillTpMultiple) { + const InstanceInfo local_info = make_info(1, {0, 1, 2}); + const InstanceInfo remote_info = make_info(1, {3, 4}); + + const PdTopoResult result = + check_pd_topo(local_info, remote_info, "PUSH", false); + EXPECT_EQ(result.status, PdTopoStatus::DENY_HETERO); + EXPECT_EQ(result.reason, + "non-mla hetero pd requires prefill tp_size divisible by decode " + "tp_size"); +} + TEST(PdTopologyGuardTest, CheckPdTopoRejectInvalidLocalTopo) { const InstanceInfo local_info = make_info(0, {0, 1, 2, 3}); const InstanceInfo remote_info = make_info(1, {0, 1, 2, 3}); diff --git a/tests/core/framework/kv_cache_transfer/push_route_test.cpp b/tests/core/framework/kv_cache_transfer/push_route_test.cpp index df18471948..c21cc61c27 100644 --- a/tests/core/framework/kv_cache_transfer/push_route_test.cpp +++ b/tests/core/framework/kv_cache_transfer/push_route_test.cpp @@ -71,4 +71,22 @@ TEST(PushRouteTest, DstDpRankOffsetApplied) { EXPECT_EQ(dst_ranks, expect_ranks); } +TEST(PushRouteTest, DecodeTpOneLinksEveryPrefillTpRank) { + const std::vector src_tp_ranks = get_src_tp_ranks(0, 2, 1); + const std::vector expect_src_tp_ranks = {0, 1}; + EXPECT_EQ(src_tp_ranks, expect_src_tp_ranks); +} + +TEST(PushRouteTest, DecodeRankLinksMatchingPrefillOwners) { + const std::vector src_tp_ranks = get_src_tp_ranks(1, 8, 2); + const std::vector expect_src_tp_ranks = {1, 3, 5, 7}; + EXPECT_EQ(src_tp_ranks, expect_src_tp_ranks); +} + +TEST(PushRouteTest, LargerDecodeTpKeepsRoundRobinSource) { + const std::vector src_tp_ranks = get_src_tp_ranks(5, 2, 8); + const std::vector expect_src_tp_ranks = {1}; + EXPECT_EQ(src_tp_ranks, expect_src_tp_ranks); +} + } // namespace xllm diff --git a/xllm/core/distributed_runtime/comm_channel.cpp b/xllm/core/distributed_runtime/comm_channel.cpp index 7f63f59027..6af0c5ad43 100644 --- a/xllm/core/distributed_runtime/comm_channel.cpp +++ b/xllm/core/distributed_runtime/comm_channel.cpp @@ -266,6 +266,30 @@ bool CommChannel::pull_kv_blocks( return !cntl.Failed() && s.ok(); } +bool CommChannel::pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + proto::PullKVCacheRequest request; + request.set_hetero_merge(true); + ADD_VECTOR_TO_PROTO(request.mutable_src_cluster_ids(), src_cluster_ids); + ADD_VECTOR_TO_PROTO(request.mutable_src_addrs(), src_addrs); + ADD_VECTOR_TO_PROTO(request.mutable_src_blocks(), src_blocks); + ADD_VECTOR_TO_PROTO(request.mutable_dst_blocks(), dst_blocks); + ADD_VECTOR_TO_PROTO(request.mutable_src_linear_state_ids(), + src_linear_state_ids); + ADD_VECTOR_TO_PROTO(request.mutable_dst_linear_state_ids(), + dst_linear_state_ids); + + proto::Status s; + brpc::Controller cntl; + stub_->PullKVCache(&cntl, &request, &s, nullptr); + return !cntl.Failed() && s.ok(); +} + void CommChannel::execute_model_async( const ForwardInput& input, folly::Promise>& promise) { diff --git a/xllm/core/distributed_runtime/comm_channel.h b/xllm/core/distributed_runtime/comm_channel.h index a0d2575be7..353ea2fc31 100644 --- a/xllm/core/distributed_runtime/comm_channel.h +++ b/xllm/core/distributed_runtime/comm_channel.h @@ -79,6 +79,14 @@ class CommChannel { const std::vector& src_linear_state_ids = {}, const std::vector& dst_linear_state_ids = {}); + virtual bool pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}); + virtual void execute_model_async( const ForwardInput& input, folly::Promise>& promise); diff --git a/xllm/core/distributed_runtime/engine.h b/xllm/core/distributed_runtime/engine.h index ee7419b168..74c6eab183 100644 --- a/xllm/core/distributed_runtime/engine.h +++ b/xllm/core/distributed_runtime/engine.h @@ -78,6 +78,20 @@ class Engine { return false; }; + virtual bool pull_hetero_kv_blocks( + const int32_t src_dp_size, + const int32_t src_dp_rank, + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const int32_t dst_dp_rank, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}) { + NOT_IMPLEMENTED(); + return false; + }; + virtual std::vector> transfer_kv_blocks( const uint32_t dp_rank, const std::vector& block_transfer_info) { diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 7aa1f42f0d..4e1be1d390 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -45,6 +45,7 @@ limitations under the License. #include "framework/kv_cache/kv_cache_estimation.h" #include "framework/kv_cache/kv_cache_shape.h" #include "framework/kv_cache/kv_cache_utils.h" +#include "framework/kv_cache_transfer/push_route.h" #include "framework/model/model_args.h" #include "framework/model_loader.h" #include "framework/xtensor/page_allocator.h" @@ -701,6 +702,55 @@ bool LLMEngine::pull_kv_blocks( return true; } +bool LLMEngine::pull_hetero_kv_blocks( + const int32_t src_dp_size, + const int32_t src_dp_rank, + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const int32_t dst_dp_rank, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + if (src_dp_size <= 0 || src_dp_rank < 0 || src_dp_rank >= src_dp_size || + src_cluster_ids.size() != src_addrs.size() || + src_cluster_ids.size() % static_cast(src_dp_size) != 0) { + LOG(ERROR) << "Invalid heterogeneous KV pull topology."; + return false; + } + const int32_t src_tp_size = + static_cast(src_cluster_ids.size()) / src_dp_size; + const int32_t dst_tp_size = static_cast(dp_local_tp_size_); + if (src_tp_size < dst_tp_size || src_tp_size % dst_tp_size != 0) { + LOG(ERROR) << "Unsupported heterogeneous KV pull ratio: prefill_tp_size=" + << src_tp_size << ", decode_tp_size=" << dst_tp_size; + return false; + } + + std::vector results; + results.reserve(dst_tp_size); + for (int32_t dst_tp_rank = 0; dst_tp_rank < dst_tp_size; ++dst_tp_rank) { + std::vector worker_src_cluster_ids; + std::vector worker_src_addrs; + for (int32_t src_tp_rank : + get_src_tp_ranks(dst_tp_rank, src_tp_size, dst_tp_size)) { + const int32_t src_worker_rank = src_dp_rank * src_tp_size + src_tp_rank; + worker_src_cluster_ids.push_back(src_cluster_ids[src_worker_rank]); + worker_src_addrs.push_back(src_addrs[src_worker_rank]); + } + const int32_t dst_worker_rank = dst_dp_rank * dst_tp_size + dst_tp_rank; + results.push_back(worker_clients_[dst_worker_rank]->pull_hetero_kv_blocks( + worker_src_cluster_ids, + worker_src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids)); + } + return std::all_of( + results.begin(), results.end(), [](bool result) { return result; }); +} + std::vector> LLMEngine::transfer_kv_blocks( const uint32_t dp_rank, const std::vector& block_transfer_info) { @@ -786,15 +836,13 @@ bool LLMEngine::link_cluster(const std::vector& cluster_ids, const std::vector& ports, const int32_t src_dp_size, const int32_t src_kv_split_size) { - // Each D worker connects to all P workers that share the same TP rank. + const int32_t src_world_size = static_cast(cluster_ids.size()); + + // Each D worker connects to every P worker that routes KV blocks to its TP + // rank. When P TP is larger, multiple P ranks share one D-side owner. // P layout: rank = dp_i * src_cp_tp_size + split_j * src_tp_size + tp_rank - // D workers cycle through tp_rank in [0, src_tp_size) round-robin. - // Requires: D-side dp_local_tp_size_ == src_tp_size. - int32_t src_world_size = static_cast(cluster_ids.size()); int32_t src_cp_tp_size = src_world_size / src_dp_size; int32_t src_tp_size = src_cp_tp_size / src_kv_split_size; - int32_t src_dp_worker_index = 0; - std::vector> futures; futures.reserve(worker_clients_num_); for (size_t worker_rank = 0; worker_rank < worker_clients_num_; @@ -802,22 +850,29 @@ bool LLMEngine::link_cluster(const std::vector& cluster_ids, std::vector target_cluster_ids; std::vector target_addrs; std::vector target_ports; - target_cluster_ids.reserve(src_dp_size * src_kv_split_size); - target_addrs.reserve(src_dp_size * src_kv_split_size); - target_ports.reserve(src_dp_size * src_kv_split_size); + const int32_t dst_tp_rank = + static_cast(worker_rank % dp_local_tp_size_); + const std::vector src_tp_ranks = + get_src_tp_ranks(dst_tp_rank, src_tp_size, dp_local_tp_size_); + const size_t endpoint_count = static_cast(src_dp_size) * + static_cast(src_kv_split_size) * + src_tp_ranks.size(); + target_cluster_ids.reserve(endpoint_count); + target_addrs.reserve(endpoint_count); + target_ports.reserve(endpoint_count); for (int32_t dp_i = 0; dp_i < src_dp_size; ++dp_i) { for (int32_t split_j = 0; split_j < src_kv_split_size; ++split_j) { - int32_t p_idx = - dp_i * src_cp_tp_size + split_j * src_tp_size + src_dp_worker_index; - target_cluster_ids.emplace_back(cluster_ids[p_idx]); - target_addrs.emplace_back(addrs[p_idx]); - target_ports.emplace_back(ports[p_idx]); + for (int32_t src_tp_rank : src_tp_ranks) { + const int32_t p_idx = + dp_i * src_cp_tp_size + split_j * src_tp_size + src_tp_rank; + target_cluster_ids.emplace_back(cluster_ids[p_idx]); + target_addrs.emplace_back(addrs[p_idx]); + target_ports.emplace_back(ports[p_idx]); + } } } - src_dp_worker_index = (src_dp_worker_index + 1) % src_tp_size; - folly::Promise promise; auto future = promise.getSemiFuture(); link_threadpool_->schedule( @@ -848,12 +903,11 @@ bool LLMEngine::unlink_cluster(const std::vector& cluster_ids, const std::vector& ports, const int32_t src_dp_size, const int32_t src_kv_split_size) { - // Symmetric to link_cluster; uses the same rank mapping. - int32_t src_world_size = static_cast(cluster_ids.size()); + const int32_t src_world_size = static_cast(cluster_ids.size()); + + // Symmetric to link_cluster; uses the same owner mapping. int32_t src_cp_tp_size = src_world_size / src_dp_size; int32_t src_tp_size = src_cp_tp_size / src_kv_split_size; - int32_t src_dp_worker_index = 0; - std::vector> futures; futures.reserve(worker_clients_num_); for (size_t worker_rank = 0; worker_rank < worker_clients_num_; @@ -861,22 +915,29 @@ bool LLMEngine::unlink_cluster(const std::vector& cluster_ids, std::vector target_cluster_ids; std::vector target_addrs; std::vector target_ports; - target_cluster_ids.reserve(src_dp_size * src_kv_split_size); - target_addrs.reserve(src_dp_size * src_kv_split_size); - target_ports.reserve(src_dp_size * src_kv_split_size); + const int32_t dst_tp_rank = + static_cast(worker_rank % dp_local_tp_size_); + const std::vector src_tp_ranks = + get_src_tp_ranks(dst_tp_rank, src_tp_size, dp_local_tp_size_); + const size_t endpoint_count = static_cast(src_dp_size) * + static_cast(src_kv_split_size) * + src_tp_ranks.size(); + target_cluster_ids.reserve(endpoint_count); + target_addrs.reserve(endpoint_count); + target_ports.reserve(endpoint_count); for (int32_t dp_i = 0; dp_i < src_dp_size; ++dp_i) { for (int32_t split_j = 0; split_j < src_kv_split_size; ++split_j) { - int32_t p_idx = - dp_i * src_cp_tp_size + split_j * src_tp_size + src_dp_worker_index; - target_cluster_ids.emplace_back(cluster_ids[p_idx]); - target_addrs.emplace_back(addrs[p_idx]); - target_ports.emplace_back(ports[p_idx]); + for (int32_t src_tp_rank : src_tp_ranks) { + const int32_t p_idx = + dp_i * src_cp_tp_size + split_j * src_tp_size + src_tp_rank; + target_cluster_ids.emplace_back(cluster_ids[p_idx]); + target_addrs.emplace_back(addrs[p_idx]); + target_ports.emplace_back(ports[p_idx]); + } } } - src_dp_worker_index = (src_dp_worker_index + 1) % src_tp_size; - folly::Promise promise; auto future = promise.getSemiFuture(); link_threadpool_->schedule( diff --git a/xllm/core/distributed_runtime/llm_engine.h b/xllm/core/distributed_runtime/llm_engine.h index 9b7b743c6d..ab80eacade 100644 --- a/xllm/core/distributed_runtime/llm_engine.h +++ b/xllm/core/distributed_runtime/llm_engine.h @@ -75,6 +75,17 @@ class LLMEngine : public Engine { const std::vector& src_linear_state_ids = {}, const std::vector& dst_linear_state_ids = {}) override; + bool pull_hetero_kv_blocks( + const int32_t src_dp_size, + const int32_t src_dp_rank, + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const int32_t dst_dp_rank, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}) override; + std::vector> transfer_kv_blocks( const uint32_t dp_rank, const std::vector& block_transfer_info) override; diff --git a/xllm/core/distributed_runtime/remote_worker.cpp b/xllm/core/distributed_runtime/remote_worker.cpp index 26158665a9..13ae18e679 100644 --- a/xllm/core/distributed_runtime/remote_worker.cpp +++ b/xllm/core/distributed_runtime/remote_worker.cpp @@ -275,6 +275,21 @@ folly::SemiFuture RemoteWorker::pull_kv_blocks_async( return future; } +bool RemoteWorker::pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + return channel_->pull_hetero_kv_blocks(src_cluster_ids, + src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); +} + folly::SemiFuture RemoteWorker::transfer_kv_blocks( const std::vector& block_transfer_info) { folly::Promise promise; diff --git a/xllm/core/distributed_runtime/remote_worker.h b/xllm/core/distributed_runtime/remote_worker.h index d337989069..d24fa6e07c 100644 --- a/xllm/core/distributed_runtime/remote_worker.h +++ b/xllm/core/distributed_runtime/remote_worker.h @@ -79,6 +79,14 @@ class RemoteWorker : public WorkerClient { const std::vector& src_linear_state_ids = {}, const std::vector& dst_linear_state_ids = {}) override; + virtual bool pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}) override; + // prepare input request virtual ForwardInput prepare_inputs(Batch& batch) override; diff --git a/xllm/core/distributed_runtime/speculative_engine.cpp b/xllm/core/distributed_runtime/speculative_engine.cpp index c17abdc3fd..3c43a21e18 100644 --- a/xllm/core/distributed_runtime/speculative_engine.cpp +++ b/xllm/core/distributed_runtime/speculative_engine.cpp @@ -262,6 +262,27 @@ bool SpeculativeEngine::pull_kv_blocks( dst_linear_state_ids); }; +bool SpeculativeEngine::pull_hetero_kv_blocks( + const int32_t src_dp_size, + const int32_t src_dp_rank, + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const int32_t dst_dp_rank, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + return engine_->pull_hetero_kv_blocks(src_dp_size, + src_dp_rank, + src_cluster_ids, + src_addrs, + src_blocks, + dst_dp_rank, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); +} + void SpeculativeEngine::get_cache_info(std::vector& cluster_ids, std::vector& addrs, std::vector& ports) { diff --git a/xllm/core/distributed_runtime/speculative_engine.h b/xllm/core/distributed_runtime/speculative_engine.h index 58b1c79883..3555b5d38c 100644 --- a/xllm/core/distributed_runtime/speculative_engine.h +++ b/xllm/core/distributed_runtime/speculative_engine.h @@ -68,6 +68,17 @@ class SpeculativeEngine : public Engine { const std::vector& src_linear_state_ids = {}, const std::vector& dst_linear_state_ids = {}) override; + bool pull_hetero_kv_blocks( + const int32_t src_dp_size, + const int32_t src_dp_rank, + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const int32_t dst_dp_rank, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}) override; + void get_cache_info(std::vector& cluster_ids, std::vector& addrs, std::vector& ports) override; diff --git a/xllm/core/distributed_runtime/worker_service.cpp b/xllm/core/distributed_runtime/worker_service.cpp index e4a970892e..9155c65ade 100644 --- a/xllm/core/distributed_runtime/worker_service.cpp +++ b/xllm/core/distributed_runtime/worker_service.cpp @@ -487,8 +487,6 @@ void WorkerService::PullKVCache(::google::protobuf::RpcController* controller, ::google::protobuf::Closure* done) { threadpool_->schedule([this, controller, req, resp, done]() mutable { brpc::ClosureGuard done_guard(done); - uint64_t src_cluster_id = req->cluster_id(); - std::string addr = req->addr(); std::vector src_blocks(req->src_blocks().begin(), req->src_blocks().end()); std::vector dst_blocks(req->dst_blocks().begin(), @@ -497,12 +495,26 @@ void WorkerService::PullKVCache(::google::protobuf::RpcController* controller, req->src_linear_state_ids().begin(), req->src_linear_state_ids().end()); std::vector dst_linear_state_ids( req->dst_linear_state_ids().begin(), req->dst_linear_state_ids().end()); - auto future = worker_->pull_kv_blocks_async(src_cluster_id, - addr, - src_blocks, - dst_blocks, - src_linear_state_ids, - dst_linear_state_ids); + auto future = [&]() { + if (req->hetero_merge()) { + std::vector src_cluster_ids(req->src_cluster_ids().begin(), + req->src_cluster_ids().end()); + std::vector src_addrs(req->src_addrs().begin(), + req->src_addrs().end()); + return worker_->pull_hetero_kv_blocks_async(src_cluster_ids, + src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); + } + return worker_->pull_kv_blocks_async(req->cluster_id(), + req->addr(), + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); + }(); bool status = std::move(future).get(); resp->set_ok(status); }); diff --git a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h index c631c0d5b1..6adec2d306 100644 --- a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h +++ b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h @@ -136,6 +136,18 @@ class KVCacheTransfer { const std::vector& src_linear_state_ids, const std::vector& dst_linear_state_ids); + // Heterogeneous TP fallback transport: pull every source TP shard into + // temporary local buffers and merge them into the decode-side cache. + virtual bool pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + return false; + } + #if defined(USE_NPU) || defined(USE_MLU) || defined(USE_DCU) virtual folly::SemiFuture push_kv_blocks_async( const std::vector& transfer_kv_infos, diff --git a/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.cpp b/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.cpp index 64acbf9c9c..61ae35d8b5 100644 --- a/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.cpp +++ b/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.cpp @@ -19,10 +19,14 @@ limitations under the License. #include #include +#include +#include +#include #include "common/macros.h" #include "core/framework/config/disagg_pd_config.h" #include "util/net.h" +#include "util/timer.h" namespace xllm { @@ -45,6 +49,74 @@ ge::DataType dtype_to_ge_dtype(torch::ScalarType dtype) { return it->second; } +bool is_linear_state_cache(KVCacheTensorRole role) { + return role == KVCacheTensorRole::CONV || role == KVCacheTensorRole::SSM; +} + +int64_t sharded_dimension(KVCacheTensorRole role, const torch::Tensor& tensor) { + if (role == KVCacheTensorRole::KEY || role == KVCacheTensorRole::VALUE) { + return 2; + } + if (role == KVCacheTensorRole::CONV) { + return tensor.dim() - 1; + } + if (role == KVCacheTensorRole::SSM) { + return 1; + } + return -1; +} + +std::vector make_compact_ids(size_t count) { + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + return ids; +} + +std::vector expand_checkpoint_ids( + const std::vector& logical_ids, + int64_t checkpoint_stride) { + std::vector ids; + ids.reserve(logical_ids.size() * static_cast(checkpoint_stride)); + for (uint64_t logical_id : logical_ids) { + for (int64_t checkpoint = 0; checkpoint < checkpoint_stride; ++checkpoint) { + ids.push_back(logical_id * static_cast(checkpoint_stride) + + static_cast(checkpoint)); + } + } + return ids; +} + +torch::Tensor make_page_aligned_staging_tensor( + std::vector& shape, + const torch::TensorOptions& options, + int64_t element_size) { + constexpr int64_t kHcclPageSize = 2 * 1024 * 1024; + int64_t row_elements = 1; + for (size_t dim = 1; dim < shape.size(); ++dim) { + row_elements *= shape[dim]; + } + const int64_t row_bytes = row_elements * element_size; + const int64_t rows_per_alignment = + kHcclPageSize / std::gcd(kHcclPageSize, row_bytes); + shape[0] = ((shape[0] + rows_per_alignment - 1) / rows_per_alignment) * + rows_per_alignment; + + const int64_t aligned_numel = std::accumulate( + shape.begin(), shape.end(), int64_t{1}, std::multiplies()); + const int64_t padding_elements = kHcclPageSize / element_size; + torch::Tensor backing = + torch::empty({aligned_numel + padding_elements}, options); + const uintptr_t base = reinterpret_cast(backing.data_ptr()); + const uintptr_t aligned = (base + kHcclPageSize - 1) & ~(kHcclPageSize - 1); + const int64_t offset_elements = + static_cast(aligned - base) / element_size; + torch::Tensor stage = + backing.narrow(0, offset_elements, aligned_numel).view(shape); + CHECK_EQ(reinterpret_cast(stage.data_ptr()) % kHcclPageSize, 0); + CHECK_EQ(stage.numel() * element_size % kHcclPageSize, 0); + return stage; +} + LlmDataDistTransfer::LlmDataDistTransfer(const uint16_t listen_port, const InstanceRole& instance_role, bool enable_lighting_indexer) @@ -229,7 +301,8 @@ RegisteredCache LlmDataDistTransfer::register_cache_tensor( RegisteredCache registered_cache{cache_tensor.role, cache_tensor.group_id, cache_tensor.sequence_scoped, - Cache{}}; + Cache{}, + tensor}; registered_cache.cache.tensor_addrs = {tensor_addr}; CacheDesc& desc = registered_cache.cache.cache_desc; @@ -242,10 +315,572 @@ RegisteredCache LlmDataDistTransfer::register_cache_tensor( CHECK(ret == LLM_SUCCESS) << "Register " << cache_tensor.role.to_string() << " cache failed at layer " << layer_id << ", ret = " << std::hex << ret; - + VLOG(5) << "Registered KV cache: layer=" << layer_id + << ", role=" << cache_tensor.role.to_string() + << ", cache_id=" << registered_cache.cache.cache_id + << ", shape=" << tensor.sizes(); return registered_cache; } +bool LlmDataDistTransfer::pull_and_merge_sharded_caches( + const LayerRegisteredCaches& layer_registered_caches, + const LayerRegisteredCaches& staging_registered_caches, + const std::vector& src_cluster_ids, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + if (src_cluster_ids.size() != 2 || + layer_registered_caches.size() != staging_registered_caches.size() || + src_blocks.size() != dst_blocks.size() || + src_linear_state_ids.size() != dst_linear_state_ids.size()) { + LOG(ERROR) << "Invalid heterogeneous KV pull metadata: src_tp=" + << src_cluster_ids.size() << ", src_blocks=" << src_blocks.size() + << ", dst_blocks=" << dst_blocks.size() + << ", src_linear_states=" << src_linear_state_ids.size() + << ", dst_linear_states=" << dst_linear_state_ids.size(); + return false; + } + + // The staging tensors are registered once and reused by all requests. + // Serialize request-level restore/merge so concurrent FirstGeneration RPCs + // cannot overwrite each other's staging rows. + std::lock_guard hetero_pull_lock(hetero_pull_mutex_); + + const int64_t shard_count = static_cast(src_cluster_ids.size()); + Timer breakdown_total_timer; + const char* parallel_pull_env = std::getenv("XLLM_PD_PARALLEL_SHARD_PULL"); + // Parallel shard pulls are the default for heterogeneous TP. Keep an + // explicit opt-out so deployments can immediately fall back to the serial + // path without rebuilding if the transport backend rejects concurrency. + const bool parallel_shard_pull = + parallel_pull_env == nullptr || std::string(parallel_pull_env) != "0"; + double pull_seconds = 0.0; + double pull_wall_seconds = 0.0; + double merge_seconds = 0.0; + double conv_pull_seconds = 0.0; + double conv_merge_seconds = 0.0; + double ssm_pull_seconds = 0.0; + double ssm_merge_seconds = 0.0; + std::vector shard_pull_seconds(src_cluster_ids.size(), 0.0); + size_t pull_calls = 0; + size_t merge_calls = 0; + bool success = true; + for (int64_t layer_id = 0; + layer_id < static_cast(layer_registered_caches.size()); + ++layer_id) { + const auto& layer_caches = layer_registered_caches[layer_id]; + const auto& layer_staging_caches = staging_registered_caches[layer_id]; + if (layer_caches.size() != layer_staging_caches.size()) { + LOG(ERROR) << "Heterogeneous KV staging layout mismatch at layer " + << layer_id; + return false; + } + int64_t checkpoint_stride = 1; + for (const RegisteredCache& cache : layer_caches) { + if (cache.role == KVCacheTensorRole::CONV && cache.tensor.defined()) { + for (const RegisteredCache& candidate : layer_caches) { + if (candidate.role == KVCacheTensorRole::SSM && + candidate.tensor.defined()) { + CHECK_EQ(candidate.tensor.size(0) % cache.tensor.size(0), 0); + checkpoint_stride = candidate.tensor.size(0) / cache.tensor.size(0); + } + } + } + } + + for (size_t cache_index = 0; cache_index < layer_caches.size(); + ++cache_index) { + const RegisteredCache& registered_cache = layer_caches[cache_index]; + const RegisteredCache& stage_cache = layer_staging_caches[cache_index]; + const int64_t shard_dim = + sharded_dimension(registered_cache.role, registered_cache.tensor); + if (shard_dim < 0) { + LOG(ERROR) << "Unsupported heterogeneous KV tensor role: " + << registered_cache.role.to_string(); + return false; + } + CHECK_EQ(registered_cache.tensor.size(shard_dim) % shard_count, 0) + << "Cache shard dimension is not divisible by source shard count, " + << "layer=" << layer_id + << ", role=" << registered_cache.role.to_string() + << ", shape=" << registered_cache.tensor.sizes(); + + const bool linear_state_cache = + is_linear_state_cache(registered_cache.role); + std::vector remote_ids = + linear_state_cache ? src_linear_state_ids : src_blocks; + std::vector final_ids = + linear_state_cache ? dst_linear_state_ids : dst_blocks; + if (registered_cache.role == KVCacheTensorRole::SSM) { + remote_ids = expand_checkpoint_ids(remote_ids, checkpoint_stride); + final_ids = expand_checkpoint_ids(final_ids, checkpoint_stride); + } + if (remote_ids.empty()) { + continue; + } + + if (stage_cache.tensor.size(0) < + static_cast(remote_ids.size())) { + LOG(ERROR) << "Heterogeneous KV transfer exceeds staging capacity, " + << "layer=" << layer_id + << ", role=" << registered_cache.role.to_string() + << ", requested=" << remote_ids.size() + << ", capacity=" << stage_cache.tensor.size(0); + return false; + } + CHECK_EQ(stage_cache.tensor.size(0) % shard_count, 0); + const int64_t rows_per_shard = stage_cache.tensor.size(0) / shard_count; + CHECK_LE(static_cast(remote_ids.size()), rows_per_shard); + std::vector shard_tensors; + shard_tensors.reserve(src_cluster_ids.size()); + + std::vector pull_rets(src_cluster_ids.size(), + LLM_SUCCESS); + std::vector current_pull_seconds(src_cluster_ids.size(), 0.0); + auto pull_one_shard = [&](size_t source_shard) { + const uint64_t src_cluster_id = src_cluster_ids[source_shard]; + std::vector staging_ids = make_compact_ids(remote_ids.size()); + const uint64_t staging_offset = + static_cast(source_shard * rows_per_shard); + for (uint64_t& staging_id : staging_ids) { + staging_id += staging_offset; + } + CacheIndex src_cache_index{src_cluster_id, + registered_cache.cache.cache_id}; + KvCacheExtParam ext_param{}; + ext_param.src_layer_range = {0, 0}; + ext_param.dst_layer_range = {0, 0}; + ext_param.tensor_num_per_layer = 1; + Timer pull_timer; + pull_rets[source_shard] = + llm_data_dist_->PullKvBlocks(src_cache_index, + stage_cache.cache, + remote_ids, + staging_ids, + ext_param); + current_pull_seconds[source_shard] = pull_timer.elapsed_seconds(); + }; + + Timer pull_group_timer; + if (parallel_shard_pull && src_cluster_ids.size() == 2) { + TaskGroup shard_pull_group(1); + shard_pull_threadpool_.schedule( + shard_pull_group.wrap([&]() { pull_one_shard(1); })); + std::exception_ptr request_thread_exception; + try { + pull_one_shard(0); + } catch (...) { + request_thread_exception = std::current_exception(); + } + shard_pull_group.wait(); + if (request_thread_exception) { + std::rethrow_exception(request_thread_exception); + } + } else { + for (size_t source_shard = 0; source_shard < src_cluster_ids.size(); + ++source_shard) { + pull_one_shard(source_shard); + } + } + pull_wall_seconds += pull_group_timer.elapsed_seconds(); + + for (size_t source_shard = 0; source_shard < src_cluster_ids.size(); + ++source_shard) { + const uint64_t src_cluster_id = src_cluster_ids[source_shard]; + const uint64_t staging_offset = + static_cast(source_shard * rows_per_shard); + pull_seconds += current_pull_seconds[source_shard]; + shard_pull_seconds[source_shard] += current_pull_seconds[source_shard]; + ++pull_calls; + if (registered_cache.role == KVCacheTensorRole::CONV) { + conv_pull_seconds += current_pull_seconds[source_shard]; + } else if (registered_cache.role == KVCacheTensorRole::SSM) { + ssm_pull_seconds += current_pull_seconds[source_shard]; + } + if (pull_rets[source_shard] != LLM_SUCCESS) { + LOG(ERROR) << "Heterogeneous PullKvBlocks failed, layer=" << layer_id + << ", role=" << registered_cache.role.to_string() + << ", src_cluster_id=" << src_cluster_id + << ", src_cache_id=" << registered_cache.cache.cache_id + << ", ret=" << std::hex << pull_rets[source_shard]; + success = false; + break; + } + // Each Prefill shard owns a disjoint staging row range. PullKvBlocks + // is synchronous, so the merge can consume these rows directly + // without cloning the entire recurrent state between shard pulls. + shard_tensors.push_back( + stage_cache.tensor.narrow(0, staging_offset, remote_ids.size())); + } + + if (success) { + Timer merge_timer; + torch::Tensor merged; + if (registered_cache.role == KVCacheTensorRole::CONV) { + // Each TP rank stores its causal-conv state as + // [Q_rank, K_rank, V_rank]. Concatenating whole rank tensors would + // produce [Q0,K0,V0,Q1,K1,V1], while a TP1 model consumes + // [Q0,Q1,K0,K1,V0,V1]. Recover the local V width from the matching + // SSM cache, then merge each projection segment independently. + int64_t local_v_width = -1; + for (const RegisteredCache& candidate : layer_staging_caches) { + if (candidate.role == KVCacheTensorRole::SSM && + candidate.tensor.defined()) { + local_v_width = + candidate.tensor.size(1) * candidate.tensor.size(3); + break; + } + } + CHECK_GT(local_v_width, 0) + << "CONV cache requires a matching SSM cache at layer " + << layer_id; + const int64_t local_conv_width = stage_cache.tensor.size(shard_dim); + CHECK_EQ((local_conv_width - local_v_width) % 2, 0) + << "Invalid Q/K/V CONV cache layout at layer " << layer_id; + const int64_t local_qk_width = (local_conv_width - local_v_width) / 2; + std::vector q_shards; + std::vector k_shards; + std::vector v_shards; + q_shards.reserve(shard_tensors.size()); + k_shards.reserve(shard_tensors.size()); + v_shards.reserve(shard_tensors.size()); + for (const torch::Tensor& shard : shard_tensors) { + auto qkv = torch::split_with_sizes( + shard, + {local_qk_width, local_qk_width, local_v_width}, + shard_dim); + q_shards.push_back(qkv[0]); + k_shards.push_back(qkv[1]); + v_shards.push_back(qkv[2]); + } + merged = torch::cat({torch::cat(q_shards, shard_dim), + torch::cat(k_shards, shard_dim), + torch::cat(v_shards, shard_dim)}, + shard_dim); + } else { + merged = torch::cat(shard_tensors, shard_dim); + } + std::vector signed_final_ids(final_ids.begin(), + final_ids.end()); + torch::Tensor dst_indices = + torch::tensor(signed_final_ids, + torch::TensorOptions().dtype(torch::kLong)) + .to(registered_cache.tensor.device(), /*non_blocking=*/false); + registered_cache.tensor.index_copy_(0, dst_indices, merged); + const double current_merge_seconds = merge_timer.elapsed_seconds(); + merge_seconds += current_merge_seconds; + ++merge_calls; + if (registered_cache.role == KVCacheTensorRole::CONV) { + conv_merge_seconds += current_merge_seconds; + } else if (registered_cache.role == KVCacheTensorRole::SSM) { + ssm_merge_seconds += current_merge_seconds; + } + } + if (!success) { + return false; + } + } + } + LOG(INFO) << "[PD-PERF] Heterogeneous pull-merge breakdown" + << " linear_request=" << !src_linear_state_ids.empty() + << " parallel_shard_pull=" << parallel_shard_pull + << " pull_calls=" << pull_calls << " merge_calls=" << merge_calls + << " shard0_pull_ms=" << shard_pull_seconds[0] * 1000.0 + << " shard1_pull_ms=" << shard_pull_seconds[1] * 1000.0 + << " conv_pull_ms=" << conv_pull_seconds * 1000.0 + << " ssm_pull_ms=" << ssm_pull_seconds * 1000.0 + << " pull_work_ms=" << pull_seconds * 1000.0 + << " pull_ms=" << pull_wall_seconds * 1000.0 + << " conv_merge_ms=" << conv_merge_seconds * 1000.0 + << " ssm_merge_ms=" << ssm_merge_seconds * 1000.0 + << " merge_ms=" << merge_seconds * 1000.0 << " other_ms=" + << (breakdown_total_timer.elapsed_seconds() - pull_wall_seconds - + merge_seconds) * + 1000.0 + << " total_ms=" << breakdown_total_timer.elapsed_seconds() * 1000.0; + return true; +} + +bool LlmDataDistTransfer::merge_pre_pushed_sharded_caches( + const LayerRegisteredCaches& layer_registered_caches, + const LayerRegisteredCaches& staging_registered_caches, + const std::vector& dst_blocks, + const std::vector& dst_linear_state_ids, + int64_t source_shard_count) { + if (source_shard_count != 2 || + layer_registered_caches.size() != staging_registered_caches.size()) { + LOG(ERROR) << "Invalid pre-pushed heterogeneous KV layout: source_tp=" + << source_shard_count; + return false; + } + + for (size_t layer_id = 0; layer_id < layer_registered_caches.size(); + ++layer_id) { + const auto& layer_caches = layer_registered_caches[layer_id]; + const auto& layer_staging_caches = staging_registered_caches[layer_id]; + if (layer_caches.size() != layer_staging_caches.size()) { + LOG(ERROR) << "Pre-pushed KV staging layout mismatch at layer " + << layer_id; + return false; + } + int64_t checkpoint_stride = 1; + for (const RegisteredCache& cache : layer_caches) { + if (cache.role != KVCacheTensorRole::CONV || !cache.tensor.defined()) { + continue; + } + for (const RegisteredCache& candidate : layer_caches) { + if (candidate.role == KVCacheTensorRole::SSM && + candidate.tensor.defined()) { + CHECK_EQ(candidate.tensor.size(0) % cache.tensor.size(0), 0); + checkpoint_stride = candidate.tensor.size(0) / cache.tensor.size(0); + break; + } + } + } + + for (size_t cache_index = 0; cache_index < layer_caches.size(); + ++cache_index) { + const RegisteredCache& registered_cache = layer_caches[cache_index]; + const RegisteredCache& stage_cache = layer_staging_caches[cache_index]; + const int64_t shard_dim = + sharded_dimension(registered_cache.role, registered_cache.tensor); + CHECK_GE(shard_dim, 0); + std::vector final_ids = + is_linear_state_cache(registered_cache.role) ? dst_linear_state_ids + : dst_blocks; + if (registered_cache.role == KVCacheTensorRole::SSM) { + final_ids = expand_checkpoint_ids(final_ids, checkpoint_stride); + } + if (final_ids.empty()) { + continue; + } + + CHECK_EQ(stage_cache.tensor.size(0) % source_shard_count, 0); + const int64_t rows_per_shard = + stage_cache.tensor.size(0) / source_shard_count; + // Staging is a compact per-request scratch buffer, not a mirror of + // Decode's block allocator. Both sides preserve request block order, so + // each shard occupies a contiguous ordinal range. + CHECK_LE(static_cast(final_ids.size()), rows_per_shard); + std::vector shard_tensors; + shard_tensors.reserve(source_shard_count); + for (int64_t shard = 0; shard < source_shard_count; ++shard) { + shard_tensors.push_back(stage_cache.tensor.narrow( + 0, shard * rows_per_shard, final_ids.size())); + } + + torch::Tensor merged; + if (registered_cache.role == KVCacheTensorRole::CONV) { + int64_t local_v_width = -1; + for (const RegisteredCache& candidate : layer_staging_caches) { + if (candidate.role == KVCacheTensorRole::SSM && + candidate.tensor.defined()) { + local_v_width = candidate.tensor.size(1) * candidate.tensor.size(3); + break; + } + } + CHECK_GT(local_v_width, 0); + const int64_t local_conv_width = stage_cache.tensor.size(shard_dim); + CHECK_EQ((local_conv_width - local_v_width) % 2, 0); + const int64_t local_qk_width = (local_conv_width - local_v_width) / 2; + std::vector q_shards; + std::vector k_shards; + std::vector v_shards; + for (const torch::Tensor& shard : shard_tensors) { + auto qkv = torch::split_with_sizes( + shard, + {local_qk_width, local_qk_width, local_v_width}, + shard_dim); + q_shards.push_back(qkv[0]); + k_shards.push_back(qkv[1]); + v_shards.push_back(qkv[2]); + } + merged = torch::cat({torch::cat(q_shards, shard_dim), + torch::cat(k_shards, shard_dim), + torch::cat(v_shards, shard_dim)}, + shard_dim); + } else { + merged = torch::cat(shard_tensors, shard_dim); + } + + std::vector signed_final_ids(final_ids.begin(), final_ids.end()); + torch::Tensor final_indices = + torch::tensor(signed_final_ids, + torch::TensorOptions().dtype(torch::kLong)) + .to(registered_cache.tensor.device(), /*non_blocking=*/false); + registered_cache.tensor.index_copy_(0, final_indices, merged); + } + } + return true; +} + +bool LlmDataDistTransfer::push_layer_registered_caches_to_staging( + const LayerRegisteredCaches& layer_registered_caches, + const LayerRegisteredCaches& staging_registered_caches, + std::unordered_map& merged_kv_infos, + std::shared_ptr& layer_synchronizer, + int64_t source_shard_rank, + int64_t source_shard_count) { + CHECK_GE(source_shard_rank, 0); + CHECK_LT(source_shard_rank, source_shard_count); + CHECK_EQ(layer_registered_caches.size(), staging_registered_caches.size()); + std::vector keys; + keys.reserve(merged_kv_infos.size()); + for (const auto& pair : merged_kv_infos) { + keys.push_back(pair.first); + } + std::sort(keys.begin(), keys.end()); + + Timer total_timer; + double layer_wait_seconds = 0.0; + double push_seconds = 0.0; + bool success = true; + for (size_t layer_id = 0; layer_id < layer_registered_caches.size(); + ++layer_id) { + VLOG(5) << "Heterogeneous staged push waiting for layer=" << layer_id + << ", source_shard=" << source_shard_rank; + Timer layer_wait_timer; + layer_synchronizer->synchronize_layer(layer_id); + layer_wait_seconds += layer_wait_timer.elapsed_seconds(); + VLOG(5) << "Heterogeneous staged push layer ready: layer=" << layer_id + << ", source_shard=" << source_shard_rank; + const auto& layer_caches = layer_registered_caches[layer_id]; + const auto& layer_staging_caches = staging_registered_caches[layer_id]; + CHECK_EQ(layer_caches.size(), layer_staging_caches.size()); + int64_t checkpoint_stride = 1; + for (const RegisteredCache& cache : layer_caches) { + if (cache.role != KVCacheTensorRole::CONV || !cache.tensor.defined()) { + continue; + } + for (const RegisteredCache& candidate : layer_caches) { + if (candidate.role == KVCacheTensorRole::SSM && + candidate.tensor.defined()) { + CHECK_EQ(candidate.tensor.size(0) % cache.tensor.size(0), 0); + checkpoint_stride = candidate.tensor.size(0) / cache.tensor.size(0); + break; + } + } + } + + for (const std::string& key : keys) { + const KVCacheInfo& kv_info = merged_kv_infos.at(key); + for (size_t cache_index = 0; cache_index < layer_caches.size(); + ++cache_index) { + const RegisteredCache& source_cache = layer_caches[cache_index]; + const RegisteredCache& stage_cache = layer_staging_caches[cache_index]; + // Decode restores CONV/SSM with a synchronous PULL because consuming + // their pre-pushed staging rows is not correct on the heterogeneous + // Qwen3.5 path. Avoid sending the same large recurrent state twice; + // heterogeneous staging PUSH is only useful for target KEY/VALUE. + if (is_linear_state_cache(source_cache.role)) { + continue; + } + std::vector src_ids = is_linear_state_cache(source_cache.role) + ? kv_info.src_linear_state_ids + : kv_info.src_blocks; + std::vector dst_ids = is_linear_state_cache(source_cache.role) + ? kv_info.dst_linear_state_ids + : kv_info.dst_blocks; + if (source_cache.role == KVCacheTensorRole::SSM) { + src_ids = expand_checkpoint_ids(src_ids, checkpoint_stride); + dst_ids = expand_checkpoint_ids(dst_ids, checkpoint_stride); + } + if (src_ids.empty() || dst_ids.empty()) { + continue; + } + CHECK_EQ(stage_cache.tensor.size(0) % source_shard_count, 0); + const int64_t rows_per_shard = + stage_cache.tensor.size(0) / source_shard_count; + CHECK_LE(static_cast(dst_ids.size()), rows_per_shard); + // Use compact request-local staging rows. dst_ids are Decode's real + // allocator ids and may exceed the bounded staging capacity. + for (size_t ordinal = 0; ordinal < dst_ids.size(); ++ordinal) { + dst_ids[ordinal] = static_cast( + source_shard_rank * rows_per_shard + ordinal); + } + CacheIndex destination{kv_info.dst_cluster_id, + stage_cache.cache.cache_id}; + KvCacheExtParam ext_param{}; + ext_param.src_layer_range = {0, 0}; + ext_param.dst_layer_range = {0, 0}; + ext_param.tensor_num_per_layer = 1; + VLOG(5) << "Heterogeneous staged push begin: layer=" << layer_id + << ", role=" << source_cache.role.to_string() + << ", source_shard=" << source_shard_rank + << ", source_cache_id=" << source_cache.cache.cache_id + << ", destination_cache_id=" << stage_cache.cache.cache_id; + Timer push_timer; + const auto ret = llm_data_dist_->PushKvBlocks( + source_cache.cache, destination, src_ids, dst_ids, ext_param); + push_seconds += push_timer.elapsed_seconds(); + VLOG(5) << "Heterogeneous staged push end: layer=" << layer_id + << ", role=" << source_cache.role.to_string() + << ", source_shard=" << source_shard_rank + << ", ret=" << std::hex << ret; + if (ret != LLM_SUCCESS) { + LOG(ERROR) << "Heterogeneous staged PushKvBlocks failed, layer=" + << layer_id << ", role=" << source_cache.role.to_string() + << ", source_shard=" << source_shard_rank + << ", destination_cache_id=" << stage_cache.cache.cache_id + << ", ret=" << std::hex << ret; + success = false; + } + } + } + } + LOG(INFO) << "[PD-PERF] Heterogeneous staging push source_shard=" + << source_shard_rank << ", request_count=" << keys.size() + << ", layer_wait_ms=" << layer_wait_seconds * 1000.0 + << ", push_ms=" << push_seconds * 1000.0 + << ", total_ms=" << total_timer.elapsed_seconds() * 1000.0; + return success; +} + +void LlmDataDistTransfer::register_hetero_staging_caches( + const LayerRegisteredCaches& source_registered_caches, + LayerRegisteredCaches& staging_registered_caches, + int64_t source_shard_count, + bool source_is_sharded) { + CHECK_EQ(source_shard_count, 2) + << "Only Prefill TP2 to Decode TP1 staging is supported."; + staging_registered_caches.clear(); + staging_registered_caches.resize(source_registered_caches.size()); + for (size_t layer_id = 0; layer_id < source_registered_caches.size(); + ++layer_id) { + for (const RegisteredCache& source_cache : + source_registered_caches[layer_id]) { + std::vector shape = source_cache.tensor.sizes().vec(); + const int64_t shard_dim = + sharded_dimension(source_cache.role, source_cache.tensor); + CHECK_GE(shard_dim, 0); + if (!source_is_sharded) { + CHECK_EQ(shape[shard_dim] % source_shard_count, 0); + shape[shard_dim] /= source_shard_count; + } + if (source_cache.role == KVCacheTensorRole::KEY || + source_cache.role == KVCacheTensorRole::VALUE) { + // max_tokens_per_batch=32768 with block_size=128. + shape[0] = std::min(shape[0], 256) * source_shard_count; + } else if (source_cache.role == KVCacheTensorRole::SSM) { + shape[0] = std::min(shape[0], 4) * source_shard_count; + } else { + shape[0] = source_shard_count; + } + torch::Tensor stage_tensor = + make_page_aligned_staging_tensor(shape, + source_cache.tensor.options(), + source_cache.tensor.element_size()); + staging_registered_caches[layer_id].push_back( + register_cache_tensor(layer_id, + KVCacheTensor{source_cache.role, + stage_tensor, + source_cache.group_id, + source_cache.sequence_scoped})); + } + } +} + void LlmDataDistTransfer::register_layer_registered_caches( std::vector& kv_caches, LayerRegisteredCaches& layer_registered_caches) { diff --git a/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.h b/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.h index fd4a5d13a0..2b54c18dff 100644 --- a/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.h +++ b/xllm/core/framework/kv_cache_transfer/llm_data_dist_transfer.h @@ -17,6 +17,8 @@ limitations under the License. #include +#include + #include "framework/kv_cache_transfer/kv_cache_transfer.h" namespace xllm { @@ -28,6 +30,9 @@ struct RegisteredCache { int32_t group_id; bool sequence_scoped; Cache cache; + // Keep the tensor owner and layout available for heterogeneous-TP staging + // and decode-side merge. Cache only stores the raw address. + torch::Tensor tensor; }; using LayerRegisteredCaches = std::vector>; @@ -94,6 +99,36 @@ class LlmDataDistTransfer : public KVCacheTransfer { int32_t kv_split_rank = 0, int32_t kv_split_size = 1); + bool pull_and_merge_sharded_caches( + const LayerRegisteredCaches& layer_registered_caches, + const LayerRegisteredCaches& staging_registered_caches, + const std::vector& src_cluster_ids, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids); + + bool merge_pre_pushed_sharded_caches( + const LayerRegisteredCaches& layer_registered_caches, + const LayerRegisteredCaches& staging_registered_caches, + const std::vector& dst_blocks, + const std::vector& dst_linear_state_ids, + int64_t source_shard_count); + + bool push_layer_registered_caches_to_staging( + const LayerRegisteredCaches& layer_registered_caches, + const LayerRegisteredCaches& staging_registered_caches, + std::unordered_map& merged_kv_infos, + std::shared_ptr& layer_synchronizer, + int64_t source_shard_rank, + int64_t source_shard_count); + + void register_hetero_staging_caches( + const LayerRegisteredCaches& source_registered_caches, + LayerRegisteredCaches& staging_registered_caches, + int64_t source_shard_count = 2, + bool source_is_sharded = false); + protected: uint64_t cluster_id_; std::string host_ip_; @@ -106,6 +141,13 @@ class LlmDataDistTransfer : public KVCacheTransfer { std::shared_ptr llm_data_dist_; LayerRegisteredCaches layer_registered_caches_; + // Heterogeneous requests share staging caches, so only one request may + // restore and merge them at a time. The two source shards within that + // request are still pulled concurrently. + std::mutex hetero_pull_mutex_; + ThreadPool shard_pull_threadpool_{/*num_threads=*/1, + /*cpu_binding=*/false, + /*pool_name=*/"KVCacheTransfer.shard_pull"}; }; } // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/pd_topology_guard.cpp b/xllm/core/framework/kv_cache_transfer/pd_topology_guard.cpp index fd5e2f1f69..6248b0399b 100644 --- a/xllm/core/framework/kv_cache_transfer/pd_topology_guard.cpp +++ b/xllm/core/framework/kv_cache_transfer/pd_topology_guard.cpp @@ -31,18 +31,31 @@ bool fail_topo(const std::string& msg, std::string* reason) { return false; } -PdTopoResult check_hetero_pd_req(const std::string& kv_mode, bool enable_mla) { +PdTopoResult check_hetero_pd_req(const PdTopo& prefill_topo, + const PdTopo& decode_topo, + const std::string& kv_mode, + bool enable_mla) { if (kv_mode != "PUSH") { return PdTopoResult{PdTopoStatus::DENY_HETERO, "hetero pd requires kv_mode=PUSH"}; } - // Non-MLA KV cache still shards KV heads by TP. Hetero TP needs separate - // head-dimension split/merge support, so this path is limited to MLA. - if (!enable_mla) { + + if (enable_mla) { + return PdTopoResult{PdTopoStatus::ALLOW_HETERO, ""}; + } + + if (prefill_topo.dp_size != decode_topo.dp_size) { return PdTopoResult{PdTopoStatus::DENY_HETERO, - "hetero pd requires enable_mla=true"}; + "non-mla hetero pd requires equal dp_size"}; } + if (prefill_topo.tp_size < decode_topo.tp_size || + prefill_topo.tp_size % decode_topo.tp_size != 0) { + return PdTopoResult{ + PdTopoStatus::DENY_HETERO, + "non-mla hetero pd requires prefill tp_size divisible by decode " + "tp_size"}; + } return PdTopoResult{PdTopoStatus::ALLOW_HETERO, ""}; } @@ -109,7 +122,7 @@ PdTopoResult check_pd_topo(const InstanceInfo& local, return PdTopoResult{PdTopoStatus::ALLOW_HOMO, ""}; } - return check_hetero_pd_req(kv_mode, enable_mla); + return check_hetero_pd_req(local_topo, remote_topo, kv_mode, enable_mla); } } // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/push_route.cpp b/xllm/core/framework/kv_cache_transfer/push_route.cpp index 7215eeff53..3ecf501144 100644 --- a/xllm/core/framework/kv_cache_transfer/push_route.cpp +++ b/xllm/core/framework/kv_cache_transfer/push_route.cpp @@ -51,4 +51,29 @@ std::vector get_dst_ranks(const int32_t src_tp_rank, return dst_ranks; } +std::vector get_src_tp_ranks(const int32_t dst_tp_rank, + const int32_t src_tp_size, + const int32_t dst_tp_size) { + std::vector src_tp_ranks; + if (src_tp_size <= 0 || dst_tp_size <= 0 || dst_tp_rank < 0 || + dst_tp_rank >= dst_tp_size) { + return src_tp_ranks; + } + + if (src_tp_size <= dst_tp_size) { + src_tp_ranks.reserve(1); + src_tp_ranks.emplace_back(dst_tp_rank % src_tp_size); + return src_tp_ranks; + } + + const size_t src_rank_num = + static_cast((src_tp_size - 1 - dst_tp_rank) / dst_tp_size + 1); + src_tp_ranks.reserve(src_rank_num); + for (int32_t src_tp_rank = dst_tp_rank; src_tp_rank < src_tp_size; + src_tp_rank += dst_tp_size) { + src_tp_ranks.emplace_back(src_tp_rank); + } + return src_tp_ranks; +} + } // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/push_route.h b/xllm/core/framework/kv_cache_transfer/push_route.h index c04bbd1f32..2eda51d352 100644 --- a/xllm/core/framework/kv_cache_transfer/push_route.h +++ b/xllm/core/framework/kv_cache_transfer/push_route.h @@ -27,4 +27,8 @@ std::vector get_dst_ranks(int32_t src_tp_rank, int32_t dst_tp_size, int32_t dst_dp_rank); +std::vector get_src_tp_ranks(int32_t dst_tp_rank, + int32_t src_tp_size, + int32_t dst_tp_size); + } // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.cpp b/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.cpp index 58f72f96dc..9433a59d29 100644 --- a/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.cpp +++ b/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.cpp @@ -17,13 +17,86 @@ limitations under the License. #include +#include + #include "common/macros.h" +#include "util/timer.h" namespace xllm { + +namespace { + +std::optional get_remote_tp_size( + const std::vector& transfer_kv_infos) { + for (const TransferKVInfo& info : transfer_kv_infos) { + const int32_t remote_dp_size = info.remote_instance_info.dp_size; + const size_t remote_world_size = + info.remote_instance_info.cluster_ids.size(); + if (remote_dp_size <= 0 || remote_world_size == 0 || + remote_world_size % static_cast(remote_dp_size) != 0) { + continue; + } + return static_cast(remote_world_size / + static_cast(remote_dp_size)); + } + return std::nullopt; +} + +void merge_heterogeneous_kv_blocks( + std::unordered_map& merged, + const std::vector& transfer_kv_infos, + int32_t source_shard_rank) { + for (const TransferKVInfo& info : transfer_kv_infos) { + const int32_t dst_dp_size = info.remote_instance_info.dp_size; + const size_t dst_world_size = info.remote_instance_info.cluster_ids.size(); + CHECK_GT(dst_dp_size, 0); + CHECK_EQ(dst_world_size % static_cast(dst_dp_size), 0); + const int32_t dst_tp_size = + static_cast(dst_world_size / dst_dp_size); + CHECK_GT(dst_tp_size, 0); + CHECK_GE(info.dp_rank, 0); + CHECK_LT(info.dp_rank, dst_dp_size); + + // Every Prefill head shard must reach a Decode worker. The generic + // merge_kv_blocks() route is one-to-one in TP rank and drops source ranks + // whose rank is >= dst_tp_size (for TP2 -> TP1, rank 1 is dropped). Map + // source shards modulo the destination TP group instead; staging row + // ranges keep multiple source shards disjoint on the selected worker. + const int32_t dst_rank = + info.dp_rank * dst_tp_size + source_shard_rank % dst_tp_size; + CHECK_LT(static_cast(dst_rank), dst_world_size); + const uint64_t dst_cluster_id = + info.remote_instance_info.cluster_ids[dst_rank]; + const std::string& dst_addr = info.remote_instance_info.addrs[dst_rank]; + const std::string key = std::to_string(dst_cluster_id) + "_" + dst_addr; + auto& kv_info = merged[key]; + kv_info.dst_cluster_id = dst_cluster_id; + kv_info.dst_addr = dst_addr; + kv_info.src_blocks.insert(kv_info.src_blocks.end(), + info.local_blocks_ids.begin(), + info.local_blocks_ids.end()); + kv_info.dst_blocks.insert(kv_info.dst_blocks.end(), + info.remote_blocks_ids.begin(), + info.remote_blocks_ids.end()); + kv_info.src_linear_state_ids.insert(kv_info.src_linear_state_ids.end(), + info.local_linear_state_ids.begin(), + info.local_linear_state_ids.end()); + kv_info.dst_linear_state_ids.insert(kv_info.dst_linear_state_ids.end(), + info.remote_linear_state_ids.begin(), + info.remote_linear_state_ids.end()); + } +} + +} // namespace + SpecKVCacheTransfer::SpecKVCacheTransfer(const uint16_t listen_port, const InstanceRole& instance_role, - bool enable_lighting_indexer) + bool enable_lighting_indexer, + bool enable_mla, + bool draft_body_uses_tp1) : LlmDataDistTransfer(listen_port, instance_role, enable_lighting_indexer) { + enable_mla_ = enable_mla; + draft_body_uses_tp1_ = draft_body_uses_tp1; } void SpecKVCacheTransfer::register_kv_cache( @@ -42,6 +115,46 @@ void SpecKVCacheTransfer::register_kv_cache_spec( UNUSED_PARAMETER(kv_cache_shape); UNUSED_PARAMETER(dtype); register_kv_cache_internal(kv_caches, spec_layer_registered_caches_); + // Register matching staging cache IDs on both Prefill and Decode before the + // first DataDist link. Prefill pushes each local TP shard into a disjoint row + // range; Decode merges those already-local rows into its TP1 cache. + const bool source_is_sharded = role_ == LlmRole::kPrompt; + register_hetero_staging_caches(layer_registered_caches_, + hetero_staging_registered_caches_, + /*source_shard_count=*/2, + source_is_sharded); + register_hetero_staging_caches(spec_layer_registered_caches_, + spec_hetero_staging_registered_caches_, + /*source_shard_count=*/2, + source_is_sharded && !draft_body_uses_tp1_); +} + +bool SpecKVCacheTransfer::pull_replicated_spec_kv_blocks( + uint64_t src_cluster_id, + const std::vector& src_blocks, + const std::vector& dst_blocks) { + CHECK_EQ(src_blocks.size(), dst_blocks.size()); + bool success = true; + for (size_t layer_id = 0; layer_id < spec_layer_registered_caches_.size(); + ++layer_id) { + for (const RegisteredCache& cache : + spec_layer_registered_caches_[layer_id]) { + CacheIndex source{src_cluster_id, cache.cache.cache_id}; + KvCacheExtParam ext_param{}; + ext_param.src_layer_range = {0, 0}; + ext_param.dst_layer_range = {0, 0}; + ext_param.tensor_num_per_layer = 1; + const auto ret = llm_data_dist_->PullKvBlocks( + source, cache.cache, src_blocks, dst_blocks, ext_param); + if (ret != LLM_SUCCESS) { + LOG(ERROR) << "Pull replicated TP1 draft KV failed, layer=" << layer_id + << ", role=" << cache.role.to_string() + << ", ret=" << std::hex << ret; + success = false; + } + } + } + return success; } void SpecKVCacheTransfer::register_kv_cache_internal( @@ -54,6 +167,8 @@ void SpecKVCacheTransfer::free_kv_cache() { layer_registered_caches_.clear(); spec_layer_registered_caches_.clear(); has_grouped_cache_layout_ = false; + hetero_staging_registered_caches_.clear(); + spec_hetero_staging_registered_caches_.clear(); } bool SpecKVCacheTransfer::pull_kv_blocks( @@ -109,6 +224,73 @@ bool SpecKVCacheTransfer::pull_kv_blocks( return base_success && spec_success; } +bool SpecKVCacheTransfer::pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + (void)src_addrs; + Timer phase_timer; + // DataDist PUSH does not expose a Decode-side completion primitive for the + // large recurrent state tensors. Keep CONV/SSM on the established + // synchronous PULL path while KEY/VALUE are pre-pushed into staging. + const bool linear_success = + pull_and_merge_sharded_caches(layer_registered_caches_, + hetero_staging_registered_caches_, + src_cluster_ids, + /*src_blocks=*/{}, + /*dst_blocks=*/{}, + src_linear_state_ids, + dst_linear_state_ids); + if (!linear_success) { + return false; + } + const double linear_seconds = phase_timer.elapsed_seconds(); + phase_timer.reset(); + const bool target_success = + merge_pre_pushed_sharded_caches(layer_registered_caches_, + hetero_staging_registered_caches_, + dst_blocks, + /*dst_linear_state_ids=*/{}, + /*source_shard_count=*/2); + if (!target_success) { + return false; + } + const double target_merge_seconds = phase_timer.elapsed_seconds(); + phase_timer.reset(); + // Keep the one-layer MTP draft cache on the established synchronous pull + // path while validating the new layer-overlapped target-cache push. The + // draft pull is small (~1 ms), and this isolates whether its newly-added + // layer event observes the cache before the MTP prefill write is complete. + const bool draft_success = + draft_body_uses_tp1_ + ? pull_replicated_spec_kv_blocks( + src_cluster_ids.front(), src_blocks, dst_blocks) + : pull_and_merge_sharded_caches( + spec_layer_registered_caches_, + spec_hetero_staging_registered_caches_, + src_cluster_ids, + src_blocks, + dst_blocks, + /*src_linear_state_ids=*/{}, + /*dst_linear_state_ids=*/{}); + if (draft_success) { + const double draft_seconds = phase_timer.elapsed_seconds(); + LOG(INFO) << "Merged heterogeneous TP KV cache (target KV pre-pushed, " + "linear state and draft pulled): source_shards=" + << 2 << ", blocks=" << dst_blocks.size() + << ", linear_states=" << dst_linear_state_ids.size() + << ", linear_ms=" << linear_seconds * 1000.0 + << ", target_merge_ms=" << target_merge_seconds * 1000.0 + << ", draft_ms=" << draft_seconds * 1000.0 << ", total_ms=" + << (linear_seconds + target_merge_seconds + draft_seconds) * + 1000.0; + } + return draft_success; +} + bool SpecKVCacheTransfer::push_kv_blocks( std::unordered_map& merged_kv_infos, std::shared_ptr& layer_synchronizer, @@ -152,18 +334,73 @@ bool SpecKVCacheTransfer::push_kv_blocks_internal( kv_split_size); } +bool SpecKVCacheTransfer::push_kv_blocks_to_hetero_staging( + std::unordered_map& merged_kv_infos, + std::shared_ptr& layer_synchronizer, + bool is_spec_draft, + int64_t source_shard_rank, + int64_t source_shard_count) { + const LayerRegisteredCaches& source_caches = + is_spec_draft ? spec_layer_registered_caches_ : layer_registered_caches_; + const LayerRegisteredCaches& staging_caches = + is_spec_draft ? spec_hetero_staging_registered_caches_ + : hetero_staging_registered_caches_; + return push_layer_registered_caches_to_staging(source_caches, + staging_caches, + merged_kv_infos, + layer_synchronizer, + source_shard_rank, + source_shard_count); +} + folly::SemiFuture SpecKVCacheTransfer::push_kv_blocks_async( const std::vector& transfer_kv_infos, const ParallelArgs& parallel_args, std::shared_ptr layer_synchronizer, bool is_spec_draft) { + const int32_t local_dp_size = parallel_args.dp_size(); + const int32_t kv_split_size = parallel_args.kv_split_size_effective(); + const std::optional remote_tp_size = + get_remote_tp_size(transfer_kv_infos); + bool heterogeneous_non_mla = false; + int32_t local_tp_size = 1; + if (!enable_mla_ && local_dp_size > 0 && kv_split_size > 0 && + remote_tp_size.has_value()) { + local_tp_size = parallel_args.world_size() / local_dp_size / kv_split_size; + if (local_tp_size != remote_tp_size.value()) { + heterogeneous_non_mla = true; + LOG(INFO) << "Push non-MLA heterogeneous KV shards to decode staging: " + << "prefill_tp_size=" << local_tp_size + << ", decode_tp_size=" << remote_tp_size.value() + << ", is_spec_draft=" << is_spec_draft + << "; decode will only perform a local merge."; + } + } + const int64_t source_shard_rank = + heterogeneous_non_mla ? parallel_args.rank() % local_tp_size : 0; + folly::Promise promise; auto future = promise.getSemiFuture(); + // In heterogeneous non-MLA mode Decode intentionally restores the draft + // cache from the source shards with a synchronous PULL. Pushing the same + // one-layer draft cache into staging is therefore redundant: no Decode + // path consumes spec_hetero_staging_registered_caches_. The source cache + // remains alive until the synchronous FirstGeneration RPC returns, so + // skipping this PUSH does not shorten its lifetime for the later PULL. + if (heterogeneous_non_mla && is_spec_draft) { + VLOG(5) << "Skip redundant heterogeneous MTP draft staging PUSH; " + "Decode restores draft KV from source shards."; + promise.setValue(true); + return future; + } threadpool_.schedule([this, transfer_kv_infos, ¶llel_args, layer_synchronizer, is_spec_draft, + heterogeneous_non_mla, + local_tp_size, + source_shard_rank, promise = std::move(promise)]() mutable { std::unordered_map merged_kv_infos; std::vector filtered_kv_infos; @@ -172,24 +409,38 @@ folly::SemiFuture SpecKVCacheTransfer::push_kv_blocks_async( // (kv_split_size_effective > 1), filter remote_blocks_ids down to this // rank's slice. When kv_split_size==1 each rank holds the full replica and // we keep the legacy 1:1 remote_blocks_ids mapping. - const int32_t kv_split_size = parallel_args.kv_split_size_effective(); - if (kv_split_size > 1) { + const int32_t effective_kv_split_size = + parallel_args.kv_split_size_effective(); + if (effective_kv_split_size > 1) { filtered_kv_infos = filter_kv_split_infos( - parallel_args.kv_split_rank(), kv_split_size, *kv_infos); + parallel_args.kv_split_rank(), effective_kv_split_size, *kv_infos); kv_infos = &filtered_kv_infos; if (kv_infos->empty()) { promise.setValue(true); return; } } - merge_kv_blocks(merged_kv_infos, *kv_infos, parallel_args); + if (heterogeneous_non_mla) { + merge_heterogeneous_kv_blocks( + merged_kv_infos, *kv_infos, source_shard_rank); + } else { + merge_kv_blocks(merged_kv_infos, *kv_infos, parallel_args); + } bool success = true; if (!merged_kv_infos.empty()) { - success = this->push_kv_blocks(merged_kv_infos, - layer_synchronizer, - is_spec_draft, - parallel_args.kv_split_rank(), - parallel_args.kv_split_size_effective()); + if (heterogeneous_non_mla) { + success = this->push_kv_blocks_to_hetero_staging(merged_kv_infos, + layer_synchronizer, + is_spec_draft, + source_shard_rank, + local_tp_size); + } else { + success = this->push_kv_blocks(merged_kv_infos, + layer_synchronizer, + is_spec_draft, + parallel_args.kv_split_rank(), + parallel_args.kv_split_size_effective()); + } } promise.setValue(success); }); diff --git a/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.h b/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.h index e2fc963abd..2f3e4c0fad 100644 --- a/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.h +++ b/xllm/core/framework/kv_cache_transfer/spec_kv_cache_transfer.h @@ -26,7 +26,9 @@ class SpecKVCacheTransfer : public LlmDataDistTransfer { public: SpecKVCacheTransfer(const uint16_t listen_port, const InstanceRole& instance_role, - bool enable_lighting_indexer = false); + bool enable_lighting_indexer = false, + bool enable_mla = false, + bool draft_body_uses_tp1 = false); virtual ~SpecKVCacheTransfer() = default; @@ -52,6 +54,14 @@ class SpecKVCacheTransfer : public LlmDataDistTransfer { const std::vector& src_linear_state_ids, const std::vector& dst_linear_state_ids) override; + bool pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) override; + folly::SemiFuture push_kv_blocks_async( const std::vector& transfer_kv_infos, const ParallelArgs& parallel_args, @@ -78,7 +88,21 @@ class SpecKVCacheTransfer : public LlmDataDistTransfer { int32_t kv_split_rank = 0, int32_t kv_split_size = 1); + bool push_kv_blocks_to_hetero_staging( + std::unordered_map& merged_kv_infos, + std::shared_ptr& layer_synchronizer, + bool is_spec_draft, + int64_t source_shard_rank, + int64_t source_shard_count); + private: + bool pull_replicated_spec_kv_blocks(uint64_t src_cluster_id, + const std::vector& src_blocks, + const std::vector& dst_blocks); + + bool draft_body_uses_tp1_ = false; + LayerRegisteredCaches hetero_staging_registered_caches_; + LayerRegisteredCaches spec_hetero_staging_registered_caches_; LayerRegisteredCaches spec_layer_registered_caches_; }; diff --git a/xllm/core/runtime/dflash_worker_impl.cpp b/xllm/core/runtime/dflash_worker_impl.cpp index a0f3d10e84..7168a63552 100644 --- a/xllm/core/runtime/dflash_worker_impl.cpp +++ b/xllm/core/runtime/dflash_worker_impl.cpp @@ -406,7 +406,8 @@ bool DFlashWorkerImpl::allocate_kv_cache_with_transfer( kv_cache_transfer_ = std::make_shared( options_.transfer_listen_port(), options_.instance_role(), - context_.get_model_args().index_n_heads() > 0); + context_.get_model_args().index_n_heads() > 0, + context_.get_model_args().enable_mla()); #elif defined(USE_MLU) CHECK_EQ(::xllm::DisaggPDConfig::get_instance().kv_cache_transfer_type(), "Mooncake") diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index 47bf3786a4..9d365d8cac 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -796,7 +796,9 @@ bool MTPWorkerImpl::allocate_kv_cache_with_transfer( kv_cache_transfer_ = std::make_shared( options_.transfer_listen_port(), options_.instance_role(), - context_.get_model_args().index_n_heads() > 0); + context_.get_model_args().index_n_heads() > 0, + context_.get_model_args().enable_mla(), + options_.enable_mtp_draft_body_tp1()); #elif defined(USE_MLU) CHECK_EQ(::xllm::DisaggPDConfig::get_instance().kv_cache_transfer_type(), "Mooncake") diff --git a/xllm/core/runtime/speculative_worker_impl.h b/xllm/core/runtime/speculative_worker_impl.h index 7ddac65593..ecde027539 100644 --- a/xllm/core/runtime/speculative_worker_impl.h +++ b/xllm/core/runtime/speculative_worker_impl.h @@ -111,6 +111,21 @@ class SpeculativeWorkerImpl : public WorkerImpl { dst_linear_state_ids); }; + folly::SemiFuture pull_hetero_kv_blocks_async( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}) override { + return impl_->pull_hetero_kv_blocks_async(src_cluster_ids, + src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); + }; + protected: // Algorithm-specific virtual methods for subclasses to implement virtual std::optional step_prefill( diff --git a/xllm/core/runtime/worker.cpp b/xllm/core/runtime/worker.cpp index f5c6c41b75..d5d24833d5 100644 --- a/xllm/core/runtime/worker.cpp +++ b/xllm/core/runtime/worker.cpp @@ -179,6 +179,21 @@ folly::SemiFuture Worker::pull_kv_blocks_async( dst_linear_state_ids); } +folly::SemiFuture Worker::pull_hetero_kv_blocks_async( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + return impl_->pull_hetero_kv_blocks_async(src_cluster_ids, + src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); +} + uint32_t Worker::transfer_kv_blocks( const uint64_t batch_id, const std::vector& block_transfer_info) { diff --git a/xllm/core/runtime/worker.h b/xllm/core/runtime/worker.h index 5b7e30fc40..88a9b13a07 100644 --- a/xllm/core/runtime/worker.h +++ b/xllm/core/runtime/worker.h @@ -116,6 +116,14 @@ class Worker { const std::vector& src_linear_state_ids = {}, const std::vector& dst_linear_state_ids = {}); + virtual folly::SemiFuture pull_hetero_kv_blocks_async( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}); + virtual uint32_t transfer_kv_blocks( const uint64_t batch_id, const std::vector& block_transfer_info); diff --git a/xllm/core/runtime/worker_client.cpp b/xllm/core/runtime/worker_client.cpp index 0759598e4e..3acaaa3234 100644 --- a/xllm/core/runtime/worker_client.cpp +++ b/xllm/core/runtime/worker_client.cpp @@ -88,6 +88,22 @@ bool WorkerClient::pull_kv_blocks( return std::move(future).get(); } +bool WorkerClient::pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + auto future = worker_->pull_hetero_kv_blocks_async(src_cluster_ids, + src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); + return std::move(future).get(); +} + ForwardInput WorkerClient::prepare_inputs(Batch& batch) { return worker_->prepare_inputs(batch); } diff --git a/xllm/core/runtime/worker_client.h b/xllm/core/runtime/worker_client.h index 8186527ccc..0255f36409 100644 --- a/xllm/core/runtime/worker_client.h +++ b/xllm/core/runtime/worker_client.h @@ -86,6 +86,14 @@ class WorkerClient { const std::vector& src_linear_state_ids = {}, const std::vector& dst_linear_state_ids = {}); + virtual bool pull_hetero_kv_blocks( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}); + // prepare input for execution virtual ForwardInput prepare_inputs(Batch& batch); diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 203e62e6ea..11540506f6 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -1763,6 +1763,55 @@ folly::SemiFuture WorkerImpl::pull_kv_blocks_async( return false; } +folly::SemiFuture WorkerImpl::pull_hetero_kv_blocks_async( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids, + const std::vector& dst_linear_state_ids) { + folly::Promise promise; + auto future = promise.getSemiFuture(); +#if defined(USE_NPU) + threadpool_.schedule([this, + src_cluster_ids, + src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids, + promise = std::move(promise)]() mutable { + const bool success = + kv_cache_transfer_->pull_hetero_kv_blocks(src_cluster_ids, + src_addrs, + src_blocks, + dst_blocks, + src_linear_state_ids, + dst_linear_state_ids); + if (success) { + const int ret = device_.synchronize_default_stream(); + if (ret != 0) { + LOG(ERROR) << "synchronize_default_stream after heterogeneous KV " + "merge failed, ret=" + << ret; + promise.setValue(false); + return; + } + } + promise.setValue(success); + }); +#else + (void)src_cluster_ids; + (void)src_addrs; + (void)src_blocks; + (void)dst_blocks; + (void)src_linear_state_ids; + (void)dst_linear_state_ids; + promise.setValue(false); +#endif + return future; +} + uint32_t WorkerImpl::transfer_kv_blocks( const uint64_t batch_id, const std::vector& block_transfer_info) { diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index 62f8b97266..309f6b6db4 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -162,6 +162,14 @@ class WorkerImpl { const std::vector& src_linear_state_ids = {}, const std::vector& dst_linear_state_ids = {}); + virtual folly::SemiFuture pull_hetero_kv_blocks_async( + const std::vector& src_cluster_ids, + const std::vector& src_addrs, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& src_linear_state_ids = {}, + const std::vector& dst_linear_state_ids = {}); + virtual uint32_t transfer_kv_blocks( const uint64_t batch_id, const std::vector& block_transfer_info); diff --git a/xllm/core/scheduler/async_response_processor.cpp b/xllm/core/scheduler/async_response_processor.cpp index 635006fc2c..f09344e6ac 100644 --- a/xllm/core/scheduler/async_response_processor.cpp +++ b/xllm/core/scheduler/async_response_processor.cpp @@ -289,6 +289,7 @@ void AsyncResponseProcessor::batch_process_stream_requests( num_tokens = std::move(num_tokens), req_output = &request_outputs[i]]() mutable { AUTO_COUNTER(responsing_latency_seconds_stream); + const absl::Time response_start_time = absl::Now(); // RequestOutput req_output; req_output->request_id = request->request_id(); @@ -310,6 +311,13 @@ void AsyncResponseProcessor::batch_process_stream_requests( req_output->finished_on_prefill_instance = true; } } + if (req_output->finished_on_prefill_instance) { + LOG(INFO) << "[PD-PERF] Prefill response generation request_id=" + << request->request_id() << ", response_thread_id=" + << request->state().response_thread_id << ", total_ms=" + << absl::ToDoubleMilliseconds(absl::Now() - + response_start_time); + } counter->decrement_count(); }; if (request->state().response_thread_id < 0) { @@ -327,8 +335,19 @@ void AsyncResponseProcessor::batch_process_stream_requests( requests = std::move(requests), request_outputs = std::move(request_outputs)]() mutable { auto& resp_callback = requests[0]->state().outputs_func; + const absl::Time wait_start_time = absl::Now(); counter->wait(); + const double wait_ms = + absl::ToDoubleMilliseconds(absl::Now() - wait_start_time); + const absl::Time rpc_start_time = absl::Now(); std::vector status_set = resp_callback(request_outputs); + if (!request_outputs.empty() && + request_outputs[0].finished_on_prefill_instance) { + LOG(INFO) << "[PD-PERF] Prefill response RPC request_id=" + << requests[0]->request_id() + << ", response_wait_ms=" << wait_ms << ", rpc_ms=" + << absl::ToDoubleMilliseconds(absl::Now() - rpc_start_time); + } for (size_t i = 0; i < requests.size(); ++i) { if (!status_set[i]) { cancel_request(requests[i]); diff --git a/xllm/core/scheduler/disagg_pd_scheduler.cpp b/xllm/core/scheduler/disagg_pd_scheduler.cpp index ba3cad45d6..85df2e32d7 100644 --- a/xllm/core/scheduler/disagg_pd_scheduler.cpp +++ b/xllm/core/scheduler/disagg_pd_scheduler.cpp @@ -43,6 +43,7 @@ limitations under the License. #include "scheduler/chunked_prefill_scheduler.h" #include "scheduler/continuous_scheduler.h" #include "util/env_var.h" +#include "util/timer.h" #include "util/utils.h" namespace xllm { @@ -441,6 +442,7 @@ void DisaggPDScheduler::dispatch_requests() { // TODO: send the request to the selected D instance // Send 'DisaggRequests' and recv 'DisaggResponses' + Timer decode_allocation_timer; xllm::proto::DisaggRequests reqs; xllm::proto::DisaggResponses resps; // prefill name (ID) @@ -509,10 +511,15 @@ void DisaggPDScheduler::dispatch_requests() { reqs.mutable_cluster_infos()->mutable_ports()->Add( instance_info_.ports.begin(), instance_info_.ports.end()); reqs.mutable_cluster_infos()->set_dp_size(options_.dp_size()); + const double allocation_prepare_seconds = + decode_allocation_timer.elapsed_seconds(); // TODO: sync rpc here currently brpc::Controller cntl; + Timer allocation_rpc_timer; stub->AddNewRequests(&cntl, &reqs, &resps, nullptr); + const double allocation_rpc_seconds = + allocation_rpc_timer.elapsed_seconds(); if (cntl.Failed()) { LOG(ERROR) << "Failed to add new requests to decode instance : " << selected_instance << ", error text : " << cntl.ErrorText(); @@ -609,6 +616,12 @@ void DisaggPDScheduler::dispatch_requests() { request_queue_.write(requests[i]); } } + LOG(INFO) << "[PD-PERF] Prefill Decode allocation request_id=" + << requests.front()->request_id() + << ", request_count=" << requests.size() + << ", prepare_ms=" << allocation_prepare_seconds * 1000.0 + << ", rpc_ms=" << allocation_rpc_seconds * 1000.0 << ", total_ms=" + << decode_allocation_timer.elapsed_seconds() * 1000.0; } } @@ -660,6 +673,7 @@ void DisaggPDScheduler::prefill_send_first_generation() { kv_cache_manager_->deallocate(request.get()); }; for (auto& request : requests) { + Timer first_generation_timer; // TODO: support batch request later proto::DisaggGenerationsRequests gens; auto gen = gens.mutable_multi_gens()->Add(); @@ -692,24 +706,23 @@ void DisaggPDScheduler::prefill_send_first_generation() { request->sequences()[0]->first_token().value().token_top_logprobs); } gen->set_kv_cache_transfer_mode(options_.kv_cache_transfer_mode()); - if (options_.kv_cache_transfer_mode() == "PULL") { - ADD_VECTOR_TO_PROTO(gen->mutable_cluster_ids(), - instance_info_.cluster_ids); - ADD_VECTOR_TO_PROTO(gen->mutable_addrs(), instance_info_.addrs); - - const auto blocks = - request->sequences()[0]->kv_state().blocks(BlockType::KV); - std::vector block_ids; - block_ids.reserve(blocks.size()); - for (const auto& block : blocks) { - block_ids.push_back(block.id()); - } - ADD_VECTOR_TO_PROTO(gen->mutable_block_ids(), block_ids); - gen->set_linear_state_id( - request->sequences()[0]->get_single_block_id()); - gen->set_dp_size(instance_info_.dp_size); - gen->set_dp_rank(request->sequences()[0]->dp_rank()); + // Heterogeneous TP uses decode-side PULL+merge even when the configured + // transport mode is PUSH, so source topology and cache ids are always + // included in the first-generation metadata. + ADD_VECTOR_TO_PROTO(gen->mutable_cluster_ids(), + instance_info_.cluster_ids); + ADD_VECTOR_TO_PROTO(gen->mutable_addrs(), instance_info_.addrs); + gen->set_dp_size(instance_info_.dp_size); + gen->set_dp_rank(request->sequences()[0]->dp_rank()); + const auto blocks = + request->sequences()[0]->kv_state().blocks(BlockType::KV); + std::vector block_ids; + block_ids.reserve(blocks.size()); + for (const auto& block : blocks) { + block_ids.push_back(block.id()); } + ADD_VECTOR_TO_PROTO(gen->mutable_block_ids(), block_ids); + gen->set_linear_state_id(request->sequences()[0]->get_single_block_id()); if (options_.num_speculative_tokens() > 0) { torch::Tensor embedding = request->sequences()[0]->get_mtp_bootstrap_embedding(); @@ -733,6 +746,7 @@ void DisaggPDScheduler::prefill_send_first_generation() { } request->sequences()[0]->clear_mtp_bootstrap_embedding(); } + const double prepare_seconds = first_generation_timer.elapsed_seconds(); // send first gens to remote instance proto::DisaggPDService_Stub* stub = nullptr; @@ -745,7 +759,14 @@ void DisaggPDScheduler::prefill_send_first_generation() { // TODO: Async call later proto::Status resp; brpc::Controller cntl; + Timer rpc_timer; stub->FirstGeneration(&cntl, &gens, &resp, nullptr); + const double rpc_seconds = rpc_timer.elapsed_seconds(); + LOG(INFO) << "[PD-PERF] Prefill first-generation request_id=" + << request->request_id() + << ", prepare_ms=" << prepare_seconds * 1000.0 + << ", rpc_ms=" << rpc_seconds * 1000.0 << ", total_ms=" + << first_generation_timer.elapsed_seconds() * 1000.0; const bool sent_first_generation = !cntl.Failed() && resp.ok(); if (!sent_first_generation) { @@ -809,6 +830,7 @@ bool DisaggPDScheduler::decode_recv_first_generation( int32_t src_dp_size, int32_t src_dp_rank, torch::Tensor mtp_bootstrap_embedding) { + Timer receive_timer; // push to request_queue_, and will be executed by engine. std::shared_ptr request = nullptr; { @@ -838,6 +860,19 @@ bool DisaggPDScheduler::decode_recv_first_generation( return false; } Sequence* sequence = request->sequences()[0].get(); + bool hetero_kv_pull = false; + if (kv_cache_transfer_mode == "PUSH" && !engine_->model_args().enable_mla() && + src_dp_size > 0 && instance_info_.dp_size > 0 && + src_cluster_ids.size() % static_cast(src_dp_size) == 0 && + instance_info_.cluster_ids.size() % + static_cast(instance_info_.dp_size) == + 0) { + const size_t src_tp_size = + src_cluster_ids.size() / static_cast(src_dp_size); + const size_t dst_tp_size = instance_info_.cluster_ids.size() / + static_cast(instance_info_.dp_size); + hetero_kv_pull = src_tp_size != dst_tp_size; + } const bool need_mtp_bootstrap = options_.num_speculative_tokens() > 0; if (need_mtp_bootstrap) { const int32_t slot_id = sequence->get_single_block_id(); @@ -887,6 +922,10 @@ bool DisaggPDScheduler::decode_recv_first_generation( sequence->tbt(absl::Now()); // TODO: we only support one sequence for currently. + // Heterogeneous TP now restores the Prefill KV/linear state before Decode, + // so its first token follows the same overlap protocol as every other PD + // transfer path. Appending the real token here bypasses last_step_token and + // makes the first Decode step consume misaligned sequence state. if (enable_schedule_overlap()) { Token fake_token(-1); sequence->append_token(fake_token); @@ -894,9 +933,13 @@ bool DisaggPDScheduler::decode_recv_first_generation( } else { sequence->append_token(first_token); } + const double prepare_seconds = receive_timer.elapsed_seconds(); - // pull kv cache - if (kv_cache_transfer_mode == "PULL") { + // Pull KV cache in native PULL mode. For a non-MLA heterogeneous TP PUSH + // deployment, pull every P-side shard into temporary D-side caches and + // concatenate the sharded tensor dimensions before decode starts. + if (kv_cache_transfer_mode == "PULL" || hetero_kv_pull) { + Timer pull_timer; const auto blocks = sequence->kv_state().blocks(BlockType::KV); std::vector dst_block_ids; dst_block_ids.reserve(blocks.size()); @@ -911,27 +954,47 @@ bool DisaggPDScheduler::decode_recv_first_generation( } int32_t dst_dp_rank = sequence->dp_rank(); - const bool pulled = engine_->pull_kv_blocks(src_dp_size, - src_dp_rank, - src_cluster_ids, - src_addrs, - src_block_ids, - dst_dp_rank, - dst_block_ids, - src_linear_state_ids, - dst_linear_state_ids); + const bool pulled = + hetero_kv_pull ? engine_->pull_hetero_kv_blocks(src_dp_size, + src_dp_rank, + src_cluster_ids, + src_addrs, + src_block_ids, + dst_dp_rank, + dst_block_ids, + src_linear_state_ids, + dst_linear_state_ids) + : engine_->pull_kv_blocks(src_dp_size, + src_dp_rank, + src_cluster_ids, + src_addrs, + src_block_ids, + dst_dp_rank, + dst_block_ids, + src_linear_state_ids, + dst_linear_state_ids); if (!pulled) { - LOG(ERROR) << "Failed to pull KV blocks, request_id: " << req_id; + LOG(ERROR) << "Failed to pull" + << (hetero_kv_pull ? " and merge heterogeneous" : "") + << " KV blocks, request_id: " << req_id; kv_cache_manager_->deallocate(request.get()); return false; } + LOG(INFO) << "[PD-PERF] Decode KV restore request_id=" << req_id + << ", hetero=" << hetero_kv_pull + << ", pull_ms=" << pull_timer.elapsed_seconds() * 1000.0; } + Timer enqueue_timer; if (!request_queue_.write(request)) { LOG(ERROR) << "Failed to enqueue decode request, request_id: " << req_id; kv_cache_manager_->deallocate(request.get()); return false; } + LOG(INFO) << "[PD-PERF] Decode first-generation request_id=" << req_id + << ", prepare_ms=" << prepare_seconds * 1000.0 + << ", enqueue_ms=" << enqueue_timer.elapsed_seconds() * 1000.0 + << ", total_ms=" << receive_timer.elapsed_seconds() * 1000.0; return true; } diff --git a/xllm/models/llm/qwen3_5_mtp.h b/xllm/models/llm/qwen3_5_mtp.h index f026b430c0..cfbde669cb 100644 --- a/xllm/models/llm/qwen3_5_mtp.h +++ b/xllm/models/llm/qwen3_5_mtp.h @@ -144,6 +144,13 @@ class Qwen3_5MtpModelImpl : public Qwen3HybridModelImplBase { kv_caches[i], input_params, mrope_cos_sin); +#if defined(USE_NPU) + if (input_params.parallel.layer_synchronizer != nullptr && + !input_params.parallel.layer_synchronizer->record_event( + static_cast(i), device_.index())) { + return ModelOutput(); + } +#endif } auto [new_mtp_hidden, new_res] = norm_->forward(mtp_hidden, residual); mtp_hidden = new_mtp_hidden; diff --git a/xllm/proto/worker.proto b/xllm/proto/worker.proto index 3dca764b11..a30d7239d2 100644 --- a/xllm/proto/worker.proto +++ b/xllm/proto/worker.proto @@ -99,6 +99,9 @@ message PullKVCacheRequest { repeated uint64 dst_blocks = 5; repeated uint64 src_linear_state_ids = 7; repeated uint64 dst_linear_state_ids = 8; + repeated uint64 src_cluster_ids = 9; + repeated string src_addrs = 10; + bool hetero_merge = 11; } enum TransferType {