Skip to content

refactor: eliminate DisaggPDChunkedPrefillScheduler. - #2017

Open
weizhehuang0827 wants to merge 1 commit into
xLLM-AI:mainfrom
weizhehuang0827:feat/eliminate-pd-chunked-scheduler
Open

refactor: eliminate DisaggPDChunkedPrefillScheduler.#2017
weizhehuang0827 wants to merge 1 commit into
xLLM-AI:mainfrom
weizhehuang0827:feat/eliminate-pd-chunked-scheduler

Conversation

@weizhehuang0827

@weizhehuang0827 weizhehuang0827 commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

Motivation

After #1992 unified scheduling logic into SchedulerPolicy, DisaggPDChunkedPrefillScheduler duplicated the same scheduling work in its own prepare_batch(). The key insight is that PrefillFirstPolicy already schedules chunk_queue (in-flight) before prefill_queue (fresh), which naturally prevents the hold-and-wait deadlock that the old class was designed to solve.

Changes

  1. Delete disagg_pd_chunked_prefill_scheduler.cpp/.h and its test
  2. Remove DisaggPDScheduler::prepare_batch() override — defer to ContinuousScheduler::prepare_batch() + policy
  3. Remove DISAGG_PD_CHUNKED_PREFILL from SchedulerKind enum and factory
  4. Add full-footprint admission gate in schedule_prefill_from_queue: fresh requests check if their total block footprint fits before admission (defense-in-depth against deadlock)
  5. PD PREFILL forces exclusive batch via resolve_batch_mode (enable_mix_batch = false)
  6. Skip expand_sequences for PD PREFILL in drain_request_queue (best_of_n deferred to decode)
  7. Remove redundant external deallocateallocate() internally handles cleanup on failure

Test Plan

  • Unit tests: 42 pass (including new FullFootprintAdmissionGate and FullFootprintAdmitsBothWhenFits)
  • PD + chunked_prefill=true + prefix_cache=false: shared-prefix 50 reqs, sharegpt 100 reqs, random 50 reqs — all PASS
  • PD + chunked_prefill=false + prefix_cache=false: same tests — all PASS
  • PD + prefix_cache=true: see bugfix: skip prefix-cached blocks during PD KV transfer. #2040 for a separate bugfix needed (pre-existing bug unrelated to this refactor)

Change Type

  • Refactoring (no behavior change)

Reviewer Notes

The prefix cache + PD crash (batch_input_builder.cpp:269: remote block coverage shortage) is a pre-existing bug that also reproduces on ca6a5fd7e (before the scheduler refactor). Fix submitted separately in #2040.

@weizhehuang0827 weizhehuang0827 changed the title refactor(scheduler): eliminate DisaggPDChunkedPrefillScheduler. refactor: eliminate DisaggPDChunkedPrefillScheduler. Jul 23, 2026
JimHsiung
JimHsiung previously approved these changes Jul 23, 2026
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants