Refactor PyTorch scheduler ownership and API boundaries

[grimoire]

Motivation

The PyTorch scheduler had accumulated several responsibilities with different
owners and lifecycles:

• prefill ordering and admission
• tentative prefix-cache matching and rollback
• external KV lookup/load admission
• sequence state transitions and paging cleanup
• eviction and state-checkpoint coordination
• engine and connector integration

These responsibilities were coupled through broad Scheduler references,
nested control flow, implicit state, and cross-module resource reach-through.
This made scheduling behavior difficult to understand and risky to extend,
especially after the SSM prefix-cache and Mooncake KV-transfer integrations.

What changed

• Extract prefill ordering and admission into prefill_scheduler.py.
• Represent prefill outcomes, prefix-match phases, long-prefill policy, and
  deferred KV-load cleanup with explicit states.
• Move external KV-load planning, capacity admission, allocation, binding,
  rollback, and asynchronous ownership into KVLoadCoordinator.
• Centralize sequence registration, transitions, paging cleanup, and connector
  notifications in SequenceLifecycle.
• Move checkpoint restore/save planning and publication to
  StateCheckpointLifecycle.
• Simplify eviction by removing the unused base hierarchy and making resource
  dependencies explicit.
• Replace engine and connector reach-through with complete scheduler operations
  or direct injection of the actual narrow owner.
• Introduce connector-owned KVConnectorStepInput while keeping
  SchedulerOutput compatible.
• Align lifecycle and command/query naming across scheduler collaborators.
• Split scheduler tests by behavior owner: facade/decode, prefill, SSM,
  external KV transfer, and state management.

The scheduler cross-module reach-through audit was reduced from 48 findings to
zero.

Behavior and compatibility

This is intended to be a behavior-preserving refactor:

• Prefill ordering, token budgets, TTFT policy, prefix-cache behavior,
  long-context chunking, SSM checkpoints, migration, and KV-transfer semantics
  are preserved.
• Public engine configuration is unchanged.
• SchedulerOutput remains compatible.

Two intentional compatibility notes:

• schedule(is_prefill=False) now rejects the obsolete decode path explicitly;
  production decode already uses schedule_running().
• Connector extensions now receive the shape-compatible
  KVConnectorStepInput instead of depending on SchedulerOutput.

[Copilot]

Pull request overview

Refactors the PyTorch paging scheduler into narrower ownership components (prefill admission, external KV-load/save coordination, and sequence lifecycle/state transitions) while keeping engine-facing behavior and SchedulerOutput compatibility.

Changes:

• Introduces a dedicated prefill_scheduler.py with explicit prefill turn policies, tentative prefix-match rollback, long-prefill gating, and external KV-load admission integration.
• Centralizes request/session lifecycle and paging cleanup in SequenceLifecycle, and moves external KV-load/save ownership into KVLoadCoordinator / KVSaveCoordinator.
• Updates engine/connector interfaces and expands tests to cover prefill ordering, SSM checkpoint lifecycle, and async KV-transfer behaviors.

Reviewed changes

Copilot reviewed 37 out of 37 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
lmdeploy/pytorch/paging/prefill_scheduler.py	New prefill admission/ordering owner with explicit policies and rollback states.
lmdeploy/pytorch/paging/seq_states/states.py	Adds SequenceLifecycle; replaces direct scheduler reach-through with lifecycle-owned transitions/cleanup.
lmdeploy/pytorch/paging/seq_states/init.py	Exports updated lifecycle/state symbols.
lmdeploy/pytorch/paging/kv_load_coordinator.py	Refactors external KV-load planning/admission and deferred cleanup handling.
lmdeploy/pytorch/paging/kv_save_coordinator.py	Refactors save lease ownership to depend on BaseBlockManager (not scheduler).
lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py	Makes eviction helper depend on explicit owners; aligns eviction with load tracking and lifecycle cleanup.
lmdeploy/pytorch/paging/eviction_helper/init.py	Updates eviction helper construction to inject explicit dependencies.
lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py	Removes unused base hierarchy.
lmdeploy/pytorch/paging/block_trie/trie.py	Adds finalize_match() for shared cached-token accounting across local matches and external loads.
lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py	Adds CheckpointCopyPlan + batch APIs and forward-dispatch pin/unpin coordination.
lmdeploy/pytorch/paging/block_trie/README.md	Updates docs to point at the new prefill owner and responsibilities.
lmdeploy/pytorch/messages.py	Renames/clarifies SequenceManager APIs; wires sessions through SequenceLifecycle; adds SchedulerSequence.activate/finish.
lmdeploy/pytorch/kv_connector/base.py	Introduces connector-owned KVConnectorStepInput interface.
lmdeploy/pytorch/kv_connector/init.py	Re-exports KVConnectorStepInput.
lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py	Updates connector scheduler to consume KVConnectorStepInput.
lmdeploy/pytorch/kv_connector/mooncake/store/connector.py	Updates connector wrapper to accept KVConnectorStepInput.
lmdeploy/pytorch/engine/inputs_maker.py	Routes checkpoint planning through StateCheckpointLifecycle; resumes completed migrations before selecting work; adjusts connector detection.
lmdeploy/pytorch/engine/engine_loop.py	Uses injected StateCheckpointLifecycle; updates migration scheduling and checkpoint pin/unpin flow.
lmdeploy/pytorch/engine/engine.py	Switches session access to Scheduler.get_session/get_sessions; updates finish/sleep transfer drain call.
lmdeploy/pytorch/disagg/conn/engine_conn.py	Updates session-end semantics to use boolean engine.end_session result.
lmdeploy/pytorch/strategies/ar/sequence.py	Uses msg.finish() convenience wrapper.
lmdeploy/pytorch/strategies/ar_spec/sequence.py	Uses msg.finish() convenience wrapper.
lmdeploy/pytorch/strategies/dllm/sequence.py	Uses msg.finish() convenience wrapper.
tests/pytorch/paging/test_state_manager.py	Adds direct StateManager behavior tests for reserved/runtime/checkpoint slot semantics.
tests/pytorch/paging/test_prefill_scheduler.py	New coverage for prefill ordering, token budgets, prefix-hit rollback, and long-context chunk admission.
tests/pytorch/paging/test_scheduler_ssm.py	New coverage for SSM runtime/checkpoint slot interactions and rollback behavior.
tests/pytorch/paging/test_scheduler_kv_transfer.py	New coverage for async external KV lookup/load/save lifecycle and cleanup.
tests/pytorch/paging/test_block_trie/test_trie.py	Updates API usage to release_paging_resources().
tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py	Renames tests and adds coverage for new batch copy-plan APIs and dispatch pinning.
tests/pytorch/kv_connector/test_mooncake_store_scheduler.py	Updates tests to use KVConnectorStepInput instead of SchedulerOutput-shaped namespaces.
tests/pytorch/kv_connector/test_mooncake_store_connector.py	Updates tests to use KVConnectorStepInput and new scheduler delegation signatures.
tests/pytorch/engine/test_inputs_maker.py	Updates tests for new checkpoint lifecycle batch APIs and migration-resume ordering.
tests/pytorch/engine/test_engine_sleep.py	Updates fake scheduler/session APIs used by sleep flow.
tests/pytorch/engine/test_abort_stopped_seqs.py	Adds get_session() to fake scheduler to match new engine expectations.
tests/pytorch/engine/test_kv_connector_wiring.py	Removes outdated shutdown test; keeps engine loop shutdown ordering coverage.
lmdeploy/pytorch/paging/scheduler.py	Public facade updated to compose new owners and expose narrow helper APIs.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.