refactor: unify scheduler subclasses into ContinuousScheduler + SchedulerPolicy.

[weizhehuang0827]

Description

Unify 4 scheduler subclasses (ContinuousScheduler, ChunkedPrefillScheduler, PrefillOnlyScheduler, MixScheduler) into a single ContinuousScheduler driven by a SchedulerPolicy class hierarchy. Net -1775 lines of code while preserving all scheduling behaviors.

Motivation

The old design had 4 scheduler classes with heavily duplicated scheduling logic (block allocation, preemption, decode scheduling, latency-aware budget, prefix cache matching, etc.) spread across each subclass. Any bug fix or new feature had to be applied to all 4 implementations independently. This refactor extracts the common logic into a shared SchedulerPolicy base class, and lets each policy subclass only define batch assembly order (how prefill and decode requests are prioritized within a step).

Architecture (After)

Scheduler (abstract)
└── ContinuousScheduler                (unified lifecycle: step loop, batch creation, request queue drain)
    ├── policy_: SchedulerPolicy*      (batch assembly strategy, selected at construction)
    │   ├── PrefillFirstPolicy         (exclusive batch: prefill and decode never share a batch)
    │   ├── DecodeFirstPolicy          (mixed batch: decode fills first, remaining budget goes to prefill)
    │   └── UnifiedPolicy              (mixed batch: single urgency-sorted queue, multi-SLO/priority)
    ├── FixedStepsScheduler            (for generative recommendation; completely overrides prepare_batch)
    ├── ZeroEvictionScheduler          (never preempts, rejects when full)
    ├── DisaggPDScheduler              (PD disaggregation base, inherits ContinuousScheduler directly)
    │   ├── DisaggPDChunkedPrefillScheduler  (PD + chunked prefill with full-footprint block reservation)
    │   └── PDOOCScheduler             (online-offline co-location, own scheduling with separate queues)

Note: FixedStepsScheduler is designed for generative recommendation workloads and completely overrides prepare_batch() with its own logic. It does not use SchedulerPolicy — it inherits ContinuousScheduler for lifecycle management only.

Three-Queue Design

The refactored ContinuousScheduler maintains three explicit queues (replacing the old single waiting_priority_queue_ + implicit running_queue_ approach):

Queue	Contents	Used by
prefill_queue_	New requests awaiting first prefill (kv_cache_tokens == 0)	All policies
chunk_queue_	Chunked prefill continuations (has partial KV, preemptable)	PrefillFirst, DecodeFirst
decode_queue_	Requests in decode stage	All policies
unified_queue_	All requests merged into one list (sorted by urgency)	UnifiedPolicy only

This separation makes the scheduling invariants explicit:

• chunk_queue_ requests are preemptable (their partial KV can be evicted to make room for decode)
• decode_queue_ requests have committed KV and should not be preempted unless absolutely necessary
• prefill_queue_ holds fresh arrivals that haven't allocated any blocks yet

BatchMode and Policy Selection

resolve_batch_mode() determines the BatchMode struct from user configuration, then create_scheduler_policy() selects the concrete policy:

BatchMode resolve_batch_mode(const Options& options) {
  mode.enable_mix_batch = SchedulerConfig::enable_mix_batch();  // gflag, default true

  // Hard constraints that override user setting:
  if (multi_slo_and_prio) mode.enable_chunked_prefill = true;
  if (cp_or_mtp)          mode.enable_mix_batch = false;
  if (!chunked_prefill)   mode.enable_mix_batch = false;
}

enable_mix_batch (resolved)	priority_strategy	→ Policy
false	any	PrefillFirstPolicy
true	fcfs / priority / deadline	DecodeFirstPolicy
true	multi_slo_and_prio	UnifiedPolicy

The new --enable_mix_batch gflag (default true) replaces the old --use_mix_scheduler flag. Users can explicitly set --enable_mix_batch=false to force exclusive batch mode regardless of other settings.

SchedulerState Boundary

Policies operate solely through the SchedulerState struct — they never access ContinuousScheduler members directly. This provides a clean boundary:

struct SchedulerState {
  RequestPriorityQueue& prefill_queue, chunk_queue, decode_queue;
  std::list<Request>& unified_queue;
  std::vector<Request>& running_requests;
  std::vector<Sequence*>& running_sequences;
  KVCacheManager* kv_cache_manager;
  ProfileManager* profile_manager;
  // ... configuration flags
};

Common Scheduling Primitives (in SchedulerPolicy base)

The following operations are shared across all policies, eliminating the code duplication:

