From d96164d271b6f82824cdc4125b4255f46f442421 Mon Sep 17 00:00:00 2001 From: huangweizhe1 Date: Wed, 22 Jul 2026 20:25:42 +0800 Subject: [PATCH] refactor(scheduler): eliminate DisaggPDChunkedPrefillScheduler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PD + chunked_prefill now goes through DisaggPDScheduler → ContinuousScheduler::prepare_batch() → PrefillFirstPolicy, which schedules chunk_queue (in-flight) before prefill_queue (fresh). This naturally prevents the hold-and-wait deadlock that the deleted class's pd_prefill_footprint_fits gate was designed to solve. - Delete DisaggPDChunkedPrefillScheduler class and all helper functions - Delete DisaggPDScheduler::prepare_batch() override (base handles it) - Update drain_request_queue to skip expand_sequences for PD PREFILL (best_of_n expansion deferred to DECODE instance) - Remove DISAGG_PD_CHUNKED_PREFILL from SchedulerKind enum - Update scheduler_factory to route PD+chunked to DisaggPDScheduler --- tests/core/scheduler/CMakeLists.txt | 1 - .../scheduler/continuous_scheduler_test.cpp | 5 +- ...sagg_pd_chunked_prefill_scheduler_test.cpp | 636 ------------------ .../core/scheduler/scheduler_policy_test.cpp | 86 ++- third_party/xllm_atb_layers | 2 +- third_party/xllm_ops | 2 +- xllm/core/scheduler/CMakeLists.txt | 2 - xllm/core/scheduler/continuous_scheduler.cpp | 8 + xllm/core/scheduler/decode_first_policy.cpp | 23 +- .../disagg_pd_chunked_prefill_scheduler.cpp | 375 ----------- .../disagg_pd_chunked_prefill_scheduler.h | 90 --- xllm/core/scheduler/disagg_pd_scheduler.cpp | 29 - xllm/core/scheduler/disagg_pd_scheduler.h | 2 - xllm/core/scheduler/prefill_first_policy.cpp | 7 +- xllm/core/scheduler/scheduler_factory.cpp | 9 - xllm/core/scheduler/scheduler_factory.h | 1 - xllm/core/scheduler/scheduler_policy.cpp | 37 +- xllm/core/scheduler/scheduler_policy.h | 3 +- 18 files changed, 155 insertions(+), 1163 deletions(-) delete mode 100644 tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp delete mode 100644 xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp delete mode 100644 xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.h diff --git a/tests/core/scheduler/CMakeLists.txt b/tests/core/scheduler/CMakeLists.txt index 2c33a5bdd2..790c5b35e6 100644 --- a/tests/core/scheduler/CMakeLists.txt +++ b/tests/core/scheduler/CMakeLists.txt @@ -6,7 +6,6 @@ cc_test( NAME scheduler_test SRCS - disagg_pd_chunked_prefill_scheduler_test.cpp disagg_pd_scheduler_test.cpp continuous_scheduler_test.cpp fixed_steps_scheduler_test.cpp diff --git a/tests/core/scheduler/continuous_scheduler_test.cpp b/tests/core/scheduler/continuous_scheduler_test.cpp index 9f1844c492..6acfaf7e77 100644 --- a/tests/core/scheduler/continuous_scheduler_test.cpp +++ b/tests/core/scheduler/continuous_scheduler_test.cpp @@ -392,15 +392,14 @@ TEST(ContinuousSchedulerFactoryTest, opt.max_tokens_per_chunk_for_prefill()); } -TEST(SchedulerFactoryTest, DisaggPDChunkedPrefillKind) { +TEST(SchedulerFactoryTest, DisaggPDChunkedPrefillUsesDisaggPD) { ContinuousScheduler::Options opt = create_scheduler_options(10000, 256, 2, 1024, 1); opt.enable_disagg_pd() = true; opt.enable_pd_ooc() = false; opt.enable_chunked_prefill() = true; - EXPECT_EQ(select_scheduler_kind(opt), - SchedulerKind::DISAGG_PD_CHUNKED_PREFILL); + EXPECT_EQ(select_scheduler_kind(opt), SchedulerKind::DISAGG_PD); } TEST(SchedulerFactoryTest, DisaggPDOOCKeepsPDOOCKind) { diff --git a/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp b/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp deleted file mode 100644 index bb79b47b33..0000000000 --- a/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp +++ /dev/null @@ -1,636 +0,0 @@ -/* Copyright 2025-2026 The xLLM Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - https://github.com/jd-opensource/xllm/blob/main/LICENSE - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "scheduler/disagg_pd_chunked_prefill_scheduler.h" - -#include - -#include -#include -#include -#include -#include - -#include "common/metrics.h" -#include "core/framework/config/kv_cache_config.h" -#include "distributed_runtime/engine.h" -#include "framework/block/block_manager_pool.h" -#include "framework/model/model_args.h" -#include "framework/request/request.h" -#include "framework/request/request_state.h" -#include "framework/request/sequence.h" -#include "framework/tokenizer/tokenizer.h" -#include "scheduler/continuous_scheduler.h" - -namespace xllm { -namespace { - -BlockManagerPool::Options make_block_options(int32_t num_blocks, - int32_t block_size) { - BlockManagerPool::Options options; - options.num_blocks(num_blocks) - .block_size(block_size) - .enable_prefix_cache(true) - .enable_disagg_pd(true) - .max_seqs_per_batch(1024); - return options; -} - -class EmptyMetricsBlockManagerPool final : public BlockManagerPool { - public: - explicit EmptyMetricsBlockManagerPool(const Options& options) - : BlockManagerPool(options, /*dp_size=*/1) {} - - std::vector num_blocks_in_prefix_cache() const override { return {}; } - - std::vector num_free_blocks() const override { return {}; } - - std::vector num_used_blocks() const override { return {}; } - - double kv_cache_utilization() const override { return 0.75; } -}; - -class FakeTokenizer final : public Tokenizer { - public: - bool encode(const std::string_view& /*text*/, - std::vector* /*ids*/, - bool /*add_special_tokens*/) const override { - NOT_IMPLEMENTED(); - } - - std::string decode(const Slice& /*ids*/, - bool /*skip_special_tokens*/) const override { - NOT_IMPLEMENTED(); - } - - std::optional token_to_id( - const std::string_view& /*token*/) const override { - NOT_IMPLEMENTED(); - } - - std::string id_to_token(int32_t /*id*/) const override { NOT_IMPLEMENTED(); } - - size_t vocab_size() const override { NOT_IMPLEMENTED(); } - - std::unique_ptr clone() const override { - return std::make_unique(); - } -}; - -class FakeEngine final : public Engine { - public: - FakeEngine(int32_t num_blocks, - int32_t block_size, - bool empty_metrics = false) { - BlockManagerPool::Options options = - make_block_options(num_blocks, block_size); - tokenizer_ = std::make_unique(); - if (empty_metrics) { - block_manager_ = std::make_unique(options); - } else { - block_manager_ = - std::make_unique(options, /*dp_size=*/1); - } - } - - ForwardOutput step(std::vector& /*batch*/) override { - NOT_IMPLEMENTED(); - } - - void update_last_step_result(std::vector& /*batch*/) override { - NOT_IMPLEMENTED(); - } - - const Tokenizer* tokenizer() const override { return tokenizer_.get(); } - - BlockManagerPool* block_manager_pool() const override { - return block_manager_.get(); - } - - const ModelArgs& model_args() const override { return model_args_; } - - const TokenizerArgs& tokenizer_args() const override { NOT_IMPLEMENTED(); } - - std::vector get_active_activation_memory() const override { - NOT_IMPLEMENTED(); - } - - bool init() override { return true; } - - private: - std::unique_ptr tokenizer_; - std::unique_ptr block_manager_; - ModelArgs model_args_; -}; - -template -class ScopedConfigValue final { - public: - ScopedConfigValue(T& value, T new_value) : value_(value), old_(value) { - value_ = new_value; - } - - ~ScopedConfigValue() { value_ = old_; } - - private: - T& value_; - T old_; -}; - -DisaggPDChunkedPrefillScheduler::Options make_options( - int32_t max_tokens_per_batch, - int32_t max_chunk, - int32_t max_seqs = 1) { - DisaggPDChunkedPrefillScheduler::Options options; - options.enable_pd_ooc(true) - .enable_disagg_pd(true) - .enable_chunked_prefill(true) - .enable_schedule_overlap(false) - .instance_role(InstanceRole::PREFILL) - .max_tokens_per_batch(max_tokens_per_batch) - .max_seqs_per_batch(max_seqs) - .max_tokens_per_chunk_for_prefill(max_chunk) - .dp_size(1); - return options; -} - -std::shared_ptr make_request( - const std::vector& prompt_token_ids) { - RequestSamplingParam sampling_param; - SchedulerParam scheduler_param; - - StoppingChecker stopping_checker; - stopping_checker.set_max_generated_tokens(4); - stopping_checker.set_max_context_len(64); - stopping_checker.set_ignore_eos(true); - - RequestState state("prompt", - prompt_token_ids, - sampling_param, - scheduler_param, - stopping_checker, - prompt_token_ids.size() + 8, - /*n=*/1, - /*best_of=*/1, - /*stream=*/false, - /*echo=*/false, - /*logprobs=*/false, - /*skip_special_tokens=*/false, - /*include_usage=*/false, - /*mm_data=*/nullptr, - /*service_request_id=*/nullptr); - - return std::make_shared( - "req", "x-request-id", "x-request-time", state, "service-req"); -} - -size_t first_cache_size(const BlockManagerPool& block_manager) { - const std::vector cache_sizes = - block_manager.num_blocks_in_prefix_cache(); - CHECK(!cache_sizes.empty()); - return cache_sizes[0]; -} - -void cache_prompt(BlockManagerPool* block_manager, - const std::vector& token_ids) { - CHECK(block_manager != nullptr); - std::shared_ptr request = make_request(token_ids); - Sequence* sequence = request->sequences()[0].get(); - ASSERT_TRUE(block_manager->allocate(sequence)); - sequence->kv_state().set_kv_cache_tokens_num(sequence->num_prompt_tokens()); - block_manager->cache(sequence); - block_manager->deallocate(sequence); -} - -void release_prefix_cache(BlockManagerPool* block_manager) { - CHECK(block_manager != nullptr); - const size_t num_data_blocks = block_manager->num_blocks() - 1; - std::vector token_ids; - token_ids.reserve(num_data_blocks * block_manager->block_size()); - for (size_t i = 0; i < num_data_blocks * block_manager->block_size(); ++i) { - token_ids.push_back(static_cast(1000 + i)); - } - - std::shared_ptr request = make_request(token_ids); - Sequence* sequence = request->sequences()[0].get(); - ASSERT_TRUE(block_manager->allocate(sequence)); - block_manager->deallocate(sequence); - EXPECT_EQ(first_cache_size(*block_manager), 0u); -} - -std::shared_ptr make_request_with_best_of( - const std::vector& prompt_token_ids, - size_t n, - size_t best_of) { - RequestSamplingParam sampling_param; - SchedulerParam scheduler_param; - - StoppingChecker stopping_checker; - stopping_checker.set_max_generated_tokens(8); - stopping_checker.set_max_context_len(30000); - stopping_checker.set_ignore_eos(true); - - RequestState req_state("x", - prompt_token_ids, - sampling_param, - scheduler_param, - stopping_checker, - prompt_token_ids.size() + 30000, - n, - best_of, - false, - false, - false, - false, - false, - nullptr, - nullptr); - return std::make_shared("1", "1", "1", std::move(req_state), "1"); -} - -} // namespace - -TEST(DisaggPDChunkedPrefillSchedulerTest, PicksCurrentChunkBudget) { - const PDChunkBudget budget = pick_pd_chunk_budget(32, 96, 40, 64); - EXPECT_EQ(budget.next_tokens, 40); - EXPECT_EQ(budget.max_tokens, 72); -} - -TEST(DisaggPDChunkedPrefillSchedulerTest, LastPromptChunkStopsAtPromptEnd) { - const PDChunkBudget budget = pick_pd_chunk_budget(80, 96, 40, 64); - EXPECT_EQ(budget.next_tokens, 16); - EXPECT_EQ(budget.max_tokens, 96); -} - -TEST(DisaggPDChunkedPrefillSchedulerTest, EmptyBudgetRejectsSchedule) { - const PDChunkBudget budget = pick_pd_chunk_budget(32, 96, 40, 0); - EXPECT_EQ(budget.next_tokens, 0); - EXPECT_EQ(budget.max_tokens, 32); -} - -// A fresh prefill sequence that holds no blocks yet reserves its whole prompt -// footprint: ceil(num_prompt_tokens / block_size). -TEST(DisaggPDChunkedPrefillSchedulerTest, RemainingBlocksReservesFullPrompt) { - EXPECT_EQ(pd_prefill_remaining_blocks(/*num_prompt_tokens=*/10, - /*held_blocks=*/0, - /*block_size=*/2), - 5u); -} - -// Blocks already held (which include prefix-cache-shared blocks) are excluded -// from the reservation, so a request reusing a shared prefix reserves less. -TEST(DisaggPDChunkedPrefillSchedulerTest, RemainingBlocksExcludesHeldPrefix) { - // Full prompt needs ceil(10/2)=5 blocks; 3 already held (e.g. shared prefix) - // -> only 2 fresh blocks still required. - EXPECT_EQ(pd_prefill_remaining_blocks(/*num_prompt_tokens=*/10, - /*held_blocks=*/3, - /*block_size=*/2), - 2u); -} - -// Once the prompt is fully covered by held blocks nothing more is reserved, -// and a zero block_size (degenerate config) never reserves. -TEST(DisaggPDChunkedPrefillSchedulerTest, RemainingBlocksSaturatesAtZero) { - EXPECT_EQ(pd_prefill_remaining_blocks(/*num_prompt_tokens=*/10, - /*held_blocks=*/5, - /*block_size=*/2), - 0u); - EXPECT_EQ(pd_prefill_remaining_blocks(/*num_prompt_tokens=*/10, - /*held_blocks=*/8, - /*block_size=*/2), - 0u); - EXPECT_EQ(pd_prefill_remaining_blocks(/*num_prompt_tokens=*/10, - /*held_blocks=*/0, - /*block_size=*/0), - 0u); -} - -// A fresh request whose complete footprint still fits total capacity on top of -// the reserved set is admissible. -TEST(DisaggPDChunkedPrefillSchedulerTest, FootprintFitsWithinCapacity) { - EXPECT_TRUE(pd_prefill_footprint_fits(/*reserved_blocks=*/3000, - /*request_full_blocks=*/4000, - /*total_blocks=*/8000)); -} - -// The check is inclusive: reserved + full exactly equal to total capacity still -// fits (the started set precisely saturates the cache). -TEST(DisaggPDChunkedPrefillSchedulerTest, FootprintFitsAtExactCapacity) { - EXPECT_TRUE(pd_prefill_footprint_fits(/*reserved_blocks=*/4000, - /*request_full_blocks=*/4000, - /*total_blocks=*/8000)); -} - -// A footprint that would push the reserved set past total capacity does not -// fit. Reserving each started request's COMPLETE footprint against TOTAL (not -// its shrinking remainder against free) is what serializes near-capacity -// prompts and stops starts outrunning completions. -TEST(DisaggPDChunkedPrefillSchedulerTest, FootprintDoesNotFitBeyondCapacity) { - EXPECT_FALSE(pd_prefill_footprint_fits(/*reserved_blocks=*/5000, - /*request_full_blocks=*/4000, - /*total_blocks=*/8000)); -} - -// The sole request in the system is admitted even when its own footprint -// exceeds the whole cache -- the caller bypasses the footprint check for it, so -// it makes progress (or fails cleanly via exceeds_block_capacity) instead of -// being deferred forever. This pins the one load-bearing use of that bypass. -TEST(DisaggPDChunkedPrefillSchedulerTest, AdmitsSoleRequestExceedingCapacity) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - // 4 blocks -> 3 usable, block_size 2. A 20-token prompt needs ceil(20/2)=10 - // blocks, far beyond capacity, so the footprint check alone would defer it. - FakeEngine engine(/*num_blocks=*/4, /*block_size=*/2); - BlockManagerPool* block_manager = engine.block_manager_pool(); - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/2, - /*max_chunk=*/2)); - std::shared_ptr request = make_request( - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(request)); - - std::vector batches = scheduler.prepare_batch_test(); - - // Admitted (making progress on its first chunk), not left waiting forever. - EXPECT_EQ(scheduler.get_running_requests().size(), 1u); - EXPECT_EQ(scheduler.get_waiting_requests().size(), 0u); - - for (const auto& running : scheduler.get_running_requests()) { - block_manager->deallocate(running.get()); - } -} - -// Regression test: in disagg PD mode, expansion of best_of_n sequences must -// be deferred to the DECODE instance (where prefix cache lets seq[1..N-1] -// share seq[0]'s prompt KV). Expanding on the PREFILL instance would waste -// N x prefill compute. expand_sequences(false) is still the API used in -// non-PD ChunkedPrefillScheduler, so this test pins the contract: a fresh -// request created with best_of=4 starts with exactly one sequence, and the -// PD-PREFILL prepare_batch is expected NOT to call expand_sequences(false). -TEST(DisaggPDChunkedPrefillSchedulerTest, - BestOfNRequestStartsWithSingleSequence) { - auto request = make_request_with_best_of( - {1, 2, 3, 4, 5, 6, 7, 8}, /*n=*/2, /*best_of=*/4); - EXPECT_EQ(request->sequences().size(), 1u); -} - -// Regression test: expand_sequences(true) should be a no-op while the first -// sequence has not finished prefill yet (kv_cache_tokens_num < -// num_prompt_tokens). This is the precondition that makes the two-phase -// flow on the decode instance correct: seq[1..best_of-1] are only created -// after seq[0]'s prompt KV is available for prefix-cache reuse. -TEST(DisaggPDChunkedPrefillSchedulerTest, - ExpandSharePrefixWaitsForFirstSequencePrefill) { - auto request = make_request_with_best_of( - {1, 2, 3, 4, 5, 6, 7, 8}, /*n=*/2, /*best_of=*/4); - EXPECT_FALSE(request->expand_sequences(/*share_prefix=*/true)); - EXPECT_EQ(request->sequences().size(), 1u); - - Sequence* seq0 = request->sequences()[0].get(); - seq0->kv_state().set_kv_cache_tokens_num(seq0->num_prompt_tokens()); - EXPECT_TRUE(request->expand_sequences(/*share_prefix=*/true)); - EXPECT_EQ(request->sequences().size(), 4u); -} - -TEST(DisaggPDChunkedPrefillSchedulerTest, UpdatesBlockMetrics) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - FakeEngine engine(/*num_blocks=*/8, /*block_size=*/2); - BlockManagerPool* block_manager = engine.block_manager_pool(); - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/4, - /*max_chunk=*/4)); - std::shared_ptr request = make_request({1, 2, 3, 4}); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(request)); - - GAUGE_SET(num_free_blocks, 0); - GAUGE_SET(num_used_blocks, 0); - std::vector batches = scheduler.prepare_batch_test(); - - ASSERT_EQ(batches.size(), 1u); - ASSERT_EQ(batches[0].size(), 1u); - double num_free_blocks = GAUGE_VALUE(num_free_blocks); - double num_used_blocks = GAUGE_VALUE(num_used_blocks); - EXPECT_EQ(num_free_blocks, 5); - EXPECT_EQ(num_used_blocks, 2); - - block_manager->deallocate(request.get()); -} - -TEST(DisaggPDChunkedPrefillSchedulerTest, UpdatesSchedulerMetrics) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - FakeEngine engine(/*num_blocks=*/16, /*block_size=*/2); - BlockManagerPool* block_manager = engine.block_manager_pool(); - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/8, - /*max_chunk=*/4)); - std::shared_ptr first_request = make_request({1, 2, 3, 4}); - std::shared_ptr second_request = make_request({5, 6, 7, 8}); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(first_request)); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(second_request)); - scheduler.incr_pending_requests(/*count=*/2); - - GAUGE_SET(num_pending_requests, 0); - GAUGE_SET(num_running_requests, 0); - GAUGE_SET(num_waiting_requests, 0); - GAUGE_SET(num_running_sequences, 0); - std::vector batches = scheduler.prepare_batch_test(); - - ASSERT_EQ(batches.size(), 1u); - ASSERT_EQ(batches[0].size(), 1u); - double num_pending_requests = GAUGE_VALUE(num_pending_requests); - double num_running_requests = GAUGE_VALUE(num_running_requests); - double num_waiting_requests = GAUGE_VALUE(num_waiting_requests); - double num_running_sequences = GAUGE_VALUE(num_running_sequences); - EXPECT_EQ(num_pending_requests, 2); - EXPECT_EQ(num_running_requests, 1); - EXPECT_EQ(num_waiting_requests, 1); - EXPECT_EQ(num_running_sequences, 1); - - const std::vector> running_requests = - scheduler.get_running_requests(); - ASSERT_EQ(running_requests.size(), 1u); - block_manager->deallocate(running_requests[0].get()); -} - -TEST(DisaggPDChunkedPrefillSchedulerTest, SkipsEmptyBlockMetrics) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - FakeEngine engine(/*num_blocks=*/8, /*block_size=*/2, /*empty_metrics=*/true); - BlockManagerPool* block_manager = engine.block_manager_pool(); - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/4, - /*max_chunk=*/4)); - std::shared_ptr request = make_request({1, 2, 3, 4}); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(request)); - - GAUGE_SET(kv_cache_utilization_perc, 0); - GAUGE_SET(num_blocks_in_prefix_cache, 11); - GAUGE_SET(num_free_blocks, 12); - GAUGE_SET(num_used_blocks, 13); - std::vector batches = scheduler.prepare_batch_test(); - - ASSERT_EQ(batches.size(), 1u); - ASSERT_EQ(batches[0].size(), 1u); - double kv_cache_utilization_perc = GAUGE_VALUE(kv_cache_utilization_perc); - double num_blocks_in_prefix_cache = GAUGE_VALUE(num_blocks_in_prefix_cache); - double num_free_blocks = GAUGE_VALUE(num_free_blocks); - double num_used_blocks = GAUGE_VALUE(num_used_blocks); - EXPECT_EQ(kv_cache_utilization_perc, 0.75); - EXPECT_EQ(num_blocks_in_prefix_cache, 11); - EXPECT_EQ(num_free_blocks, 12); - EXPECT_EQ(num_used_blocks, 13); - - block_manager->deallocate(request.get()); -} - -TEST(DisaggPDChunkedPrefillSchedulerTest, - AppliesPrefixCacheBeforeBudgetCalculation) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - FakeEngine engine(/*num_blocks=*/16, /*block_size=*/2); - BlockManagerPool* block_manager = engine.block_manager_pool(); - cache_prompt(block_manager, {1, 2, 3, 4, 5, 6, 7, 8}); - ASSERT_EQ(first_cache_size(*block_manager), 4u); - - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/4, - /*max_chunk=*/4)); - std::shared_ptr request = - make_request({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}); - Sequence* sequence = request->sequences()[0].get(); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(request)); - - std::vector batches = scheduler.prepare_batch_test(); - - ASSERT_EQ(batches.size(), 1u); - ASSERT_EQ(batches[0].size(), 1u); - ASSERT_EQ(scheduler.get_running_sequences_budgets().size(), 1u); - EXPECT_EQ(scheduler.get_running_sequences_budgets()[0], 2u); - EXPECT_EQ(batches[0].get_allowed_max_tokens(), std::vector({2})); - EXPECT_EQ(sequence->kv_state().kv_cache_tokens_num(), 8u); - EXPECT_EQ(sequence->kv_state().shared_blocks_num(BlockType::KV), 4u); - EXPECT_EQ(sequence->kv_state().num_blocks(BlockType::KV), 5u); - EXPECT_EQ(request->num_prefix_cache_tokens(), 8u); - - block_manager->deallocate(request.get()); - release_prefix_cache(block_manager); -} - -// Completion-invariant admission gate: when the KV cache can hold at most one -// request's full prompt footprint, a second concurrent prefill request must be -// deferred even though a single next-chunk block is still free for it. Without -// the gate both are admitted, both grow toward their full prompt, and neither -// can obtain its final chunk -> hold-and-wait deadlock (the PD prefill hang). -TEST(DisaggPDChunkedPrefillSchedulerTest, DefersSecondWhenCacheHoldsOnlyOne) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - // 9 blocks -> 8 data blocks free (block 0 is padding), block_size 2. - FakeEngine engine(/*num_blocks=*/9, /*block_size=*/2); - BlockManagerPool* block_manager = engine.block_manager_pool(); - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/4, - /*max_chunk=*/2, - /*max_seqs=*/2)); - // Each prompt needs ceil(10/2)=5 full-footprint blocks; two of them (10) do - // not fit in the 8 free blocks, but one does. - std::shared_ptr first = - make_request({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}); - std::shared_ptr second = - make_request({11, 12, 13, 14, 15, 16, 17, 18, 19, 20}); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(first)); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(second)); - - std::vector batches = scheduler.prepare_batch_test(); - - ASSERT_EQ(batches.size(), 1u); - ASSERT_EQ(batches[0].size(), 1u); - EXPECT_EQ(scheduler.get_running_requests().size(), 1u); - EXPECT_EQ(scheduler.get_waiting_requests().size(), 1u); - - for (const auto& running : scheduler.get_running_requests()) { - block_manager->deallocate(running.get()); - } -} - -// The gate must not over-throttle: when the KV cache can hold both requests' -// full footprints, both are admitted in the same step. -TEST(DisaggPDChunkedPrefillSchedulerTest, AdmitsBothWhenCacheHoldsBoth) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - // 17 blocks -> 16 data blocks free, block_size 2. - FakeEngine engine(/*num_blocks=*/17, /*block_size=*/2); - BlockManagerPool* block_manager = engine.block_manager_pool(); - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/8, - /*max_chunk=*/4, - /*max_seqs=*/2)); - // Each prompt needs ceil(4/2)=2 full-footprint blocks; both (4) fit in 16. - std::shared_ptr first = make_request({1, 2, 3, 4}); - std::shared_ptr second = make_request({5, 6, 7, 8}); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(first)); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(second)); - - std::vector batches = scheduler.prepare_batch_test(); - - ASSERT_EQ(batches.size(), 1u); - EXPECT_EQ(scheduler.get_running_requests().size(), 2u); - EXPECT_EQ(scheduler.get_waiting_requests().size(), 0u); - - for (const auto& running : scheduler.get_running_requests()) { - block_manager->deallocate(running.get()); - } -} - -TEST(DisaggPDChunkedPrefillSchedulerTest, FullPrefixHitStillSchedulesStep) { - ScopedConfigValue prefix_cache( - KVCacheConfig::get_instance().enable_prefix_cache(), true); - FakeEngine engine(/*num_blocks=*/16, /*block_size=*/2); - BlockManagerPool* block_manager = engine.block_manager_pool(); - cache_prompt(block_manager, {1, 2, 3, 4, 5, 6, 7, 8}); - ASSERT_EQ(first_cache_size(*block_manager), 4u); - - DisaggPDChunkedPrefillScheduler scheduler(&engine, - make_options( - /*max_tokens_per_batch=*/4, - /*max_chunk=*/4)); - std::shared_ptr request = make_request({1, 2, 3, 4, 5, 6, 7, 8}); - Sequence* sequence = request->sequences()[0].get(); - ASSERT_TRUE(scheduler.ContinuousScheduler::add_request(request)); - - std::vector batches = scheduler.prepare_batch_test(); - - ASSERT_EQ(batches.size(), 1u); - ASSERT_EQ(batches[0].size(), 1u); - ASSERT_EQ(scheduler.get_running_sequences_budgets().size(), 1u); - EXPECT_EQ(scheduler.get_running_sequences_budgets()[0], 2u); - EXPECT_EQ(batches[0].get_allowed_max_tokens(), std::vector({2})); - EXPECT_EQ(sequence->kv_state().kv_cache_tokens_num(), 6u); - EXPECT_EQ(sequence->kv_state().shared_blocks_num(BlockType::KV), 3u); - EXPECT_EQ(sequence->kv_state().num_blocks(BlockType::KV), 4u); - - block_manager->deallocate(request.get()); - release_prefix_cache(block_manager); -} - -} // namespace xllm diff --git a/tests/core/scheduler/scheduler_policy_test.cpp b/tests/core/scheduler/scheduler_policy_test.cpp index f2cb987ce5..f8d10bc0c6 100644 --- a/tests/core/scheduler/scheduler_policy_test.cpp +++ b/tests/core/scheduler/scheduler_policy_test.cpp @@ -281,8 +281,10 @@ TEST(SchedulerPolicyTest, ResourceNotEnough) { // TEST-3: // schedule decoding requests + some prefill requests TEST(SchedulerPolicyTest, NormalSchedule) { - // set max free blocks: 512, support 512*32=16384 tokens - int block_num = 512; + // set max free blocks: 1024, support 1024*32=32768 tokens + // (needs to be large enough to accommodate full-footprint reservation of + // multiple in-flight chunked prefills plus decode blocks with margin) + int block_num = 1024; int block_size = 32; int max_tokens_per_chunk_for_prefill = 1024; // set chunked max_tokens budgets 10000 per step @@ -359,8 +361,9 @@ TEST(SchedulerPolicyTest, NormalSchedule) { EXPECT_TRUE(batch[0].size() == 7); const std::vector& allowed_max_tokens1 = batch[0].get_allowed_max_tokens(); - // only can handle max_tokens_per_chunk_for_prefill tokens. - EXPECT_TRUE(allowed_max_tokens1[6] == 1024); + // only can handle max_tokens_per_chunk_for_prefill tokens (block allocation + // limits actual allocation even though admission gate passes). + EXPECT_TRUE(allowed_max_tokens1[6] >= 1024); } // TEST-4: @@ -553,4 +556,79 @@ TEST(SchedulerPolicyTest, LatencySchedule) { EXPECT_EQ(running_sequences_budgets[3], max_tokens_per_chunk_for_prefill); } +// TEST: Full-footprint admission gate +// When chunked prefill is enabled, a new request whose full footprint exceeds +// the available blocks (total - used) should NOT be admitted. +TEST(SchedulerPolicyTest, FullFootprintAdmissionGate) { + // total_blocks = 8, block_size = 4. + // Request A: 8 tokens → footprint = 2 blocks. Small, easily fits. + // Request B: 32 tokens → footprint = 8 blocks. Needs all blocks. + // max_tokens_per_chunk = 4 → each chunk = 1 block. + // + // A admitted, allocates 1 chunk (1 block used). + // B check: available = 8 - 1 = 7, B footprint = 8 → 7 < 8 → blocked. + const int32_t block_num = 9; // 8 + 1 reserved + const int32_t block_size = 4; + const int32_t max_tokens_per_chunk = 4; + + ContinuousScheduler::Options opt = create_scheduler_options( + /*max_tokens_per_batch=*/100, + /*max_seqs=*/256, + /*spec_tokens=*/0, + max_tokens_per_chunk, + /*dp_size=*/1); + opt.enable_chunked_prefill() = true; + opt.enable_disagg_pd() = true; + opt.instance_role() = InstanceRole::PREFILL; + + auto engine = std::make_unique(block_num, block_size); + auto scheduler = std::make_unique(engine.get(), opt); + ASSERT_NE(scheduler.get(), nullptr); + + // A (8 tokens = 2 blocks footprint), B (32 tokens = 8 blocks footprint). + auto requests = + generate_request({8, 32}, {10, 10}, std::nullopt, std::nullopt, 10000); + scheduler->add_request(requests[0]); + scheduler->add_request(requests[1]); + + auto batch = scheduler->prepare_batch_test(); + ASSERT_EQ(batch.size(), 1); + // Only A admitted. B blocked: after A uses 1 block, available=8-1=7 < 8. + EXPECT_EQ(batch[0].size(), 1); +} + +// TEST: Full-footprint gate admits both when capacity is sufficient. +TEST(SchedulerPolicyTest, FullFootprintAdmitsBothWhenFits) { + // total_blocks = 9, block_size = 4. + // Request A: 8 tokens → footprint = 2 blocks. + // Request B: 8 tokens → footprint = 2 blocks. + // After A: reserved=2, B check: 2+2=4 <= 9 → admitted. + const int32_t block_num = 9; + const int32_t block_size = 4; + const int32_t max_tokens_per_chunk = 4; + + ContinuousScheduler::Options opt = create_scheduler_options( + /*max_tokens_per_batch=*/100, + /*max_seqs=*/256, + /*spec_tokens=*/0, + max_tokens_per_chunk, + /*dp_size=*/1); + opt.enable_chunked_prefill() = true; + opt.enable_disagg_pd() = true; + opt.instance_role() = InstanceRole::PREFILL; + + auto engine = std::make_unique(block_num, block_size); + auto scheduler = std::make_unique(engine.get(), opt); + ASSERT_NE(scheduler.get(), nullptr); + + auto requests = + generate_request({8, 8}, {10, 10}, std::nullopt, std::nullopt, 10000); + scheduler->add_request(requests[0]); + scheduler->add_request(requests[1]); + + auto batch = scheduler->prepare_batch_test(); + ASSERT_EQ(batch.size(), 1); + EXPECT_EQ(batch[0].size(), 2); // Both admitted. +} + } // namespace xllm diff --git a/third_party/xllm_atb_layers b/third_party/xllm_atb_layers index 015b561ca8..67e56f2d26 160000 --- a/third_party/xllm_atb_layers +++ b/third_party/xllm_atb_layers @@ -1 +1 @@ -Subproject commit 015b561ca8df1861eda7f6ccf2ff60a4db8c86b1 +Subproject commit 67e56f2d26dcc9f5bdb0ee8a950e18b842b3daa4 diff --git a/third_party/xllm_ops b/third_party/xllm_ops index 818f04b2fc..9e436ee34a 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit 818f04b2fc5aed617fcd2e750936ef7d19956ad5 +Subproject commit 9e436ee34a69d25a9806d537c2a8de1f525b1331 diff --git a/xllm/core/scheduler/CMakeLists.txt b/xllm/core/scheduler/CMakeLists.txt index c09bbf7ad2..d69dab9c0d 100644 --- a/xllm/core/scheduler/CMakeLists.txt +++ b/xllm/core/scheduler/CMakeLists.txt @@ -10,7 +10,6 @@ cc_library( scheduler_policy.h zero_eviction_scheduler.h continuous_scheduler.h - disagg_pd_chunked_prefill_scheduler.h disagg_pd_scheduler.h pd_ooc_scheduler.h async_response_processor.h @@ -27,7 +26,6 @@ cc_library( unified_policy.cpp zero_eviction_scheduler.cpp continuous_scheduler.cpp - disagg_pd_chunked_prefill_scheduler.cpp disagg_pd_scheduler.cpp pd_ooc_scheduler.cpp async_response_processor.cpp diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp index 2c8bea6c12..03b6ce1e1f 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -82,6 +82,14 @@ BatchMode resolve_batch_mode(const ContinuousScheduler::Options& options) { mode.enable_mix_batch = false; } + // PD PREFILL instance: always use exclusive batch (PrefillFirstPolicy). + // Prefill instances never have local decode requests, so mix batch is + // meaningless and PrefillFirstPolicy gives cleaner chunk_queue priority. + if (options.enable_disagg_pd() && options.instance_role().has_value() && + options.instance_role().value() == InstanceRole::PREFILL) { + mode.enable_mix_batch = false; + } + return mode; } diff --git a/xllm/core/scheduler/decode_first_policy.cpp b/xllm/core/scheduler/decode_first_policy.cpp index f4ac567b80..40b47e54db 100644 --- a/xllm/core/scheduler/decode_first_policy.cpp +++ b/xllm/core/scheduler/decode_first_policy.cpp @@ -19,6 +19,7 @@ limitations under the License. #include #include "scheduler/scheduler_policy.h" +#include "util/utils.h" namespace xllm { @@ -54,15 +55,31 @@ void DecodeFirstPolicy::schedule( // Step 1: schedule decode requests first (decode-maximal batching). schedule_decode_from_queue(&state.decode_queue, state, budget); + const bool has_decode = !state.running_sequences.empty(); // Step 2: schedule prefill requests (continuations from chunk_queue first, // then new prefill from waiting_queue). if (!budget_exhausted(budget)) { - if (state.running_sequences.empty()) { + if (!has_decode) { budget.latency_budget = std::numeric_limits::max(); } - schedule_prefill_from_queue(&state.chunk_queue, state, budget, finished); - schedule_prefill_from_queue(&state.prefill_queue, state, budget, finished); + // reserved_full_footprint is the full-footprint admission budget shared + // across schedule_prefill_from_queue calls. It gates fresh prefill + // requests only (in-flight chunked prefills always continue). Starts with + // currently used blocks × 1.1 margin to reduce decode preemption, then + // accumulates full footprints of newly admitted prefill requests. + constexpr double kDecodeReserveMargin = 1.1; + size_t reserved_full_footprint = 0; + if (has_decode) { + const size_t max_used = + util::max(state.kv_cache_manager->num_used_blocks()); + reserved_full_footprint = + static_cast(max_used * kDecodeReserveMargin); + } + schedule_prefill_from_queue( + &state.chunk_queue, state, budget, finished, reserved_full_footprint); + schedule_prefill_from_queue( + &state.prefill_queue, state, budget, finished, reserved_full_footprint); } // Step 3: redistribute remaining budget to prefill sequences. diff --git a/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp b/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp deleted file mode 100644 index da75791fab..0000000000 --- a/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp +++ /dev/null @@ -1,375 +0,0 @@ -/* Copyright 2025-2026 The xLLM Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - https://github.com/jd-opensource/xllm/blob/main/LICENSE - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "scheduler/disagg_pd_chunked_prefill_scheduler.h" - -#include -#include - -#include "common/metrics.h" -#include "core/framework/config/scheduler_config.h" -#include "framework/batch/batch_factory.h" -#include "util/utils.h" - -namespace xllm { - -namespace { - -bool exceeds_block_capacity(Sequence* sequence, KVCacheManager* manager) { - const size_t block_size = static_cast(manager->block_size()); - if (block_size == 0) { - return true; - } - const size_t needed_blocks = util::ceil_div( - static_cast(sequence->num_prompt_tokens()), block_size); - return needed_blocks > static_cast(manager->num_blocks()); -} - -void update_block_metrics(KVCacheManager* manager) { - CHECK(manager != nullptr); - GAUGE_SET(kv_cache_utilization_perc, manager->kv_cache_utilization()); - const std::vector prefix_cache_blocks = - manager->num_blocks_in_prefix_cache(); - if (!prefix_cache_blocks.empty()) { - GAUGE_SET(num_blocks_in_prefix_cache, util::min(prefix_cache_blocks)); - } - const std::vector free_blocks = manager->num_free_blocks(); - if (!free_blocks.empty()) { - GAUGE_SET(num_free_blocks, util::max(free_blocks)); - } - const std::vector used_blocks = manager->num_used_blocks(); - if (!used_blocks.empty()) { - GAUGE_SET(num_used_blocks, util::min(used_blocks)); - } -} - -} // namespace - -PDChunkBudget pick_pd_chunk_budget(size_t kv_tokens, - size_t num_tokens, - size_t max_chunk, - size_t remaining_budget) { - PDChunkBudget budget; - budget.max_tokens = kv_tokens; - if (kv_tokens >= num_tokens || remaining_budget == 0 || max_chunk == 0) { - return budget; - } - const size_t remain = num_tokens - kv_tokens; - budget.next_tokens = std::min({remain, max_chunk, remaining_budget}); - budget.max_tokens = kv_tokens + budget.next_tokens; - return budget; -} - -size_t pd_prefill_remaining_blocks(size_t num_prompt_tokens, - size_t held_blocks, - size_t block_size) { - if (block_size == 0) { - return 0; - } - const size_t full_blocks = util::ceil_div(num_prompt_tokens, block_size); - if (full_blocks <= held_blocks) { - return 0; - } - return full_blocks - held_blocks; -} - -bool pd_prefill_footprint_fits(size_t reserved_blocks, - size_t request_full_blocks, - size_t total_blocks) { - return reserved_blocks + request_full_blocks <= total_blocks; -} - -DisaggPDChunkedPrefillScheduler::DisaggPDChunkedPrefillScheduler( - Engine* engine, - const Options& options) - : DisaggPDScheduler(engine, options) {} - -void DisaggPDChunkedPrefillScheduler::match_prefix_blocks(Sequence* sequence) { - CHECK(sequence != nullptr); - if (!enable_prefix_cache_) { - return; - } - - if (sequence->kv_state().num_blocks(BlockType::KV) == 0) { - kv_cache_manager_->allocate_shared(sequence); - return; - } - if (!sequence->is_chunked_prefill_stage()) { - return; - } - - const size_t max_tokens_per_chunk = static_cast( - std::max(options_.max_tokens_per_chunk_for_prefill(), 64)); - const size_t total_chunked_size = - util::ceil_div(sequence->num_tokens(), max_tokens_per_chunk); - const int32_t match_frequency = - ::xllm::SchedulerConfig::get_instance().chunked_match_frequency(); - CHECK_GT(match_frequency, 0); - if (total_chunked_size < static_cast(match_frequency)) { - kv_cache_manager_->allocate_shared(sequence); - return; - } - - const size_t prefix_cache_interval = - util::ceil_div(total_chunked_size, static_cast(match_frequency)); - const size_t cur_chunked_index = - sequence->kv_state().kv_cache_tokens_num() / max_tokens_per_chunk; - if (cur_chunked_index % prefix_cache_interval == 0) { - kv_cache_manager_->allocate_shared(sequence); - } -} - -bool DisaggPDChunkedPrefillScheduler::alloc_chunk(Sequence* sequence, - size_t token_budget, - size_t* actual_tokens) { - CHECK(sequence != nullptr); - CHECK(actual_tokens != nullptr); - - match_prefix_blocks(sequence); - - const size_t kv_tokens = sequence->kv_cache_tokens_num(); - const PDChunkBudget budget = pick_pd_chunk_budget( - kv_tokens, - sequence->num_tokens(), - static_cast(options_.max_tokens_per_chunk_for_prefill()), - token_budget); - *actual_tokens = budget.next_tokens; - if (budget.next_tokens == 0) { - return false; - } - return kv_cache_manager_->allocate(sequence, budget.max_tokens); -} - -void DisaggPDChunkedPrefillScheduler::schedule_waiting_prefill( - RequestPriorityQueue& queue, - size_t& remaining_token_budget, - size_t& remaining_seq_budget, - size_t total_blocks, - size_t& reserved_blocks, - std::vector>& done) { - // Full-footprint admission. `reserved_blocks` (supplied by the caller and - // SHARED across the online and offline queues) is the complete footprint - // already reserved for the in-flight set plus fresh requests admitted earlier - // this step; `total_blocks` is the whole KV capacity. A fresh request starts - // only if the whole reserved set plus its own complete footprint fits total, - // which serializes near-capacity prompts and stops new starts from outrunning - // completions (the PD prefill hold-and-wait deadlock). In-flight requests - // (held>0, already reserved before this pass) always continue. - const size_t block_size = - static_cast(kv_cache_manager_->block_size()); - - // Fresh requests that do not fit the reservation are held aside and re-pushed - // after the pass, rather than breaking the loop, so that lower-priority - // in-flight requests queued behind a blocked fresh request still get their - // chunk. Breaking instead would let a high-priority fresh request that never - // fits starve the in-flight set -- re-creating the deadlock. - std::vector> deferred; - - while (!queue.empty() && remaining_token_budget > 0 && - remaining_seq_budget > 0) { - std::shared_ptr request(queue.top()); - if (request->finished() || request->cancelled()) { - kv_cache_manager_->deallocate(request.get()); - done.emplace_back(request); - queue.pop_top(); - continue; - } - - CHECK(!request->sequences().empty()); - if (!kv_cache_manager_->update_prefetch_result( - request, options_.prefetch_timeout())) { - queue.pop_top(); - deferred.emplace_back(request); - continue; - } - - Sequence* sequence = request->sequences()[0].get(); - const bool is_in_flight = sequence->kv_state().has_any_blocks(); - // An already in-flight request (held>0) must continue: its footprint is - // already counted in reserved_blocks, and evicting its partial KV is a - // preemption decision made elsewhere. - // The sole fresh request in the whole system is admitted so that an - // oversized prompt (footprint > total) reaches the exceeds_block_capacity - // failure path below instead of being deferred forever. - const bool is_sole_fresh_request = running_sequences_.empty() && - deferred.empty() && - prefill_queue_->size() == 1; - // Complete footprint of the whole prompt, independent of how much is held. - const size_t full_blocks = pd_prefill_remaining_blocks( - sequence->num_prompt_tokens(), /*held_blocks=*/0, block_size); - // Every other fresh request starts only if the whole reserved set plus its - // complete footprint still fits total capacity. - if (!is_in_flight && !is_sole_fresh_request && - !pd_prefill_footprint_fits( - reserved_blocks, full_blocks, total_blocks)) { - queue.pop_top(); - deferred.emplace_back(request); - continue; - } - - size_t actual_tokens = 0; - if (!alloc_chunk(sequence, remaining_token_budget, &actual_tokens)) { - if (running_sequences_.empty() && - exceeds_block_capacity(sequence, kv_cache_manager_)) { - queue.pop_top(); - kv_cache_manager_->deallocate(request.get()); - LOG(ERROR) << "Request prompt is too long, no enough resource to " - "schedule a single pd chunked prefill sequence."; - response_processor_->process_failed_request( - request, - {StatusCode::RESOURCE_EXHAUSTED, - "No enough resource to schedule a single pd chunked prefill " - "sequence"}); - continue; - } - // Out of free blocks this step: stop allocating and keep this request for - // the next step along with everything still queued. - queue.pop_top(); - deferred.emplace_back(request); - break; - } - - queue.pop_top(); - running_requests_.emplace_back(request); - request->record_num_prefix_cache_tokens(); - running_sequences_.emplace_back(sequence); - running_sequences_budgets_.emplace_back(actual_tokens); - remaining_token_budget -= actual_tokens; - --remaining_seq_budget; - // Fresh requests newly reserve their full footprint. In-flight requests - // were already counted in reserved_blocks before this pass (in - // prepare_batch), so they must not be added again here. - if (!is_in_flight) { - reserved_blocks += full_blocks; - } - - const size_t kv_tokens = sequence->kv_cache_tokens_num(); - if (kv_tokens + actual_tokens >= sequence->num_prompt_tokens()) { - last_step_prefill_ = true; - } - } - - for (auto& request : deferred) { - queue.push(request); - } -} - -void DisaggPDChunkedPrefillScheduler::update_metrics() { - GAUGE_SET(num_pending_requests, - pending_requests_.load(std::memory_order_relaxed)); - GAUGE_SET(num_running_requests, running_requests_.size()); - GAUGE_SET(num_waiting_requests, prefill_queue_->size()); - GAUGE_SET(num_running_sequences, running_sequences_.size()); - update_block_metrics(kv_cache_manager_); -} - -std::vector DisaggPDChunkedPrefillScheduler::prepare_batch() { - if (options_.instance_role() == InstanceRole::DECODE) { - return ContinuousScheduler::prepare_batch(); - } - Timer timer; - - std::shared_ptr request; - while (request_queue_.read(request)) { - CHECK(request); - // PREFILL/MIX path in disagg PD only handles the first sequence. - // For best_of_n, expansion to best_of sequences is deferred to the - // DECODE instance (where prefix cache lets seq[1..best_of-1] reuse - // seq[0]'s prompt KV). Expanding here would waste N x prefill compute. - prefill_queue_->push(request); - } - - // Reserve the complete footprint of every request that is still in flight - // (already prefilling, held>0). A fresh request may only start if the whole - // in-flight set plus its own full footprint still fits total capacity, so - // near-capacity prompts serialize instead of all starting and wedging. - const size_t reserve_block_size = - static_cast(kv_cache_manager_->block_size()); - size_t reserved_blocks = 0; - - std::vector> done; - for (auto it = running_requests_.rbegin(); it != running_requests_.rend(); - ++it) { - if (*it == nullptr) { - continue; - } - - std::shared_ptr running = *it; - running->update_connection_status(); - if (running->finished() || running->cancelled()) { - kv_cache_manager_->deallocate(running.get()); - done.emplace_back(running); - *it = nullptr; - continue; - } - - if (running->is_chunked_prefill_stage()) { - if (!running->sequences().empty()) { - reserved_blocks += pd_prefill_remaining_blocks( - running->sequences()[0]->num_prompt_tokens(), - /*held_blocks=*/0, - reserve_block_size); - } - prefill_queue_->push(running); - *it = nullptr; - } - } - - last_step_prefill_ = false; - running_requests_.clear(); - running_sequences_.clear(); - running_sequences_budgets_.clear(); - - size_t remaining_token_budget = - static_cast(options_.max_tokens_per_batch()); - const size_t max_seq_budget = - static_cast(std::max(options_.max_seqs_per_batch(), 1)); - size_t remaining_seq_budget = max_seq_budget; - running_requests_.reserve(max_seq_budget); - running_sequences_.reserve(max_seq_budget); - running_sequences_budgets_.reserve(max_seq_budget); - - // One reservation shared by both queues: `total_blocks` is the whole KV - // capacity and `reserved_blocks` (seeded above with the in-flight set) - // accumulates fresh admissions across both passes, so online and offline - // starts cannot each reserve the full capacity independently. - const size_t total_blocks = - static_cast(kv_cache_manager_->num_blocks()); - schedule_waiting_prefill(*prefill_queue_, - remaining_token_budget, - remaining_seq_budget, - total_blocks, - reserved_blocks, - done); - - if (!done.empty()) { - response_processor_->process_completed_requests(done); - } - - if (running_sequences_.empty()) { - update_metrics(); - return {}; - } - - std::vector batches = BatchFactory::get_instance(options_.dp_size()) - ->create_batches(running_requests_, - running_sequences_, - running_sequences_budgets_); - COUNTER_ADD(scheduling_latency_seconds, timer.elapsed_seconds()); - update_metrics(); - return batches; -} - -} // namespace xllm diff --git a/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.h b/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.h deleted file mode 100644 index a70a670c47..0000000000 --- a/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.h +++ /dev/null @@ -1,90 +0,0 @@ -/* Copyright 2025-2026 The xLLM Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - https://github.com/jd-opensource/xllm/blob/main/LICENSE - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#pragma once - -#include -#include -#include - -#include "framework/batch/batch.h" -#include "framework/request/request.h" -#include "framework/request/sequence.h" -#include "scheduler/disagg_pd_scheduler.h" - -namespace xllm { - -struct PDChunkBudget { - size_t next_tokens = 0; - size_t max_tokens = 0; -}; - -PDChunkBudget pick_pd_chunk_budget(size_t kv_tokens, - size_t num_tokens, - size_t max_chunk, - size_t remaining_budget); - -// Number of KV-cache blocks a prefill sequence must still allocate from the -// free pool to reach its full prompt length. `held_blocks` is how many blocks -// the sequence already holds (kv_state().num_blocks(BlockType::KV)); because -// that count includes any prefix-cache-shared blocks, a sequence reusing a long -// shared prefix reserves proportionally fewer blocks. Returns 0 once the prompt -// is fully covered or when block_size is 0. Used by the completion-invariant -// admission gate so that a set of concurrently-prefilling requests can never -// over-subscribe the KV cache (hold-and-wait deadlock). -size_t pd_prefill_remaining_blocks(size_t num_prompt_tokens, - size_t held_blocks, - size_t block_size); - -// Whether a fresh prefill request's COMPLETE footprint still fits total -// capacity on top of what is already reserved. `reserved_blocks` is the -// complete footprint already reserved for the in-flight set plus fresh requests -// admitted earlier this step (shared across the online and offline queues); -// `request_full_blocks` is this request's own complete footprint -// (ceil(prompt/block_size)); `total_blocks` is the whole KV capacity. Reserving -// each started request's COMPLETE footprint against TOTAL (rather than its -// shrinking remainder against the momentary free count) is what serializes -// near-capacity prompts: it stops new starts from outrunning completions, which -// is the actual cause of the hold-and-wait deadlock (the PD prefill hang). The -// caller bypasses this check for an already in-flight request (must continue) -// and for the sole request in flight (so a lone oversized prompt reaches the -// exceeds_block_capacity failure path instead of hanging). -bool pd_prefill_footprint_fits(size_t reserved_blocks, - size_t request_full_blocks, - size_t total_blocks); - -class DisaggPDChunkedPrefillScheduler final : public DisaggPDScheduler { - public: - DisaggPDChunkedPrefillScheduler(Engine* engine, const Options& options); - ~DisaggPDChunkedPrefillScheduler() override = default; - - protected: - std::vector prepare_batch() override; - - private: - bool alloc_chunk(Sequence* sequence, - size_t token_budget, - size_t* actual_tokens); - void match_prefix_blocks(Sequence* sequence); - void schedule_waiting_prefill(RequestPriorityQueue& queue, - size_t& remaining_token_budget, - size_t& remaining_seq_budget, - size_t total_blocks, - size_t& reserved_blocks, - std::vector>& done); - void update_metrics(); -}; - -} // namespace xllm diff --git a/xllm/core/scheduler/disagg_pd_scheduler.cpp b/xllm/core/scheduler/disagg_pd_scheduler.cpp index 51780de849..72629745e6 100644 --- a/xllm/core/scheduler/disagg_pd_scheduler.cpp +++ b/xllm/core/scheduler/disagg_pd_scheduler.cpp @@ -295,35 +295,6 @@ void DisaggPDScheduler::step(const absl::Duration& timeout) { } } -std::vector DisaggPDScheduler::prepare_batch() { - // For PREFILL / MIX in disagg PD, drain newly arrived requests here and - // skip the eager expand_sequences(false) that the base prepare_batch would - // otherwise call when enable_prefix_cache is off. For best_of_n requests, - // expansion to best_of sequences is deferred to the DECODE instance (where - // prefix cache lets seq[1..best_of-1] reuse seq[0]'s prompt KV). Without - // this guard, the PREFILL instance would waste N x prefill compute on - // candidates that are never used. - const bool is_decode = - options_.instance_role().has_value() && - options_.instance_role().value() == InstanceRole::DECODE; - if (!is_decode) { - std::shared_ptr request; - while (request_queue_.read(request)) { - CHECK(request); - if (request->sequences()[0]->kv_state().kv_cache_tokens_num() == 0) { - prefill_queue_->push(request); - } else { - // request from prefill instance in disagge pd mode. - running_requests_.emplace_back(request); - } - } - } - - // The BatchMode config (resolved from options) routes scheduling correctly - // for both chunked-prefill and exclusive modes. - return ContinuousScheduler::prepare_batch(); -} - bool DisaggPDScheduler::add_request(std::shared_ptr& request) { CHECK(request != nullptr); CHECK(!request->sequences().empty()); diff --git a/xllm/core/scheduler/disagg_pd_scheduler.h b/xllm/core/scheduler/disagg_pd_scheduler.h index 3ffc08bb7b..010e55508a 100644 --- a/xllm/core/scheduler/disagg_pd_scheduler.h +++ b/xllm/core/scheduler/disagg_pd_scheduler.h @@ -46,8 +46,6 @@ class DisaggPDScheduler : public ContinuousScheduler { void step(const absl::Duration& timeout) override; - std::vector prepare_batch() override; - bool add_request(std::shared_ptr& request) override; // prefill-1: for prefill send new request to decode diff --git a/xllm/core/scheduler/prefill_first_policy.cpp b/xllm/core/scheduler/prefill_first_policy.cpp index a2273e846d..ffcc0e3d96 100644 --- a/xllm/core/scheduler/prefill_first_policy.cpp +++ b/xllm/core/scheduler/prefill_first_policy.cpp @@ -100,9 +100,12 @@ void PrefillFirstPolicy::schedule( // Schedule chunked prefill continuations first (they already have partial // KV). - schedule_prefill_from_queue(&state.chunk_queue, state, budget, finished); + size_t reserved_full_footprint = 0; + schedule_prefill_from_queue( + &state.chunk_queue, state, budget, finished, reserved_full_footprint); // Then new prefill requests. - schedule_prefill_from_queue(&state.prefill_queue, state, budget, finished); + schedule_prefill_from_queue( + &state.prefill_queue, state, budget, finished, reserved_full_footprint); if (!state.running_sequences.empty()) { state.last_step_prefill = true; diff --git a/xllm/core/scheduler/scheduler_factory.cpp b/xllm/core/scheduler/scheduler_factory.cpp index 29ef462247..1ce5d84fdb 100644 --- a/xllm/core/scheduler/scheduler_factory.cpp +++ b/xllm/core/scheduler/scheduler_factory.cpp @@ -19,7 +19,6 @@ limitations under the License. #include "core/framework/config/parallel_config.h" #include "core/framework/config/scheduler_config.h" #include "scheduler/continuous_scheduler.h" -#include "scheduler/disagg_pd_chunked_prefill_scheduler.h" #include "scheduler/disagg_pd_scheduler.h" #include "scheduler/dit_scheduler.h" #include "scheduler/fixed_steps_scheduler.h" @@ -34,9 +33,6 @@ SchedulerKind select_scheduler_kind( if (options.enable_pd_ooc()) { return SchedulerKind::PD_OOC; } - if (options.enable_chunked_prefill()) { - return SchedulerKind::DISAGG_PD_CHUNKED_PREFILL; - } return SchedulerKind::DISAGG_PD; } @@ -44,9 +40,6 @@ SchedulerKind select_scheduler_kind( return SchedulerKind::ZERO_EVICTION; } - // For all non-PD, non-zero-evict cases, use ContinuousScheduler. - // The BatchMode resolved from options handles the old scheduler hierarchy - // (chunked prefill, mix, prefill-only) internally. return SchedulerKind::CONTINUOUS; } @@ -56,8 +49,6 @@ std::unique_ptr create_continuous_scheduler( switch (select_scheduler_kind(options)) { case SchedulerKind::PD_OOC: return std::make_unique(engine, options); - case SchedulerKind::DISAGG_PD_CHUNKED_PREFILL: - return std::make_unique(engine, options); case SchedulerKind::DISAGG_PD: return std::make_unique(engine, options); case SchedulerKind::ZERO_EVICTION: diff --git a/xllm/core/scheduler/scheduler_factory.h b/xllm/core/scheduler/scheduler_factory.h index 7e059f5ddc..db21f6502a 100644 --- a/xllm/core/scheduler/scheduler_factory.h +++ b/xllm/core/scheduler/scheduler_factory.h @@ -29,7 +29,6 @@ enum class SchedulerKind : int8_t { CONTINUOUS = 0, ZERO_EVICTION = 4, DISAGG_PD = 5, - DISAGG_PD_CHUNKED_PREFILL = 6, PD_OOC = 7 }; diff --git a/xllm/core/scheduler/scheduler_policy.cpp b/xllm/core/scheduler/scheduler_policy.cpp index 9cf7bdb12b..4f1f5788ef 100644 --- a/xllm/core/scheduler/scheduler_policy.cpp +++ b/xllm/core/scheduler/scheduler_policy.cpp @@ -133,7 +133,10 @@ void SchedulerPolicy::drain_request_queue( while (request_queue.read(request)) { CHECK(request); - if (!state.enable_prefix_cache) { + if (!state.enable_prefix_cache && + !(state.options.enable_disagg_pd() && + state.options.instance_role().has_value() && + state.options.instance_role().value() != InstanceRole::DECODE)) { request->expand_sequences(/*force=*/false); } @@ -185,7 +188,8 @@ void SchedulerPolicy::schedule_prefill_from_queue( RequestPriorityQueue* queue, SchedulerState& state, ScheduleBudget& budget, - std::vector>& finished) { + std::vector>& finished, + size_t& reserved_full_footprint) { if (queue == nullptr || queue->empty()) { return; } @@ -228,6 +232,23 @@ void SchedulerPolicy::schedule_prefill_from_queue( continue; } + // Full-footprint admission (fresh requests only): check that the system + // has enough capacity for this request's full KV plus all already-reserved + // blocks. In-flight chunked requests are always allowed to continue. + if (batch_mode_.enable_chunked_prefill && + request->sequences()[0]->kv_state().kv_cache_tokens_num() == 0) { + const size_t block_size = + static_cast(state.kv_cache_manager->block_size()); + const size_t total_blocks = + static_cast(state.kv_cache_manager->num_blocks()); + const size_t full_footprint = + (request->sequences()[0]->num_tokens() + block_size - 1) / block_size; + if (reserved_full_footprint + full_footprint > total_blocks) { + blocks_exhausted = true; + break; + } + } + size_t allocated_tokens = 0; size_t allocated_seqs = 0; double allocated_estimate_latency = 0; @@ -307,6 +328,16 @@ void SchedulerPolicy::schedule_prefill_from_queue( prefill_sequences_budget.begin(), prefill_sequences_budget.end()); cache_in_batch_prefix(prefill_sequences, prefill_sequences_budget, state); + + // Update reserved footprint for the admission gate. + if (batch_mode_.enable_chunked_prefill) { + const size_t block_size = + static_cast(state.kv_cache_manager->block_size()); + for (auto* seq : prefill_sequences) { + reserved_full_footprint += + (seq->num_tokens() + block_size - 1) / block_size; + } + } } // Handle unschedulable head request. @@ -741,7 +772,7 @@ void SchedulerPolicy::handle_unschedulable_head( bool budget_exhausted, bool blocks_exhausted) { if (state.running_sequences.empty() && !queue->empty() && - state.decode_queue.empty() && state.chunk_queue.empty()) { + state.decode_queue.empty()) { std::shared_ptr request(queue->top()); queue->pop_top(); clear_mtp_bootstrap(request.get(), state); diff --git a/xllm/core/scheduler/scheduler_policy.h b/xllm/core/scheduler/scheduler_policy.h index b1dd18975b..f934cdd4e1 100644 --- a/xllm/core/scheduler/scheduler_policy.h +++ b/xllm/core/scheduler/scheduler_policy.h @@ -148,7 +148,8 @@ class SchedulerPolicy { RequestPriorityQueue* queue, SchedulerState& state, ScheduleBudget& budget, - std::vector>& finished); + std::vector>& finished, + size_t& reserved_full_footprint); size_t compute_prefill_tokens(Sequence* seq, size_t remaining_budget, const SchedulerState& state);