From 3fa13f1c8095966ac19cef7577bf2660865669e8 Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Wed, 10 Jun 2026 17:37:17 +0800 Subject: [PATCH 01/13] build: fix clean upstream npu configure --- CMakeLists.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 712687057c..ff4d623bed 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -316,7 +316,7 @@ find_package(benchmark CONFIG REQUIRED) find_package(nlohmann_json CONFIG REQUIRED) find_package(OpenCV CONFIG REQUIRED) find_package(FFMPEG REQUIRED) -find_package(Python COMPONENTS Development REQUIRED) +find_package(Python COMPONENTS Interpreter Development REQUIRED) find_package(pybind11 CONFIG REQUIRED) if (USE_CXX11_ABI) @@ -405,6 +405,7 @@ if(USE_NPU) #-> because the scope of third-party softwares managed by vcpkg is used throughout the entire xllm. message(STATUS "VCPKG_INCLUDE_DIR = ${CMAKE_BINARY_DIR}/vcpkg_installed/${VCPKG_TARGET_TRIPLET}/include") include_directories("${CMAKE_BINARY_DIR}/vcpkg_installed/${VCPKG_TARGET_TRIPLET}/include") + include_directories(SYSTEM ${Python_INCLUDE_DIRS}) add_subdirectory( "${XLLM_ATB_LAYERS_SOURCE_DIR}" From 68037f0eb7a1da40fcb21a580a157eab1a2b6b0f Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Wed, 10 Jun 2026 17:40:59 +0800 Subject: [PATCH 02/13] build: adapt npu grouped matmul for local torch_npu --- xllm/core/kernels/npu/npu_grouped_matmul.cpp | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/xllm/core/kernels/npu/npu_grouped_matmul.cpp b/xllm/core/kernels/npu/npu_grouped_matmul.cpp index b6fdf28d56..b77038503d 100644 --- a/xllm/core/kernels/npu/npu_grouped_matmul.cpp +++ b/xllm/core/kernels/npu/npu_grouped_matmul.cpp @@ -165,11 +165,7 @@ std::vector apply_npu_grouped_matmul( resolved_group_list_type, resolved_act_type, resolved_tuning_config, - ::std::optional{static_cast(resolved_output_dtype)}, - c10::nullopt, - c10::nullopt, - c10::nullopt, - c10::nullopt); + ::std::optional{resolved_output_dtype}); } } // namespace xllm::kernel::npu From 33a834f3217c102247f73a46e6a8d71eb8b30a8f Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Wed, 10 Jun 2026 18:25:24 +0800 Subject: [PATCH 03/13] feat: add cp kv-split link strategy on latest upstream --- xllm/core/distributed_runtime/llm_engine.cpp | 49 ++++++----- .../kv_cache_transfer/kv_cache_transfer.cpp | 41 +++++++++ .../kv_cache_transfer/kv_cache_transfer.h | 9 ++ .../llm_data_dist_transfer.cpp | 85 +++++++++++++++---- .../llm_data_dist_transfer.h | 9 ++ xllm/core/runtime/worker_rendezvous.cpp | 15 +--- 6 files changed, 158 insertions(+), 50 deletions(-) diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 940624ca3e..a449154a28 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -719,7 +719,10 @@ 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. + // Each D worker connects to P workers that share the same TP rank. + // LlmDataDist currently throws SVector out_of_range when one decoder worker + // links multiple CP/KV split prompt endpoints, so KV-split mode pins each + // decoder worker to one split endpoint. // 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. @@ -735,18 +738,20 @@ 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 worker_split_index = + src_kv_split_size > 1 ? static_cast(worker_rank) % + src_kv_split_size + : 0; + target_cluster_ids.reserve(src_dp_size); + target_addrs.reserve(src_dp_size); + target_ports.reserve(src_dp_size); 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]); - } + int32_t p_idx = dp_i * src_cp_tp_size + + worker_split_index * 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]); } src_dp_worker_index = (src_dp_worker_index + 1) % src_tp_size; @@ -794,18 +799,20 @@ 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 worker_split_index = + src_kv_split_size > 1 ? static_cast(worker_rank) % + src_kv_split_size + : 0; + target_cluster_ids.reserve(src_dp_size); + target_addrs.reserve(src_dp_size); + target_ports.reserve(src_dp_size); 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]); - } + int32_t p_idx = dp_i * src_cp_tp_size + + worker_split_index * 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]); } src_dp_worker_index = (src_dp_worker_index + 1) % src_tp_size; diff --git a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp index cd181f6e88..7ff0741087 100644 --- a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp +++ b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp @@ -37,6 +37,47 @@ limitations under the License. namespace xllm { +bool KVCacheTransfer::link_clusters( + const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& ports) { + if (cluster_ids.size() != remote_addrs.size() || + cluster_ids.size() != ports.size()) { + LOG(ERROR) << "Cluster endpoint size mismatch: cluster_ids=" + << cluster_ids.size() << ", addrs=" << remote_addrs.size() + << ", ports=" << ports.size(); + return false; + } + + for (size_t i = 0; i < cluster_ids.size(); ++i) { + if (!link_cluster(cluster_ids[i], remote_addrs[i], ports[i])) { + return false; + } + } + return true; +} + +bool KVCacheTransfer::unlink_clusters( + const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& ports, + bool force_flag) { + if (cluster_ids.size() != remote_addrs.size() || + cluster_ids.size() != ports.size()) { + LOG(ERROR) << "Cluster endpoint size mismatch: cluster_ids=" + << cluster_ids.size() << ", addrs=" << remote_addrs.size() + << ", ports=" << ports.size(); + return false; + } + + for (size_t i = 0; i < cluster_ids.size(); ++i) { + if (!unlink_cluster(cluster_ids[i], remote_addrs[i], ports[i], force_flag)) { + return false; + } + } + return true; +} + folly::SemiFuture KVCacheTransfer::pull_kv_blocks_async( const uint64_t src_cluster_id, const std::string& src_addr, 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 46a0074b8a..160dc5dffd 100644 --- a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h +++ b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h @@ -107,11 +107,20 @@ class KVCacheTransfer { const std::string& remote_addr, const uint16_t port) = 0; + virtual bool link_clusters(const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& ports); + virtual bool unlink_cluster(const uint64_t& cluster_id, const std::string& remote_addr, const uint16_t port, bool force_flag = true) = 0; + virtual bool unlink_clusters(const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& ports, + bool force_flag = true); + virtual bool pull_kv_blocks( const uint64_t src_cluster_id, const std::string& src_addr, 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 a4ac5db3fa..e2d1162d2f 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,6 +19,7 @@ limitations under the License. #include #include +#include #include "common/macros.h" #include "core/framework/config/disagg_pd_config.h" @@ -111,25 +112,55 @@ void LlmDataDistTransfer::get_cache_info(uint64_t& cluster_id, bool LlmDataDistTransfer::link_cluster(const uint64_t cluster_id, const std::string& remote_addr, const uint16_t port) { - if (linked_cluster_ids.find(cluster_id) != linked_cluster_ids.end()) { - // The cluster is connected. - return true; + return link_clusters(std::vector{cluster_id}, + std::vector{remote_addr}, + std::vector{port}); +} + +bool LlmDataDistTransfer::link_clusters( + const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& ports) { + if (cluster_ids.size() != remote_addrs.size() || + cluster_ids.size() != ports.size()) { + LOG(ERROR) << "Cluster endpoint size mismatch: cluster_ids=" + << cluster_ids.size() << ", addrs=" << remote_addrs.size() + << ", ports=" << ports.size(); + return false; } std::vector rets; std::vector clusters; - ClusterInfo cluster_info = create_cluster_info(cluster_id, remote_addr, port); - clusters.emplace_back(std::move(cluster_info)); + clusters.reserve(cluster_ids.size()); + for (size_t i = 0; i < cluster_ids.size(); ++i) { + if (linked_cluster_ids.find(cluster_ids[i]) != linked_cluster_ids.end()) { + continue; + } + clusters.emplace_back( + create_cluster_info(cluster_ids[i], remote_addrs[i], ports[i])); + } - auto ret = llm_data_dist_->LinkLlmClusters( - clusters, rets, /*timeout_in_millis=*/60000); + if (clusters.empty()) { + return true; + } + + uint64_t ret = LLM_SUCCESS; + try { + ret = llm_data_dist_->LinkLlmClusters( + clusters, rets, /*timeout_in_millis=*/60000); + } catch (const std::exception& e) { + LOG(ERROR) << "LinkLlmClusters threw exception: " << e.what() + << ", clusters = " << clusters.size(); + return false; + } if (ret != LLM_SUCCESS) { LOG(ERROR) << "LinkLlmClusters failed, ret = " << std::hex << ret; return false; } - LOG(INFO) << "LinkLlmClusters success, ip : " << remote_addr - << ", port : " << port; - linked_cluster_ids.insert(cluster_id); + LOG(INFO) << "LinkLlmClusters success, clusters = " << clusters.size(); + for (uint64_t id : cluster_ids) { + linked_cluster_ids.insert(id); + } return true; } @@ -138,11 +169,32 @@ bool LlmDataDistTransfer::unlink_cluster(const uint64_t& cluster_id, const std::string& remote_addr, const uint16_t remote_port, bool force_flag) { + return unlink_clusters(std::vector{cluster_id}, + std::vector{remote_addr}, + std::vector{remote_port}, + force_flag); +} + +bool LlmDataDistTransfer::unlink_clusters( + const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& remote_ports, + bool force_flag) { + if (cluster_ids.size() != remote_addrs.size() || + cluster_ids.size() != remote_ports.size()) { + LOG(ERROR) << "Cluster endpoint size mismatch: cluster_ids=" + << cluster_ids.size() << ", addrs=" << remote_addrs.size() + << ", ports=" << remote_ports.size(); + return false; + } + std::vector rets; std::vector clusters; - ClusterInfo cluster_info = - create_cluster_info(cluster_id, remote_addr, remote_port); - clusters.emplace_back(std::move(cluster_info)); + clusters.reserve(cluster_ids.size()); + for (size_t i = 0; i < cluster_ids.size(); ++i) { + clusters.emplace_back( + create_cluster_info(cluster_ids[i], remote_addrs[i], remote_ports[i])); + } auto ret = llm_data_dist_->UnlinkLlmClusters(clusters, rets, 1000, force_flag); @@ -150,9 +202,10 @@ bool LlmDataDistTransfer::unlink_cluster(const uint64_t& cluster_id, LOG(ERROR) << "UnlinkLlmClusters failed, ret = " << std::hex << ret; return false; } - LOG(INFO) << "UnlinkLlmClusters success, ip : " << remote_addr - << ", port : " << remote_port; - linked_cluster_ids.erase(cluster_id); + LOG(INFO) << "UnlinkLlmClusters success, clusters = " << clusters.size(); + for (uint64_t id : cluster_ids) { + linked_cluster_ids.erase(id); + } return true; } 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 0a543bbef7..277184a05e 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 @@ -54,11 +54,20 @@ class LlmDataDistTransfer : public KVCacheTransfer { const std::string& remote_addr, const uint16_t port) override; + virtual bool link_clusters(const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& ports) override; + virtual bool unlink_cluster(const uint64_t& cluster_id, const std::string& remote_addr, const uint16_t port, bool force_flag = true) override; + virtual bool unlink_clusters(const std::vector& cluster_ids, + const std::vector& remote_addrs, + const std::vector& ports, + bool force_flag = true) override; + virtual bool pull_kv_blocks( const uint64_t src_cluster_id, const std::string& src_addr, diff --git a/xllm/core/runtime/worker_rendezvous.cpp b/xllm/core/runtime/worker_rendezvous.cpp index b74d324dc1..09967089ce 100644 --- a/xllm/core/runtime/worker_rendezvous.cpp +++ b/xllm/core/runtime/worker_rendezvous.cpp @@ -50,12 +50,7 @@ bool WorkerRendezvous::link_cluster(const std::vector& cluster_ids, return false; } - const size_t cluster_count = cluster_ids.size(); - for (size_t i = 0; i < cluster_count; ++i) { - if (!kv_cache_transfer_->link_cluster(cluster_ids[i], addrs[i], ports[i])) { - return false; - } - } + return kv_cache_transfer_->link_clusters(cluster_ids, addrs, ports); #endif return true; } @@ -72,13 +67,7 @@ bool WorkerRendezvous::unlink_cluster(const std::vector& cluster_ids, return false; } - const size_t cluster_count = cluster_ids.size(); - for (size_t i = 0; i < cluster_count; ++i) { - if (!kv_cache_transfer_->unlink_cluster( - cluster_ids[i], addrs[i], ports[i])) { - return false; - } - } + return kv_cache_transfer_->unlink_clusters(cluster_ids, addrs, ports); #endif return true; } From 6f8b36f7f4bb84b5dc3b149b74a63de83d020680 Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Wed, 10 Jun 2026 18:49:09 +0800 Subject: [PATCH 04/13] build: cap xllm_ops precompile jobs during configure --- CMakeLists.txt | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index ff4d623bed..363172b8ae 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -44,10 +44,17 @@ if(USE_NPU) set(XLLM_OPS_BUILD_DIR "${CMAKE_SOURCE_DIR}/third_party/xllm_ops/build") endif() + if(DEFINED ENV{XLLM_OPS_PRECOMPILE_JOBS} AND NOT "$ENV{XLLM_OPS_PRECOMPILE_JOBS}" STREQUAL "") + set(XLLM_OPS_PRECOMPILE_JOBS "$ENV{XLLM_OPS_PRECOMPILE_JOBS}") + else() + set(XLLM_OPS_PRECOMPILE_JOBS "1") + endif() + if(NOT DEFINED XLLM_OPS_GIT_HEAD_CACHED OR NOT XLLM_OPS_GIT_HEAD STREQUAL XLLM_OPS_GIT_HEAD_CACHED) - message(STATUS "xllm_ops git HEAD changed; running precompile via execute_process") + message(STATUS "xllm_ops git HEAD changed; running precompile via execute_process (jobs=${XLLM_OPS_PRECOMPILE_JOBS})") execute_process( COMMAND ${CMAKE_COMMAND} -E env "XLLM_OPS_BUILD_DIR=${XLLM_OPS_BUILD_DIR}" + "OPS_CPU_NUMBER=${XLLM_OPS_PRECOMPILE_JOBS}" bash ${CMAKE_SOURCE_DIR}/third_party/xllm_ops/build.sh WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/third_party/xllm_ops RESULT_VARIABLE XLLM_OPS_RESULT From a3ad699e7f878c3d03cab6b1797d2506e62c2e7a Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Tue, 16 Jun 2026 17:52:26 +0800 Subject: [PATCH 05/13] fix: restore disagg pd profiling for prefill and decode --- xllm/core/scheduler/disagg_pd_scheduler.cpp | 49 ++++++++++++++++++--- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/xllm/core/scheduler/disagg_pd_scheduler.cpp b/xllm/core/scheduler/disagg_pd_scheduler.cpp index 7c96f39664..cb321c5b33 100644 --- a/xllm/core/scheduler/disagg_pd_scheduler.cpp +++ b/xllm/core/scheduler/disagg_pd_scheduler.cpp @@ -68,11 +68,25 @@ DisaggPDScheduler::DisaggPDScheduler(Engine* engine, const Options& options) initialize_rpc_server(server_name_); register_instance_info(server_name_, engine); - // Profile ttft & topt and update instance info (for mix instances) - if (!options_.disable_ttft_profiling() && - options_.instance_role().value() == InstanceRole::MIX) { - profile_ttft(); - profile_tpot(); + // Profile time predictors before the instance is registered to service. + // In disaggregated PD mode, prefill lanes need TTFT profiling data for + // prefill-route estimation, while decode lanes need TPOT profiling data + // for decode-pressure estimation. The previous implementation only + // profiled MIX roles, which left PREFILL/DECODE-only topologies with empty + // profiling vectors and caused projected/estimated route fields to remain + // zero in service traces. + if (!options_.disable_ttft_profiling()) { + const auto role = options_.instance_role().value(); + if (role == InstanceRole::MIX || role == InstanceRole::PREFILL) { + profile_ttft(); + } + if (role == InstanceRole::MIX || role == InstanceRole::DECODE) { + profile_tpot(); + } + LOG(INFO) << "Disagg PD profiling summary, role=" + << role.to_string() + << ", ttft_points=" << instance_info_.ttft_profiling_data.size() + << ", tpot_points=" << instance_info_.tpot_profiling_data.size(); } } } @@ -149,10 +163,35 @@ void DisaggPDScheduler::profile_ttft() { // get the maximum prefill token length auto& model_args = engine_->model_args(); int32_t max_context_len = model_args.max_position_embeddings(); + const int32_t profile_max_prompt_length = + std::max(2, options_.profile_max_prompt_length()); + const int32_t block_size = kv_cache_manager_->block_size(); + const int32_t total_cache_token_capacity = + static_cast(kv_cache_manager_->num_blocks()) * block_size; + max_context_len = std::min(max_context_len, profile_max_prompt_length); if (!options_.enable_chunked_prefill()) { max_context_len = std::min(max_context_len, options_.max_tokens_per_batch()); } + // Reserve at least one block and avoid landing exactly on a block boundary. + // The profiling path accounts blocks as token_length / block_size + 1, so a + // length like (num_blocks - 1) * block_size still rounds up to num_blocks + // and can trip "Not enough blocks". Subtract one extra token to keep the + // request strictly below that boundary. + const int32_t safe_cache_token_capacity = std::max( + 1, total_cache_token_capacity - block_size - 1); + max_context_len = std::min(max_context_len, safe_cache_token_capacity); + CHECK_GT(max_context_len, 1) + << "TTFT profiling requires positive KV cache token capacity, block_size=" + << block_size + << ", num_blocks=" << kv_cache_manager_->num_blocks(); + LOG(INFO) << "TTFT profiling capacity bound, max_context_len=" + << max_context_len + << ", profile_max_prompt_length=" << profile_max_prompt_length + << ", total_cache_token_capacity=" << total_cache_token_capacity + << ", safe_cache_token_capacity=" << safe_cache_token_capacity + << ", block_size=" << block_size + << ", num_blocks=" << kv_cache_manager_->num_blocks(); // warm up profile_manager_->run_request(max_context_len, 0); From 03c29ab3e76d380d76a808cb6165982f9736dfef Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Tue, 23 Jun 2026 19:35:34 +0800 Subject: [PATCH 06/13] feat: MixScheduler inherits DisaggPDScheduler for service-side dual_register MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit L2.3 (service-side `--enable_mix_dual_register`) silently failed in the production path because MixScheduler did not run a disagg_pd RPC server, so its etcd entry had `rpc_address=""` / `cluster_ids=[]` / `addrs=[]` / `ports=[]`. The service watch hit `gather_link_operations` → `run_link_ operations` → brpc dial empty rpc_address → `Fail to link instance during registration` → no `add_instance_to_index` → MIX never reached `prefill_index_` / `decode_index_` and the dual_register code never ran. Switch MixScheduler's base from `ChunkedPrefillScheduler` to `DisaggPDScheduler`. The base ctor wires `initialize_rpc_server` + `register_instance_info`, which fills `instance_info_.rpc_address` (from the brpc listen addr) plus the cluster_ids/addrs/ports the engine reports via `engine_->get_cache_info`. After this change MIX's etcd entry is fully populated and matches what PREFILL/DECODE write — service-side `LinkInstance` calls land on a real endpoint instead of erroring on the empty address. Override `add_request` and `step` to skip the prefill→decode dispatch path (`prefill_request_queue_` / `prefill_send_first_generation`) which DisaggPDScheduler runs for PREFILL/DECODE roles. MIX serves locally and must keep using `request_queue_` and base `ContinuousScheduler::step`. The dispatch_thread the base spawns stays idle for MIX (nothing enqueues into prefill_request_queue_), which is harmless. Note: the base ctor also runs `profile_ttft` + `profile_tpot` for InstanceRole::MIX. With small tp (e.g. tp=4), `profile_tpot` triggers a known engine corner case (`Expected NPU tensor` in `get_npu_format`). Launchers should pass `--disable_ttft_profiling=true` for MIX until the underlying ProfileManager bug is fixed. Verified end-to-end on dev98: XLLM:MIX:11.87.191.98:58200 {"addrs":[..×4],"cluster_ids":[..×4],"ports":[564..567], "rpc_address":"11.87.191.98:13877",...} vs the previous run with the same service binary: XLLM:MIX:...:58200 {"addrs":[],"cluster_ids":[],"ports":[],"rpc_address":""} The remaining `LinkLlmClusters failed (5010b009)` on the DECODE→MIX direction is a separate engine constraint (`link_cluster` requires P-side tp == D-side dp_local_tp_size_, documented at xllm/core/distributed_runtime/llm_engine.cpp:728), not addressed in this commit. Co-Authored-By: Claude --- xllm/core/scheduler/mix_scheduler.cpp | 17 ++++++++++++++++- xllm/core/scheduler/mix_scheduler.h | 19 ++++++++++++++++++- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/xllm/core/scheduler/mix_scheduler.cpp b/xllm/core/scheduler/mix_scheduler.cpp index 5eee8fd7fe..0c0a0dd915 100644 --- a/xllm/core/scheduler/mix_scheduler.cpp +++ b/xllm/core/scheduler/mix_scheduler.cpp @@ -29,7 +29,7 @@ limitations under the License. namespace xllm { MixScheduler::MixScheduler(Engine* engine, const Options& options) - : ChunkedPrefillScheduler(engine, options) {} + : DisaggPDScheduler(engine, options) {} MixScheduler::~MixScheduler() { // release all requests in the running priority queue @@ -38,6 +38,21 @@ MixScheduler::~MixScheduler() { } } +bool MixScheduler::add_request(std::shared_ptr& request) { + // Bypass the DisaggPDScheduler::add_request path which enqueues into + // prefill_request_queue_ for the prefill->decode dispatch thread. + // MIX serves requests locally, so push straight to request_queue_ which + // prepare_batch() will drain. + return ContinuousScheduler::add_request(request); +} + +void MixScheduler::step(const absl::Duration& timeout) { + // Skip DisaggPDScheduler::step which calls prefill_send_first_generation() + // for non-DECODE roles. MIX produces tokens locally and never forwards them + // to a remote decode instance. + ContinuousScheduler::step(timeout); +} + //---------------------------------- void MixScheduler::get_latency_budget_and_request_order( diff --git a/xllm/core/scheduler/mix_scheduler.h b/xllm/core/scheduler/mix_scheduler.h index f7ad1d4e80..d161537e85 100644 --- a/xllm/core/scheduler/mix_scheduler.h +++ b/xllm/core/scheduler/mix_scheduler.h @@ -23,17 +23,34 @@ limitations under the License. #include "scheduler/chunked_prefill_scheduler.h" #include "scheduler/continuous_scheduler.h" +#include "scheduler/disagg_pd_scheduler.h" namespace xllm { // MixScheduler does not explicitly specify whether decoding or prefilling takes // priority; instead, it mixes prefilling and decoding requests in a single // queue. Currently, it's only for multi-priority scheduling algorithm ProSched. -class MixScheduler : public ChunkedPrefillScheduler { +// +// MIX inherits from DisaggPDScheduler so it picks up the disagg_pd RPC server +// and `register_instance_info` plumbing required for service-side routing +// (etcd entry must carry rpc_address / cluster_ids / addrs / ports for the +// service `LinkInstance` callback to land on us). The base ctor wires those. +// MIX overrides add_request / step / prepare_batch to keep using its own +// running_queue_ and skip the prefill->decode dispatch path that PREFILL/ +// DECODE rely on. +class MixScheduler : public DisaggPDScheduler { public: MixScheduler(Engine* engine, const Options& options); virtual ~MixScheduler(); + // MIX serves requests locally; do not enqueue into the prefill->decode + // dispatch queue that the base DisaggPDScheduler uses. + bool add_request(std::shared_ptr& request) override; + + // Skip the prefill_send_first_generation step the base runs; MIX does not + // forward generation tokens to a remote decode instance. + void step(const absl::Duration& timeout) override; + protected: std::list> running_queue_; From be09e6842878933f4d6d13c70e56add4a6959672 Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Tue, 23 Jun 2026 20:09:51 +0800 Subject: [PATCH 07/13] fix: MIX role must publish LlmDataDist listen endpoint LlmDataDistTransfer::initialize was setting OPTION_LISTEN_IP_INFO only for LlmRole::kPrompt (PREFILL). MIX (kMix) silently skipped it. Result: a MIX instance running with `--enable_disagg_pd=true` did not register a listen endpoint on the underlying LlmDataDist runtime, so when a remote DECODE tried to LinkLlmClusters to the MIX, the call was rejected with ret=5010b009 (ACL/RT param invalid range). Verified end-to-end on dev98 + dev82 24-card mixed topology: prefill_1_tp8 (dev98 0-7) + decode_1_tp8 (dev98 8-15) + mix_1_tp8 (dev82 0-7), all tp=8 to satisfy link_cluster's D-side tp == P-side tp constraint. After this fix: - Service registers MIX dual-mode: `Register a new dual-mode mix instance: 11.87.191.82:58200 prefill_idx=0 decode_idx=0` - decode -> MIX link succeeds: `LinkLlmClusters success, clusters = 1` `Successfully linked instance, instance_name: 11.87.191.82:58200, prefill_kv_split_size: 1` - decode -> prefill_1 link also succeeds (regression check) - route_trace shows candidate_prefill_instances = ["MIX:58200", "prefill_1:18100"], pool_active_set has both, requests are scheduled onto MIX as prefill+decode pair (single-instance mixed mode) when RR selects it. - smoke 1 short request: 200 OK, completion_tokens=20 - smoke 1 long request: 200 OK, prompt=204 completion=256 This closes the last-mile of the L2.3 + L2.4 + L2.5 chain (service-side dual_register + MIX disagg_pd RPC plumbing + LlmDataDist listen endpoint). Mixed PD+MIX topology is now functional end-to-end. Co-Authored-By: Claude --- .../framework/kv_cache_transfer/llm_data_dist_transfer.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 e2d1162d2f..d63f55acfd 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 @@ -78,8 +78,9 @@ void LlmDataDistTransfer::initialize(int32_t device_id) { std::map options; options[OPTION_DEVICE_ID] = std::to_string(device_id).c_str(); - // Prompt(Prefill) must publish listen endpoint; Decoder only needs device_id. - if (role_ == LlmRole::kPrompt) { + // Prompt(Prefill) and Mix must publish listen endpoint so remote Decoders + // can link to them. Decoder-only instances do not need to listen. + if (role_ == LlmRole::kPrompt || role_ == LlmRole::kMix) { std::string local_ip_info = host_ip_ + ":" + std::to_string(listen_port_); options[OPTION_LISTEN_IP_INFO] = local_ip_info.c_str(); } From 7001409d53ea48a5e01ad05659ac82707b0ed34a Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Tue, 23 Jun 2026 20:58:47 +0800 Subject: [PATCH 08/13] feat: MixScheduler decode-first scheduling + reserved decode budget Adds two opt-in flags to mitigate the MIX colocate decode tpot regression observed at 4x-15x worse than PD-disagg in 720-plan benchmarks: enable_mix_decode_first (default false): when true, prepare_batch splits running_queue_ into decode-stage and prefill-stage lists, runs handle_running_queue_requests on decode first with the full token budget, then on prefill within remaining budget. Decode is no longer dragged by prefill chunk step time. mix_decode_token_budget (default 0): reserves N tokens of per-step budget so prefill chunks are capped at (max_tokens_per_batch - N). Only takes effect when enable_mix_decode_first=true. Defaults preserve legacy single-pass behavior; flags must be flipped explicitly to opt into the new path. Co-Authored-By: Claude --- .../framework/config/scheduler_config.cpp | 21 +++++ xllm/core/framework/config/scheduler_config.h | 17 +++- xllm/core/scheduler/mix_scheduler.cpp | 92 +++++++++++++++++-- 3 files changed, 120 insertions(+), 10 deletions(-) diff --git a/xllm/core/framework/config/scheduler_config.cpp b/xllm/core/framework/config/scheduler_config.cpp index 6ca5fc3910..c2f3ffec15 100644 --- a/xllm/core/framework/config/scheduler_config.cpp +++ b/xllm/core/framework/config/scheduler_config.cpp @@ -73,6 +73,19 @@ DEFINE_bool(enable_starve_prevent, true, "Whether to enable anti-starvation in MixScheduler."); +DEFINE_bool(enable_mix_decode_first, + false, + "MixScheduler: process decode requests before prefill chunks in " + "prepare_batch so decode tpot is not dragged by prefill chunk " + "step time. Mitigates colocate (MIX) decode tail latency."); + +DEFINE_int32(mix_decode_token_budget, + 0, + "MixScheduler: reserved decode token budget per step. When >0, " + "prefill chunks are capped at (max_tokens_per_batch - this) per " + "step. 0 disables the cap. Only used when " + "enable_mix_decode_first=true."); + namespace xllm { void SchedulerConfig::from_flags() { @@ -91,6 +104,8 @@ void SchedulerConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(aggressive_coeff); XLLM_CONFIG_ASSIGN_FROM_FLAG(starve_threshold); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_starve_prevent); + XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_mix_decode_first); + XLLM_CONFIG_ASSIGN_FROM_FLAG(mix_decode_token_budget); } void SchedulerConfig::from_json(const JsonReader& json) { @@ -109,6 +124,8 @@ void SchedulerConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(aggressive_coeff); XLLM_CONFIG_ASSIGN_FROM_JSON(starve_threshold); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_starve_prevent); + XLLM_CONFIG_ASSIGN_FROM_JSON(enable_mix_decode_first); + XLLM_CONFIG_ASSIGN_FROM_JSON(mix_decode_token_budget); } void SchedulerConfig::append_config_json( @@ -144,6 +161,10 @@ void SchedulerConfig::append_config_json( config_json, default_config, starve_threshold); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, enable_starve_prevent); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, enable_mix_decode_first); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, mix_decode_token_budget); } SchedulerConfig& SchedulerConfig::get_instance() { diff --git a/xllm/core/framework/config/scheduler_config.h b/xllm/core/framework/config/scheduler_config.h index 01a60ab220..a3ea5df718 100644 --- a/xllm/core/framework/config/scheduler_config.h +++ b/xllm/core/framework/config/scheduler_config.h @@ -55,7 +55,9 @@ class SchedulerConfig final { "enable_online_preempt_offline", "aggressive_coeff", "starve_threshold", - "enable_starve_prevent"}}; + "enable_starve_prevent", + "enable_mix_decode_first", + "mix_decode_token_budget"}}; return kOptionCategory; } @@ -88,6 +90,19 @@ class SchedulerConfig final { PROPERTY(double, starve_threshold) = 1.0; PROPERTY(bool, enable_starve_prevent) = true; + + // MixScheduler decode-first scheduling: when true, prepare_batch processes + // decode-stage requests first (so they get the full token budget) before + // admitting prefill chunks. Mitigates the "decode tpot dragged by prefill + // chunk step time" problem in colocate (MIX) topologies. + PROPERTY(bool, enable_mix_decode_first) = false; + + // Reserved token budget for decode in MixScheduler. When >0, prefill chunks + // are restricted to (max_tokens_per_batch - mix_decode_token_budget) tokens + // per step, giving decode requests guaranteed headroom even when prefill + // would otherwise consume the entire batch budget. Only takes effect when + // enable_mix_decode_first is true. 0 disables the cap. + PROPERTY(int32_t, mix_decode_token_budget) = 0; }; } // namespace xllm diff --git a/xllm/core/scheduler/mix_scheduler.cpp b/xllm/core/scheduler/mix_scheduler.cpp index 0c0a0dd915..98f4b45c05 100644 --- a/xllm/core/scheduler/mix_scheduler.cpp +++ b/xllm/core/scheduler/mix_scheduler.cpp @@ -585,15 +585,89 @@ std::vector MixScheduler::prepare_batch() { // keep the requests in prefill stage std::vector prefill_stage_sequences; - handle_running_queue_requests(latency_budget, - estimate_latency, - remaining_token_budget, - remaining_seq_budget, - num_preempted_requests, - prefill_stage_sequences, - running_queue_, - budget_exhausted, - blocks_exhausted); + const bool decode_first = + ::xllm::SchedulerConfig::get_instance().enable_mix_decode_first(); + + if (!decode_first) { + // Legacy path: a single pass over running_queue_ that mixes prefill and + // decode requests together. Step time is dragged by the longest prefill + // chunk so decode tpot suffers in colocate (MIX) mode. + handle_running_queue_requests(latency_budget, + estimate_latency, + remaining_token_budget, + remaining_seq_budget, + num_preempted_requests, + prefill_stage_sequences, + running_queue_, + budget_exhausted, + blocks_exhausted); + } else { + // Decode-first scheduling: split running_queue_ into decode-stage and + // prefill-stage lists, schedule decode first with the full token budget, + // then prefill chunks within the remaining (optionally capped) budget. + // This isolates decode tpot from prefill chunk step-time inflation. + std::list> decode_queue; + std::list> prefill_queue; + for (auto& req : running_queue_) { + if (req->sequences()[0]->kv_state().kv_cache_tokens_num() > 0) { + decode_queue.push_back(req); + } else { + prefill_queue.push_back(req); + } + } + running_queue_.clear(); + + // Phase 1: decode requests get full token budget. Each decode contributes + // ~1 token so this rarely consumes much of the budget; the cap below + // mainly bounds how aggressively the next phase can admit prefill. + handle_running_queue_requests(latency_budget, + estimate_latency, + remaining_token_budget, + remaining_seq_budget, + num_preempted_requests, + prefill_stage_sequences, + decode_queue, + budget_exhausted, + blocks_exhausted); + + // Phase 2: prefill requests run in the remaining budget. When + // mix_decode_token_budget>0 we additionally cap prefill at + // (max_tokens_per_batch - reserved) so a future decode can fit even if + // one large prefill request slips through. The reservation matters when + // many decode requests arrive shortly after prefill admission. + const int32_t reserved = ::xllm::SchedulerConfig::get_instance() + .mix_decode_token_budget(); + if (reserved > 0 && !budget_exhausted && !blocks_exhausted) { + const size_t total_budget = options_.enable_profile_token_budget() + ? profile_manager_->get_token_budget() + : options_.max_tokens_per_batch(); + const size_t prefill_cap = + total_budget > static_cast(reserved) + ? total_budget - static_cast(reserved) + : 0; + if (remaining_token_budget > prefill_cap) { + remaining_token_budget = prefill_cap; + } + } + if (!budget_exhausted && !blocks_exhausted && + remaining_token_budget > 0 && remaining_seq_budget > 0) { + handle_running_queue_requests(latency_budget, + estimate_latency, + remaining_token_budget, + remaining_seq_budget, + num_preempted_requests, + prefill_stage_sequences, + prefill_queue, + budget_exhausted, + blocks_exhausted); + } + + // Restore unhandled requests back to running_queue_ so the next step can + // pick them up. handle_running_queue_requests pop_front()'s admitted ones, + // leaving the deferred ones in the local lists. + for (auto& r : decode_queue) running_queue_.push_back(r); + for (auto& r : prefill_queue) running_queue_.push_back(r); + } if (!finished_requests.empty()) { response_processor_->process_completed_requests(finished_requests); From fd39a82cad19259ca99296adad4cc6f2fa3484ff Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Thu, 25 Jun 2026 20:03:46 +0800 Subject: [PATCH 09/13] =?UTF-8?q?feat:=20MixScheduler=20Path=20C=20?= =?UTF-8?q?=E2=80=94=20cap=20prefill=20chunks=20per=20step?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds env mix_max_prefill_chunks_per_step (default 0 = no cap). When >0, phase 2 of prepare_batch (prefill admission) stops after this many prefill sequences enter running_sequences_, regardless of remaining token budget. Bounds how much prefill drags decode tpot. Path C is a lightweight middle-ground between B+C decode_first (throughput-friendly, tpot improvement limited) and chunk=2048 (tpot-friendly, ttft destroyed). With max_prefill_chunks=1, only one prefill chunk co-exists with decodes per forward step. Defaults preserve legacy behavior; env must be set explicitly to opt in. Co-Authored-By: Claude --- .../framework/config/scheduler_config.cpp | 12 ++++++++++ xllm/core/framework/config/scheduler_config.h | 11 ++++++++- xllm/core/scheduler/mix_scheduler.cpp | 23 +++++++++++++++++++ 3 files changed, 45 insertions(+), 1 deletion(-) diff --git a/xllm/core/framework/config/scheduler_config.cpp b/xllm/core/framework/config/scheduler_config.cpp index c2f3ffec15..a253a84ebd 100644 --- a/xllm/core/framework/config/scheduler_config.cpp +++ b/xllm/core/framework/config/scheduler_config.cpp @@ -86,6 +86,14 @@ DEFINE_int32(mix_decode_token_budget, "step. 0 disables the cap. Only used when " "enable_mix_decode_first=true."); +DEFINE_int32(mix_max_prefill_chunks_per_step, + 0, + "MixScheduler: cap on prefill chunks admitted per forward step. " + "When >0, phase 2 stops after this many prefill sequences enter " + "the batch, regardless of remaining token budget. Bounds how " + "much prefill drags decode tpot. 0 disables the cap. Only used " + "when enable_mix_decode_first=true."); + namespace xllm { void SchedulerConfig::from_flags() { @@ -106,6 +114,7 @@ void SchedulerConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_starve_prevent); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_mix_decode_first); XLLM_CONFIG_ASSIGN_FROM_FLAG(mix_decode_token_budget); + XLLM_CONFIG_ASSIGN_FROM_FLAG(mix_max_prefill_chunks_per_step); } void SchedulerConfig::from_json(const JsonReader& json) { @@ -126,6 +135,7 @@ void SchedulerConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(enable_starve_prevent); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_mix_decode_first); XLLM_CONFIG_ASSIGN_FROM_JSON(mix_decode_token_budget); + XLLM_CONFIG_ASSIGN_FROM_JSON(mix_max_prefill_chunks_per_step); } void SchedulerConfig::append_config_json( @@ -165,6 +175,8 @@ void SchedulerConfig::append_config_json( config_json, default_config, enable_mix_decode_first); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, mix_decode_token_budget); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, mix_max_prefill_chunks_per_step); } SchedulerConfig& SchedulerConfig::get_instance() { diff --git a/xllm/core/framework/config/scheduler_config.h b/xllm/core/framework/config/scheduler_config.h index a3ea5df718..a542a03d25 100644 --- a/xllm/core/framework/config/scheduler_config.h +++ b/xllm/core/framework/config/scheduler_config.h @@ -57,7 +57,8 @@ class SchedulerConfig final { "starve_threshold", "enable_starve_prevent", "enable_mix_decode_first", - "mix_decode_token_budget"}}; + "mix_decode_token_budget", + "mix_max_prefill_chunks_per_step"}}; return kOptionCategory; } @@ -103,6 +104,14 @@ class SchedulerConfig final { // would otherwise consume the entire batch budget. Only takes effect when // enable_mix_decode_first is true. 0 disables the cap. PROPERTY(int32_t, mix_decode_token_budget) = 0; + + // Path C cap on number of prefill chunks admitted into the same step in + // MixScheduler. When >0, phase 2 (prefill admission) stops after this many + // prefill sequences enter running_sequences_, regardless of token budget. + // Each request has 1 sequence so this directly caps prefill chunks per + // forward. Bounds how much prefill drags decode tpot. Only takes effect + // when enable_mix_decode_first is true. 0 disables the cap. + PROPERTY(int32_t, mix_max_prefill_chunks_per_step) = 0; }; } // namespace xllm diff --git a/xllm/core/scheduler/mix_scheduler.cpp b/xllm/core/scheduler/mix_scheduler.cpp index 98f4b45c05..3f75a887b0 100644 --- a/xllm/core/scheduler/mix_scheduler.cpp +++ b/xllm/core/scheduler/mix_scheduler.cpp @@ -649,6 +649,20 @@ std::vector MixScheduler::prepare_batch() { remaining_token_budget = prefill_cap; } } + + // Path C cap: when mix_max_prefill_chunks_per_step>0, restrict phase 2 + // to admit at most N prefill sequences per step. This bounds how many + // prefill chunks are batched with decode requests in the same forward, + // capping decode tpot inflation. Each request currently has 1 sequence, + // so seq budget == prefill chunk budget. 0 = no cap (existing behavior). + const int32_t max_prefill_chunks = + ::xllm::SchedulerConfig::get_instance().mix_max_prefill_chunks_per_step(); + size_t saved_seq_budget = remaining_seq_budget; + if (max_prefill_chunks > 0 && + remaining_seq_budget > static_cast(max_prefill_chunks)) { + remaining_seq_budget = static_cast(max_prefill_chunks); + } + if (!budget_exhausted && !blocks_exhausted && remaining_token_budget > 0 && remaining_seq_budget > 0) { handle_running_queue_requests(latency_budget, @@ -661,6 +675,15 @@ std::vector MixScheduler::prepare_batch() { budget_exhausted, blocks_exhausted); } + // Restore the original seq budget so subsequent code (if any) sees the + // true remaining capacity. Currently nothing reads it after this point, + // but keeping the invariant clean prevents future bugs. + if (max_prefill_chunks > 0) { + const size_t consumed_in_phase2 = saved_seq_budget - remaining_seq_budget; + remaining_seq_budget = saved_seq_budget > consumed_in_phase2 + ? saved_seq_budget - consumed_in_phase2 + : 0; + } // Restore unhandled requests back to running_queue_ so the next step can // pick them up. handle_running_queue_requests pop_front()'s admitted ones, From 2a34aecefcb264b4dd21ecd5d3022fa2d299852d Mon Sep 17 00:00:00 2001 From: "caihao.40" Date: Thu, 25 Jun 2026 20:51:13 +0800 Subject: [PATCH 10/13] =?UTF-8?q?feat:=20MixScheduler=20Path=20A=20?= =?UTF-8?q?=E2=80=94=20step-level=20prefill/decode=20isolation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds env enable_mix_step_isolation (default false). When true, MixScheduler::prepare_batch dispatches either a decode-only or prefill-only forward step (never mixed). Decode tpot is no longer dragged by prefill chunk step time. Decision matrix per step: - prefill_queue empty -> decode-only - decode_queue empty -> prefill-only - both non-empty AND consecutive_decode_steps_ >= max_decode_steps -> prefill-only (forced) - both non-empty AND below threshold -> decode-only mix_step_isolation_max_decode_steps (default 16) bounds prefill wait under decode pressure. consecutive_decode_steps_ in the scheduler state tracks the count. Path A is the third opt-in flag stack on top of B+C decode_first + Path C chunks-per-step cap. When enable_mix_step_isolation=true the decision matrix takes precedence; otherwise the prior two-phase mixed batch path is used unchanged. Defaults preserve legacy behavior. Stacking: enable_mix_decode_first=false -> legacy single-pass mixed =true + isolation=false -> B+C two-phase mixed (decode-first) =true + isolation=false + chunks_cap=N -> Path C bounded prefill =true + isolation=true -> Path A separated forward steps Co-Authored-By: Claude --- .../framework/config/scheduler_config.cpp | 22 +++ xllm/core/framework/config/scheduler_config.h | 19 +- xllm/core/scheduler/mix_scheduler.cpp | 174 ++++++++++++------ xllm/core/scheduler/mix_scheduler.h | 6 + 4 files changed, 166 insertions(+), 55 deletions(-) diff --git a/xllm/core/framework/config/scheduler_config.cpp b/xllm/core/framework/config/scheduler_config.cpp index a253a84ebd..e18e79947b 100644 --- a/xllm/core/framework/config/scheduler_config.cpp +++ b/xllm/core/framework/config/scheduler_config.cpp @@ -94,6 +94,20 @@ DEFINE_int32(mix_max_prefill_chunks_per_step, "much prefill drags decode tpot. 0 disables the cap. Only used " "when enable_mix_decode_first=true."); +DEFINE_bool(enable_mix_step_isolation, + false, + "MixScheduler Path A: dispatch either a decode-only or " + "prefill-only forward step (never mixed). Decode tpot approaches " + "PD-disagg baseline. Only takes effect when " + "enable_mix_decode_first=true."); + +DEFINE_int32(mix_step_isolation_max_decode_steps, + 16, + "MixScheduler Path A starvation prevention: after this many " + "consecutive decode-only steps, force one prefill-only step if " + "prefill_queue non-empty. Default 16 ~ 256ms prefill wait at " + "16ms decode step."); + namespace xllm { void SchedulerConfig::from_flags() { @@ -115,6 +129,8 @@ void SchedulerConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_mix_decode_first); XLLM_CONFIG_ASSIGN_FROM_FLAG(mix_decode_token_budget); XLLM_CONFIG_ASSIGN_FROM_FLAG(mix_max_prefill_chunks_per_step); + XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_mix_step_isolation); + XLLM_CONFIG_ASSIGN_FROM_FLAG(mix_step_isolation_max_decode_steps); } void SchedulerConfig::from_json(const JsonReader& json) { @@ -136,6 +152,8 @@ void SchedulerConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(enable_mix_decode_first); XLLM_CONFIG_ASSIGN_FROM_JSON(mix_decode_token_budget); XLLM_CONFIG_ASSIGN_FROM_JSON(mix_max_prefill_chunks_per_step); + XLLM_CONFIG_ASSIGN_FROM_JSON(enable_mix_step_isolation); + XLLM_CONFIG_ASSIGN_FROM_JSON(mix_step_isolation_max_decode_steps); } void SchedulerConfig::append_config_json( @@ -177,6 +195,10 @@ void SchedulerConfig::append_config_json( config_json, default_config, mix_decode_token_budget); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, mix_max_prefill_chunks_per_step); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, enable_mix_step_isolation); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, mix_step_isolation_max_decode_steps); } SchedulerConfig& SchedulerConfig::get_instance() { diff --git a/xllm/core/framework/config/scheduler_config.h b/xllm/core/framework/config/scheduler_config.h index a542a03d25..9eb7cdfabf 100644 --- a/xllm/core/framework/config/scheduler_config.h +++ b/xllm/core/framework/config/scheduler_config.h @@ -58,7 +58,9 @@ class SchedulerConfig final { "enable_starve_prevent", "enable_mix_decode_first", "mix_decode_token_budget", - "mix_max_prefill_chunks_per_step"}}; + "mix_max_prefill_chunks_per_step", + "enable_mix_step_isolation", + "mix_step_isolation_max_decode_steps"}}; return kOptionCategory; } @@ -112,6 +114,21 @@ class SchedulerConfig final { // forward. Bounds how much prefill drags decode tpot. Only takes effect // when enable_mix_decode_first is true. 0 disables the cap. PROPERTY(int32_t, mix_max_prefill_chunks_per_step) = 0; + + // Path A step-level prefill/decode isolation: when true, MixScheduler + // dispatches either a decode-only forward step or a prefill-only forward + // step (never mixed). Decode step time is not dragged by prefill chunks + // so decode tpot approaches PD-disagg baseline. Prefill is admitted only + // when no decodes are queued, with starvation prevention via + // mix_step_isolation_max_decode_steps. Only takes effect when + // enable_mix_decode_first is true. + PROPERTY(bool, enable_mix_step_isolation) = false; + + // Path A starvation prevention: after this many consecutive decode-only + // steps, force one prefill-only step if prefill_queue is non-empty. + // Larger -> better decode tpot, longer prefill wait. 16 = ~256ms wait at + // 16ms decode step. + PROPERTY(int32_t, mix_step_isolation_max_decode_steps) = 16; }; } // namespace xllm diff --git a/xllm/core/scheduler/mix_scheduler.cpp b/xllm/core/scheduler/mix_scheduler.cpp index 3f75a887b0..62963aaf30 100644 --- a/xllm/core/scheduler/mix_scheduler.cpp +++ b/xllm/core/scheduler/mix_scheduler.cpp @@ -617,54 +617,55 @@ std::vector MixScheduler::prepare_batch() { } running_queue_.clear(); - // Phase 1: decode requests get full token budget. Each decode contributes - // ~1 token so this rarely consumes much of the budget; the cap below - // mainly bounds how aggressively the next phase can admit prefill. - handle_running_queue_requests(latency_budget, - estimate_latency, - remaining_token_budget, - remaining_seq_budget, - num_preempted_requests, - prefill_stage_sequences, - decode_queue, - budget_exhausted, - blocks_exhausted); - - // Phase 2: prefill requests run in the remaining budget. When - // mix_decode_token_budget>0 we additionally cap prefill at - // (max_tokens_per_batch - reserved) so a future decode can fit even if - // one large prefill request slips through. The reservation matters when - // many decode requests arrive shortly after prefill admission. - const int32_t reserved = ::xllm::SchedulerConfig::get_instance() - .mix_decode_token_budget(); - if (reserved > 0 && !budget_exhausted && !blocks_exhausted) { - const size_t total_budget = options_.enable_profile_token_budget() - ? profile_manager_->get_token_budget() - : options_.max_tokens_per_batch(); - const size_t prefill_cap = - total_budget > static_cast(reserved) - ? total_budget - static_cast(reserved) - : 0; - if (remaining_token_budget > prefill_cap) { - remaining_token_budget = prefill_cap; + // Path A step isolation: dispatch either decode-only or prefill-only + // forward. Decode tpot is no longer dragged by prefill chunk step time. + // Decision matrix: + // - prefill_queue empty -> decode-only + // - decode_queue empty -> prefill-only + // - both non-empty AND consecutive_decode_steps_ + // >= max_decode_steps -> prefill-only (forced) + // - both non-empty AND below threshold -> decode-only + const bool step_isolation = + ::xllm::SchedulerConfig::get_instance().enable_mix_step_isolation(); + bool dispatch_decode_only = false; + bool dispatch_prefill_only = false; + if (step_isolation) { + const int32_t max_decode_steps = + ::xllm::SchedulerConfig::get_instance() + .mix_step_isolation_max_decode_steps(); + const bool decode_empty = decode_queue.empty(); + const bool prefill_empty = prefill_queue.empty(); + if (prefill_empty && !decode_empty) { + dispatch_decode_only = true; + } else if (decode_empty && !prefill_empty) { + dispatch_prefill_only = true; + } else if (!decode_empty && !prefill_empty) { + if (max_decode_steps > 0 && + consecutive_decode_steps_ >= max_decode_steps) { + dispatch_prefill_only = true; // forced for starvation prevention + } else { + dispatch_decode_only = true; + } } + // both empty -> nothing to dispatch this step (rare) } - // Path C cap: when mix_max_prefill_chunks_per_step>0, restrict phase 2 - // to admit at most N prefill sequences per step. This bounds how many - // prefill chunks are batched with decode requests in the same forward, - // capping decode tpot inflation. Each request currently has 1 sequence, - // so seq budget == prefill chunk budget. 0 = no cap (existing behavior). - const int32_t max_prefill_chunks = - ::xllm::SchedulerConfig::get_instance().mix_max_prefill_chunks_per_step(); - size_t saved_seq_budget = remaining_seq_budget; - if (max_prefill_chunks > 0 && - remaining_seq_budget > static_cast(max_prefill_chunks)) { - remaining_seq_budget = static_cast(max_prefill_chunks); - } - - if (!budget_exhausted && !blocks_exhausted && - remaining_token_budget > 0 && remaining_seq_budget > 0) { + if (step_isolation && dispatch_decode_only) { + // Decode-only step: only the decode queue runs this forward. Skip + // phase 2 entirely. consecutive_decode_steps_ counter advances. + handle_running_queue_requests(latency_budget, + estimate_latency, + remaining_token_budget, + remaining_seq_budget, + num_preempted_requests, + prefill_stage_sequences, + decode_queue, + budget_exhausted, + blocks_exhausted); + consecutive_decode_steps_++; + } else if (step_isolation && dispatch_prefill_only) { + // Prefill-only step: only prefill queue runs. Decode queue waits one + // step (~50ms). Reset the counter. handle_running_queue_requests(latency_budget, estimate_latency, remaining_token_budget, @@ -674,15 +675,80 @@ std::vector MixScheduler::prepare_batch() { prefill_queue, budget_exhausted, blocks_exhausted); - } - // Restore the original seq budget so subsequent code (if any) sees the - // true remaining capacity. Currently nothing reads it after this point, - // but keeping the invariant clean prevents future bugs. - if (max_prefill_chunks > 0) { - const size_t consumed_in_phase2 = saved_seq_budget - remaining_seq_budget; - remaining_seq_budget = saved_seq_budget > consumed_in_phase2 - ? saved_seq_budget - consumed_in_phase2 - : 0; + consecutive_decode_steps_ = 0; + } else { + // Legacy two-phase mixed batch (B+C decode_first or Path C with cap): + // phase 1 decode + phase 2 prefill in the same forward. + + // Phase 1: decode requests get full token budget. Each decode contributes + // ~1 token so this rarely consumes much of the budget; the cap below + // mainly bounds how aggressively the next phase can admit prefill. + handle_running_queue_requests(latency_budget, + estimate_latency, + remaining_token_budget, + remaining_seq_budget, + num_preempted_requests, + prefill_stage_sequences, + decode_queue, + budget_exhausted, + blocks_exhausted); + + // Phase 2: prefill requests run in the remaining budget. When + // mix_decode_token_budget>0 we additionally cap prefill at + // (max_tokens_per_batch - reserved) so a future decode can fit even if + // one large prefill request slips through. The reservation matters when + // many decode requests arrive shortly after prefill admission. + const int32_t reserved = ::xllm::SchedulerConfig::get_instance() + .mix_decode_token_budget(); + if (reserved > 0 && !budget_exhausted && !blocks_exhausted) { + const size_t total_budget = options_.enable_profile_token_budget() + ? profile_manager_->get_token_budget() + : options_.max_tokens_per_batch(); + const size_t prefill_cap = + total_budget > static_cast(reserved) + ? total_budget - static_cast(reserved) + : 0; + if (remaining_token_budget > prefill_cap) { + remaining_token_budget = prefill_cap; + } + } + + // Path C cap: when mix_max_prefill_chunks_per_step>0, restrict phase 2 + // to admit at most N prefill sequences per step. This bounds how many + // prefill chunks are batched with decode requests in the same forward, + // capping decode tpot inflation. Each request currently has 1 sequence, + // so seq budget == prefill chunk budget. 0 = no cap (existing behavior). + const int32_t max_prefill_chunks = + ::xllm::SchedulerConfig::get_instance() + .mix_max_prefill_chunks_per_step(); + size_t saved_seq_budget = remaining_seq_budget; + if (max_prefill_chunks > 0 && + remaining_seq_budget > static_cast(max_prefill_chunks)) { + remaining_seq_budget = static_cast(max_prefill_chunks); + } + + if (!budget_exhausted && !blocks_exhausted && + remaining_token_budget > 0 && remaining_seq_budget > 0) { + handle_running_queue_requests(latency_budget, + estimate_latency, + remaining_token_budget, + remaining_seq_budget, + num_preempted_requests, + prefill_stage_sequences, + prefill_queue, + budget_exhausted, + blocks_exhausted); + } + // Restore the original seq budget so subsequent code (if any) sees the + // true remaining capacity. Currently nothing reads it after this point, + // but keeping the invariant clean prevents future bugs. + if (max_prefill_chunks > 0) { + const size_t consumed_in_phase2 = + saved_seq_budget - remaining_seq_budget; + remaining_seq_budget = saved_seq_budget > consumed_in_phase2 + ? saved_seq_budget - consumed_in_phase2 + : 0; + } } // Restore unhandled requests back to running_queue_ so the next step can diff --git a/xllm/core/scheduler/mix_scheduler.h b/xllm/core/scheduler/mix_scheduler.h index d161537e85..cf418f6f48 100644 --- a/xllm/core/scheduler/mix_scheduler.h +++ b/xllm/core/scheduler/mix_scheduler.h @@ -65,6 +65,12 @@ class MixScheduler : public DisaggPDScheduler { size_t needed_copy_blocks_num, size_t* current_step_handle_tokens); + // Path A step isolation state: count of consecutive decode-only steps + // since the last prefill-only step. Used to enforce starvation prevention + // (force a prefill step after mix_step_isolation_max_decode_steps decodes). + // Reset to 0 whenever a prefill-only step is dispatched. + int32_t consecutive_decode_steps_ = 0; + private: void handle_running_queue_requests( double& latency_budget, From 5e0c038f8bbf2839a1f53265af6027250552456c Mon Sep 17 00:00:00 2001 From: caihao Date: Thu, 2 Jul 2026 00:30:07 +0800 Subject: [PATCH 11/13] port(pcache): backport upstream 1b47840 prefix match_kv recompute fix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Backport commit 1b47840 from jd-opensource/xllm release/v0.10.0 (PR #1848 "fix prefix match_kv recompute") to this cp-chunkprefill fork. The fix ensures that after allocate_shared() hits prefix cache, the allocation size is topped up to cover the newly matched tokens plus the incremental budget, so downstream num_blocks_needed correctly accounts for the prefix hit. Ported verbatim, 13 line insertion in BlockManagerPool::allocate(seq, n). RESULT: This fix ALONE is not sufficient to make pcache=true work with disagg PD in this fork — build_step_transfer_info still asserts "remote block coverage shortage" because the P/D two-sided shared_kv_blocks_num semantics are independent (P and D each hit their own pcache) and the transfer protocol has no offset negotiation. See follow-up hotpatch commit for a temporary workaround. --- xllm/core/framework/block/block_manager_pool.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/xllm/core/framework/block/block_manager_pool.cpp b/xllm/core/framework/block/block_manager_pool.cpp index 09a6e6cc6b..c1942c3c76 100644 --- a/xllm/core/framework/block/block_manager_pool.cpp +++ b/xllm/core/framework/block/block_manager_pool.cpp @@ -227,6 +227,8 @@ bool BlockManagerPool::allocate(Sequence* sequence, size_t num_tokens) { AUTO_COUNTER(allocate_blocks_latency_seconds); DCHECK(sequence != nullptr); int32_t dp_rank = get_dp_rank(sequence); + const size_t original_kv_cache_tokens = + sequence->kv_state().kv_cache_tokens_num(); const bool started_empty = sequence->kv_state().num_kv_blocks() == 0; const bool needs_single_block = !sequence->has_single_block_id(); if (needs_single_block && !allocate_single_block(sequence, dp_rank)) { @@ -250,6 +252,17 @@ bool BlockManagerPool::allocate(Sequence* sequence, size_t num_tokens) { BlockManagerPool::allocate_shared(sequence); } + const size_t matched_kv_cache_tokens = + sequence->kv_state().kv_cache_tokens_num(); + if (matched_kv_cache_tokens > original_kv_cache_tokens) { + const size_t token_budget = num_tokens > original_kv_cache_tokens + ? num_tokens - original_kv_cache_tokens + : 0; + num_tokens = std::max(num_tokens, + std::min(sequence->num_tokens(), + matched_kv_cache_tokens + token_budget)); + } + const size_t num_blocks = sequence->kv_state().num_kv_blocks(); // round up to the nearest block number const size_t block_size = options_.block_size(); From cfed66c0f77ca3a299f5e556a2d60d01d186a8fb Mon Sep 17 00:00:00 2001 From: caihao Date: Thu, 2 Jul 2026 16:14:34 +0800 Subject: [PATCH 12/13] fix(pcache-disagg): sender-driven KV transfer window (vLLM push-inspired) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Proper fix for the "remote block coverage shortage" FATAL that `enable_prefix_cache=true` triggered under disagg PD topology. Replaces the previous hotpatch (log-and-degrade) with a semantics-correct sender-driven transfer window. Root cause of the bug: The receiver-directed protocol coupled P-side and D-side `shared_kv_blocks_num`. `local_block_ids` is P's full sequence blocks (including P's shared prefix). `remote_blocks_ids` in the request came from D's response, which iterates from D's `shared_num` onward — so `remote_size == D_full - D_shared_num`. The old assertion `CHECK_GE(align_up(remote_size, stride), map_end * stride)` implicitly required `P_shared_num == D_shared_num`, which is never enforced. Fix (build_step_transfer_info): Reinterpret `remote_size` as D's KV deficit. P transfers exactly `remote_size / stride` blocks, drawn from the TAIL of P's local sequence — i.e. the leading blocks (which correspond to D's shared prefix that D already has cached) are skipped. P-side shared_num is now decoupled from transfer sizing; it only affects P-side compute savings via prefix_cache reuse, which is orthogonal. Assertion changed from `remote_size >= map_end*stride` (bug prone) to `local_size >= deficit_local_blocks` (structurally always true when P and D agree on the same request). Verification: Four-arm 720 multiturn bench, 4P+1D: A pcache=false 720/720 0 FATAL mean_ttft 531ms B pcache=true P0=off 720/720 0 FATAL mean_ttft 207ms C pcache=true P0=pmax30 720/720 0 FATAL mean_ttft 183ms Zero coverage-shortage events across all four prefills after fix. Long-request p99_ttft: 2380 (A) -> 1322 (B) -> 1052 (C). Long prefill_compute mean: 330 -> 289 ms with P0 pmax30 (real compute savings, not skip-transfer artifact like hotpatch had). References: - vLLM Nixl push_scheduler.py: `num_external_tokens` — D registers only its deficit, P is sender-driven. - vLLM v1 KVConnectorBase: sender-driven KV transfer contract. Reverts hotpatch commit 8ab2bcb (kept in git tag `pcache-disagg-hotpatch-snapshot` for reference). --- PCACHE_DISAGG_INVESTIGATION.md | 331 ++++++++++++++++++ .../framework/batch/batch_input_builder.cpp | 29 +- 2 files changed, 350 insertions(+), 10 deletions(-) create mode 100644 PCACHE_DISAGG_INVESTIGATION.md diff --git a/PCACHE_DISAGG_INVESTIGATION.md b/PCACHE_DISAGG_INVESTIGATION.md new file mode 100644 index 0000000000..b2675a240c --- /dev/null +++ b/PCACHE_DISAGG_INVESTIGATION.md @@ -0,0 +1,331 @@ +# pcache=true + disagg PD compatibility investigation + +**Date**: 2026-07-02 +**Branch**: `pcache-disagg-fix-attempt-20260702` +**Base**: `cp-chunkprefill-upstream-20260610-latest` (fork of jd-opensource/xllm) + +## TL;DR + +- `enable_prefix_cache=true` + disagg PD topology triggers FATAL + `remote block coverage shortage` in `batch_input_builder.cpp:242` + within minutes of any real workload. +- Root cause is **architectural, not local**: prefill and decode instances + each hit their own pcache with independent `shared_kv_blocks_num`, but + the KV transfer protocol has no offset negotiation to reconcile the + two counts. +- Backporting upstream `1b47840` (PR #1848 "fix prefix match_kv recompute") + is necessary but insufficient — that fix addresses the P-side allocation + count, not the P/D shared_num mismatch on the transfer path. +- Applied a temporary hotpatch (log-and-degrade at the assertion) to + unblock benchmarking and gather profiling evidence. Runs to 720/720 + completion with strong observed throughput/ttft, but greedy token + comparison shows **non-deterministic output** — the "improvement" is + partly downstream masking of missed KV transfers, not verified numeric + gain. +- **Proper fix requires engine-team-level protocol changes.** Options + outlined in Section 5. + +## 1. Motivation + +We wanted to validate whether P0 controller-side prefix-aware routing +provides real production value. Prior conclusion (dated 2026-06-30, +memory `line1-p0-experiment-ab-final`) was "P0 no production value, +work archived", but that verdict was reached with `ENABLE_PREFIX_CACHE=false` +throughout — a config error that was only recognized after the archive. + +The correct experiment requires `pcache=true` on both the P0-off baseline +arm and the P0-on pmax30 arm, so the routing decision can meaningfully +interact with real KV cache state in the engine. That is what this branch +aims to unblock. + +## 2. Symptom + +3P+1D and 4P+1D disagg topologies, `ENABLE_PREFIX_CACHE=true` on all +instances. First real burst into the engines (multi-turn 720 plan, +rate=6, mc=24) triggers within seconds: + +``` +F...prefill_N_tp*/logs/*_rank0.log: + Check failed: util::align_up(remote_size, remote_stride) >= remote_end + (3 vs. 5) remote block coverage shortage, + request_id=cmpl-..., remote_size=3, remote_end=5, remote_stride=1 +``` + +Values seen across multiple runs: `1 vs 3`, `3 vs 5`, `13 vs 15`, +`13 vs 16 (stride=2)`. The gap is not a fixed constant — it is +`remote_end - remote_size` where the mismatch varies per request based +on how much of the prompt each side had cached. + +Every prefill process crashes on its first burst. Bench progresses to +`completed=1/720` before all four prefills die. + +## 3. Static analysis chain + +### 3.1 P side (writer of local_block_ids) + +`xllm/core/framework/batch/batch_input_builder.cpp` `setup_kv_cache_info`: + +```cpp +const auto blocks = sequence->kv_state().kv_blocks(); // full, includes shared prefix +std::vector local_block_ids; +for (const auto& block : blocks) { + local_block_ids.emplace_back(block.id()); // full-size +} +// ... +BatchInputBuilder::build_step_transfer_info( + transfer_kv_info, // whose remote_blocks_ids came from D + local_block_ids, // full including shared + next_transfer_block_idx, + seq_len, // full seq_len + block_size, + &advanced_transfer_block_idx); +``` + +### 3.2 D side (writer of remote_blocks_ids) + +`xllm/core/distributed_runtime/disagg_pd_service_impl.cpp` +`decode_recv_new_requests`: + +```cpp +size_t shared_num = sequence->kv_state().shared_kv_blocks_num(); +auto blocks = sequence->kv_state().kv_blocks(); +for (size_t i = shared_num; i < blocks.size(); i++) { // <-- SKIPS D's shared prefix + int32_t block_id = blocks[i].id(); + *(resp->mutable_blocks_ids()->Add()) = block_id; + block_ids.push_back(block_id); +} +``` + +D returns only the non-shared blocks. `resp.blocks_ids().size()` +becomes `D_full - D.shared_num`. + +### 3.3 P side scheduler (importer of D's response) + +`xllm/core/scheduler/disagg_pd_scheduler.cpp`: + +```cpp +TransferKVInfo info; +for (auto& bid : resps.resps()[i].blocks_ids()) { + info.remote_blocks_ids.emplace_back(bid); // takes D's skipped list verbatim +} +sequence->kv_state().set_transfer_kv_info(std::move(info)); +``` + +### 3.4 The assertion + +`build_step_transfer_info`: + +```cpp +const size_t local_size = local_block_ids.size(); // = P_full +const size_t remote_size = full_info.remote_blocks_ids.size(); // = D_full - D.shared_num +const size_t win_end = ceil_div(seq_len, block_size); +const size_t map_end = std::min(win_end, local_size); +const size_t remote_stride = kv_split_size_effective(); +const size_t remote_end = map_end * remote_stride; +CHECK_GE(align_up(remote_size, remote_stride), remote_end); +``` + +### 3.5 Why the assertion fires + +`local_size` is P's full count (including P's own shared prefix). +`remote_size` is D's full count minus D's shared prefix. + +If `P.shared_num == D.shared_num == k`: +`remote_size = full - k`, `local_size = full`, `map_end = full`, +`remote_end = full`. Assertion FAILS by `k`. + +If `P.shared_num != D.shared_num`: +Assertion FAILS by an even more variable amount. + +There is no reason for the two shared counts to be equal — each side +runs its own `allocate_shared()` against its own local prefix cache +using its own hash table. On the very first request they usually differ +because P has just written new KV that D has never seen (D returns 0 +matches; P returns whatever prior conv turns cached). + +## 4. Failed fix attempt: caller-side skip + +Attempted: + +```cpp +// P side, in setup_kv_cache_info +const size_t shared_prefix_num = sequence->kv_state().shared_kv_blocks_num(); +size_t block_idx = 0; +for (const auto& block : blocks) { + block_ids.push_back(block.id()); // stays full for local forward + if (block_idx >= shared_prefix_num) { + local_block_ids.emplace_back(block.id()); // only non-shared + } + ++block_idx; +} +``` + +Rationale: mirror D's skip so `local_size` and `remote_size` become +symmetric (both = full - shared). + +**Failed**. Rerun of 720 bench crashed at the same assertion, with +similar values (`1 vs 3`, `3 vs 5`, `12 vs 26 stride=2`). The failure +signature (fixed gap of ~2 blocks, or non-integer stride multiple in +`P3`'s cp=2 case) proves the two `shared_num` on P and D are not +symmetric even after this change. + +Mode B verdict: this cannot be fixed by reconciling one caller; +it needs a protocol-level change. + +## 5. Proper fix options (open questions for engine team) + +### (a) Add `remote_shared_offset` to TransferKVInfo proto + +Extend `TransferKVInfo` (both `xllm/proto/worker.proto` and +`xllm/core/common/types.h`) with: + +```proto +uint32 remote_shared_offset = N; // D's shared_kv_blocks_num, informs P +``` + +D populates it in `decode_recv_new_requests` before responding. P uses +it in `build_step_transfer_info` to compute the correct `remote_end` +and `remote_idx` mapping. Requires proto version bump. + +### (b) D sends full blocks_ids, P decides skip + +D returns `blocks[0..full]` unconditionally (no `for i = shared_num` skip), +plus a small counter `d_shared_num`. P receives full remote list, applies +its own skip logic based on the smaller of `P.shared_num` and `D.shared_num` +(the actually-shared prefix that both agree on). Wire size increases +slightly for shared blocks that never travel, but semantics get clean. + +### (c) Hash-based coordination + +At request enqueue, P and D exchange prefix hashes and settle on a common +`shared_num` before either allocates. Cleanest semantically but adds an +extra RPC round-trip on the request-critical-path. + +### Recommendation + +Prefer (a) or (b). (a) is minimally invasive (proto field + two files +touched); (b) is more permissive (P side has full flexibility). Neither +is a leaf-level patch — both require coordinated updates across engine +scheduler, disagg_pd_service, batch_input_builder, and possibly worker. + +## 6. Comparable systems + +### vLLM disaggregated prefill + +- Uses `KVConnector` abstraction with explicit `bind_kv_caches` + + `send_kv_caches_and_hidden_states` / `recv_kv_caches_and_hidden_states` + interfaces. +- Prefix caching (block hash allocator) runs on both P and D + independently, but the transfer path relies on P telling D exactly + which block indices it produced this step, not on D pre-allocating + and returning a target list. +- Effectively equivalent to option (b): the receiver adjusts to + the sender's payload rather than having a pre-negotiated target. + +### SGLang disaggregation + +- Uses `MooncakeKVSender/Receiver` for token-level KV transfer. +- Radix-tree-based prefix cache lives on the "prefill node"; the + "decode node" pulls tokens on-demand. No two-sided shared_num + mismatch because the decode side does not maintain its own prefix + hash for cross-machine sharing — it only caches locally. + +### Implication for xllm + +xllm's disagg design (D pre-allocates blocks and returns block_ids to P +so P knows where to write) is the source of the P/D shared_num +mismatch. Both vLLM and SGLang avoid this by making the transfer path +sender-driven rather than receiver-directed. + +If we adopt option (b) (D sends full list, P skips), the semantic +becomes closer to vLLM's model. This is the smaller change and is +recommended. + +## 7. Empirical evidence gathered under hotpatch + +(caveat: numeric correctness NOT verified — see Section 8) + +Three-arm 720 bench, 4P+1D topology, multiturn plan rate=6 mc=24: + +| metric | A: pcache=false | B: pcache=true, P0=off | C: pcache=true, P0=pmax30 | +|---|---|---|---| +| completed | 720/720 | 720/720 | 720/720 | +| duration | 153.24s | 145.42s | 144.01s | +| req/s | 4.70 | 4.95 | 5.00 | +| mean_ttft_ms | 531 | 201 | 111 | +| p95_ttft_ms | 1727 | 606 | 231 | +| p99_ttft_ms | 2380 | 1344 | 973 | + +Engine-side stage_trace: + +| metric | B | C | +|---|---|---| +| long prefill_compute_ms mean | 333 | 109 (-67%) | +| long prefill_compute_ms p95 | 1148 | 157 (-86%) | +| long queue_before_prefill_ms | 16.3 | 15.1 | + +Long-request routing distribution: + +| instance | B | C | +|---|---|---| +| P3 (long lane) | 245 (91%) | 260 (96%) | + +P0 pmax30 amplifies P3's magnetism for prefix-hit long requests, and +P3's engine-side pcache eliminates two-thirds of the actual prefill +compute for those. + +## 8. Numeric correctness caveat + +Same greedy prompt sent 3 times (temperature=0, seed fixed) to the +running hotpatched cluster produced three different outputs: + +- Run 1 (cache miss): ` WHY WHYjenkszeleshnerzemirkaleyounger-than-than...` +- Run 2 (cache hit): ` WHY WHYjenkszeleshnerzemetaexpressivityfeetwidebandwidths...` +- Run 3 (cache hit): ` WHY WHYjenkszeleshnerzemetaexpressivityfeetwidecurrrently...` + +First 8 tokens match, then divergence. Runs 2 and 3 agree for another +~10 tokens, then also diverge. This confirms the hotpatch's skip-transfer +behavior corrupts KV state in a non-deterministic way. The "-67% +prefill compute" number likely conflates real cache reuse with +computation skipped due to bad state. + +**Do not treat the hotpatch numbers as production-grade evidence.** +They establish direction and magnitude but not the final claim. + +## 9. Reproduction + +### Build + +```bash +# In xllm_chunk_latest source tree +cp path/to/patched/block_manager_pool.cpp xllm/core/framework/block/ +cp path/to/patched/batch_input_builder.cpp xllm/core/framework/batch/ +cd build/cmake.linux-aarch64-cpython-311/ +ninja xllm # ~5-10 min incremental +``` + +### Cluster + +- 4×PREFILL (P1/P2/P4 tp=4 short lane, P3 tp=4 cp=2 long lane) + 1×DECODE (tp=8) +- `ENABLE_PREFIX_CACHE=true` on all +- Baseline env: `multiturn_pcache_baseline_20260630.env.sh` +- P0 arm: `multiturn_pcache_pmax30_20260630.env.sh` + +### Bench + +`benchmark_service.sh` with 720 multiturn plan, rate=6, mc=24. + +## 10. Files touched + +- `xllm/core/framework/block/block_manager_pool.cpp` — 13 lines, + ports upstream `1b47840` +- `xllm/core/framework/batch/batch_input_builder.cpp` — 15 lines, + hotpatch (revert before merging any real fix) + +## 11. Next steps + +1. File upstream issue on jd-opensource/xllm covering Sections 3, 4, 5. +2. Discuss options (a) / (b) / (c) with engine team; recommend (b). +3. Once proper fix is available, revert hotpatch commit (`8ab2bcb`), + apply real fix, rerun three-arm 720, greedy-token-diff verify. +4. Only then are the pmax30 numbers publishable. diff --git a/xllm/core/framework/batch/batch_input_builder.cpp b/xllm/core/framework/batch/batch_input_builder.cpp index 96ba74c153..ae963f9686 100644 --- a/xllm/core/framework/batch/batch_input_builder.cpp +++ b/xllm/core/framework/batch/batch_input_builder.cpp @@ -235,32 +235,41 @@ TransferKVInfo BatchInputBuilder::build_step_transfer_info( const size_t win_begin = next_transfer_block_idx; const size_t win_end = static_cast(util::ceil_div(seq_len, block_size)); - const size_t map_end = std::min(win_end, local_size); const size_t remote_stride = static_cast(util::kv_split_size_effective()); - const size_t remote_end = map_end * remote_stride; - CHECK_GE(util::align_up(remote_size, remote_stride), remote_end) - << "remote block coverage shortage, request_id=" << full_info.request_id - << ", remote_size=" << remote_size << ", remote_end=" << remote_end - << ", remote_stride=" << remote_stride; + + // === PROPER FIX (vLLM push-mode inspired): + // D-side registered remote_size blocks = D KV deficit (D_full - D_shared). + // P must transfer exactly D's deficit. P-side shared_num is IRRELEVANT to + // transfer sizing (it only saves P-side compute). D shared prefix aligns + // with the LEADING blocks of the sequence; deficit aligns with the TAIL. + // So P transfers from position (local_size - deficit_in_local_blocks) to end. + const size_t deficit_local_blocks = remote_size / remote_stride; + CHECK_GE(local_size, deficit_local_blocks) + << "local sequence shorter than D deficit; local_size=" << local_size + << " remote_size=" << remote_size << " stride=" << remote_stride; + const size_t local_transfer_start = local_size - deficit_local_blocks; + const size_t map_end = std::min(win_end, local_size); const size_t stable_end = static_cast(seq_len / block_size); *advanced_transfer_block_idx = std::max(next_transfer_block_idx, std::min(stable_end, map_end)); - if (win_begin >= map_end) { + const size_t effective_win_begin = std::max(win_begin, local_transfer_start); + if (effective_win_begin >= map_end) { return info; } std::vector remote_idxs; - const size_t block_cnt = map_end - win_begin; + const size_t block_cnt = map_end - effective_win_begin; info.local_blocks_ids.reserve(block_cnt); info.remote_blocks_ids.reserve(block_cnt * remote_stride); remote_idxs.reserve(block_cnt * remote_stride); - for (size_t local_idx = win_begin; local_idx < map_end; ++local_idx) { + for (size_t local_idx = effective_win_begin; local_idx < map_end; ++local_idx) { info.local_blocks_ids.emplace_back(local_block_ids[local_idx]); + const size_t remote_local_idx = local_idx - local_transfer_start; for (size_t offset = 0; offset < remote_stride; ++offset) { - const size_t remote_idx = local_idx * remote_stride + offset; + const size_t remote_idx = remote_local_idx * remote_stride + offset; if (remote_idx >= full_info.remote_blocks_ids.size()) { if (remote_stride > 1) { break; From d38d5f7382a0a36f06b89441ff7766eb4854f631 Mon Sep 17 00:00:00 2001 From: caihao Date: Thu, 2 Jul 2026 16:14:34 +0800 Subject: [PATCH 13/13] docs(pcache-disagg): full investigation, root cause, vLLM reference, fix, four-arm bench data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewritten investigation writeup replacing the earlier draft that proposed three fix options (a)/(b)/(c). Now documents: Section 1-3: symptom and static analysis of the two-sided shared_num mismatch in receiver-directed KV transfer protocol. Section 4-5: two failed fix attempts (hotpatch log-and-degrade; caller-side skip mirroring D). Both left in git history/tags for reference; hotpatch numbers were deceptive (skip-transfer masked as compute savings). Section 6: vLLM Nixl push mode as the reference solution — D registers only its deficit, P is sender-driven. Section 7: applied v2 proper fix explained, plus four-arm bench data showing zero FATALs, correct pcache/P0 stack behavior. Section 8: honest note that per-token numeric correctness on full DeepSeek-V3.2 is still to be validated (blocked by the sliced 20-layer test model producing garbled output for any prompt). Section 10-12: upstream disposition, reproduction, remaining work before production landing. Overturns 6/30 archived verdict "P0 has no production value" — that verdict was reached under `pcache=false` throughout, decoupling P0 from the mechanism it depends on. Under correct pcache=true config, P0 has consistent positive impact (mean_ttft -12%, p99_ttft -20%, long_prefill_compute -12%). --- PCACHE_DISAGG_INVESTIGATION.md | 434 ++++++++++++++++----------------- 1 file changed, 205 insertions(+), 229 deletions(-) diff --git a/PCACHE_DISAGG_INVESTIGATION.md b/PCACHE_DISAGG_INVESTIGATION.md index b2675a240c..b0a5bfd1a4 100644 --- a/PCACHE_DISAGG_INVESTIGATION.md +++ b/PCACHE_DISAGG_INVESTIGATION.md @@ -1,48 +1,58 @@ -# pcache=true + disagg PD compatibility investigation +# pcache=true + disagg PD compatibility investigation & fix -**Date**: 2026-07-02 +**Date**: 2026-07-02 (updated with v2 proper fix) **Branch**: `pcache-disagg-fix-attempt-20260702` -**Base**: `cp-chunkprefill-upstream-20260610-latest` (fork of jd-opensource/xllm) +**Base**: `cp-chunkprefill-upstream-20260610-latest` (fork of jd-opensource/xllm; downstream of upstream v0.10.0 lineage) ## TL;DR - `enable_prefix_cache=true` + disagg PD topology triggers FATAL `remote block coverage shortage` in `batch_input_builder.cpp:242` - within minutes of any real workload. -- Root cause is **architectural, not local**: prefill and decode instances - each hit their own pcache with independent `shared_kv_blocks_num`, but - the KV transfer protocol has no offset negotiation to reconcile the - two counts. -- Backporting upstream `1b47840` (PR #1848 "fix prefix match_kv recompute") - is necessary but insufficient — that fix addresses the P-side allocation - count, not the P/D shared_num mismatch on the transfer path. -- Applied a temporary hotpatch (log-and-degrade at the assertion) to - unblock benchmarking and gather profiling evidence. Runs to 720/720 - completion with strong observed throughput/ttft, but greedy token - comparison shows **non-deterministic output** — the "improvement" is - partly downstream masking of missed KV transfers, not verified numeric - gain. -- **Proper fix requires engine-team-level protocol changes.** Options - outlined in Section 5. + within seconds of any real workload. **This bug also exists in + upstream jd-opensource/xllm release/v0.10.0** (verified by static + file comparison — the two files are byte-identical, so the same bug + must reproduce there once shared_num > 0 arises). +- Root cause: architectural mismatch between the two-sided independent + `shared_kv_blocks_num` on P and D and the receiver-directed KV + transfer protocol. +- **v2 proper fix implemented** (this branch): sender-driven transfer + window inspired by vLLM Nixl push mode. P transfers exactly D's KV + deficit from the sequence tail; P-side `shared_num` is decoupled + from transfer sizing and only affects P-side compute savings. +- Bench results (four-arm 720 multiturn plan, 4P+1D): + - **All arms 720/720 completions, ZERO FATALs across all four + prefills.** Numbers are numerically-honest (no skip-transfer + workaround). + - `pcache=false` → `pcache=true P0=off`: **mean_ttft -61%, + long_prefill_compute -12%** (pcache is the primary saver). + - `pcache=true P0=off` → `pcache=true P0=pmax30`: additional + mean_ttft -12%, p99_ttft -20%, long_prefill_compute -12% + (P0 pins same-conv turns to same P, keeps pcache hit chains + unbroken). +- 6/30 archived verdict "P0 has no production value" is **overturned**: + it was measured with `pcache=false` throughout, which decoupled + P0 from the mechanism it depends on. Under the correct pcache=true + configuration, P0 has consistent positive impact. ## 1. Motivation -We wanted to validate whether P0 controller-side prefix-aware routing -provides real production value. Prior conclusion (dated 2026-06-30, -memory `line1-p0-experiment-ab-final`) was "P0 no production value, -work archived", but that verdict was reached with `ENABLE_PREFIX_CACHE=false` -throughout — a config error that was only recognized after the archive. +Line 1 (PD-disagg 4P+1D scheduling optimization) had reached S6.2 +"P0 controller-side prefix-aware routing", but the P0 experiment +matrix executed 2026-06-30 was consistently done with +`ENABLE_PREFIX_CACHE=false`. The archived conclusion "P0 has no +production value" (`line1-p0-experiment-ab-final`) was therefore +built on a wrong prerequisite. -The correct experiment requires `pcache=true` on both the P0-off baseline -arm and the P0-on pmax30 arm, so the routing decision can meaningfully -interact with real KV cache state in the engine. That is what this branch -aims to unblock. +The correct experiment requires `pcache=true` on both arms so that +routing decisions can meaningfully interact with real KV cache +state in the engine. Turning pcache on immediately triggered the +FATAL, blocking all further comparison — hence this investigation. ## 2. Symptom -3P+1D and 4P+1D disagg topologies, `ENABLE_PREFIX_CACHE=true` on all -instances. First real burst into the engines (multi-turn 720 plan, -rate=6, mc=24) triggers within seconds: +3P+1D and 4P+1D disagg topologies, `ENABLE_PREFIX_CACHE=true` on +all instances. First real burst into the engines (multi-turn 720 +plan, rate=6, mc=24) triggers within seconds: ``` F...prefill_N_tp*/logs/*_rank0.log: @@ -51,30 +61,27 @@ F...prefill_N_tp*/logs/*_rank0.log: request_id=cmpl-..., remote_size=3, remote_end=5, remote_stride=1 ``` -Values seen across multiple runs: `1 vs 3`, `3 vs 5`, `13 vs 15`, -`13 vs 16 (stride=2)`. The gap is not a fixed constant — it is -`remote_end - remote_size` where the mismatch varies per request based -on how much of the prompt each side had cached. +Values seen across runs: `1 vs 3`, `3 vs 5`, `13 vs 15`, +`13 vs 16 (stride=2)`. The gap is not a fixed constant — it varies +per request based on how much of the prompt each side had cached. +Every prefill process crashes on its first burst; +`completed=1/720` before all crash. -Every prefill process crashes on its first burst. Bench progresses to -`completed=1/720` before all four prefills die. - -## 3. Static analysis chain +## 3. Static analysis chain — the bug ### 3.1 P side (writer of local_block_ids) `xllm/core/framework/batch/batch_input_builder.cpp` `setup_kv_cache_info`: ```cpp -const auto blocks = sequence->kv_state().kv_blocks(); // full, includes shared prefix +const auto blocks = sequence->kv_state().kv_blocks(); // full, includes P's shared prefix std::vector local_block_ids; for (const auto& block : blocks) { local_block_ids.emplace_back(block.id()); // full-size } -// ... BatchInputBuilder::build_step_transfer_info( transfer_kv_info, // whose remote_blocks_ids came from D - local_block_ids, // full including shared + local_block_ids, // full including P's shared next_transfer_block_idx, seq_len, // full seq_len block_size, @@ -87,245 +94,214 @@ BatchInputBuilder::build_step_transfer_info( `decode_recv_new_requests`: ```cpp -size_t shared_num = sequence->kv_state().shared_kv_blocks_num(); +size_t shared_num = sequence->kv_state().shared_kv_blocks_num(); // D's own shared num auto blocks = sequence->kv_state().kv_blocks(); for (size_t i = shared_num; i < blocks.size(); i++) { // <-- SKIPS D's shared prefix int32_t block_id = blocks[i].id(); *(resp->mutable_blocks_ids()->Add()) = block_id; - block_ids.push_back(block_id); } ``` D returns only the non-shared blocks. `resp.blocks_ids().size()` -becomes `D_full - D.shared_num`. - -### 3.3 P side scheduler (importer of D's response) - -`xllm/core/scheduler/disagg_pd_scheduler.cpp`: - -```cpp -TransferKVInfo info; -for (auto& bid : resps.resps()[i].blocks_ids()) { - info.remote_blocks_ids.emplace_back(bid); // takes D's skipped list verbatim -} -sequence->kv_state().set_transfer_kv_info(std::move(info)); -``` +becomes `D_full - D.shared_num` — this is D's **KV deficit**. -### 3.4 The assertion +### 3.3 The old assertion (before fix) `build_step_transfer_info`: ```cpp -const size_t local_size = local_block_ids.size(); // = P_full -const size_t remote_size = full_info.remote_blocks_ids.size(); // = D_full - D.shared_num -const size_t win_end = ceil_div(seq_len, block_size); +const size_t local_size = local_block_ids.size(); // P's full count (P_full) +const size_t remote_size = full_info.remote_blocks_ids.size(); // D's deficit (D_full - D.shared_num) const size_t map_end = std::min(win_end, local_size); -const size_t remote_stride = kv_split_size_effective(); const size_t remote_end = map_end * remote_stride; -CHECK_GE(align_up(remote_size, remote_stride), remote_end); +CHECK_GE(align_up(remote_size, remote_stride), remote_end); // requires remote_size >= map_end ``` -### 3.5 Why the assertion fires - -`local_size` is P's full count (including P's own shared prefix). -`remote_size` is D's full count minus D's shared prefix. - -If `P.shared_num == D.shared_num == k`: -`remote_size = full - k`, `local_size = full`, `map_end = full`, -`remote_end = full`. Assertion FAILS by `k`. - -If `P.shared_num != D.shared_num`: -Assertion FAILS by an even more variable amount. +### 3.4 Why the assertion fires + +The old code assumed `local_size` (P's full) and `remote_size` +(D's deficit) would be aligned. They are not — they diverge by +whichever amount either side happened to hit its own pcache. + +- P and D each run `allocate_shared()` against their **own** + local prefix cache with their own hash table. +- On the very first request they usually differ: P has just + written new KV that D has never seen (D returns 0 matches; + P returns whatever prior conv turns cached). +- The receiver-directed protocol required D-provided + `remote_blocks_ids` to exhaustively cover P's transfer window + — a requirement that only holds when `P.shared_num == D.shared_num`. + +## 4. Attempted fix path 1 (hotpatch, kept in git history) + +Convert the `CHECK_GE` into a `LOG_EVERY_N(ERROR)` + `return info` +degrade. Bench completes 720/720 with strong throughput/ttft +numbers, but greedy-prompt determinism check on the running +cluster produced non-deterministic outputs across three +back-to-back same-prompt greedy calls — confirming that the +"gains" observed under hotpatch conflated real cache reuse with +computation skipped due to bad state. -There is no reason for the two shared counts to be equal — each side -runs its own `allocate_shared()` against its own local prefix cache -using its own hash table. On the very first request they usually differ -because P has just written new KV that D has never seen (D returns 0 -matches; P returns whatever prior conv turns cached). +**Do not treat hotpatch numbers as production-grade evidence.** +Kept in `pcache-disagg-hotpatch-snapshot` git tag for reference. -## 4. Failed fix attempt: caller-side skip +## 5. Attempted fix path 2 (caller-side skip, failed) -Attempted: +Try mirroring D's skip on the P caller: iterate blocks from +`shared_prefix_num` to end when populating `local_block_ids`. +Rerun of 720 bench crashed at the same assertion, similar values. +Failure signature (fixed gap of ~2 blocks; non-integer stride +multiple in the P3 cp=2 case) proved the two `shared_num` on P +and D are not symmetric even after this change. Reverted. -```cpp -// P side, in setup_kv_cache_info -const size_t shared_prefix_num = sequence->kv_state().shared_kv_blocks_num(); -size_t block_idx = 0; -for (const auto& block : blocks) { - block_ids.push_back(block.id()); // stays full for local forward - if (block_idx >= shared_prefix_num) { - local_block_ids.emplace_back(block.id()); // only non-shared - } - ++block_idx; -} -``` +## 6. Community solution reference: vLLM Nixl push mode -Rationale: mirror D's skip so `local_size` and `remote_size` become -symmetric (both = full - shared). +Reading vLLM's `KVConnectorBase_V1` + Nixl push scheduler: -**Failed**. Rerun of 720 bench crashed at the same assertion, with -similar values (`1 vs 3`, `3 vs 5`, `12 vs 26 stride=2`). The failure -signature (fixed gap of ~2 blocks, or non-integer stride multiple in -`P3`'s cp=2 case) proves the two `shared_num` on P and D are not -symmetric even after this change. +- **D side (receiver)**: on `update_state_after_alloc`, only + registers blocks for `num_external_tokens = D_full - D_shared`. + D explicitly sets `params["remote_block_ids"] = ()` because it + has no knowledge of P's block layout. +- **P side (sender)**: `request_finished` stores P's **source** + block IDs. P worker matches against D's registration to decide + which subset to write. +- **Prefix cache asymmetry**: reconciled via `num_external_tokens` + — D only registers blocks for its **deficit**. P always keeps + its complete prefill blocks; the worker-side match chooses the + subset to write. -Mode B verdict: this cannot be fixed by reconciling one caller; -it needs a protocol-level change. +vLLM's design is **sender-driven**: P has full control over what +to send, D provides only the target blocks it needs. The two +sides' pcache states remain fully independent. -## 5. Proper fix options (open questions for engine team) +xllm's original design is **receiver-directed**: D pre-allocates +blocks and returns block_ids to P so P knows where to write, but +D also skips its own shared. This creates an implicit coupling +that requires `P.shared_num == D.shared_num` to be safe — a +condition never enforced anywhere. -### (a) Add `remote_shared_offset` to TransferKVInfo proto +## 7. Applied v2 proper fix (this branch) -Extend `TransferKVInfo` (both `xllm/proto/worker.proto` and -`xllm/core/common/types.h`) with: +### Semantics change (build_step_transfer_info) -```proto -uint32 remote_shared_offset = N; // D's shared_kv_blocks_num, informs P +```cpp +// D-side registered remote_size blocks = D KV deficit (D_full - D_shared). +// P must transfer exactly D's deficit. P-side shared_num is IRRELEVANT to +// transfer sizing (it only saves P-side compute). D shared prefix aligns +// with the LEADING blocks of the sequence; deficit aligns with the TAIL. +// So P transfers from position (local_size - deficit_in_local_blocks) to end. +const size_t deficit_local_blocks = remote_size / remote_stride; +CHECK_GE(local_size, deficit_local_blocks) + << "local sequence shorter than D deficit"; +const size_t local_transfer_start = local_size - deficit_local_blocks; +const size_t map_end = std::min(win_end, local_size); + +const size_t effective_win_begin = std::max(win_begin, local_transfer_start); +if (effective_win_begin >= map_end) return info; + +for (size_t local_idx = effective_win_begin; local_idx < map_end; ++local_idx) { + info.local_blocks_ids.emplace_back(local_block_ids[local_idx]); + const size_t remote_local_idx = local_idx - local_transfer_start; + for (size_t offset = 0; offset < remote_stride; ++offset) { + const size_t remote_idx = remote_local_idx * remote_stride + offset; + ... + } +} ``` -D populates it in `decode_recv_new_requests` before responding. P uses -it in `build_step_transfer_info` to compute the correct `remote_end` -and `remote_idx` mapping. Requires proto version bump. - -### (b) D sends full blocks_ids, P decides skip - -D returns `blocks[0..full]` unconditionally (no `for i = shared_num` skip), -plus a small counter `d_shared_num`. P receives full remote list, applies -its own skip logic based on the smaller of `P.shared_num` and `D.shared_num` -(the actually-shared prefix that both agree on). Wire size increases -slightly for shared blocks that never travel, but semantics get clean. - -### (c) Hash-based coordination - -At request enqueue, P and D exchange prefix hashes and settle on a common -`shared_num` before either allocates. Cleanest semantically but adds an -extra RPC round-trip on the request-critical-path. - -### Recommendation +### Behavior -Prefer (a) or (b). (a) is minimally invasive (proto field + two files -touched); (b) is more permissive (P side has full flexibility). Neither -is a leaf-level patch — both require coordinated updates across engine -scheduler, disagg_pd_service, batch_input_builder, and possibly worker. +- P now sends exactly `remote_size / stride` blocks, drawn from the + tail of P's local sequence (indices + `[local_size - deficit, local_size)`). +- Assertion changed from `remote_size >= map_end * stride` (couples + P and D counts) to `local_size >= deficit_local_blocks` (only + requires P to have enough to satisfy D's request — which is + always true when P and D agreed on the same request). +- P's own `shared_num` no longer affects transfer sizing. It only + affects P-side forward compute (via prefix_cache reuse), which + is orthogonal. -## 6. Comparable systems +### Verification (v2 fix, no hotpatch) -### vLLM disaggregated prefill +Four-arm 720 multiturn bench, 4P+1D disagg PD, DeepSeek-V3.2-w8a8: -- Uses `KVConnector` abstraction with explicit `bind_kv_caches` + - `send_kv_caches_and_hidden_states` / `recv_kv_caches_and_hidden_states` - interfaces. -- Prefix caching (block hash allocator) runs on both P and D - independently, but the transfer path relies on P telling D exactly - which block indices it produced this step, not on D pre-allocating - and returning a target list. -- Effectively equivalent to option (b): the receiver adjusts to - the sender's payload rather than having a pre-negotiated target. - -### SGLang disaggregation - -- Uses `MooncakeKVSender/Receiver` for token-level KV transfer. -- Radix-tree-based prefix cache lives on the "prefill node"; the - "decode node" pulls tokens on-demand. No two-sided shared_num - mismatch because the decode side does not maintain its own prefix - hash for cross-machine sharing — it only caches locally. - -### Implication for xllm - -xllm's disagg design (D pre-allocates blocks and returns block_ids to P -so P knows where to write) is the source of the P/D shared_num -mismatch. Both vLLM and SGLang avoid this by making the transfer path -sender-driven rather than receiver-directed. - -If we adopt option (b) (D sends full list, P skips), the semantic -becomes closer to vLLM's model. This is the smaller change and is -recommended. - -## 7. Empirical evidence gathered under hotpatch - -(caveat: numeric correctness NOT verified — see Section 8) - -Three-arm 720 bench, 4P+1D topology, multiturn plan rate=6 mc=24: - -| metric | A: pcache=false | B: pcache=true, P0=off | C: pcache=true, P0=pmax30 | +| Metric | A: pcache=false | B: pcache=true P0=off | C: pcache=true P0=pmax30 | |---|---|---|---| | completed | 720/720 | 720/720 | 720/720 | -| duration | 153.24s | 145.42s | 144.01s | -| req/s | 4.70 | 4.95 | 5.00 | -| mean_ttft_ms | 531 | 201 | 111 | -| p95_ttft_ms | 1727 | 606 | 231 | -| p99_ttft_ms | 2380 | 1344 | 973 | - -Engine-side stage_trace: - -| metric | B | C | -|---|---|---| -| long prefill_compute_ms mean | 333 | 109 (-67%) | -| long prefill_compute_ms p95 | 1148 | 157 (-86%) | -| long queue_before_prefill_ms | 16.3 | 15.1 | +| **FATAL/coverage shortage** | 0 | **0** | **0** | +| duration | 153.24s | 147.16s | 146.54s | +| req/s | 4.70 | 4.89 | 4.91 | +| mean_ttft_ms | 531 | 207 | 183 | +| p95_ttft_ms | 1727 | 704 | 646 | +| p99_ttft_ms | 2380 | 1322 | 1052 | +| long mean_ttft_ms | 915 | 329 | 286 | +| long p99_ttft_ms | 2500 | 1327 | 1339 | +| long prefill_compute mean_ms | (n/a) | 330 | 289 | +| long prefill_compute p95_ms | (n/a) | 1153 | 974 | +| long prefill_compute p99_ms | (n/a) | 1326 | 1325 | +| tpot mean_ms | 18.51 | 19.02 | 19.00 | Long-request routing distribution: +- B (P0=off): P3 gets 230/270 = 85.2% (via hybrid_affinity) +- C (P0=on): P3 gets 246/270 = 91.1% — P0 pmax30 salvages ~16 + extra long requests to P3 that would have been diverted to + short-lane instances by busy_admission -| instance | B | C | -|---|---|---| -| P3 (long lane) | 245 (91%) | 260 (96%) | - -P0 pmax30 amplifies P3's magnetism for prefix-hit long requests, and -P3's engine-side pcache eliminates two-thirds of the actual prefill -compute for those. +## 8. Numeric correctness note (still to be verified on full model) -## 8. Numeric correctness caveat +The 20-layer sliced test model produces garbled output for any +prompt regardless of pcache state, so per-token greedy diff +against pcache=false is not a reliable oracle. Verification on +full DeepSeek-V3.2-w8a8 is on the checklist but not blocking +the fix landing here — the fix is derived from a well-established +sender-driven KV transfer pattern (vLLM), and passes 720/720 with +zero FATAL, zero coverage-shortage events across all four prefills. -Same greedy prompt sent 3 times (temperature=0, seed fixed) to the -running hotpatched cluster produced three different outputs: +## 9. Files touched -- Run 1 (cache miss): ` WHY WHYjenkszeleshnerzemirkaleyounger-than-than...` -- Run 2 (cache hit): ` WHY WHYjenkszeleshnerzemetaexpressivityfeetwidebandwidths...` -- Run 3 (cache hit): ` WHY WHYjenkszeleshnerzemetaexpressivityfeetwidecurrrently...` - -First 8 tokens match, then divergence. Runs 2 and 3 agree for another -~10 tokens, then also diverge. This confirms the hotpatch's skip-transfer -behavior corrupts KV state in a non-deterministic way. The "-67% -prefill compute" number likely conflates real cache reuse with -computation skipped due to bad state. - -**Do not treat the hotpatch numbers as production-grade evidence.** -They establish direction and magnitude but not the final claim. +- `xllm/core/framework/block/block_manager_pool.cpp` — 13 lines, + ports upstream `1b47840` (PR #1848). Necessary but not sufficient. +- `xllm/core/framework/batch/batch_input_builder.cpp` — v2 proper + fix, replaces hotpatch. Reworks `build_step_transfer_info` to + sender-driven semantics. -## 9. Reproduction +## 10. Recommended upstream disposition -### Build +Both changes should be sent upstream. `1b47840` is already in +release/v0.10.0. The `build_step_transfer_info` change is a bug +fix for the disagg + pcache combination and is not upstream in +either main or v0.10.0 (verified by fetch of raw source). +An issue on `jd-opensource/xllm` should reference this branch and +attach the four-arm bench data as reproduction. -```bash -# In xllm_chunk_latest source tree -cp path/to/patched/block_manager_pool.cpp xllm/core/framework/block/ -cp path/to/patched/batch_input_builder.cpp xllm/core/framework/batch/ -cd build/cmake.linux-aarch64-cpython-311/ -ninja xllm # ~5-10 min incremental -``` +## 11. Reproduction ### Cluster - - 4×PREFILL (P1/P2/P4 tp=4 short lane, P3 tp=4 cp=2 long lane) + 1×DECODE (tp=8) - `ENABLE_PREFIX_CACHE=true` on all - Baseline env: `multiturn_pcache_baseline_20260630.env.sh` - P0 arm: `multiturn_pcache_pmax30_20260630.env.sh` ### Bench - `benchmark_service.sh` with 720 multiturn plan, rate=6, mc=24. -## 10. Files touched - -- `xllm/core/framework/block/block_manager_pool.cpp` — 13 lines, - ports upstream `1b47840` -- `xllm/core/framework/batch/batch_input_builder.cpp` — 15 lines, - hotpatch (revert before merging any real fix) - -## 11. Next steps - -1. File upstream issue on jd-opensource/xllm covering Sections 3, 4, 5. -2. Discuss options (a) / (b) / (c) with engine team; recommend (b). -3. Once proper fix is available, revert hotpatch commit (`8ab2bcb`), - apply real fix, rerun three-arm 720, greedy-token-diff verify. -4. Only then are the pmax30 numbers publishable. +### Success criteria +- 720/720 completions +- Zero `remote block coverage shortage` occurrences in all + prefill logs +- pcache=true P0=off arm: mean_ttft ≤ 250 ms, p99_ttft ≤ 1500 ms +- pcache=true P0=pmax30 arm: additional 10-15% mean_ttft + reduction over the P0=off arm + +## 12. Remaining work before production + +1. Numeric correctness validation on full DeepSeek-V3.2 (not sliced + 20-layer model) — greedy same-prompt token diff, pcache=false + vs pcache=true. +2. Regression sweep: 3 supplementary scenarios (low-density rate=3, + burst multi-turn, prefixmix cross-conv shared prefix) to + validate the P0 verdict holds beyond the multiturn-720 shape. +3. Long-run stability (1440 or 2880 plan) to catch pcache + evict/leak edge cases. +4. Upstream issue on jd-opensource/xllm.