• schedule_prefill_from_queue — iterates a prefill/chunk queue, computes token budget per sequence, allocates blocks, handles beam-search expansion, publishes in-batch prefix cache
• schedule_decode_from_queue — iterates decode queue, allocates decode blocks, handles latency-aware budget checking, preemption when blocks exhausted
• compute_prefill_tokens — determines per-sequence token count respecting max_tokens_per_chunk, CP alignment (maybe_align_cp_chunk_tokens), and remaining budget
• allocate_for_prefill — allocates KV blocks for a sequence, handling full-prefill vs chunked-prefill paths, GDN linear-state block alignment, prefix cache matching
• allocate_shared_blocks_for — prefix cache matching at chunk boundaries (with special every-chunk re-match for linear-state models)

Key Invariants Preserved

• GDN linear-state block alignment: When has_linear_attention_layers && enable_prefix_cache, chunk boundaries are rounded down to block_size so linear-state checkpoints land at recoverable positions. Every chunk re-matches prefix cache.
• Platform::uses_model_cp_partition() guard: When the model handles CP partition internally, the scheduler skips CP alignment (treats effective cp_size=1).
• Host-device KV swap: UnifiedPolicy uses the 3-arg kv_cache_manager->allocate(seq, tokens, copy_blocks) when host_blocks_factor > 1.0, accounting for H2D copy block budget.
• PD-OOC online/offline batch exclusivity: PDOOCScheduler maintains separate prefill_queue_offline_ and decode_queue_offline_ queues. A batch consists entirely of online OR entirely of offline requests.
• Speculative decoding token reservation: Decode scheduling loop requires remaining_token_budget > min_speculative_tokens_required before scheduling a decode request.
• transfer_blocks call: After batch creation, kv_cache_manager->transfer_blocks(batches) is called to drain pending H2D/D2H KV transfers (host prefix cache framework).

Future Work

• The Disaggregated PD schedulers (DisaggPDScheduler, DisaggPDChunkedPrefillScheduler, PDOOCScheduler) still contain substantial standalone scheduling logic bypassing the SchedulerPolicy abstraction. A follow-up PR by @JimHsiung will refactor these.
• SchedulerState is currently passed by reference to every policy method. A future cleanup will make it a policy member (set once per step) to simplify all signatures.
• Urgency/SLO computation is partially duplicated between PrefillFirstPolicy::adjust_latency_budget_and_reorder and UnifiedPolicy::get_latency_budget_and_request_order — to be extracted into a shared base method.

Related Issues

N/A (architecture refactoring)

Change Type

• Bug fix
• New feature
• Performance improvement
• Refactor
• Documentation
• Test
• Build or CI

Pull Request Checklist

PR Title and Commit Messages

• The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit or an equivalent command.
• I have installed the hooks with pre-commit install.
• I have run pre-commit run --all-files and fixed any reported issues.

Self Review

• I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
• I have rebased this PR onto the latest main branch.

Build and Test Coverage

• Tests have been added or updated as needed.
• CUDA: python setup.py build test has passed on a CUDA machine.
• NPU: python setup.py build test has passed on an NPU machine.
• MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

Test Setup

Base commit: d5ce3007f (before 1d0056b23 which introduces an unrelated batch_input_builder regression)

Unit Tests (all PASSED):

• scheduler_test (57 tests) — covers all policy paths, PD chunked prefill, disagg PD scheduler

Serving Tests (TP=2, DeepSeek-V3.2-w8a8-5layer, Ascend 910B npu:12+npu:13):

Each case: smoke test (1 curl request) + load test (ShareGPT dataset, 50 requests, max-concurrency 10, request_rate=inf)

Case	Config	Policy	Result
1	chunked_prefill=false	PrefillFirstPolicy	PASSED
2	chunked_prefill=true	DecodeFirstPolicy	PASSED
3	chunked_prefill=true, MTP=on	PrefillFirstPolicy	PASSED
4	chunked_prefill=true, multi_slo_and_prio	UnifiedPolicy	PASSED
5	chunked_prefill=true, MTP=on, multi_slo_and_prio	PrefillFirstPolicy	PASSED
6	chunked_prefill=true, enable_mix_batch=false	PrefillFirstPolicy	PASSED
7	chunked_prefill=false, multi_slo_and_prio	UnifiedPolicy	PASSED

Additional validation (Qwen2-7B, TP=2, max_tokens_per_batch=2500):

• Random dataset (input=1000, output=100) confirms chunked prefill actually splits prompts into multiple chunks — all 4 non-MTP cases PASSED

PD Disaggregation E2E Test (etcd + xllm_service + prefill + decode):

• Config: enable_disagg_pd=true, enable_chunked_prefill=true
• Smoke: PASSED | Load (50 reqs, concurrency 10): PASSED
• Performance: TTFT median 234ms, TPOT median 10.6ms

Known Gaps

• PDOOCScheduler bypasses SchedulerPolicy (uses own handle_prefill_requests_impl/handle_decode_requests) — to be addressed in follow-up by @JimHsiung
• enable_chunked_prefill=false crashes on commits after 1d0056b23 due to an unrelated batch_input_builder regression (confirmed on main without this PR)