diff --git a/CMakeLists.txt b/CMakeLists.txt index 712687057c..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 @@ -316,7 +323,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 +412,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}" diff --git a/PCACHE_DISAGG_INVESTIGATION.md b/PCACHE_DISAGG_INVESTIGATION.md new file mode 100644 index 0000000000..b0a5bfd1a4 --- /dev/null +++ b/PCACHE_DISAGG_INVESTIGATION.md @@ -0,0 +1,307 @@ +# pcache=true + disagg PD compatibility investigation & fix + +**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; 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 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 + +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 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: + +``` +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 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. + +## 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 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 P's 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(); // 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; +} +``` + +D returns only the non-shared blocks. `resp.blocks_ids().size()` +becomes `D_full - D.shared_num` — this is D's **KV deficit**. + +### 3.3 The old assertion (before fix) + +`build_step_transfer_info`: + +```cpp +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_end = map_end * remote_stride; +CHECK_GE(align_up(remote_size, remote_stride), remote_end); // requires remote_size >= map_end +``` + +### 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. + +**Do not treat hotpatch numbers as production-grade evidence.** +Kept in `pcache-disagg-hotpatch-snapshot` git tag for reference. + +## 5. Attempted fix path 2 (caller-side skip, failed) + +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. + +## 6. Community solution reference: vLLM Nixl push mode + +Reading vLLM's `KVConnectorBase_V1` + Nixl push scheduler: + +- **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. + +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. + +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. + +## 7. Applied v2 proper fix (this branch) + +### Semantics change (build_step_transfer_info) + +```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; + ... + } +} +``` + +### Behavior + +- 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. + +### Verification (v2 fix, no hotpatch) + +Four-arm 720 multiturn bench, 4P+1D disagg PD, DeepSeek-V3.2-w8a8: + +| Metric | A: pcache=false | B: pcache=true P0=off | C: pcache=true P0=pmax30 | +|---|---|---|---| +| completed | 720/720 | 720/720 | 720/720 | +| **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 + +## 8. Numeric correctness note (still to be verified on full model) + +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. + +## 9. Files touched + +- `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. + +## 10. Recommended upstream disposition + +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. + +## 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. + +### 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. 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/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; 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(); diff --git a/xllm/core/framework/config/scheduler_config.cpp b/xllm/core/framework/config/scheduler_config.cpp index 6ca5fc3910..e18e79947b 100644 --- a/xllm/core/framework/config/scheduler_config.cpp +++ b/xllm/core/framework/config/scheduler_config.cpp @@ -73,6 +73,41 @@ 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."); + +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."); + +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() { @@ -91,6 +126,11 @@ 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); + 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) { @@ -109,6 +149,11 @@ 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); + 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( @@ -144,6 +189,16 @@ 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); + 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 01a60ab220..9eb7cdfabf 100644 --- a/xllm/core/framework/config/scheduler_config.h +++ b/xllm/core/framework/config/scheduler_config.h @@ -55,7 +55,12 @@ 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", + "mix_max_prefill_chunks_per_step", + "enable_mix_step_isolation", + "mix_step_isolation_max_decode_steps"}}; return kOptionCategory; } @@ -88,6 +93,42 @@ 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; + + // 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; + + // 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/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..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 @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include "common/macros.h" #include "core/framework/config/disagg_pd_config.h" @@ -77,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(); } @@ -111,25 +113,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 +170,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 +203,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/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 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; } 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); diff --git a/xllm/core/scheduler/mix_scheduler.cpp b/xllm/core/scheduler/mix_scheduler.cpp index 5eee8fd7fe..62963aaf30 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( @@ -570,15 +585,178 @@ 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(); + + // 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) + } + + 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, + remaining_seq_budget, + num_preempted_requests, + prefill_stage_sequences, + prefill_queue, + budget_exhausted, + blocks_exhausted); + 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 + // 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); diff --git a/xllm/core/scheduler/mix_scheduler.h b/xllm/core/scheduler/mix_scheduler.h index f7ad1d4e80..cf418f6f48 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_; @@ -48,6 +65,12 @@ class MixScheduler : public ChunkedPrefillScheduler { 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,