From cfcb5ce68a297ae2c288cbf1a18ddb7c78ffabc1 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 12:16:46 +0800
Subject: [PATCH 01/22] test: characterize scheduler admission rollback

---
 tests/pytorch/paging/test_scheduler.py | 67 ++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index 4275c80366..5de4388369 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -1029,6 +1029,43 @@ def test_async_load_requires_capacity_for_the_complete_prefill():
     assert connector.allocations == []
 
 
+def test_async_load_capacity_failure_restores_tentative_local_prefix(monkeypatch):
+    connector = _AsyncLookupConnector([(4, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        num_gpu_blocks=3,
+    )
+    tokens = torch.arange(13)
+    cached = scheduler.add_session(76).add_sequence(tokens[:5])
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached.state.stop()
+    cached_block = cached.logical_blocks.get_real_blocks()[:1]
+    ref_count = scheduler.block_manager.allocator.get_ref_count(cached_block).copy()
+    scheduler.block_trie.stats.reset()
+
+    seq = scheduler.add_session(77).add_sequence(tokens)
+    evict_for_seq = Mock(return_value=False)
+    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_history_ids == 0
+    assert seq.num_blocks == 0
+    assert seq.kv_token_limit is None
+    assert seq.cached_tokens == 0
+    assert seq.prefix_cache.trie_cursor is None
+    assert seq.prefix_cache.match_start_step == -1
+    assert connector.lookup_calls == [(seq.seq_id, 4)]
+    assert connector.allocations == []
+    assert evict_for_seq.call_count == 1
+    assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
 def test_async_load_does_not_consume_model_batch_slot():
     connector = _AsyncLookupConnector([(8, True), (0, False)])
     scheduler = _make_async_lookup_scheduler(
@@ -1432,6 +1469,36 @@ def test_scheduler_reorder_cache_stays_order_only_after_prefix_hit():
     assert normal.status == MessageStatus.READY
 
 
+def test_scheduler_resource_rejection_rolls_back_tentative_prefix_match(monkeypatch):
+    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=1)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2])
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+    cached_block = cached.logical_blocks.get_real_blocks()[:1]
+    ref_count = scheduler.block_manager.allocator.get_ref_count(cached_block).copy()
+    scheduler.block_trie.stats.reset()
+
+    seq = scheduler.add_session(1).add_sequence([1] * block_size + [3])
+    evict_for_seq = Mock(return_value=False)
+    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_history_ids == 0
+    assert seq.num_blocks == 0
+    assert seq.kv_token_limit is None
+    assert seq.cached_tokens == 0
+    assert seq.prefix_cache.trie_cursor is None
+    assert seq.prefix_cache.match_start_step == -1
+    assert evict_for_seq.call_count == 1
+    assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
 def test_scheduler_rolls_back_prefix_match_for_prefill_gate_when_tail_still_exceeds_budget():
     scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
 

From 14931e8d58ac31325a1e216f4303f0f4df00402b Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 14:28:02 +0800
Subject: [PATCH 02/22] refactor: make prefill admission outcome explicit

---
 lmdeploy/pytorch/paging/scheduler.py | 64 +++++++++++-----------------
 1 file changed, 26 insertions(+), 38 deletions(-)

diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index d5f2450ddc..9298c8e2cc 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -56,6 +56,7 @@
   TP rank reports terminal progress or worker queues are drained.
 """
 
+import enum
 import time
 from collections import Counter, OrderedDict
 from contextlib import contextmanager
@@ -246,43 +247,45 @@ def _reorder_for_short_turn(self, waiting: SeqList):
         return self._sort_normal_prefills(normal_waiting) + long_waiting
 
 
+class _PrefillAdmissionAction(enum.Enum):
+    ADMIT = enum.auto()
+    SKIP = enum.auto()
+    STOP = enum.auto()
+    LOAD_STARTED = enum.auto()
+
+
 @dataclass(frozen=True)
 class _PrefillAdmissionResult:
     """Outcome from trying to admit one waiting prefill request.
 
     The outer loop distinguishes four outcomes:
 
-    * ``admitted``: include the request in this tick's model batch.
-    * ``should_skip``: leave it waiting but continue trying later candidates.
-    * ``should_stop``: resource pressure ends this prefill admission turn.
-    * ``load_started``: no model work was selected, but the request left the
+    * ``ADMIT``: include the request in this tick's model batch.
+    * ``SKIP``: leave it waiting but continue trying later candidates.
+    * ``STOP``: resource pressure ends this prefill admission turn.
+    * ``LOAD_STARTED``: no model work was selected, but the request left the
       waiting queue for asynchronous KV load.
     """
 
-    admitted: bool
+    action: _PrefillAdmissionAction
     prefill_token_count: int = 0
-    should_skip: bool = False
-    load_started: bool = False
 
     @classmethod
     def admit(cls, prefill_token_count: int):
-        return cls(admitted=True, prefill_token_count=prefill_token_count)
+        return cls(action=_PrefillAdmissionAction.ADMIT,
+                   prefill_token_count=prefill_token_count)
 
     @classmethod
     def skip(cls):
-        return cls(admitted=False, should_skip=True)
+        return cls(action=_PrefillAdmissionAction.SKIP)
 
     @classmethod
     def stop(cls):
-        return cls(admitted=False)
+        return cls(action=_PrefillAdmissionAction.STOP)
 
     @classmethod
     def load(cls):
-        return cls(admitted=False, load_started=True)
-
-    @property
-    def should_stop(self):
-        return not self.admitted and not self.should_skip and not self.load_started
+        return cls(action=_PrefillAdmissionAction.LOAD_STARTED)
 
 
 @dataclass(frozen=True)
@@ -365,18 +368,18 @@ def run(self):
         """
         if self.scheduler._external_lookup_enabled and not self._remote_ready:
             if self._lookup_is_pending():
-                return self._check_result(_PrefillAdmissionResult.skip())
+                return _PrefillAdmissionResult.skip()
             self._capture_prefix_match_baseline()
 
         gate_result = self._check_prefill_admission_gates()
         if gate_result is not None:
-            return self._check_result(gate_result)
+            return gate_result
 
         resource_result = self._admit_resources()
         if resource_result is not None:
-            return self._check_result(resource_result)
+            return resource_result
 
-        return self._check_result(self._finish_admission())
+        return self._finish_admission()
 
     def _rollback_gate(self, stats_snapshot, reason: str):
         """Rollback a tentative prefix hit and return any gate-only rejection.
@@ -389,22 +392,6 @@ def _rollback_gate(self, stats_snapshot, reason: str):
         self._rollback_prefix_match(stats_snapshot, reason)
         return self._gate_match_rollback_result
 
-    def _check_result(self, result: _PrefillAdmissionResult):
-        if result.admitted and result.should_skip:
-            self._warn_unexpected_state(
-                f'admission result both admits and skips: prefill_token_count={result.prefill_token_count}')
-        if not result.admitted and result.prefill_token_count != 0:
-            self._warn_unexpected_state(
-                f'rejected admission result carries token count: prefill_token_count={result.prefill_token_count}')
-        if result.load_started and (result.admitted or result.should_skip):
-            self._warn_unexpected_state('external load result has conflicting admission flags')
-        return result
-
-    def _warn_unexpected_state(self, message: str):
-        seq = self.seq
-        logger.warning('Unexpected prefill admission state: session_id=%s seq_id=%s %s',
-                       seq.session_id, seq.seq_id, message)
-
     def _lookup_is_pending(self) -> bool:
         """Skip without touching local prefix state while lookup is running.
 
@@ -1206,19 +1193,20 @@ def _to_running(seq: SchedulerSequence, prefill_token_count: int):
                 allow_long_prefill=allow_long_prefill,
             ).run()
 
-            if admission.load_started:
+            if admission.action is _PrefillAdmissionAction.LOAD_STARTED:
                 # start_load already moved the sequence out of WAITING. It must
                 # not join running or skipped_waiting, both of which permit
                 # ordinary model/paging operations on the sequence. Since no
                 # model work was admitted, continue without consuming a batch
                 # slot or token budget.
                 continue
-            if admission.should_skip:
+            if admission.action is _PrefillAdmissionAction.SKIP:
                 skipped_waiting.append(seq)
                 continue
-            if admission.should_stop:
+            if admission.action is _PrefillAdmissionAction.STOP:
                 break
 
+            assert admission.action is _PrefillAdmissionAction.ADMIT
             _to_running(seq, admission.prefill_token_count)
 
             seq.record_event(EventType.SCHEDULED)

From 24e5dcf1938a9cd7d7d510a0dde6eb82d00096ba Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 14:45:22 +0800
Subject: [PATCH 03/22] refactor: centralize tentative prefix match lifecycle

---
 lmdeploy/pytorch/paging/scheduler.py   | 369 +++++++++++++------------
 tests/pytorch/paging/test_scheduler.py |  46 ++-
 2 files changed, 238 insertions(+), 177 deletions(-)

diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 9298c8e2cc..2582fd88e8 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -288,9 +288,9 @@ def load(cls):
         return cls(action=_PrefillAdmissionAction.LOAD_STARTED)
 
 
-@dataclass(frozen=True)
-class _PrefixMatchBaseline:
-    """Exact pre-attempt state used to undo a tentative local trie match.
+@dataclass(frozen=True, slots=True)
+class _PrefixMatchStateSnapshot:
+    """Exact sequence state captured before a tentative local trie match.
 
     External lookup itself does not mutate sequence paging state. The scheduler
     may, however, run ``block_trie.match`` first so the connector queries only
@@ -321,8 +321,160 @@ class _PrefixMatchBaseline:
     trie_block_map: dict[int, int]
     # Model state that must remain aligned with the committed history step.
     model_meta: Any
-    # Matching mutates global hit statistics, so those are transactional too.
-    stats_snapshot: Any
+
+    @classmethod
+    def capture(cls, seq: SchedulerSequence):
+        overlap = seq.prefix_cache.recompute_overlap
+        return cls(
+            num_history_ids=int(seq.num_history_ids),
+            num_blocks=int(seq.num_blocks),
+            trie_cursor=seq.prefix_cache.trie_cursor,
+            match_start_step=int(seq.prefix_cache.match_start_step),
+            cached_tokens=int(seq.cached_tokens),
+            kv_token_limit=seq.kv_token_limit,
+            fresh_block_range=overlap.fresh_block_range,
+            trie_block_map=dict(overlap.trie_block_map),
+            model_meta=seq.model_meta,
+        )
+
+
+class _TentativePrefixMatch:
+    """Request-local transaction around ``BlockTrie.match`` side effects.
+
+    Ordinary and SSM admission preserve the historical fallback to an unmatched request. External lookup instead needs
+    an exact pre-match snapshot because a multi-turn request may already own committed progress. Both contracts share
+    one stats snapshot, restore-pin boundary, and explicit commit/rollback lifecycle without changing their rollback
+    semantics.
+    """
+
+    __slots__ = (
+        'seq',
+        'block_trie',
+        'block_manager',
+        'is_ssm',
+        '_preserve_existing_state',
+        '_stats_snapshot',
+        '_state_snapshot',
+        '_rejection_on_rollback',
+        '_started',
+        'matched',
+    )
+
+    def __init__(self,
+                 seq: SchedulerSequence,
+                 block_trie: BlockTrie,
+                 block_manager,
+                 *,
+                 is_ssm: bool,
+                 preserve_existing_state: bool):
+        self.seq = seq
+        self.block_trie = block_trie
+        self.block_manager = block_manager
+        self.is_ssm = is_ssm
+        self._preserve_existing_state = preserve_existing_state
+        self._stats_snapshot = None
+        self._state_snapshot: _PrefixMatchStateSnapshot | None = None
+        self._rejection_on_rollback: _PrefillAdmissionResult | None = None
+        self._started = False
+        self.matched = False
+
+    def begin(self) -> None:
+        """Start the transaction before gates can mutate exact external state.
+
+        Ordinary admission starts lazily from ``match``. External admission
+        starts before gates so rollback can restore existing request state even
+        when a private partial block prevents another trie match.
+        """
+        if self._started or not self.block_trie.enabled:
+            return
+        self._stats_snapshot = self.block_trie.stats.snapshot()
+        if self._preserve_existing_state:
+            self._state_snapshot = _PrefixMatchStateSnapshot.capture(self.seq)
+        self._started = True
+
+    def match(self) -> None:
+        """Apply one tentative match after capturing its rollback boundary."""
+        assert not self.matched
+        self.begin()
+        self.block_trie.match(self.seq)
+        self.matched = True
+
+    def retain_for_admission(self, rejection_on_rollback: _PrefillAdmissionResult) -> None:
+        """Keep a gate-enabling match and remember its original rejection."""
+        assert self.matched
+        self._rejection_on_rollback = rejection_on_rollback
+
+    def pin_restore(self) -> bool:
+        """Pin an SSM restore selected by this tentative match."""
+        restore = self.seq.prefix_cache.restore
+        if not self.is_ssm or not restore.is_selected:
+            return True
+        return self.block_trie.state_checkpoints.pin_restore(self.seq)
+
+    def commit(self) -> None:
+        """Accept the match and discard request-local rollback state."""
+        self._clear()
+
+    def rollback(self, reason: str):
+        """Undo the transaction and return any gate-defined rejection."""
+        rejection = self._rejection_on_rollback
+        if not self._started:
+            return rejection
+
+        seq = self.seq
+        logger.debug('Rollback tentative prefix-cache match: session_id=%s seq_id=%s reason=%s '
+                     'num_history_ids=%s restore_state=%s', seq.session_id, seq.seq_id, reason, seq.num_history_ids,
+                     seq.prefix_cache.restore.slot)
+        self.block_trie.stats.restore(self._stats_snapshot)
+        snapshot = self._state_snapshot
+        if snapshot is None:
+            self._reset_to_unmatched()
+        else:
+            self._restore_snapshot(snapshot)
+        self._clear()
+        return rejection
+
+    def _restore_snapshot(self, snapshot: _PrefixMatchStateSnapshot) -> None:
+        seq = self.seq
+        if seq.num_blocks < snapshot.num_blocks:
+            raise RuntimeError(
+                'tentative prefix match removed sequence-owned baseline blocks')
+        if seq.num_blocks > snapshot.num_blocks:
+            self.block_manager.truncate(seq, snapshot.num_blocks)
+        seq.set_step(snapshot.num_history_ids)
+        seq.model_meta = snapshot.model_meta
+        seq.kv_token_limit = snapshot.kv_token_limit
+        prefix_cache = seq.prefix_cache
+        prefix_cache.trie_cursor = snapshot.trie_cursor
+        prefix_cache.match_start_step = snapshot.match_start_step
+        overlap = prefix_cache.recompute_overlap
+        overlap.fresh_block_range = snapshot.fresh_block_range
+        overlap.trie_block_map.clear()
+        overlap.trie_block_map.update(snapshot.trie_block_map)
+        seq.cached_tokens = snapshot.cached_tokens
+
+    def _reset_to_unmatched(self) -> None:
+        seq = self.seq
+        if self.is_ssm:
+            self.block_trie.state_checkpoints.unpin_restore(seq)
+        if seq.num_blocks > 0 or seq.logical_state >= 0:
+            seq.state.free()
+        elif seq.num_history_ids > 0:
+            seq.set_step(0)
+        seq.kv_token_limit = None
+        prefix_cache = seq.prefix_cache
+        prefix_cache.trie_cursor = None
+        prefix_cache.restore.clear()
+        prefix_cache.match_start_step = -1
+        prefix_cache.recompute_overlap.clear_tracking()
+        seq.cached_tokens = 0
+
+    def _clear(self) -> None:
+        self._stats_snapshot = None
+        self._state_snapshot = None
+        self._rejection_on_rollback = None
+        self._started = False
+        self.matched = False
 
 
 class _PrefillAdmissionAttempt:
@@ -351,9 +503,14 @@ def __init__(self,
         self._remote_ready = scheduler.kv_load_coordinator.is_remote_ready(seq)
         self.allow_long_prefill = allow_long_prefill
         self._alloc_size = prealloc_size
-        self._gate_match_stats_snapshot = None
-        self._gate_match_rollback_result = None
-        self._prefix_match_baseline: _PrefixMatchBaseline | None = None
+        self._prefix_match = _TentativePrefixMatch(
+            seq,
+            scheduler.block_trie,
+            scheduler.block_manager,
+            is_ssm=scheduler.is_ssm,
+            preserve_existing_state=(
+                scheduler._external_lookup_enabled and not self._remote_ready),
+        )
 
     def run(self):
         """Run the admission route for one waiting prefill.
@@ -369,7 +526,7 @@ def run(self):
         if self.scheduler._external_lookup_enabled and not self._remote_ready:
             if self._lookup_is_pending():
                 return _PrefillAdmissionResult.skip()
-            self._capture_prefix_match_baseline()
+            self._prefix_match.begin()
 
         gate_result = self._check_prefill_admission_gates()
         if gate_result is not None:
@@ -381,17 +538,6 @@ def run(self):
 
         return self._finish_admission()
 
-    def _rollback_gate(self, stats_snapshot, reason: str):
-        """Rollback a tentative prefix hit and return any gate-only rejection.
-
-        A prefill gate may do a tentative prefix-cache match before resource
-        admission. If that match is rolled back, the candidate should follow
-        the gate's original skip/stop result. Matches created after the gate
-        return ``None`` so the resource branch keeps its own retry/stop behavior.
-        """
-        self._rollback_prefix_match(stats_snapshot, reason)
-        return self._gate_match_rollback_result
-
     def _lookup_is_pending(self) -> bool:
         """Skip without touching local prefix state while lookup is running.
 
@@ -406,42 +552,6 @@ def _lookup_is_pending(self) -> bool:
         scheduler.last_schedule_had_pending_lookup = True
         return True
 
-    def _capture_prefix_match_baseline(self) -> None:
-        """Snapshot state before gates or resource admission may match locally.
-
-        The snapshot must precede prefill gates because a gate may call
-        ``block_trie.match`` to see whether a local hit makes a long prefill or
-        token-budget rejection schedulable. Resource admission may also match
-        before polling the external prefix. Both paths immediately advance
-        history, attach shared blocks, and mutate trie/overlap/statistics state.
-
-        If the following external poll returns ``None``, this tick skips the
-        request while the connector Future remains pending. ``_query_external_prefix``
-        uses the snapshot to remove only those tentative local-match side
-        effects, leaving any pre-existing multi-turn KV intact. No snapshot is
-        needed when the local trie is disabled because lookup polling has not
-        mutated sequence paging state. Once ``start_load`` succeeds, this
-        baseline is no longer the failure boundary; ``KVLoadCoordinator`` owns
-        rollback for potentially written destination blocks.
-        """
-        scheduler = self.scheduler
-        if not scheduler.block_trie.enabled:
-            return
-        seq = self.seq
-        overlap = seq.prefix_cache.recompute_overlap
-        self._prefix_match_baseline = _PrefixMatchBaseline(
-            num_history_ids=int(seq.num_history_ids),
-            num_blocks=int(seq.num_blocks),
-            trie_cursor=seq.prefix_cache.trie_cursor,
-            match_start_step=int(seq.prefix_cache.match_start_step),
-            cached_tokens=int(seq.cached_tokens),
-            kv_token_limit=seq.kv_token_limit,
-            fresh_block_range=overlap.fresh_block_range,
-            trie_block_map=dict(overlap.trie_block_map),
-            model_meta=seq.model_meta,
-            stats_snapshot=scheduler.block_trie.stats.snapshot(),
-        )
-
     def _admit_resources(self):
         if self.scheduler.block_trie.enabled:
             return self._admit_prefix_cache_resources()
@@ -469,14 +579,12 @@ def _admit_prefix_cache_resources(self):
         """
         scheduler = self.scheduler
         seq = self.seq
-        stats_snapshot = self._gate_match_stats_snapshot
-        if stats_snapshot is None:
-            stats_snapshot = scheduler.block_trie.stats.snapshot()
+        if not self._prefix_match.matched:
             # A completed external load has already published the accepted
             # prefix interval. Matching again would restart accounting at the
             # remote step and drop the restored tokens from request metrics.
             if not self._remote_ready and not self._has_private_local_tail():
-                scheduler.block_trie.match(seq)
+                self._prefix_match.match()
 
         if scheduler._external_lookup_enabled:
             lookup_result = self._query_external_prefix()
@@ -484,14 +592,15 @@ def _admit_prefix_cache_resources(self):
                 return lookup_result
 
         had_ssm_restore = scheduler.is_ssm and seq.prefix_cache.restore.is_selected
-        if not scheduler._pin_ssm_restore_if_needed(seq):
-            result = self._rollback_gate(stats_snapshot, 'failed to pin SSM restore checkpoint')
+        if not self._prefix_match.pin_restore():
+            result = self._prefix_match.rollback(
+                'failed to pin SSM restore checkpoint')
             if result is not None:
                 return result
 
         if not self._prepare_and_evict():
             if not had_ssm_restore:
-                result = self._rollback_gate(stats_snapshot, 'eviction failed')
+                result = self._prefix_match.rollback('eviction failed')
                 if result is not None:
                     return result
                 return _PrefillAdmissionResult.stop()
@@ -499,14 +608,16 @@ def _admit_prefix_cache_resources(self):
             # A matched SSM restore may be pinning the only checkpoint state
             # that eviction would otherwise free. Roll it back once and retry
             # eviction before declaring the sequence unschedulable.
-            result = self._rollback_gate(stats_snapshot, 'eviction failed with pinned SSM restore')
+            result = self._prefix_match.rollback(
+                'eviction failed with pinned SSM restore')
             if result is not None:
                 return result
             if not self._prepare_and_evict():
                 return _PrefillAdmissionResult.stop()
 
         if scheduler.is_ssm and not scheduler._ensure_runtime_state_available():
-            result = self._rollback_gate(stats_snapshot, 'no runtime SSM state available')
+            result = self._prefix_match.rollback(
+                'no runtime SSM state available')
             if result is not None:
                 return result
             if not self._prepare_and_evict():
@@ -548,12 +659,7 @@ def _query_external_prefix(self):
                 return self._start_external_load(int(num_external_tokens))
             return None
 
-        baseline = self._prefix_match_baseline
-        if baseline is not None:
-            scheduler._rollback_unscheduled_prefix_match(
-                self.seq,
-                baseline=baseline,
-            )
+        self._prefix_match.rollback('external lookup pending')
         scheduler.last_schedule_had_pending_lookup = True
         return _PrefillAdmissionResult.skip()
 
@@ -610,14 +716,12 @@ def _start_external_load(self, num_external_tokens: int):
             # No worker has seen the destination yet, so local match state can
             # still be restored exactly and the request can retry later.
             seq.kv_token_limit = old_kv_token_limit
-            baseline = self._prefix_match_baseline
-            if baseline is not None:
-                reason = (
-                    'full prefill capacity unavailable'
-                    if not full_prefill_fits
-                    else 'soft prefill budget unavailable'
-                )
-                self._rollback_prefix_match(baseline.stats_snapshot, reason)
+            reason = (
+                'full prefill capacity unavailable'
+                if not full_prefill_fits
+                else 'soft prefill budget unavailable'
+            )
+            self._prefix_match.rollback(reason)
             return _PrefillAdmissionResult.stop()
 
         original_num_blocks = seq.num_blocks
@@ -655,6 +759,7 @@ def _start_external_load(self, num_external_tokens: int):
             seq.kv_token_limit = old_kv_token_limit
             raise
         seq.kv_token_limit = None
+        self._prefix_match.commit()
         return _PrefillAdmissionResult.load()
 
     def _match_prefix_for_prefill_gate(self):
@@ -663,9 +768,8 @@ def _match_prefix_for_prefill_gate(self):
         if (self._remote_ready or not scheduler.block_trie.enabled
                 or self._has_private_local_tail()):
             return None
-        stats_snapshot = scheduler.block_trie.stats.snapshot()
-        scheduler.block_trie.match(self.seq)
-        return stats_snapshot
+        self._prefix_match.match()
+        return True
 
     def _has_private_local_tail(self) -> bool:
         """Whether blocks exist beyond the full-block part of local history.
@@ -708,11 +812,6 @@ def _has_private_local_tail(self) -> bool:
         seq = self.seq
         return seq.num_blocks > int(seq.num_history_ids) // seq.block_size
 
-    def _keep_gate_prefix_match(self, stats_snapshot, rollback_result: _PrefillAdmissionResult):
-        """Keep a gate-enabling match for the following resource admission."""
-        self._gate_match_stats_snapshot = stats_snapshot
-        self._gate_match_rollback_result = rollback_result
-
     def _token_budget_rejection(self):
         if self.allow_long_prefill:
             return _PrefillAdmissionResult.stop()
@@ -727,29 +826,31 @@ def _check_prefill_admission_gates(self):
         is_nonfinal_long_prefill = scheduler._prefill_kv_token_limit(seq) is not None
 
         if is_nonfinal_long_prefill and not self.allow_long_prefill:
-            stats_snapshot = self._match_prefix_for_prefill_gate()
-            if stats_snapshot is None:
+            matched = self._match_prefix_for_prefill_gate()
+            if matched is None:
                 return _PrefillAdmissionResult.skip()
             if scheduler._prefill_kv_token_limit(seq) is not None:
-                self._rollback_prefix_match(stats_snapshot, 'still non-final long prefill on short turn')
+                self._prefix_match.rollback('still non-final long prefill on short turn')
                 return _PrefillAdmissionResult.skip()
-            self._keep_gate_prefix_match(stats_snapshot, _PrefillAdmissionResult.skip())
+            self._prefix_match.retain_for_admission(
+                _PrefillAdmissionResult.skip())
             prefill_token_count = scheduler._prefill_admission_token_count(seq)
 
         exceeds_token_budget = self.has_admitted and self.token_count + prefill_token_count > token_budget
         if not exceeds_token_budget:
             return None
 
-        if self._gate_match_stats_snapshot is None:
-            stats_snapshot = self._match_prefix_for_prefill_gate()
-            if stats_snapshot is not None:
+        if not self._prefix_match.matched:
+            matched = self._match_prefix_for_prefill_gate()
+            if matched is not None:
                 prefill_token_count = scheduler._prefill_admission_token_count(seq)
                 if self.token_count + prefill_token_count <= token_budget:
-                    self._keep_gate_prefix_match(stats_snapshot, self._token_budget_rejection())
+                    self._prefix_match.retain_for_admission(
+                        self._token_budget_rejection())
                     return None
-                self._rollback_prefix_match(stats_snapshot, 'still exceeds prefill token budget')
+                self._prefix_match.rollback('still exceeds prefill token budget')
         else:
-            self._rollback_prefix_match(self._gate_match_stats_snapshot, 'still exceeds prefill token budget')
+            self._prefix_match.rollback('still exceeds prefill token budget')
         return self._token_budget_rejection()
 
     def _prepare_and_evict(self):
@@ -772,17 +873,6 @@ def _evict_for_seq(self, alloc_size: int):
         evictable = list(chain(hanging, waiting))
         return scheduler.eviction_helper.evict_for_seq(self.seq, evictable, alloc_size)
 
-    def _rollback_prefix_match(self, stats_snapshot, reason: str):
-        seq = self.seq
-        logger.debug('Rollback tentative prefix-cache match: session_id=%s seq_id=%s reason=%s '
-                     'num_history_ids=%s restore_state=%s', seq.session_id, seq.seq_id, reason, seq.num_history_ids,
-                     seq.prefix_cache.restore.slot)
-        self.scheduler._rollback_unscheduled_prefix_match(
-            seq,
-            stats_snapshot,
-            baseline=self._prefix_match_baseline,
-        )
-
     def _finish_admission(self):
         scheduler = self.scheduler
         seq = self.seq
@@ -807,6 +897,7 @@ def _finish_admission(self):
             # reservation can be released only after model output advances the
             # sequence to input_end_pos.
             scheduler.kv_load_coordinator.mark_scheduled(seq)
+        self._prefix_match.commit()
         return _PrefillAdmissionResult.admit(prefill_token_count)
 
 
@@ -899,68 +990,6 @@ def _ensure_runtime_state_available(self):
         self.block_trie.state_checkpoints.evict(1)
         return self.state_manager.get_num_free_runtime() > 0
 
-    def _pin_ssm_restore_if_needed(self, seq: SchedulerSequence):
-        """Pin a matched SSM checkpoint before scheduler-side eviction."""
-        if not self.is_ssm or not seq.prefix_cache.restore.is_selected:
-            return True
-        return self.block_trie.state_checkpoints.pin_restore(seq)
-
-    def _rollback_unscheduled_prefix_match(
-        self,
-        seq: SchedulerSequence,
-        stats_snapshot=None,
-        *,
-        baseline: _PrefixMatchBaseline | None = None,
-    ):
-        """Drop a tentative prefix match that will not be used now.
-
-        ``block_trie.match()`` mutates sequence state immediately: it advances
-        the history step, appends shared blocks, and may pin a restore node.
-        If later eviction or state allocation fails, undo those side effects so
-        the waiting sequence can be scheduled cleanly in a later round.
-
-        ``baseline`` selects precise multi-turn rollback for external lookup.
-        Without it, this is the legacy new-request/SSM rollback that releases
-        all tentative ownership and returns the sequence to an unmatched state.
-        """
-        if baseline is not None:
-            # A tentative local match may only append shared blocks. Losing a
-            # block that existed in the baseline would mean it released
-            # sequence-owned state and cannot be repaired by truncation.
-            self.block_trie.stats.restore(baseline.stats_snapshot)
-            if seq.num_blocks < baseline.num_blocks:
-                raise RuntimeError(
-                    'tentative prefix match removed sequence-owned baseline blocks')
-            if seq.num_blocks > baseline.num_blocks:
-                self.block_manager.truncate(seq, baseline.num_blocks)
-            seq.set_step(baseline.num_history_ids)
-            seq.model_meta = baseline.model_meta
-            seq.kv_token_limit = baseline.kv_token_limit
-            prefix_cache = seq.prefix_cache
-            prefix_cache.trie_cursor = baseline.trie_cursor
-            prefix_cache.match_start_step = baseline.match_start_step
-            overlap = prefix_cache.recompute_overlap
-            overlap.fresh_block_range = baseline.fresh_block_range
-            overlap.trie_block_map.clear()
-            overlap.trie_block_map.update(baseline.trie_block_map)
-            seq.cached_tokens = baseline.cached_tokens
-            return
-
-        self.block_trie.stats.restore(stats_snapshot)
-        if self.is_ssm:
-            self.block_trie.state_checkpoints.unpin_restore(seq)
-        if seq.num_blocks > 0 or seq.logical_state >= 0:
-            seq.state.free()
-        elif seq.num_history_ids > 0:
-            seq.set_step(0)
-        seq.kv_token_limit = None
-        prefix_cache = seq.prefix_cache
-        prefix_cache.trie_cursor = None
-        prefix_cache.restore.clear()
-        prefix_cache.match_start_step = -1
-        prefix_cache.recompute_overlap.clear_tracking()
-        seq.cached_tokens = 0
-
     @staticmethod
     def _finalize_prefix_cache_match(seq: SchedulerSequence):
         """Publish accepted cached-token count within the current prompt."""
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index 5de4388369..262cf6ed1a 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -683,7 +683,7 @@ def test_scheduler_ar_spec_prefix_hit_recomputes_overlap_block():
     assert scheduler.block_trie.stats.num_hit_tokens == block_size * 2
 
 
-def test_scheduler_prefix_match_rollback_clears_recompute_overlap_window():
+def test_scheduler_prefix_match_rollback_clears_recompute_overlap_window(monkeypatch):
     from lmdeploy.pytorch.strategies.ar_spec.sequence import ARSpecSequenceStrategy
     block_size = 16
     seq_meta = SequenceMeta(block_size, strategy=ARSpecSequenceStrategy())
@@ -702,16 +702,15 @@ def test_scheduler_prefix_match_rollback_clears_recompute_overlap_window():
     cached = scheduler.add_session(0).add_sequence(token_ids)
     scheduler.block_manager.allocate(cached)
     scheduler.block_trie.allocate(cached)
+    cached.state.stop()
 
     seq = scheduler.add_session(1).add_sequence(token_ids)
-    stats_snapshot = scheduler.block_trie.stats.snapshot()
-    scheduler.block_trie.match(seq)
-
-    assert seq.num_history_ids == block_size * 2
-    assert seq.prefix_cache.recompute_overlap.fresh_block_range == range(2, 3)
+    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', Mock(return_value=False))
+    scheduler.block_trie.stats.reset()
 
-    scheduler._rollback_unscheduled_prefix_match(seq, stats_snapshot)
+    output = scheduler.schedule(is_prefill=True)
 
+    assert output.running == []
     assert seq.num_history_ids == 0
     assert seq.num_token_ids == len(token_ids)
     assert seq.cached_tokens == 0
@@ -906,6 +905,39 @@ def test_async_lookup_precisely_restores_a_multiturn_local_prefix():
     assert fourth.running == [seq]
 
 
+def test_async_lookup_pending_preserves_private_partial_prefix():
+    connector = _AsyncLookupConnector([(None, False)])
+    scheduler = _make_async_lookup_scheduler(connector)
+    tokens = torch.arange(13)
+    seq = scheduler.add_session(72).add_sequence(tokens)
+
+    seq.kv_token_limit = 5
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    seq.set_step(5)
+    seq.kv_token_limit = 7
+    seq.cached_tokens = 3
+    seq.model_meta = {'state': 'keep'}
+    baseline_blocks = seq.logical_blocks.get_real_blocks().copy()
+    baseline_cursor = seq.prefix_cache.trie_cursor
+    scheduler.block_trie.stats.reset()
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert scheduler.last_schedule_had_pending_lookup
+    assert connector.lookup_calls == [(seq.seq_id, 5)]
+    assert seq.num_history_ids == 5
+    assert torch.equal(torch.from_numpy(seq.logical_blocks.get_real_blocks()),
+                       torch.from_numpy(baseline_blocks))
+    assert seq.prefix_cache.trie_cursor is baseline_cursor
+    assert seq.prefix_cache.match_start_step == -1
+    assert seq.cached_tokens == 3
+    assert seq.kv_token_limit == 7
+    assert seq.model_meta == {'state': 'keep'}
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+
+
 def test_external_cached_tokens_survive_remote_ready_admission():
     connector = _AsyncLookupConnector([(8, True)])
     scheduler = _make_async_lookup_scheduler(connector)

From 185dd43767a7b7e1739f6c9eee4412a92ba77591 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 15:07:43 +0800
Subject: [PATCH 04/22] refactor: move external load admission to coordinator

---
 lmdeploy/pytorch/paging/block_trie/trie.py    |  25 ++
 .../pytorch/paging/kv_load_coordinator.py     | 194 ++++++++++++--
 lmdeploy/pytorch/paging/scheduler.py          | 236 +++++-------------
 .../engine/test_kv_connector_wiring.py        |   2 +-
 4 files changed, 255 insertions(+), 202 deletions(-)

diff --git a/lmdeploy/pytorch/paging/block_trie/trie.py b/lmdeploy/pytorch/paging/block_trie/trie.py
index 6ca9eb88dd..cddbe35538 100644
--- a/lmdeploy/pytorch/paging/block_trie/trie.py
+++ b/lmdeploy/pytorch/paging/block_trie/trie.py
@@ -530,6 +530,31 @@ def match(self, seq: SchedulerSequence):
 
         self._match_block_prefix(seq)
 
+    @staticmethod
+    def finalize_match(seq: SchedulerSequence) -> None:
+        """Publish accepted current-prompt cache reuse for a sequence.
+
+        Local trie matches and completed external loads share this final accounting step. Recompute-preemption matches
+        remain usable internally but deliberately suppress public cached-token statistics.
+        """
+        prefix_cache = seq.prefix_cache
+        if prefix_cache.suppress_match_stats:
+            seq.cached_tokens = 0
+            prefix_cache.suppress_match_stats = False
+            return
+
+        match_start = prefix_cache.match_start_step
+        if match_start < 0:
+            seq.cached_tokens = 0
+            return
+        cached_end = seq.num_history_ids
+        prompt_start = seq.input_start_pos
+        prompt_end = seq.input_end_pos
+        seq.cached_tokens = max(
+            0,
+            min(cached_end, prompt_end) - max(match_start, prompt_start),
+        )
+
     def _ensure_attached_allocation_cursor(self, seq: SchedulerSequence):
         """Return an attached cursor, resetting a stale sequence cursor."""
         node = seq.prefix_cache.trie_cursor
diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index 6bc485d610..4341ff9087 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -1,11 +1,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 """Paging ownership for asynchronous external KV-cache loads.
 
-The connector owns lookup keys, remote I/O, and worker progress, while the
-scheduler owns sequences and GPU blocks.  This coordinator bridges those two
-lifetimes without performing I/O itself:
+The connector owns lookup keys, remote I/O, and worker progress. This
+coordinator owns destination admission and bridges connector progress to
+sequence and GPU-block lifetimes without performing worker I/O itself:
 
-1. The scheduler allocates destination blocks and calls :meth:`start_load`.
+1. :meth:`try_load` polls lookup, admits the complete prefill, allocates exact
+   destinations, binds connector metadata, and calls :meth:`start_load`.
 2. While workers may write those blocks, the sequence stays in
    ``WAITING_FOR_REMOTE_KVS`` and cannot be evicted or removed.
 3. :meth:`update` publishes a successful load or rolls a failed/cancelled load
@@ -23,14 +24,29 @@
 from __future__ import annotations
 
 import enum
+from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from lmdeploy.pytorch.kv_connector import KVLoadResult
-from lmdeploy.pytorch.messages import SchedulerSequence
+from lmdeploy.pytorch.messages import SchedulerSequence, SchedulerSession
 
 if TYPE_CHECKING:
-    from .scheduler import Scheduler
+    from lmdeploy.pytorch.kv_connector import KVConnectorBase
+
+    from .block_manager.base_block_manager import BaseBlockManager
+    from .block_trie import BlockTrie
+    from .eviction_helper.base_eviction_helper import BaseEvictionHelper
+
+
+class KVLoadAdmission(enum.Enum):
+    """Narrow external-load result interpreted by prefill queue policy."""
+
+    NO_LOAD = enum.auto()
+    PENDING = enum.auto()
+    FULL_PREFILL_UNAVAILABLE = enum.auto()
+    SOFT_BUDGET_UNAVAILABLE = enum.auto()
+    STARTED = enum.auto()
 
 
 class _LoadPhase(enum.Enum):
@@ -84,8 +100,22 @@ class KVLoadCoordinator:
       completion/preemption of the remaining prefill.
     """
 
-    def __init__(self, scheduler: Scheduler) -> None:
-        self.scheduler = scheduler
+    def __init__(
+        self,
+        *,
+        lookup_enabled: bool,
+        connector: KVConnectorBase | None,
+        block_manager: BaseBlockManager,
+        block_trie: BlockTrie,
+        eviction_helper: BaseEvictionHelper,
+        sessions: dict[int, SchedulerSession],
+    ) -> None:
+        self.lookup_enabled = lookup_enabled
+        self.connector = connector
+        self.block_manager = block_manager
+        self.block_trie = block_trie
+        self.eviction_helper = eviction_helper
+        self.sessions = sessions
         # Active load lifecycle. Records survive the LOADING -> READY ->
         # PREFILLING transitions so stop/end and preemption can find the owner.
         self._loads: dict[int, _LoadRecord] = {}
@@ -104,7 +134,7 @@ def prefill_target_blocks(
         remote-hit allocation may stop earlier, but load admission must know
         whether the eventual full prefill has a path to completion.
         """
-        block_size = self.scheduler.cache_config.block_size
+        block_size = seq.block_size
         target_tokens = int(seq.num_all_ids) + max(0, int(prealloc_size))
         return (target_tokens + block_size - 1) // block_size
 
@@ -122,7 +152,7 @@ def track_prefill(
         both paths in one table prevents the new load from consuming capacity
         required by work that the scheduler has already accepted.
         """
-        if not self.scheduler._external_lookup_enabled:
+        if not self.lookup_enabled:
             return
         if target_blocks is None:
             target_blocks = self.prefill_target_blocks(seq, prealloc_size)
@@ -153,9 +183,121 @@ def can_admit_load(
         """
         missing_blocks = max(0, int(target_blocks) - int(seq.num_blocks))
         soft_reserved = self.soft_reserved_blocks(exclude_seq=seq)
-        free_blocks = self.scheduler.block_manager.get_num_free_gpu_blocks()
+        free_blocks = self.block_manager.get_num_free_gpu_blocks()
         return missing_blocks + soft_reserved <= free_blocks
 
+    def is_lookup_pending(self, seq: SchedulerSequence) -> bool:
+        """Whether this request already has an asynchronous lookup in
+        flight."""
+        connector = self.connector
+        if not self.lookup_enabled or connector is None:
+            return False
+        return connector.is_lookup_pending(seq.seq_id)
+
+    def try_load(
+        self,
+        seq: SchedulerSequence,
+        *,
+        prealloc_size: int,
+        evictable_seqs: Iterable[SchedulerSequence],
+    ) -> KVLoadAdmission:
+        """Poll and admit one external prefix without choosing queue policy.
+
+        The return value describes only the connector/paging result. Prefill admission remains responsible for mapping
+        it to skip, stop, continue, or load-started and for committing or rolling back its local match.
+        """
+        connector = self.connector
+        if not self.lookup_enabled or connector is None:
+            return KVLoadAdmission.NO_LOAD
+
+        num_external_tokens, _ = connector.get_num_new_matched_tokens(
+            seq,
+            seq.num_history_ids,
+        )
+        if num_external_tokens is None:
+            return KVLoadAdmission.PENDING
+        if num_external_tokens <= 0:
+            return KVLoadAdmission.NO_LOAD
+        return self._admit_load(
+            seq,
+            num_external_tokens=int(num_external_tokens),
+            prealloc_size=prealloc_size,
+            evictable_seqs=evictable_seqs,
+        )
+
+    def _admit_load(
+        self,
+        seq: SchedulerSequence,
+        *,
+        num_external_tokens: int,
+        prealloc_size: int,
+        evictable_seqs: Iterable[SchedulerSequence],
+    ) -> KVLoadAdmission:
+        """Admit the complete prefill, then allocate the remote interval."""
+        connector = self.connector
+        assert connector is not None
+        block_size = seq.block_size
+        local_step = int(seq.num_history_ids)
+        # Transfers are block-granular. Reuse a private partial boundary block,
+        # publish only full loaded blocks, and leave the final token to compute.
+        fallback_step = local_step // block_size * block_size
+        remote_step = local_step + num_external_tokens
+        remote_step = min(remote_step, int(seq.get_prefix_cache_max_match_step()))
+        remote_step = remote_step // block_size * block_size
+        if remote_step <= fallback_step:
+            return KVLoadAdmission.NO_LOAD
+
+        target_blocks = self.prefill_target_blocks(seq, prealloc_size)
+        old_kv_token_limit = seq.kv_token_limit
+        # Only the remote hit is allocated now, but admission guarantees the
+        # complete prefill can finish beside every existing soft reservation.
+        seq.kv_token_limit = None
+        full_prefill_fits = self.eviction_helper.evict_for_seq(
+            seq,
+            list(evictable_seqs),
+            prealloc_size,
+        )
+        if not full_prefill_fits:
+            seq.kv_token_limit = old_kv_token_limit
+            return KVLoadAdmission.FULL_PREFILL_UNAVAILABLE
+        if not self.can_admit_load(seq, target_blocks):
+            seq.kv_token_limit = old_kv_token_limit
+            return KVLoadAdmission.SOFT_BUDGET_UNAVAILABLE
+
+        original_num_blocks = seq.num_blocks
+        try:
+            # Allocate only the checked remote interval. The unallocated local
+            # tail remains represented by the soft target above.
+            seq.kv_token_limit = remote_step
+            self.block_manager.allocate(seq)
+            block_table = self.block_manager.get_block_table(seq)
+            fallback_block = fallback_step // block_size
+            remote_block = remote_step // block_size
+            load_block_ids = tuple(
+                int(block_id)
+                for block_id in block_table[fallback_block:remote_block]
+            )
+            connector.update_state_after_alloc(
+                seq,
+                load_block_ids,
+                remote_step - fallback_step,
+            )
+            # From start_load onward, cleanup must retain destinations until
+            # workers report terminal progress or their queues are drained.
+            self.start_load(
+                seq,
+                fallback_step=fallback_step,
+                remote_step=remote_step,
+                target_blocks=target_blocks,
+            )
+        except Exception:
+            if seq.num_blocks > original_num_blocks:
+                self.block_manager.truncate(seq, original_num_blocks)
+            seq.kv_token_limit = old_kv_token_limit
+            raise
+        seq.kv_token_limit = None
+        return KVLoadAdmission.STARTED
+
     def start_load(
         self,
         seq: SchedulerSequence,
@@ -166,10 +308,9 @@ def start_load(
     ) -> None:
         """Take paging ownership after destinations have been allocated.
 
-        ``Scheduler._start_external_load`` must first allocate the exact
-        destination range and bind it to connector metadata. Only then does
-        this method move the sequence out of the normal waiting queue and make
-        the asynchronous write visible to paging cleanup paths.
+        :meth:`try_load` first allocates the exact destination range and binds
+        connector metadata. Only then does this method move the sequence out of
+        normal waiting and expose the asynchronous write to cleanup paths.
         """
         request_id = int(seq.seq_id)
         if request_id in self._loads:
@@ -217,20 +358,19 @@ def _publish(self, record: _LoadRecord) -> None:
         The blocks were private destinations while loading. Publishing inserts their full prefix into the local trie,
         advances sequence history, and exposes cached-token metrics only after all ranks have valid contents.
         """
-        scheduler = self.scheduler
         seq = record.seq
         # Limit trie publication to the successfully loaded prefix. The request
         # may already own preallocated blocks after remote_step.
         seq.kv_token_limit = record.remote_step
-        if scheduler.block_trie.enabled:
-            scheduler.block_trie.allocate(seq)
+        if self.block_trie.enabled:
+            self.block_trie.allocate(seq)
         seq.set_step(record.remote_step)
         seq.kv_token_limit = None
         if seq.prefix_cache.match_start_step < 0:
             # With no preceding local trie hit, the block-aligned load start is
             # the beginning of this request's externally cached interval.
             seq.prefix_cache.match_start_step = record.fallback_step
-        scheduler._finish_prefix_cache_schedule(seq)
+        self.block_trie.finalize_match(seq)
         seq.state.finish_remote_load()
         record.phase = _LoadPhase.READY
 
@@ -242,11 +382,10 @@ def _rollback(self, record: _LoadRecord) -> None:
         block-aligned ``fallback_step`` and move the trie cursor to an ancestor
         that refers only to retained blocks.
         """
-        scheduler = self.scheduler
         seq = record.seq
         fallback_blocks = record.fallback_step // seq.block_size
         if seq.num_blocks > fallback_blocks:
-            scheduler.block_manager.truncate(seq, fallback_blocks)
+            self.block_manager.truncate(seq, fallback_blocks)
         seq.set_step(record.fallback_step)
         seq.kv_token_limit = None
 
@@ -256,7 +395,7 @@ def _rollback(self, record: _LoadRecord) -> None:
         seq.prefix_cache.trie_cursor = cursor
         if seq.prefix_cache.match_start_step > record.fallback_step:
             seq.prefix_cache.match_start_step = -1
-        scheduler._finish_prefix_cache_schedule(seq)
+        self.block_trie.finalize_match(seq)
 
     def _finish_cancelled_or_failed(self, record: _LoadRecord) -> None:
         """Release accounting and honor cleanup deferred during ``LOADING``."""
@@ -343,15 +482,20 @@ def finish_deferred_loads_after_worker_drain(self) -> None:
             self._finish_cancelled_or_failed(record)
 
     def _remove_sequence(self, seq: SchedulerSequence) -> None:
-        scheduler = self.scheduler
-        connector = scheduler.kv_connector
+        connector = self.connector
         if connector is not None:
             connector.request_finished(seq)
 
         session = seq.session
         session.remove_sequence(seq)
         if not session.sequences:
-            scheduler.sessions.pop(session.session_id, None)
+            self.sessions.pop(session.session_id, None)
+
+    def disable(self) -> None:
+        """Stop new lookup admission and discard scheduler-side ownership."""
+        self.lookup_enabled = False
+        self.clear()
+        self.connector = None
 
     def clear(self) -> None:
         self._loads.clear()
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 2582fd88e8..fcbe84d5b0 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -75,7 +75,7 @@
 from .block_manager import build_block_manager
 from .block_trie import BlockTrie
 from .eviction_helper import build_eviction_helper
-from .kv_load_coordinator import KVLoadCoordinator
+from .kv_load_coordinator import KVLoadAdmission, KVLoadCoordinator
 from .kv_save_coordinator import KVSaveCoordinator
 from .state_manager import build_state_manager
 
@@ -500,7 +500,8 @@ def __init__(self,
         self.prealloc_size = prealloc_size
         self.token_count = token_count
         self.has_admitted = has_admitted
-        self._remote_ready = scheduler.kv_load_coordinator.is_remote_ready(seq)
+        self.load_coordinator = scheduler.kv_load_coordinator
+        self._remote_ready = self.load_coordinator.is_remote_ready(seq)
         self.allow_long_prefill = allow_long_prefill
         self._alloc_size = prealloc_size
         self._prefix_match = _TentativePrefixMatch(
@@ -509,7 +510,7 @@ def __init__(self,
             scheduler.block_manager,
             is_ssm=scheduler.is_ssm,
             preserve_existing_state=(
-                scheduler._external_lookup_enabled and not self._remote_ready),
+                self.load_coordinator.lookup_enabled and not self._remote_ready),
         )
 
     def run(self):
@@ -523,7 +524,7 @@ def run(self):
         5. Admit KV/state resources or roll the tentative match back precisely.
         6. On success, allocate blocks/states and publish the accepted hit.
         """
-        if self.scheduler._external_lookup_enabled and not self._remote_ready:
+        if self.load_coordinator.lookup_enabled and not self._remote_ready:
             if self._lookup_is_pending():
                 return _PrefillAdmissionResult.skip()
             self._prefix_match.begin()
@@ -545,9 +546,7 @@ def _lookup_is_pending(self) -> bool:
         poll delay instead of diagnosing an empty batch as GPU-cache pressure.
         """
         scheduler = self.scheduler
-        connector = scheduler.kv_connector
-        assert connector is not None
-        if not connector.is_lookup_pending(self.seq.seq_id):
+        if not self.load_coordinator.is_lookup_pending(self.seq):
             return False
         scheduler.last_schedule_had_pending_lookup = True
         return True
@@ -555,7 +554,7 @@ def _lookup_is_pending(self) -> bool:
     def _admit_resources(self):
         if self.scheduler.block_trie.enabled:
             return self._admit_prefix_cache_resources()
-        if self.scheduler._external_lookup_enabled:
+        if self.load_coordinator.lookup_enabled:
             lookup_result = self._query_external_prefix()
             if lookup_result is not None:
                 return lookup_result
@@ -586,7 +585,7 @@ def _admit_prefix_cache_resources(self):
             if not self._remote_ready and not self._has_private_local_tail():
                 self._prefix_match.match()
 
-        if scheduler._external_lookup_enabled:
+        if self.load_coordinator.lookup_enabled:
             lookup_result = self._query_external_prefix()
             if lookup_result is not None:
                 return lookup_result
@@ -629,138 +628,39 @@ def _admit_prefix_cache_resources(self):
         return None
 
     def _query_external_prefix(self):
-        """Poll the external prefix after prioritizing the local trie.
-
-        Return meanings from the connector are intentionally different:
-
-        * ``None``: lookup RPC is still pending. Restore the pre-match baseline
-          and skip this candidate so later waiters can still run.
-        * ``0``: lookup completed with no extension; continue normal local
-          allocation using any accepted trie hit.
-        * ``> 0``: allocate the block-aligned external interval and move the
-          request to the load coordinator instead of this model batch.
-
-        A coordinator ``READY`` request bypasses lookup because its accepted
-        remote boundary has already been published.
-        """
+        """Map connector/paging admission to prefill queue policy."""
         scheduler = self.scheduler
         if self._remote_ready:
             return None
-        connector = scheduler.kv_connector
-        assert connector is not None
-        # Local matching runs first, so this step asks the connector only for
-        # the remote extension beyond KV that is already resident on this node.
-        num_external_tokens, _ = connector.get_num_new_matched_tokens(
+        admission = self.load_coordinator.try_load(
             self.seq,
-            self.seq.num_history_ids,
+            prealloc_size=self.prealloc_size,
+            evictable_seqs=self._evictable_sequences(),
         )
-        if num_external_tokens is not None:
-            if num_external_tokens > 0:
-                return self._start_external_load(int(num_external_tokens))
+        if admission is KVLoadAdmission.NO_LOAD:
             return None
-
-        self._prefix_match.rollback('external lookup pending')
-        scheduler.last_schedule_had_pending_lookup = True
-        return _PrefillAdmissionResult.skip()
-
-    def _start_external_load(self, num_external_tokens: int):
-        """Admit against soft prefill budgets, then allocate the remote hit.
-
-        Lookup can start inside a partially computed block, but connector
-        transfer and local trie publication are block-granular. The load range
-        is therefore expanded down to ``fallback_step`` and truncated to the
-        deepest safe full-block ``remote_step``. A failure later recomputes from
-        fallback because an asynchronous writer may have overwritten that
-        boundary block partially.
-
-        Only the remote interval is physically allocated now. Admission still
-        checks capacity for the complete prefill and every existing soft
-        reservation before handing any destination to workers; otherwise
-        several prefix loads could occupy all blocks and leave no capacity for
-        their remaining local tails.
-        """
+        if admission is KVLoadAdmission.PENDING:
+            self._prefix_match.rollback('external lookup pending')
+            scheduler.last_schedule_had_pending_lookup = True
+            return _PrefillAdmissionResult.skip()
+        if admission is KVLoadAdmission.STARTED:
+            self._prefix_match.commit()
+            return _PrefillAdmissionResult.load()
+        if admission is KVLoadAdmission.FULL_PREFILL_UNAVAILABLE:
+            reason = 'full prefill capacity unavailable'
+        else:
+            assert admission is KVLoadAdmission.SOFT_BUDGET_UNAVAILABLE
+            reason = 'soft prefill budget unavailable'
+        # No worker has seen a destination on rejected admission, so the
+        # request-local prefix transaction remains exactly reversible.
+        self._prefix_match.rollback(reason)
+        return _PrefillAdmissionResult.stop()
+
+    def _evictable_sequences(self):
+        """Iterate queue-owned eviction candidates in historical order."""
         scheduler = self.scheduler
-        connector = scheduler.kv_connector
-        assert connector is not None
-        seq = self.seq
-        block_size = seq.block_size
-        local_step = int(seq.num_history_ids)
-        # Transfers are block-granular. Reuse the sequence's private partial
-        # block at the boundary, but only publish complete remotely loaded
-        # blocks and never match through the prompt's final token.
-        fallback_step = local_step // block_size * block_size
-        remote_step = local_step + num_external_tokens
-        remote_step = min(remote_step, int(seq.get_prefix_cache_max_match_step()))
-        remote_step = remote_step // block_size * block_size
-        if remote_step <= fallback_step:
-            return None
-
-        target_blocks = scheduler.kv_load_coordinator.prefill_target_blocks(
-            seq,
-            self.prealloc_size,
-        )
-        old_kv_token_limit = seq.kv_token_limit
-        # The load allocates only the remote hit now, but it is admitted only
-        # when the whole prefill can eventually finish. Otherwise concurrent
-        # loads could each pin a prefix and deadlock on their remaining tails.
-        seq.kv_token_limit = None
-        full_prefill_fits = self._evict_for_seq(self.prealloc_size)
-        load_admitted = (
-            full_prefill_fits
-            and scheduler.kv_load_coordinator.can_admit_load(
-                seq,
-                target_blocks,
-            )
-        )
-        if not load_admitted:
-            # No worker has seen the destination yet, so local match state can
-            # still be restored exactly and the request can retry later.
-            seq.kv_token_limit = old_kv_token_limit
-            reason = (
-                'full prefill capacity unavailable'
-                if not full_prefill_fits
-                else 'soft prefill budget unavailable'
-            )
-            self._prefix_match.rollback(reason)
-            return _PrefillAdmissionResult.stop()
-
-        original_num_blocks = seq.num_blocks
-        try:
-            # kv_token_limit prevents allocate() from reserving the unchecked
-            # local tail; can_admit_load() accounts for that tail softly.
-            seq.kv_token_limit = remote_step
-            scheduler.block_manager.allocate(seq)
-            block_table = scheduler.block_manager.get_block_table(seq)
-            fallback_block = fallback_step // block_size
-            remote_block = remote_step // block_size
-            load_block_ids = tuple(
-                int(block_id)
-                for block_id in block_table[fallback_block:remote_block]
-            )
-            connector.update_state_after_alloc(
-                seq,
-                load_block_ids,
-                remote_step - fallback_step,
-            )
-            # Register paging ownership only after connector state references
-            # concrete destinations. From start_load onward, stop/end may not
-            # free these blocks until workers report completion or are drained.
-            scheduler.kv_load_coordinator.start_load(
-                seq,
-                fallback_step=fallback_step,
-                remote_step=remote_step,
-                target_blocks=target_blocks,
-            )
-        except Exception:
-            # Allocation/binding failed synchronously, before device writes are
-            # in flight. Remove only blocks added by this attempt.
-            if seq.num_blocks > original_num_blocks:
-                scheduler.block_manager.truncate(seq, original_num_blocks)
-            seq.kv_token_limit = old_kv_token_limit
-            raise
-        seq.kv_token_limit = None
-        self._prefix_match.commit()
-        return _PrefillAdmissionResult.load()
+        yield from reversed(scheduler.hanging)
+        yield from reversed(self.evictable_waiting)
 
     def _match_prefix_for_prefill_gate(self):
         """Tentatively match once so a request can be rechecked by a gate."""
@@ -806,8 +706,7 @@ def _has_private_local_tail(self) -> bool:
         model forwards, preemption/resume, or a continued chat session; it is
         not specific to multi-turn conversation.
         """
-        scheduler = self.scheduler
-        if not scheduler._external_lookup_enabled:
+        if not self.load_coordinator.lookup_enabled:
             return False
         seq = self.seq
         return seq.num_blocks > int(seq.num_history_ids) // seq.block_size
@@ -866,12 +765,12 @@ def _prepare_and_evict(self):
 
     def _evict_for_seq(self, alloc_size: int):
         """Evict stopped or skipped waiters until this sequence can run."""
-        from itertools import chain
         scheduler = self.scheduler
-        hanging = reversed(scheduler.hanging)
-        waiting = reversed(self.evictable_waiting)
-        evictable = list(chain(hanging, waiting))
-        return scheduler.eviction_helper.evict_for_seq(self.seq, evictable, alloc_size)
+        return scheduler.eviction_helper.evict_for_seq(
+            self.seq,
+            list(self._evictable_sequences()),
+            alloc_size,
+        )
 
     def _finish_admission(self):
         scheduler = self.scheduler
@@ -887,8 +786,8 @@ def _finish_admission(self):
         if scheduler.is_ssm:
             scheduler.state_manager.allocate(seq)
         if scheduler.block_trie.enabled:
-            scheduler._finish_prefix_cache_schedule(seq)
-        scheduler.kv_load_coordinator.track_prefill(
+            scheduler.block_trie.finalize_match(seq)
+        self.load_coordinator.track_prefill(
             seq,
             prealloc_size=self.prealloc_size,
         )
@@ -896,7 +795,7 @@ def _finish_admission(self):
             # Preserve the load record through the remaining prefill so its
             # reservation can be released only after model output advances the
             # sequence to input_end_pos.
-            scheduler.kv_load_coordinator.mark_scheduled(seq)
+            self.load_coordinator.mark_scheduled(seq)
         self._prefix_match.commit()
         return _PrefillAdmissionResult.admit(prefill_token_count)
 
@@ -937,13 +836,6 @@ def __init__(
             and transfer_config.is_kv_consumer
             and not self.is_ssm
         )
-        # Keep call sites uniform even when a role is disabled. Each coordinator
-        # is a no-op until scheduler/connector metadata starts its lifecycle.
-        self.kv_load_coordinator = KVLoadCoordinator(self)
-        self.kv_save_coordinator = KVSaveCoordinator(self)
-        # Per-tick signal consumed by EngineLoop to distinguish asynchronous
-        # lookup latency from actual cache-allocation pressure.
-        self.last_schedule_had_pending_lookup = False
         checkpoint_state_manager = self.state_manager if self.is_ssm else None
         self.block_trie = BlockTrie(allocator=self.block_manager.allocator,
                                    block_size=self.cache_config.block_size,
@@ -951,6 +843,21 @@ def __init__(
                                    checkpoint_state_manager=checkpoint_state_manager)
 
         self.eviction_helper = build_eviction_helper(self, self.scheduler_config.eviction_type)
+        # Load admission receives only paging owners plus request-local queue
+        # candidates from its caller; it does not reach back through Scheduler.
+        self.kv_load_coordinator = KVLoadCoordinator(
+            lookup_enabled=self._external_lookup_enabled,
+            connector=kv_connector,
+            block_manager=self.block_manager,
+            block_trie=self.block_trie,
+            eviction_helper=self.eviction_helper,
+            sessions=self.sessions,
+        )
+        # Keep save call sites uniform even when the producer role is disabled.
+        self.kv_save_coordinator = KVSaveCoordinator(self)
+        # Per-tick signal consumed by EngineLoop to distinguish asynchronous
+        # lookup latency from actual cache-allocation pressure.
+        self.last_schedule_had_pending_lookup = False
 
         seq_meta = seq_meta or SequenceMeta(self.cache_config.block_size)
         self.seq_meta = seq_meta
@@ -972,7 +879,7 @@ def shutdown(self) -> None:
         connector = self.kv_connector
         self.kv_connector = None
         self._external_lookup_enabled = False
-        self.kv_load_coordinator.clear()
+        self.kv_load_coordinator.disable()
         self.kv_save_coordinator.clear()
         if connector is not None:
             connector.shutdown()
@@ -990,29 +897,6 @@ def _ensure_runtime_state_available(self):
         self.block_trie.state_checkpoints.evict(1)
         return self.state_manager.get_num_free_runtime() > 0
 
-    @staticmethod
-    def _finalize_prefix_cache_match(seq: SchedulerSequence):
-        """Publish accepted cached-token count within the current prompt."""
-        match_start = seq.prefix_cache.match_start_step
-        if match_start < 0:
-            seq.cached_tokens = 0
-            return
-        cached_start = match_start
-        cached_end = seq.num_history_ids
-        prompt_start = seq.input_start_pos
-        prompt_end = seq.input_end_pos
-        seq.cached_tokens = max(0, min(cached_end, prompt_end) - max(cached_start, prompt_start))
-
-    @staticmethod
-    def _finish_prefix_cache_schedule(seq: SchedulerSequence):
-        """Publish match side effects after the sequence is accepted to run."""
-        prefix_cache = seq.prefix_cache
-        if prefix_cache.suppress_match_stats:
-            seq.cached_tokens = 0
-            prefix_cache.suppress_match_stats = False
-            return
-        Scheduler._finalize_prefix_cache_match(seq)
-
     def _long_context_chunk_limit(self, seq: SchedulerSequence):
         """Return the token budget for one long-context chunk."""
         return get_long_context_chunk_limit(seq, self.cache_config.max_prefill_token_num)
@@ -1175,7 +1059,7 @@ def _reorder_migrating():
 
             # allocate session memory
             self.block_manager.allocate(seq)
-            self._finish_prefix_cache_schedule(seq)
+            self.block_trie.finalize_match(seq)
             _to_running(seq)
 
         return migration_ready
diff --git a/tests/pytorch/engine/test_kv_connector_wiring.py b/tests/pytorch/engine/test_kv_connector_wiring.py
index aa7e1eb7e7..26af606bae 100644
--- a/tests/pytorch/engine/test_kv_connector_wiring.py
+++ b/tests/pytorch/engine/test_kv_connector_wiring.py
@@ -227,7 +227,7 @@ def test_scheduler_shutdown_releases_injected_connector_once():
     scheduler.shutdown()
 
     connector.shutdown.assert_called_once_with()
-    assert scheduler.kv_load_coordinator.clear.call_count == 2
+    assert scheduler.kv_load_coordinator.disable.call_count == 2
     assert scheduler.kv_save_coordinator.clear.call_count == 2
     assert scheduler.kv_connector is None
     assert not scheduler._external_lookup_enabled

From f17a6faeddc10961f104a5bb08d09708feb956a6 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 15:22:53 +0800
Subject: [PATCH 05/22] refactor: isolate prefill scheduling ownership

---
 lmdeploy/pytorch/paging/scheduler.py   | 484 +++++++++++++++----------
 tests/pytorch/paging/test_scheduler.py |  15 +-
 2 files changed, 305 insertions(+), 194 deletions(-)

diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index fcbe84d5b0..d7e33d7b01 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -77,11 +77,14 @@
 from .eviction_helper import build_eviction_helper
 from .kv_load_coordinator import KVLoadAdmission, KVLoadCoordinator
 from .kv_save_coordinator import KVSaveCoordinator
-from .state_manager import build_state_manager
+from .state_manager import StateManager, build_state_manager
 
 if TYPE_CHECKING:
     from lmdeploy.pytorch.kv_connector.base import KVConnectorBase
 
+    from .block_manager.base_block_manager import BaseBlockManager
+    from .eviction_helper.base_eviction_helper import BaseEvictionHelper
+
 logger = get_logger('lmdeploy')
 
 MapType = dict[int, int]
@@ -120,8 +123,8 @@ class _PrefillReorderInfo:
 class _PrefillReorderer:
     """Order waiting prefills without applying scheduler side effects."""
 
-    def __init__(self, scheduler: 'Scheduler'):
-        self.scheduler = scheduler
+    def __init__(self, prefill_scheduler: '_PrefillScheduler'):
+        self.prefill_scheduler = prefill_scheduler
         self._info_cache: dict[int, _PrefillReorderInfo] = {}
 
     def reorder(self,
@@ -134,11 +137,11 @@ def reorder(self,
         # reservation. Admit it first to shorten that ownership window.
         remote_ready = [
             seq for seq in waiting
-            if self.scheduler.kv_load_coordinator.is_remote_ready(seq)
+            if self.prefill_scheduler.load_coordinator.is_remote_ready(seq)
         ]
         waiting = [
             seq for seq in waiting
-            if not self.scheduler.kv_load_coordinator.is_remote_ready(seq)
+            if not self.prefill_scheduler.load_coordinator.is_remote_ready(seq)
         ]
         if prefer_long_prefill:
             # Long-work turns choose one long waiter first. The size policy only
@@ -182,14 +185,14 @@ def _get_reorder_info(self, seq: SchedulerSequence):
         if info is not None:
             return info
 
-        scheduler = self.scheduler
-        chunk_limit = scheduler._long_context_chunk_limit(seq)
+        prefill = self.prefill_scheduler
+        chunk_limit = prefill._long_context_chunk_limit(seq)
         if seq.num_token_ids <= chunk_limit:
             info = _PrefillReorderInfo(prefill_token_count=seq.num_token_ids,
                                        is_nonfinal_long_prefill=False,
                                        estimated_long_chunks=1)
         else:
-            kv_token_limit = scheduler._next_long_context_chunk_end(seq, chunk_limit)
+            kv_token_limit = prefill._next_long_context_chunk_end(seq, chunk_limit)
             safe_chunk_limit = max(1, chunk_limit)
             info = _PrefillReorderInfo(
                 prefill_token_count=max(0, kv_token_limit - seq.num_history_ids),
@@ -201,10 +204,10 @@ def _get_reorder_info(self, seq: SchedulerSequence):
 
     def _long_priority_key(self, seq: SchedulerSequence, now: float):
         """Prefer smaller long prompts, with age credit to avoid starvation."""
-        scheduler = self.scheduler
+        prefill = self.prefill_scheduler
         info = self._get_reorder_info(seq)
         wait_age = max(0.0, now - seq.arrive_time)
-        age_credit = int(wait_age // scheduler._long_prefill_aging_seconds_per_chunk)
+        age_credit = int(wait_age // prefill._long_prefill_aging_seconds_per_chunk)
         age_adjusted_chunks = info.estimated_long_chunks - age_credit
         return age_adjusted_chunks, info.estimated_long_chunks, seq.arrive_time
 
@@ -225,8 +228,8 @@ def _sort_normal_prefills(self, waiting: SeqList):
                       key=lambda seq: (self._get_reorder_info(seq).prefill_token_count, seq.arrive_time))
 
     def _sort_long_prefills(self, waiting: SeqList):
-        scheduler = self.scheduler
-        if scheduler._long_prefill_policy != 'size':
+        prefill = self.prefill_scheduler
+        if prefill._long_prefill_policy != 'size':
             return waiting
         now = time.perf_counter()
         return sorted(waiting, key=lambda seq: self._long_priority_key(seq, now))
@@ -482,33 +485,35 @@ class _PrefillAdmissionAttempt:
 
     The attempt owns all tentative prefix-cache side effects for the sequence:
     match, SSM restore pinning, eviction, runtime-state checks, allocation, and
-    rollback. The outer scheduler loop still owns queue traversal and decides
+    rollback. The outer prefill loop still owns queue traversal and decides
     whether a rejected candidate is skipped or ends the current prefill turn.
     """
 
     def __init__(self,
-                 scheduler: 'Scheduler',
+                 prefill_scheduler: '_PrefillScheduler',
                  seq: SchedulerSequence,
+                 hanging: SeqList,
                  evictable_waiting: SeqList,
                  prealloc_size: int,
                  token_count: int,
                  has_admitted: bool,
                  allow_long_prefill: bool):
-        self.scheduler = scheduler
+        self.prefill_scheduler = prefill_scheduler
         self.seq = seq
+        self.hanging = hanging
         self.evictable_waiting = evictable_waiting
         self.prealloc_size = prealloc_size
         self.token_count = token_count
         self.has_admitted = has_admitted
-        self.load_coordinator = scheduler.kv_load_coordinator
+        self.load_coordinator = prefill_scheduler.load_coordinator
         self._remote_ready = self.load_coordinator.is_remote_ready(seq)
         self.allow_long_prefill = allow_long_prefill
         self._alloc_size = prealloc_size
         self._prefix_match = _TentativePrefixMatch(
             seq,
-            scheduler.block_trie,
-            scheduler.block_manager,
-            is_ssm=scheduler.is_ssm,
+            prefill_scheduler.block_trie,
+            prefill_scheduler.block_manager,
+            is_ssm=prefill_scheduler.is_ssm,
             preserve_existing_state=(
                 self.load_coordinator.lookup_enabled and not self._remote_ready),
         )
@@ -542,17 +547,17 @@ def run(self):
     def _lookup_is_pending(self) -> bool:
         """Skip without touching local prefix state while lookup is running.
 
-        The connector owns the Future and deduplicates polls. Marking the scheduler tick lets EngineLoop use a short I/O
-        poll delay instead of diagnosing an empty batch as GPU-cache pressure.
+        The connector owns the Future and deduplicates polls. Marking this turn lets EngineLoop use a short I/O poll
+        delay instead of diagnosing an empty batch as GPU-cache pressure.
         """
-        scheduler = self.scheduler
+        prefill = self.prefill_scheduler
         if not self.load_coordinator.is_lookup_pending(self.seq):
             return False
-        scheduler.last_schedule_had_pending_lookup = True
+        prefill.last_schedule_had_pending_lookup = True
         return True
 
     def _admit_resources(self):
-        if self.scheduler.block_trie.enabled:
+        if self.prefill_scheduler.block_trie.enabled:
             return self._admit_prefix_cache_resources()
         if self.load_coordinator.lookup_enabled:
             lookup_result = self._query_external_prefix()
@@ -576,7 +581,7 @@ def _admit_prefix_cache_resources(self):
         a prefill gate returns that gate's skip/stop result after rollback;
         normal resource failures keep their local retry/stop behavior here.
         """
-        scheduler = self.scheduler
+        prefill = self.prefill_scheduler
         seq = self.seq
         if not self._prefix_match.matched:
             # A completed external load has already published the accepted
@@ -590,7 +595,7 @@ def _admit_prefix_cache_resources(self):
             if lookup_result is not None:
                 return lookup_result
 
-        had_ssm_restore = scheduler.is_ssm and seq.prefix_cache.restore.is_selected
+        had_ssm_restore = prefill.is_ssm and seq.prefix_cache.restore.is_selected
         if not self._prefix_match.pin_restore():
             result = self._prefix_match.rollback(
                 'failed to pin SSM restore checkpoint')
@@ -614,14 +619,14 @@ def _admit_prefix_cache_resources(self):
             if not self._prepare_and_evict():
                 return _PrefillAdmissionResult.stop()
 
-        if scheduler.is_ssm and not scheduler._ensure_runtime_state_available():
+        if prefill.is_ssm and not prefill._ensure_runtime_state_available():
             result = self._prefix_match.rollback(
                 'no runtime SSM state available')
             if result is not None:
                 return result
             if not self._prepare_and_evict():
                 return _PrefillAdmissionResult.stop()
-            if not scheduler._ensure_runtime_state_available():
+            if not prefill._ensure_runtime_state_available():
                 seq.kv_token_limit = None
                 return _PrefillAdmissionResult.stop()
 
@@ -629,7 +634,7 @@ def _admit_prefix_cache_resources(self):
 
     def _query_external_prefix(self):
         """Map connector/paging admission to prefill queue policy."""
-        scheduler = self.scheduler
+        prefill = self.prefill_scheduler
         if self._remote_ready:
             return None
         admission = self.load_coordinator.try_load(
@@ -641,7 +646,7 @@ def _query_external_prefix(self):
             return None
         if admission is KVLoadAdmission.PENDING:
             self._prefix_match.rollback('external lookup pending')
-            scheduler.last_schedule_had_pending_lookup = True
+            prefill.last_schedule_had_pending_lookup = True
             return _PrefillAdmissionResult.skip()
         if admission is KVLoadAdmission.STARTED:
             self._prefix_match.commit()
@@ -658,14 +663,13 @@ def _query_external_prefix(self):
 
     def _evictable_sequences(self):
         """Iterate queue-owned eviction candidates in historical order."""
-        scheduler = self.scheduler
-        yield from reversed(scheduler.hanging)
+        yield from reversed(self.hanging)
         yield from reversed(self.evictable_waiting)
 
     def _match_prefix_for_prefill_gate(self):
         """Tentatively match once so a request can be rechecked by a gate."""
-        scheduler = self.scheduler
-        if (self._remote_ready or not scheduler.block_trie.enabled
+        prefill = self.prefill_scheduler
+        if (self._remote_ready or not prefill.block_trie.enabled
                 or self._has_private_local_tail()):
             return None
         self._prefix_match.match()
@@ -718,22 +722,22 @@ def _token_budget_rejection(self):
 
     def _check_prefill_admission_gates(self):
         """Apply prefill gates, tentatively matching only when it may help."""
-        scheduler = self.scheduler
+        prefill = self.prefill_scheduler
         seq = self.seq
-        token_budget = scheduler.cache_config.max_prefill_token_num
-        prefill_token_count = scheduler._prefill_admission_token_count(seq)
-        is_nonfinal_long_prefill = scheduler._prefill_kv_token_limit(seq) is not None
+        token_budget = prefill.cache_config.max_prefill_token_num
+        prefill_token_count = prefill._prefill_admission_token_count(seq)
+        is_nonfinal_long_prefill = prefill._prefill_kv_token_limit(seq) is not None
 
         if is_nonfinal_long_prefill and not self.allow_long_prefill:
             matched = self._match_prefix_for_prefill_gate()
             if matched is None:
                 return _PrefillAdmissionResult.skip()
-            if scheduler._prefill_kv_token_limit(seq) is not None:
+            if prefill._prefill_kv_token_limit(seq) is not None:
                 self._prefix_match.rollback('still non-final long prefill on short turn')
                 return _PrefillAdmissionResult.skip()
             self._prefix_match.retain_for_admission(
                 _PrefillAdmissionResult.skip())
-            prefill_token_count = scheduler._prefill_admission_token_count(seq)
+            prefill_token_count = prefill._prefill_admission_token_count(seq)
 
         exceeds_token_budget = self.has_admitted and self.token_count + prefill_token_count > token_budget
         if not exceeds_token_budget:
@@ -742,7 +746,7 @@ def _check_prefill_admission_gates(self):
         if not self._prefix_match.matched:
             matched = self._match_prefix_for_prefill_gate()
             if matched is not None:
-                prefill_token_count = scheduler._prefill_admission_token_count(seq)
+                prefill_token_count = prefill._prefill_admission_token_count(seq)
                 if self.token_count + prefill_token_count <= token_budget:
                     self._prefix_match.retain_for_admission(
                         self._token_budget_rejection())
@@ -754,9 +758,9 @@ def _check_prefill_admission_gates(self):
 
     def _prepare_and_evict(self):
         """Apply chunk allocation limits and evict for this prefill."""
-        scheduler = self.scheduler
+        prefill = self.prefill_scheduler
         seq = self.seq
-        alloc_size = scheduler._prepare_prefill_allocation(seq, self.prealloc_size)
+        alloc_size = prefill._prepare_prefill_allocation(seq, self.prealloc_size)
         self._alloc_size = alloc_size
         if self._evict_for_seq(alloc_size):
             return True
@@ -765,28 +769,28 @@ def _prepare_and_evict(self):
 
     def _evict_for_seq(self, alloc_size: int):
         """Evict stopped or skipped waiters until this sequence can run."""
-        scheduler = self.scheduler
-        return scheduler.eviction_helper.evict_for_seq(
+        prefill = self.prefill_scheduler
+        return prefill.eviction_helper.evict_for_seq(
             self.seq,
             list(self._evictable_sequences()),
             alloc_size,
         )
 
     def _finish_admission(self):
-        scheduler = self.scheduler
+        prefill = self.prefill_scheduler
         seq = self.seq
         # Prefix-cache matching can advance the sequence step and shrink the
         # remaining prefill tail. Charge the admitted batch with the
         # post-match/post-rollback cost, not the conservative pre-match
         # estimate used to decide whether this sequence is worth trying.
-        prefill_token_count = scheduler._prefill_admission_token_count(seq)
-        scheduler.block_manager.allocate(seq, self._alloc_size)
-        if scheduler.block_trie.enabled:
-            scheduler.block_trie.allocate(seq)
-        if scheduler.is_ssm:
-            scheduler.state_manager.allocate(seq)
-        if scheduler.block_trie.enabled:
-            scheduler.block_trie.finalize_match(seq)
+        prefill_token_count = prefill._prefill_admission_token_count(seq)
+        prefill.block_manager.allocate(seq, self._alloc_size)
+        if prefill.block_trie.enabled:
+            prefill.block_trie.allocate(seq)
+        if prefill.is_ssm:
+            prefill.state_manager.allocate(seq)
+        if prefill.block_trie.enabled:
+            prefill.block_trie.finalize_match(seq)
         self.load_coordinator.track_prefill(
             seq,
             prealloc_size=self.prealloc_size,
@@ -800,6 +804,209 @@ def _finish_admission(self):
         return _PrefillAdmissionResult.admit(prefill_token_count)
 
 
+class _PrefillScheduler:
+    """Own prefill ordering, admission, and long-context reservation.
+
+    Long-lived dependencies are the resource owners used by prefill. Queue
+    contents and active-batch counts remain request-local inputs supplied by
+    the public :class:`Scheduler` facade for each scheduling turn.
+    """
+
+    def __init__(
+        self,
+        scheduler_config: SchedulerConfig,
+        cache_config: CacheConfig,
+        *,
+        is_ssm: bool,
+        block_manager: 'BaseBlockManager',
+        block_trie: BlockTrie,
+        state_manager: StateManager,
+        eviction_helper: 'BaseEvictionHelper',
+        load_coordinator: KVLoadCoordinator,
+    ) -> None:
+        self.scheduler_config = scheduler_config
+        self.cache_config = cache_config
+        self.is_ssm = is_ssm
+        self.block_manager = block_manager
+        self.block_trie = block_trie
+        self.state_manager = state_manager
+        self.eviction_helper = eviction_helper
+        self.load_coordinator = load_coordinator
+        self.last_schedule_had_pending_lookup = False
+        self._long_prefill_policy = _envs.opt_ttft_policy
+        self._long_prefill_aging_seconds_per_chunk = max(
+            0.001,
+            _envs.opt_ttft_aging_sec,
+        )
+
+    def _ensure_runtime_state_available(self):
+        """Make one state-cache slot available for an SSM runtime state."""
+        if not self.is_ssm:
+            return True
+        if self.state_manager.get_num_free_runtime() > 0:
+            return True
+        self.block_trie.state_checkpoints.evict(1)
+        return self.state_manager.get_num_free_runtime() > 0
+
+    def _long_context_chunk_limit(self, seq: SchedulerSequence):
+        """Return the token budget for one long-context chunk."""
+        return get_long_context_chunk_limit(
+            seq,
+            self.cache_config.max_prefill_token_num,
+        )
+
+    def _next_long_context_chunk_end(
+        self,
+        seq: SchedulerSequence,
+        max_prefill_num: int | None = None,
+    ):
+        """Return the exclusive absolute token end for the next chunk."""
+        if max_prefill_num is None:
+            max_prefill_num = self._long_context_chunk_limit(seq)
+        plan = plan_long_context_chunk(
+            seq,
+            max_prefill_num,
+            include_multimodals=False,
+        )
+        return plan.chunk_end
+
+    def _prefill_kv_token_limit(self, seq: SchedulerSequence):
+        """Limit KV allocation for a non-final long-context prefill chunk."""
+        max_prefill_num = self._long_context_chunk_limit(seq)
+        if seq.num_token_ids <= max_prefill_num:
+            return None
+        return self._next_long_context_chunk_end(seq, max_prefill_num)
+
+    def _prefill_admission_token_count(self, seq: SchedulerSequence):
+        """Return token budget cost for the next prefill or chunk."""
+        kv_token_limit = self._prefill_kv_token_limit(seq)
+        if kv_token_limit is None:
+            return seq.num_token_ids
+        return max(0, kv_token_limit - seq.num_history_ids)
+
+    def _prepare_prefill_allocation(
+        self,
+        seq: SchedulerSequence,
+        prealloc_size: int,
+    ):
+        """Apply chunk KV limit and return the effective prealloc size."""
+        kv_token_limit = self._prefill_kv_token_limit(seq)
+        if kv_token_limit is None:
+            seq.kv_token_limit = None
+            return prealloc_size
+
+        seq.kv_token_limit = kv_token_limit
+        return 0
+
+    def has_waiting_long_prefill(self, waiting: SeqList):
+        """Whether a waiting request needs a non-final prefill chunk."""
+        return any(
+            self._prefill_kv_token_limit(seq) is not None
+            for seq in waiting
+        )
+
+    def reserve_long_context_chunk(
+        self,
+        seq: SchedulerSequence,
+        *,
+        hanging: SeqList,
+        waiting: SeqList,
+        chunk_size: int,
+        prealloc_size: int = 0,
+        is_last_chunk: bool = False,
+    ):
+        """Reserve KV blocks for the next chunk of a running long prefill."""
+        old_kv_token_limit = seq.kv_token_limit
+        if is_last_chunk:
+            seq.kv_token_limit = None
+        else:
+            seq.kv_token_limit = seq.num_history_ids + chunk_size
+            prealloc_size = 0
+
+        evictable = hanging + waiting
+        if not self.eviction_helper.evict_for_seq(
+            seq,
+            evictable,
+            prealloc_size,
+        ):
+            seq.kv_token_limit = old_kv_token_limit
+            return False
+
+        self.block_manager.allocate(seq, prealloc_size)
+        self.block_trie.allocate(seq)
+        return True
+
+    @record_function('schedule_prefill')
+    def schedule(
+        self,
+        *,
+        waiting: SeqList,
+        hanging: SeqList,
+        num_ready: int,
+        num_running: int,
+        prealloc_size: int = 0,
+        allow_long_prefill: bool = True,
+        prefer_long_prefill: bool = False,
+    ):
+        """Select and activate one prefill batch."""
+        self.last_schedule_had_pending_lookup = False
+        max_batches = self.scheduler_config.max_batches - num_ready - num_running
+        running: SeqList = []
+        token_count = 0
+
+        def _to_running(
+            seq: SchedulerSequence,
+            prefill_token_count: int,
+        ):
+            """Activate an admitted sequence and count its prefill tokens."""
+            seq.state.activate()
+            running.append(seq)
+            nonlocal token_count
+            token_count += prefill_token_count
+
+        if len(running) >= max_batches or len(waiting) == 0:
+            return running
+
+        waiting = _PrefillReorderer(self).reorder(
+            waiting,
+            allow_long_prefill=allow_long_prefill,
+            prefer_long_prefill=prefer_long_prefill,
+        )
+        skipped_waiting: SeqList = []
+        while len(waiting) > 0 and len(running) < max_batches:
+            seq = waiting.pop(0)
+            evictable_waiting = skipped_waiting + waiting
+            admission = _PrefillAdmissionAttempt(
+                self,
+                seq,
+                hanging=hanging,
+                evictable_waiting=evictable_waiting,
+                prealloc_size=prealloc_size,
+                token_count=token_count,
+                has_admitted=len(running) > 0,
+                allow_long_prefill=allow_long_prefill,
+            ).run()
+
+            if admission.action is _PrefillAdmissionAction.LOAD_STARTED:
+                # The request left WAITING for asynchronous load without using
+                # a model-batch slot or prefill token budget.
+                continue
+            if admission.action is _PrefillAdmissionAction.SKIP:
+                skipped_waiting.append(seq)
+                continue
+            if admission.action is _PrefillAdmissionAction.STOP:
+                break
+
+            assert admission.action is _PrefillAdmissionAction.ADMIT
+            _to_running(seq, admission.prefill_token_count)
+            seq.record_event(EventType.SCHEDULED)
+
+            if seq.kv_token_limit is not None:
+                break
+
+        return running
+
+
 class Scheduler:
     """Tools to schedule next step.
 
@@ -853,6 +1060,16 @@ def __init__(
             eviction_helper=self.eviction_helper,
             sessions=self.sessions,
         )
+        self._prefill_scheduler = _PrefillScheduler(
+            scheduler_config=self.scheduler_config,
+            cache_config=self.cache_config,
+            is_ssm=self.is_ssm,
+            block_manager=self.block_manager,
+            block_trie=self.block_trie,
+            state_manager=self.state_manager,
+            eviction_helper=self.eviction_helper,
+            load_coordinator=self.kv_load_coordinator,
+        )
         # Keep save call sites uniform even when the producer role is disabled.
         self.kv_save_coordinator = KVSaveCoordinator(self)
         # Per-tick signal consumed by EngineLoop to distinguish asynchronous
@@ -863,8 +1080,6 @@ def __init__(
         self.seq_meta = seq_meta
         self.seq_manager = SequenceManager(seq_meta)
         self.scheduler_tick = 0
-        self._long_prefill_policy = _envs.opt_ttft_policy
-        self._long_prefill_aging_seconds_per_chunk = max(0.001, _envs.opt_ttft_aging_sec)
 
     def tick(self):
         """Mark one scheduler progress step (once per forward dispatch)."""
@@ -884,57 +1099,9 @@ def shutdown(self) -> None:
         if connector is not None:
             connector.shutdown()
 
-    def _ensure_runtime_state_available(self):
-        """Make one state-cache slot available for an SSM runtime state.
-
-        Runtime states and frozen checkpoints share the same state-cache pool. Scheduling a request is more important
-        than keeping an old checkpoint, so unpinned checkpoints are evicted before we give up.
-        """
-        if not self.is_ssm:
-            return True
-        if self.state_manager.get_num_free_runtime() > 0:
-            return True
-        self.block_trie.state_checkpoints.evict(1)
-        return self.state_manager.get_num_free_runtime() > 0
-
-    def _long_context_chunk_limit(self, seq: SchedulerSequence):
-        """Return the token budget for one long-context chunk."""
-        return get_long_context_chunk_limit(seq, self.cache_config.max_prefill_token_num)
-
-    def _next_long_context_chunk_end(self, seq: SchedulerSequence, max_prefill_num: int | None = None):
-        """Return the exclusive absolute token end for the next chunk."""
-        if max_prefill_num is None:
-            max_prefill_num = self._long_context_chunk_limit(seq)
-        plan = plan_long_context_chunk(seq, max_prefill_num, include_multimodals=False)
-        return plan.chunk_end
-
-    def _prefill_kv_token_limit(self, seq: SchedulerSequence):
-        """Limit KV allocation for a non-final long-context prefill chunk."""
-        max_prefill_num = self._long_context_chunk_limit(seq)
-        if seq.num_token_ids <= max_prefill_num:
-            return None
-        return self._next_long_context_chunk_end(seq, max_prefill_num)
-
-    def _prefill_admission_token_count(self, seq: SchedulerSequence):
-        """Return token budget cost for the next prefill or chunk."""
-        kv_token_limit = self._prefill_kv_token_limit(seq)
-        if kv_token_limit is None:
-            return seq.num_token_ids
-        return max(0, kv_token_limit - seq.num_history_ids)
-
     def has_waiting_long_prefill(self):
         """Whether a waiting request would need a non-final prefill chunk."""
-        return any(self._prefill_kv_token_limit(seq) is not None for seq in self.waiting)
-
-    def _prepare_prefill_allocation(self, seq: SchedulerSequence, prealloc_size: int):
-        """Apply chunk KV limit and return the effective prealloc size."""
-        kv_token_limit = self._prefill_kv_token_limit(seq)
-        if kv_token_limit is None:
-            seq.kv_token_limit = None
-            return prealloc_size
-
-        seq.kv_token_limit = kv_token_limit
-        return 0
+        return self._prefill_scheduler.has_waiting_long_prefill(self.waiting)
 
     def reserve_long_context_chunk(self,
                                    seq: SchedulerSequence,
@@ -942,21 +1109,14 @@ def reserve_long_context_chunk(self,
                                    prealloc_size: int = 0,
                                    is_last_chunk: bool = False):
         """Reserve KV blocks for the next chunk of a running long prefill."""
-        old_kv_token_limit = seq.kv_token_limit
-        if is_last_chunk:
-            seq.kv_token_limit = None
-        else:
-            seq.kv_token_limit = seq.num_history_ids + chunk_size
-            prealloc_size = 0
-
-        evictable = self.hanging + self.waiting
-        if not self.eviction_helper.evict_for_seq(seq, evictable, prealloc_size):
-            seq.kv_token_limit = old_kv_token_limit
-            return False
-
-        self.block_manager.allocate(seq, prealloc_size)
-        self.block_trie.allocate(seq)
-        return True
+        return self._prefill_scheduler.reserve_long_context_chunk(
+            seq,
+            hanging=self.hanging,
+            waiting=self.waiting,
+            chunk_size=chunk_size,
+            prealloc_size=prealloc_size,
+            is_last_chunk=is_last_chunk,
+        )
 
     @staticmethod
     def create_status_list_property(status: MessageStatus):
@@ -1064,71 +1224,6 @@ def _reorder_migrating():
 
         return migration_ready
 
-    @record_function('schedule_prefill')
-    def _schedule_prefill(self,
-                          prealloc_size: int = 0,
-                          allow_long_prefill: bool = True,
-                          prefer_long_prefill: bool = False):
-        """Schedule for prefilling."""
-
-        max_batches = self.scheduler_config.max_batches - self.num_ready() - self.num_running()
-        swap_out_map: MapType = dict()
-        swap_in_map: MapType = dict()
-        copy_map: MapType = dict()
-        running: SeqList = []
-        token_count = 0
-
-        def _to_running(seq: SchedulerSequence, prefill_token_count: int):
-            """Activate an admitted sequence and count its prefill tokens."""
-            seq.state.activate()
-            running.append(seq)
-            nonlocal token_count
-            token_count += prefill_token_count
-
-        num_waiting = self.seq_manager.num_sequences(MessageStatus.WAITING)
-        if (len(running) >= max_batches or num_waiting == 0):
-            return running, swap_in_map, swap_out_map, copy_map
-
-        waiting = _PrefillReorderer(self).reorder(self.waiting,
-                                                 allow_long_prefill=allow_long_prefill,
-                                                 prefer_long_prefill=prefer_long_prefill)
-        skipped_waiting: SeqList = []
-        while len(waiting) > 0 and len(running) < max_batches:
-            seq = waiting.pop(0)
-            evictable_waiting = skipped_waiting + waiting
-            admission = _PrefillAdmissionAttempt(
-                self,
-                seq,
-                evictable_waiting=evictable_waiting,
-                prealloc_size=prealloc_size,
-                token_count=token_count,
-                has_admitted=len(running) > 0,
-                allow_long_prefill=allow_long_prefill,
-            ).run()
-
-            if admission.action is _PrefillAdmissionAction.LOAD_STARTED:
-                # start_load already moved the sequence out of WAITING. It must
-                # not join running or skipped_waiting, both of which permit
-                # ordinary model/paging operations on the sequence. Since no
-                # model work was admitted, continue without consuming a batch
-                # slot or token budget.
-                continue
-            if admission.action is _PrefillAdmissionAction.SKIP:
-                skipped_waiting.append(seq)
-                continue
-            if admission.action is _PrefillAdmissionAction.STOP:
-                break
-
-            assert admission.action is _PrefillAdmissionAction.ADMIT
-            _to_running(seq, admission.prefill_token_count)
-
-            seq.record_event(EventType.SCHEDULED)
-
-            if seq.kv_token_limit is not None:
-                break
-
-        return running, swap_in_map, swap_out_map, copy_map
-
     @record_function('schedule_decoding')
     def _schedule_decoding(self, prealloc_size: int = 0):
         """Schedule decoding."""
@@ -1195,10 +1290,23 @@ def schedule(self,
         """Schedule inputs for next steps."""
         self.last_schedule_had_pending_lookup = False
         if is_prefill:
-            output = self._schedule_prefill(prealloc_size, allow_long_prefill, prefer_long_prefill)
+            running = self._prefill_scheduler.schedule(
+                waiting=self.waiting,
+                hanging=self.hanging,
+                num_ready=self.num_ready(),
+                num_running=self.num_running(),
+                prealloc_size=prealloc_size,
+                allow_long_prefill=allow_long_prefill,
+                prefer_long_prefill=prefer_long_prefill,
+            )
+            self.last_schedule_had_pending_lookup = (
+                self._prefill_scheduler.last_schedule_had_pending_lookup)
+            swap_in_map: MapType = {}
+            swap_out_map: MapType = {}
+            copy_map: MapType = {}
         else:
-            output = self._schedule_decoding(prealloc_size)
-        running, swap_in_map, swap_out_map, copy_map = output
+            running, swap_in_map, swap_out_map, copy_map = self._schedule_decoding(
+                prealloc_size)
 
         return SchedulerOutput(running=running, swap_in_map=swap_in_map, swap_out_map=swap_out_map, copy_map=copy_map)
 
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index 262cf6ed1a..96dcabd580 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -1638,7 +1638,9 @@ def test_ssm_scheduler_rejects_prefix_match_for_prefill_gate_after_runtime_state
     def _ensure_runtime_state_available_once_then_succeed():
         return next(ensure_results)
 
-    monkeypatch.setattr(scheduler, '_ensure_runtime_state_available', _ensure_runtime_state_available_once_then_succeed)
+    monkeypatch.setattr(scheduler._prefill_scheduler,
+                        '_ensure_runtime_state_available',
+                        _ensure_runtime_state_available_once_then_succeed)
     scheduler.block_trie.stats.reset()
 
     cache_hit_tail = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
@@ -1927,13 +1929,13 @@ def test_scheduler_reads_opt_ttft_env(monkeypatch):
 
     scheduler, _ = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
 
-    assert scheduler._long_prefill_policy == 'fifo'
-    assert scheduler._long_prefill_aging_seconds_per_chunk == 0.25
+    assert scheduler._prefill_scheduler._long_prefill_policy == 'fifo'
+    assert scheduler._prefill_scheduler._long_prefill_aging_seconds_per_chunk == 0.25
 
 
 def test_schedule_prefill_prefer_long_fifo_policy_keeps_oldest_huge_waiter_first():
     scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-    scheduler._long_prefill_policy = 'fifo'
+    scheduler._prefill_scheduler._long_prefill_policy = 'fifo'
     now = time.perf_counter()
     huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
     huge_long.arrive_time = now - 1.0
@@ -1975,7 +1977,7 @@ def test_schedule_prefill_prefer_long_admits_smaller_long_waiter_first():
 
 def test_schedule_prefill_prefer_long_ages_huge_long_waiter():
     scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-    scheduler._long_prefill_aging_seconds_per_chunk = 0.01
+    scheduler._prefill_scheduler._long_prefill_aging_seconds_per_chunk = 0.01
     now = time.perf_counter()
     huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
     huge_long.arrive_time = now - 1.0
@@ -2002,7 +2004,8 @@ def test_schedule_prefill_reapplies_chunk_limit_after_ssm_state_rollback():
     def _ensure_runtime_state_available_once_then_succeed():
         return next(ensure_results)
 
-    scheduler._ensure_runtime_state_available = _ensure_runtime_state_available_once_then_succeed
+    scheduler._prefill_scheduler._ensure_runtime_state_available = (
+        _ensure_runtime_state_available_once_then_succeed)
 
     output = scheduler.schedule(is_prefill=True, prealloc_size=1)
 

From 1dc3f0f7b9ce725c154d1cc941569ba0cc856083 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 15:33:56 +0800
Subject: [PATCH 06/22] refactor: move prefill scheduling to dedicated module

---
 lmdeploy/pytorch/paging/block_trie/README.md |   5 +-
 lmdeploy/pytorch/paging/prefill_scheduler.py | 979 +++++++++++++++++++
 lmdeploy/pytorch/paging/scheduler.py         | 975 +-----------------
 tests/pytorch/paging/test_scheduler.py       |   8 +-
 4 files changed, 995 insertions(+), 972 deletions(-)
 create mode 100644 lmdeploy/pytorch/paging/prefill_scheduler.py

diff --git a/lmdeploy/pytorch/paging/block_trie/README.md b/lmdeploy/pytorch/paging/block_trie/README.md
index c5aed137fb..2a39b18793 100644
--- a/lmdeploy/pytorch/paging/block_trie/README.md
+++ b/lmdeploy/pytorch/paging/block_trie/README.md
@@ -16,7 +16,8 @@ asynchronous boundaries.
 1. Read this document for the ownership and lifecycle contracts.
 2. Read [`BlockTrie.match()` and `BlockTrie.allocate()`](./trie.py) for the
    public trie workflow.
-3. Read `_PrefillAdmissionAttempt` in [`scheduler.py`](../scheduler.py) for
+3. Read `_PrefillAdmissionAttempt` in
+   [`prefill_scheduler.py`](../prefill_scheduler.py) for
    tentative-match commit and rollback.
 4. For KV ownership and eviction, read [`kv_lifecycle.py`](./kv_lifecycle.py).
 5. For SSM support, read [`checkpoint.py`](./checkpoint.py) before
@@ -368,7 +369,7 @@ state that may contain stale entries.
 | Sparse checkpoint keys or exact verification           | `checkpoint.py`                                                   |
 | Checkpoint reservation, publication, pins, or eviction | `checkpoint_lifecycle.py`                                         |
 | KV references, leaf bookkeeping, or KV eviction        | `kv_lifecycle.py`                                                 |
-| Admission order or tentative-match rollback            | `../scheduler.py`                                                 |
+| Admission order or tentative-match rollback            | `../prefill_scheduler.py`                                         |
 | Per-sequence prefix-cache protocol state               | `../../prefix_cache_state.py`                                     |
 | Host restore/save copy plans                           | `../../engine/inputs_maker.py` and `../../engine/cache_inputs.py` |
 | Stream ordering around model execution                 | `../../engine/model_agent/agent.py`                               |
diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
new file mode 100644
index 0000000000..8492566e10
--- /dev/null
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -0,0 +1,979 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+"""Request scheduling and prefix-cache side-effect boundaries.
+
+The scheduler is the first owner of prefix-cache side effects.  In prefill,
+``BlockTrie.match()`` is intentionally called before eviction and allocation so
+the scheduler can account for reused KV/state.  That match is tentative:
+rollback is required if checkpoint pinning, KV eviction, or runtime state
+allocation means the request cannot safely run now.  Long-context suffixes can
+continue chunking from the accepted prefix hit.
+
+Successful prefill scheduling keeps this order:
+
+1. ``block_trie.match(seq)`` mutates sequence state to skip a cached prefix.
+2. eviction and SSM runtime-state availability are checked.
+3. ``block_manager.allocate(seq)`` allocates missing KV blocks.
+4. ``block_trie.allocate(seq)`` publishes newly allocated full blocks.
+5. For SSM, downstream input/model/engine code restores and saves checkpoint
+   states; the scheduler only owns resource decisions and rollback.
+
+SSM scheduling detail:
+
+* ``block_trie.match(seq)`` may find a published checkpoint and record
+  ``seq.prefix_cache.restore`` before the request owns a runtime state.
+  The scheduler must treat that as tentative until KV blocks and one runtime
+  state slot are guaranteed.
+* A matched restore checkpoint can be pinned before eviction so checkpoint LRU
+  cannot free the source slot.  If that pin prevents eviction from finding
+  enough resources, the scheduler rolls the match back, releases the pin, and
+  retries eviction once without the tentative hit.
+* Runtime state availability is checked after KV eviction because old unpinned
+  checkpoints may be dropped to free state-cache slots.  If no runtime slot can
+  be recovered, the tentative prefix hit is rolled back and the request waits.
+* ``state_manager.allocate(seq)`` assigns the request runtime state only after
+  ``block_manager.allocate(seq)`` and ``block_trie.allocate(seq)`` succeed.
+  Later, ``InputsMaker`` may reserve checkpoint saves for the exact produced
+  step; scheduler code does not perform state-cache tensor copies or publish
+  checkpoint readiness.
+
+External KV scheduling detail:
+
+* External lookup is enabled only for a KV consumer with a connector, and is
+  kept separate from the SSM checkpoint path. Local ``BlockTrie.match()`` runs
+  first so the connector searches only beyond KV already resident on this node.
+* Lookup is asynchronous. A pending result must leave the request schedulable
+  for a later tick without retaining a tentative local match, so multi-turn
+  sequence state is snapshotted and restored exactly.
+* A positive hit is block-aligned, allocated, and handed to
+  ``KVLoadCoordinator``. While workers may write those blocks, the sequence is
+  in ``WAITING_FOR_REMOTE_KVS`` and paging cleanup is deferred.
+* A successful load is published into the local trie and prioritized for its
+  remaining prefill. A failed or cancelled load returns to the last safe
+  block-aligned prefix because partially written destinations are untrusted.
+* Prefill saves take a physical block snapshot for workers and a logical block
+  lease for paging. ``KVSaveCoordinator`` keeps those blocks alive until every
+  TP rank reports terminal progress or worker queues are drained.
+"""
+
+import enum
+import time
+from collections import Counter
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from torch.profiler import record_function
+
+from lmdeploy.messages import EventType
+from lmdeploy.pytorch import envs as _envs
+from lmdeploy.pytorch.long_context import get_long_context_chunk_limit, plan_long_context_chunk
+from lmdeploy.utils import get_logger
+
+from ..config import CacheConfig, SchedulerConfig
+from ..messages import SchedulerSequence
+from .block_trie import BlockTrie
+from .kv_load_coordinator import KVLoadAdmission, KVLoadCoordinator
+from .state_manager import StateManager
+
+if TYPE_CHECKING:
+    from .block_manager.base_block_manager import BaseBlockManager
+    from .eviction_helper.base_eviction_helper import BaseEvictionHelper
+
+logger = get_logger('lmdeploy')
+
+SeqList = list[SchedulerSequence]
+
+
+@dataclass(frozen=True)
+class _PrefillReorderInfo:
+    """Immutable pre-admission metadata used only for waiting-list ordering."""
+
+    prefill_token_count: int
+    is_nonfinal_long_prefill: bool
+    estimated_long_chunks: int
+
+
+class _PrefillReorderer:
+    """Order waiting prefills without applying scheduler side effects."""
+
+    def __init__(self, prefill_scheduler: '_PrefillScheduler'):
+        self.prefill_scheduler = prefill_scheduler
+        self._info_cache: dict[int, _PrefillReorderInfo] = {}
+
+    def reorder(self,
+                waiting: SeqList,
+                allow_long_prefill: bool,
+                prefer_long_prefill: bool):
+        """Return waiting requests in the order the prefill loop should try."""
+        waiting = sorted(waiting, key=lambda seq: seq.arrive_time)
+        # A completed load already owns destination blocks and a soft prefill
+        # reservation. Admit it first to shorten that ownership window.
+        remote_ready = [
+            seq for seq in waiting
+            if self.prefill_scheduler.load_coordinator.is_remote_ready(seq)
+        ]
+        waiting = [
+            seq for seq in waiting
+            if not self.prefill_scheduler.load_coordinator.is_remote_ready(seq)
+        ]
+        if prefer_long_prefill:
+            # Long-work turns choose one long waiter first. The size policy only
+            # reorders this long lane; it is not global shortest-prefill-first
+            # admission.
+            long_turn_order = self._reorder_for_long_turn(waiting)
+            if long_turn_order is not None:
+                return remote_ready + self._warn_if_not_permutation(waiting, long_turn_order)
+
+        if allow_long_prefill:
+            return remote_ready + self._warn_if_not_permutation(waiting, waiting)
+
+        reordered = self._reorder_for_short_turn(waiting)
+        return remote_ready + self._warn_if_not_permutation(waiting, reordered)
+
+    def _warn_if_not_permutation(self, original: SeqList, reordered: SeqList):
+        """Warn if reorder drops, duplicates, or substitutes waiting
+        sequences."""
+        original_ids = [id(seq) for seq in original]
+        reordered_ids = [id(seq) for seq in reordered]
+        if len(original_ids) == len(reordered_ids) and Counter(original_ids) == Counter(reordered_ids):
+            return reordered
+
+        logger.warning('Unexpected prefill reorder result: original_len=%s reordered_len=%s '
+                       'original_sample=%s reordered_sample=%s',
+                       len(original), len(reordered), self._seq_id_sample(original), self._seq_id_sample(reordered))
+        return reordered
+
+    @staticmethod
+    def _seq_id_sample(seqs: SeqList):
+        return [(seq.session_id, seq.seq_id) for seq in seqs[:5]]
+
+    def _get_reorder_info(self, seq: SchedulerSequence):
+        """Return reorder-only info before prefix-cache side effects.
+
+        Prefix-cache match/rollback mutates the remaining prompt. Keep this cache confined to waiting-list ordering and
+        recompute fresh values in the admission path.
+        """
+        seq_key = id(seq)
+        info = self._info_cache.get(seq_key)
+        if info is not None:
+            return info
+
+        prefill = self.prefill_scheduler
+        chunk_limit = prefill._long_context_chunk_limit(seq)
+        if seq.num_token_ids <= chunk_limit:
+            info = _PrefillReorderInfo(prefill_token_count=seq.num_token_ids,
+                                       is_nonfinal_long_prefill=False,
+                                       estimated_long_chunks=1)
+        else:
+            kv_token_limit = prefill._next_long_context_chunk_end(seq, chunk_limit)
+            safe_chunk_limit = max(1, chunk_limit)
+            info = _PrefillReorderInfo(
+                prefill_token_count=max(0, kv_token_limit - seq.num_history_ids),
+                is_nonfinal_long_prefill=True,
+                estimated_long_chunks=max(1, (seq.num_token_ids + safe_chunk_limit - 1) // safe_chunk_limit),
+            )
+        self._info_cache[seq_key] = info
+        return info
+
+    def _long_priority_key(self, seq: SchedulerSequence, now: float):
+        """Prefer smaller long prompts, with age credit to avoid starvation."""
+        prefill = self.prefill_scheduler
+        info = self._get_reorder_info(seq)
+        wait_age = max(0.0, now - seq.arrive_time)
+        age_credit = int(wait_age // prefill._long_prefill_aging_seconds_per_chunk)
+        age_adjusted_chunks = info.estimated_long_chunks - age_credit
+        return age_adjusted_chunks, info.estimated_long_chunks, seq.arrive_time
+
+    def _split_by_prefill_kind(self, waiting: SeqList):
+        """Split waiting requests into normal/final and non-final long
+        prefill."""
+        normal_waiting: SeqList = []
+        long_waiting: SeqList = []
+        for seq in waiting:
+            if self._get_reorder_info(seq).is_nonfinal_long_prefill:
+                long_waiting.append(seq)
+            else:
+                normal_waiting.append(seq)
+        return normal_waiting, long_waiting
+
+    def _sort_normal_prefills(self, waiting: SeqList):
+        return sorted(waiting,
+                      key=lambda seq: (self._get_reorder_info(seq).prefill_token_count, seq.arrive_time))
+
+    def _sort_long_prefills(self, waiting: SeqList):
+        prefill = self.prefill_scheduler
+        if prefill._long_prefill_policy != 'size':
+            return waiting
+        now = time.perf_counter()
+        return sorted(waiting, key=lambda seq: self._long_priority_key(seq, now))
+
+    def _reorder_for_long_turn(self, waiting: SeqList):
+        """Choose one long waiter, then fill the turn with normal prefills."""
+        normal_waiting, long_waiting = self._split_by_prefill_kind(waiting)
+        if len(long_waiting) == 0:
+            return None
+
+        long_waiting = self._sort_long_prefills(long_waiting)
+        normal_waiting = self._sort_normal_prefills(normal_waiting)
+        return [long_waiting[0]] + normal_waiting + long_waiting[1:]
+
+    def _reorder_for_short_turn(self, waiting: SeqList):
+        """Prioritize normal/final prefills while preserving long waiters."""
+        normal_waiting, long_waiting = self._split_by_prefill_kind(waiting)
+        return self._sort_normal_prefills(normal_waiting) + long_waiting
+
+
+class _PrefillAdmissionAction(enum.Enum):
+    ADMIT = enum.auto()
+    SKIP = enum.auto()
+    STOP = enum.auto()
+    LOAD_STARTED = enum.auto()
+
+
+@dataclass(frozen=True)
+class _PrefillAdmissionResult:
+    """Outcome from trying to admit one waiting prefill request.
+
+    The outer loop distinguishes four outcomes:
+
+    * ``ADMIT``: include the request in this tick's model batch.
+    * ``SKIP``: leave it waiting but continue trying later candidates.
+    * ``STOP``: resource pressure ends this prefill admission turn.
+    * ``LOAD_STARTED``: no model work was selected, but the request left the
+      waiting queue for asynchronous KV load.
+    """
+
+    action: _PrefillAdmissionAction
+    prefill_token_count: int = 0
+
+    @classmethod
+    def admit(cls, prefill_token_count: int):
+        return cls(action=_PrefillAdmissionAction.ADMIT,
+                   prefill_token_count=prefill_token_count)
+
+    @classmethod
+    def skip(cls):
+        return cls(action=_PrefillAdmissionAction.SKIP)
+
+    @classmethod
+    def stop(cls):
+        return cls(action=_PrefillAdmissionAction.STOP)
+
+    @classmethod
+    def load(cls):
+        return cls(action=_PrefillAdmissionAction.LOAD_STARTED)
+
+
+@dataclass(frozen=True, slots=True)
+class _PrefixMatchStateSnapshot:
+    """Exact sequence state captured before a tentative local trie match.
+
+    External lookup itself does not mutate sequence paging state. The scheduler
+    may, however, run ``block_trie.match`` first so the connector queries only
+    beyond the locally resident prefix. If that non-blocking lookup returns
+    pending, or a positive hit cannot be admitted before worker writes start,
+    the request will not run this tick and the tentative local match must be
+    undone.
+
+    A multi-turn request may already own valid history, blocks, and model
+    metadata before this attempt. Restoring this baseline preserves that exact
+    committed state; the legacy new-request rollback to step zero would discard
+    it. This snapshot is not used after an asynchronous load starts--load
+    failure then rolls back to its block-aligned ``fallback_step`` through
+    ``KVLoadCoordinator`` because workers may have partially written KV.
+    """
+
+    # Committed sequence progress and block ownership before tentative match.
+    num_history_ids: int
+    num_blocks: int
+    # Prefix-cache cursor, public hit accounting, and temporary overlap state.
+    trie_cursor: Any
+    match_start_step: int
+    cached_tokens: int
+    # Request-local allocation limit that a multi-turn attempt may carry.
+    kv_token_limit: int | None
+    # Temporary recompute-overlap identities created by local trie matching.
+    fresh_block_range: range | None
+    trie_block_map: dict[int, int]
+    # Model state that must remain aligned with the committed history step.
+    model_meta: Any
+
+    @classmethod
+    def capture(cls, seq: SchedulerSequence):
+        overlap = seq.prefix_cache.recompute_overlap
+        return cls(
+            num_history_ids=int(seq.num_history_ids),
+            num_blocks=int(seq.num_blocks),
+            trie_cursor=seq.prefix_cache.trie_cursor,
+            match_start_step=int(seq.prefix_cache.match_start_step),
+            cached_tokens=int(seq.cached_tokens),
+            kv_token_limit=seq.kv_token_limit,
+            fresh_block_range=overlap.fresh_block_range,
+            trie_block_map=dict(overlap.trie_block_map),
+            model_meta=seq.model_meta,
+        )
+
+
+class _TentativePrefixMatch:
+    """Request-local transaction around ``BlockTrie.match`` side effects.
+
+    Ordinary and SSM admission preserve the historical fallback to an unmatched request. External lookup instead needs
+    an exact pre-match snapshot because a multi-turn request may already own committed progress. Both contracts share
+    one stats snapshot, restore-pin boundary, and explicit commit/rollback lifecycle without changing their rollback
+    semantics.
+    """
+
+    __slots__ = (
+        'seq',
+        'block_trie',
+        'block_manager',
+        'is_ssm',
+        '_preserve_existing_state',
+        '_stats_snapshot',
+        '_state_snapshot',
+        '_rejection_on_rollback',
+        '_started',
+        'matched',
+    )
+
+    def __init__(self,
+                 seq: SchedulerSequence,
+                 block_trie: BlockTrie,
+                 block_manager,
+                 *,
+                 is_ssm: bool,
+                 preserve_existing_state: bool):
+        self.seq = seq
+        self.block_trie = block_trie
+        self.block_manager = block_manager
+        self.is_ssm = is_ssm
+        self._preserve_existing_state = preserve_existing_state
+        self._stats_snapshot = None
+        self._state_snapshot: _PrefixMatchStateSnapshot | None = None
+        self._rejection_on_rollback: _PrefillAdmissionResult | None = None
+        self._started = False
+        self.matched = False
+
+    def begin(self) -> None:
+        """Start the transaction before gates can mutate exact external state.
+
+        Ordinary admission starts lazily from ``match``. External admission
+        starts before gates so rollback can restore existing request state even
+        when a private partial block prevents another trie match.
+        """
+        if self._started or not self.block_trie.enabled:
+            return
+        self._stats_snapshot = self.block_trie.stats.snapshot()
+        if self._preserve_existing_state:
+            self._state_snapshot = _PrefixMatchStateSnapshot.capture(self.seq)
+        self._started = True
+
+    def match(self) -> None:
+        """Apply one tentative match after capturing its rollback boundary."""
+        assert not self.matched
+        self.begin()
+        self.block_trie.match(self.seq)
+        self.matched = True
+
+    def retain_for_admission(self, rejection_on_rollback: _PrefillAdmissionResult) -> None:
+        """Keep a gate-enabling match and remember its original rejection."""
+        assert self.matched
+        self._rejection_on_rollback = rejection_on_rollback
+
+    def pin_restore(self) -> bool:
+        """Pin an SSM restore selected by this tentative match."""
+        restore = self.seq.prefix_cache.restore
+        if not self.is_ssm or not restore.is_selected:
+            return True
+        return self.block_trie.state_checkpoints.pin_restore(self.seq)
+
+    def commit(self) -> None:
+        """Accept the match and discard request-local rollback state."""
+        self._clear()
+
+    def rollback(self, reason: str):
+        """Undo the transaction and return any gate-defined rejection."""
+        rejection = self._rejection_on_rollback
+        if not self._started:
+            return rejection
+
+        seq = self.seq
+        logger.debug('Rollback tentative prefix-cache match: session_id=%s seq_id=%s reason=%s '
+                     'num_history_ids=%s restore_state=%s', seq.session_id, seq.seq_id, reason, seq.num_history_ids,
+                     seq.prefix_cache.restore.slot)
+        self.block_trie.stats.restore(self._stats_snapshot)
+        snapshot = self._state_snapshot
+        if snapshot is None:
+            self._reset_to_unmatched()
+        else:
+            self._restore_snapshot(snapshot)
+        self._clear()
+        return rejection
+
+    def _restore_snapshot(self, snapshot: _PrefixMatchStateSnapshot) -> None:
+        seq = self.seq
+        if seq.num_blocks < snapshot.num_blocks:
+            raise RuntimeError(
+                'tentative prefix match removed sequence-owned baseline blocks')
+        if seq.num_blocks > snapshot.num_blocks:
+            self.block_manager.truncate(seq, snapshot.num_blocks)
+        seq.set_step(snapshot.num_history_ids)
+        seq.model_meta = snapshot.model_meta
+        seq.kv_token_limit = snapshot.kv_token_limit
+        prefix_cache = seq.prefix_cache
+        prefix_cache.trie_cursor = snapshot.trie_cursor
+        prefix_cache.match_start_step = snapshot.match_start_step
+        overlap = prefix_cache.recompute_overlap
+        overlap.fresh_block_range = snapshot.fresh_block_range
+        overlap.trie_block_map.clear()
+        overlap.trie_block_map.update(snapshot.trie_block_map)
+        seq.cached_tokens = snapshot.cached_tokens
+
+    def _reset_to_unmatched(self) -> None:
+        seq = self.seq
+        if self.is_ssm:
+            self.block_trie.state_checkpoints.unpin_restore(seq)
+        if seq.num_blocks > 0 or seq.logical_state >= 0:
+            seq.state.free()
+        elif seq.num_history_ids > 0:
+            seq.set_step(0)
+        seq.kv_token_limit = None
+        prefix_cache = seq.prefix_cache
+        prefix_cache.trie_cursor = None
+        prefix_cache.restore.clear()
+        prefix_cache.match_start_step = -1
+        prefix_cache.recompute_overlap.clear_tracking()
+        seq.cached_tokens = 0
+
+    def _clear(self) -> None:
+        self._stats_snapshot = None
+        self._state_snapshot = None
+        self._rejection_on_rollback = None
+        self._started = False
+        self.matched = False
+
+
+class _PrefillAdmissionAttempt:
+    """Try to admit one waiting prefill sequence.
+
+    The attempt owns all tentative prefix-cache side effects for the sequence:
+    match, SSM restore pinning, eviction, runtime-state checks, allocation, and
+    rollback. The outer prefill loop still owns queue traversal and decides
+    whether a rejected candidate is skipped or ends the current prefill turn.
+    """
+
+    def __init__(self,
+                 prefill_scheduler: '_PrefillScheduler',
+                 seq: SchedulerSequence,
+                 hanging: SeqList,
+                 evictable_waiting: SeqList,
+                 prealloc_size: int,
+                 token_count: int,
+                 has_admitted: bool,
+                 allow_long_prefill: bool):
+        self.prefill_scheduler = prefill_scheduler
+        self.seq = seq
+        self.hanging = hanging
+        self.evictable_waiting = evictable_waiting
+        self.prealloc_size = prealloc_size
+        self.token_count = token_count
+        self.has_admitted = has_admitted
+        self.load_coordinator = prefill_scheduler.load_coordinator
+        self._remote_ready = self.load_coordinator.is_remote_ready(seq)
+        self.allow_long_prefill = allow_long_prefill
+        self._alloc_size = prealloc_size
+        self._prefix_match = _TentativePrefixMatch(
+            seq,
+            prefill_scheduler.block_trie,
+            prefill_scheduler.block_manager,
+            is_ssm=prefill_scheduler.is_ssm,
+            preserve_existing_state=(
+                self.load_coordinator.lookup_enabled and not self._remote_ready),
+        )
+
+    def run(self):
+        """Run the admission route for one waiting prefill.
+
+        1. If a previous external lookup is pending, skip without applying new
+           local prefix-cache side effects.
+        2. Snapshot multi-turn state before a local match may become tentative.
+        3. Apply long-prefill and token-budget gates.
+        4. Prefer a local trie hit, then query/load only its remote extension.
+        5. Admit KV/state resources or roll the tentative match back precisely.
+        6. On success, allocate blocks/states and publish the accepted hit.
+        """
+        if self.load_coordinator.lookup_enabled and not self._remote_ready:
+            if self._lookup_is_pending():
+                return _PrefillAdmissionResult.skip()
+            self._prefix_match.begin()
+
+        gate_result = self._check_prefill_admission_gates()
+        if gate_result is not None:
+            return gate_result
+
+        resource_result = self._admit_resources()
+        if resource_result is not None:
+            return resource_result
+
+        return self._finish_admission()
+
+    def _lookup_is_pending(self) -> bool:
+        """Skip without touching local prefix state while lookup is running.
+
+        The connector owns the Future and deduplicates polls. Marking this turn lets EngineLoop use a short I/O poll
+        delay instead of diagnosing an empty batch as GPU-cache pressure.
+        """
+        prefill = self.prefill_scheduler
+        if not self.load_coordinator.is_lookup_pending(self.seq):
+            return False
+        prefill.last_schedule_had_pending_lookup = True
+        return True
+
+    def _admit_resources(self):
+        if self.prefill_scheduler.block_trie.enabled:
+            return self._admit_prefix_cache_resources()
+        if self.load_coordinator.lookup_enabled:
+            lookup_result = self._query_external_prefix()
+            if lookup_result is not None:
+                return lookup_result
+        if not self._prepare_and_evict():
+            return _PrefillAdmissionResult.stop()
+        return None
+
+    def _admit_prefix_cache_resources(self):
+        """Admit resources for prefix-cache scheduling.
+
+        Route map:
+        1. Use or create the tentative prefix-cache match.
+        2. For external consumers, query only beyond that local match.
+        3. Pin any SSM restore state required by the match.
+        4. Prepare allocation limits and evict KV/state resources.
+        5. For SSM, verify a runtime state slot is still available.
+
+        Any failure rolls the tentative match back. A match created only to pass
+        a prefill gate returns that gate's skip/stop result after rollback;
+        normal resource failures keep their local retry/stop behavior here.
+        """
+        prefill = self.prefill_scheduler
+        seq = self.seq
+        if not self._prefix_match.matched:
+            # A completed external load has already published the accepted
+            # prefix interval. Matching again would restart accounting at the
+            # remote step and drop the restored tokens from request metrics.
+            if not self._remote_ready and not self._has_private_local_tail():
+                self._prefix_match.match()
+
+        if self.load_coordinator.lookup_enabled:
+            lookup_result = self._query_external_prefix()
+            if lookup_result is not None:
+                return lookup_result
+
+        had_ssm_restore = prefill.is_ssm and seq.prefix_cache.restore.is_selected
+        if not self._prefix_match.pin_restore():
+            result = self._prefix_match.rollback(
+                'failed to pin SSM restore checkpoint')
+            if result is not None:
+                return result
+
+        if not self._prepare_and_evict():
+            if not had_ssm_restore:
+                result = self._prefix_match.rollback('eviction failed')
+                if result is not None:
+                    return result
+                return _PrefillAdmissionResult.stop()
+
+            # A matched SSM restore may be pinning the only checkpoint state
+            # that eviction would otherwise free. Roll it back once and retry
+            # eviction before declaring the sequence unschedulable.
+            result = self._prefix_match.rollback(
+                'eviction failed with pinned SSM restore')
+            if result is not None:
+                return result
+            if not self._prepare_and_evict():
+                return _PrefillAdmissionResult.stop()
+
+        if prefill.is_ssm and not prefill._ensure_runtime_state_available():
+            result = self._prefix_match.rollback(
+                'no runtime SSM state available')
+            if result is not None:
+                return result
+            if not self._prepare_and_evict():
+                return _PrefillAdmissionResult.stop()
+            if not prefill._ensure_runtime_state_available():
+                seq.kv_token_limit = None
+                return _PrefillAdmissionResult.stop()
+
+        return None
+
+    def _query_external_prefix(self):
+        """Map connector/paging admission to prefill queue policy."""
+        prefill = self.prefill_scheduler
+        if self._remote_ready:
+            return None
+        admission = self.load_coordinator.try_load(
+            self.seq,
+            prealloc_size=self.prealloc_size,
+            evictable_seqs=self._evictable_sequences(),
+        )
+        if admission is KVLoadAdmission.NO_LOAD:
+            return None
+        if admission is KVLoadAdmission.PENDING:
+            self._prefix_match.rollback('external lookup pending')
+            prefill.last_schedule_had_pending_lookup = True
+            return _PrefillAdmissionResult.skip()
+        if admission is KVLoadAdmission.STARTED:
+            self._prefix_match.commit()
+            return _PrefillAdmissionResult.load()
+        if admission is KVLoadAdmission.FULL_PREFILL_UNAVAILABLE:
+            reason = 'full prefill capacity unavailable'
+        else:
+            assert admission is KVLoadAdmission.SOFT_BUDGET_UNAVAILABLE
+            reason = 'soft prefill budget unavailable'
+        # No worker has seen a destination on rejected admission, so the
+        # request-local prefix transaction remains exactly reversible.
+        self._prefix_match.rollback(reason)
+        return _PrefillAdmissionResult.stop()
+
+    def _evictable_sequences(self):
+        """Iterate queue-owned eviction candidates in historical order."""
+        yield from reversed(self.hanging)
+        yield from reversed(self.evictable_waiting)
+
+    def _match_prefix_for_prefill_gate(self):
+        """Tentatively match once so a request can be rechecked by a gate."""
+        prefill = self.prefill_scheduler
+        if (self._remote_ready or not prefill.block_trie.enabled
+                or self._has_private_local_tail()):
+            return None
+        self._prefix_match.match()
+        return True
+
+    def _has_private_local_tail(self) -> bool:
+        """Whether blocks exist beyond the full-block part of local history.
+
+        ``num_history_ids // block_size`` counts the completely computed
+        blocks before the current step.  A larger ``num_blocks`` means that the
+        sequence also owns the block containing a non-aligned current step, or
+        blocks preallocated after it.  Those blocks are private to this
+        sequence because their KV is partial or not computed yet, so they
+        cannot be published as complete reusable trie blocks.
+
+        For example, with block size 4, a chunked prefill may stop at step 5
+        with block table ``[P0, P1]``::
+
+            P0 -> tokens [0, 4), complete
+            P1 -> tokens [4, 8), only the KV at token 4 is valid
+
+        The trie cursor is at step 4 while ``P1`` already occupies logical
+        block index 1.  If another ``block_trie.match`` finds a shared block
+        ``S1`` for tokens [4, 8), matching appends it after ``P1`` instead of
+        filling ``P1``.  The resulting table ``[P0, P1, S1]`` is misaligned:
+        ``S1`` describes logical block 1 but resides at block-table index 2.
+
+        External lookup starts at the exact local step 5, but a block-granular
+        transfer rounds its start down to step 4.  It must therefore reuse
+        ``P1`` at index 1 as the first destination and overwrite the incomplete
+        KV there.  Skipping trie rematch keeps that destination stable until
+        the load is bound.  Without this guard, lookup may start from an
+        incorrectly advanced step or the load/model may address a block table
+        whose logical token ranges no longer match its indices.
+
+        This state means that a sequence retains local progress across
+        scheduling attempts.  It can result from chunked prefill, repeated
+        model forwards, preemption/resume, or a continued chat session; it is
+        not specific to multi-turn conversation.
+        """
+        if not self.load_coordinator.lookup_enabled:
+            return False
+        seq = self.seq
+        return seq.num_blocks > int(seq.num_history_ids) // seq.block_size
+
+    def _token_budget_rejection(self):
+        if self.allow_long_prefill:
+            return _PrefillAdmissionResult.stop()
+        return _PrefillAdmissionResult.skip()
+
+    def _check_prefill_admission_gates(self):
+        """Apply prefill gates, tentatively matching only when it may help."""
+        prefill = self.prefill_scheduler
+        seq = self.seq
+        token_budget = prefill.cache_config.max_prefill_token_num
+        prefill_token_count = prefill._prefill_admission_token_count(seq)
+        is_nonfinal_long_prefill = prefill._prefill_kv_token_limit(seq) is not None
+
+        if is_nonfinal_long_prefill and not self.allow_long_prefill:
+            matched = self._match_prefix_for_prefill_gate()
+            if matched is None:
+                return _PrefillAdmissionResult.skip()
+            if prefill._prefill_kv_token_limit(seq) is not None:
+                self._prefix_match.rollback('still non-final long prefill on short turn')
+                return _PrefillAdmissionResult.skip()
+            self._prefix_match.retain_for_admission(
+                _PrefillAdmissionResult.skip())
+            prefill_token_count = prefill._prefill_admission_token_count(seq)
+
+        exceeds_token_budget = self.has_admitted and self.token_count + prefill_token_count > token_budget
+        if not exceeds_token_budget:
+            return None
+
+        if not self._prefix_match.matched:
+            matched = self._match_prefix_for_prefill_gate()
+            if matched is not None:
+                prefill_token_count = prefill._prefill_admission_token_count(seq)
+                if self.token_count + prefill_token_count <= token_budget:
+                    self._prefix_match.retain_for_admission(
+                        self._token_budget_rejection())
+                    return None
+                self._prefix_match.rollback('still exceeds prefill token budget')
+        else:
+            self._prefix_match.rollback('still exceeds prefill token budget')
+        return self._token_budget_rejection()
+
+    def _prepare_and_evict(self):
+        """Apply chunk allocation limits and evict for this prefill."""
+        prefill = self.prefill_scheduler
+        seq = self.seq
+        alloc_size = prefill._prepare_prefill_allocation(seq, self.prealloc_size)
+        self._alloc_size = alloc_size
+        if self._evict_for_seq(alloc_size):
+            return True
+        seq.kv_token_limit = None
+        return False
+
+    def _evict_for_seq(self, alloc_size: int):
+        """Evict stopped or skipped waiters until this sequence can run."""
+        prefill = self.prefill_scheduler
+        return prefill.eviction_helper.evict_for_seq(
+            self.seq,
+            list(self._evictable_sequences()),
+            alloc_size,
+        )
+
+    def _finish_admission(self):
+        prefill = self.prefill_scheduler
+        seq = self.seq
+        # Prefix-cache matching can advance the sequence step and shrink the
+        # remaining prefill tail. Charge the admitted batch with the
+        # post-match/post-rollback cost, not the conservative pre-match
+        # estimate used to decide whether this sequence is worth trying.
+        prefill_token_count = prefill._prefill_admission_token_count(seq)
+        prefill.block_manager.allocate(seq, self._alloc_size)
+        if prefill.block_trie.enabled:
+            prefill.block_trie.allocate(seq)
+        if prefill.is_ssm:
+            prefill.state_manager.allocate(seq)
+        if prefill.block_trie.enabled:
+            prefill.block_trie.finalize_match(seq)
+        self.load_coordinator.track_prefill(
+            seq,
+            prealloc_size=self.prealloc_size,
+        )
+        if self._remote_ready:
+            # Preserve the load record through the remaining prefill so its
+            # reservation can be released only after model output advances the
+            # sequence to input_end_pos.
+            self.load_coordinator.mark_scheduled(seq)
+        self._prefix_match.commit()
+        return _PrefillAdmissionResult.admit(prefill_token_count)
+
+
+class _PrefillScheduler:
+    """Own prefill ordering, admission, and long-context reservation.
+
+    Long-lived dependencies are the resource owners used by prefill. Queue
+    contents and active-batch counts remain request-local inputs supplied by
+    the public :class:`Scheduler` facade for each scheduling turn.
+    """
+
+    def __init__(
+        self,
+        scheduler_config: SchedulerConfig,
+        cache_config: CacheConfig,
+        *,
+        is_ssm: bool,
+        block_manager: 'BaseBlockManager',
+        block_trie: BlockTrie,
+        state_manager: StateManager,
+        eviction_helper: 'BaseEvictionHelper',
+        load_coordinator: KVLoadCoordinator,
+    ) -> None:
+        self.scheduler_config = scheduler_config
+        self.cache_config = cache_config
+        self.is_ssm = is_ssm
+        self.block_manager = block_manager
+        self.block_trie = block_trie
+        self.state_manager = state_manager
+        self.eviction_helper = eviction_helper
+        self.load_coordinator = load_coordinator
+        self.last_schedule_had_pending_lookup = False
+        self._long_prefill_policy = _envs.opt_ttft_policy
+        self._long_prefill_aging_seconds_per_chunk = max(
+            0.001,
+            _envs.opt_ttft_aging_sec,
+        )
+
+    def _ensure_runtime_state_available(self):
+        """Make one state-cache slot available for an SSM runtime state."""
+        if not self.is_ssm:
+            return True
+        if self.state_manager.get_num_free_runtime() > 0:
+            return True
+        self.block_trie.state_checkpoints.evict(1)
+        return self.state_manager.get_num_free_runtime() > 0
+
+    def _long_context_chunk_limit(self, seq: SchedulerSequence):
+        """Return the token budget for one long-context chunk."""
+        return get_long_context_chunk_limit(
+            seq,
+            self.cache_config.max_prefill_token_num,
+        )
+
+    def _next_long_context_chunk_end(
+        self,
+        seq: SchedulerSequence,
+        max_prefill_num: int | None = None,
+    ):
+        """Return the exclusive absolute token end for the next chunk."""
+        if max_prefill_num is None:
+            max_prefill_num = self._long_context_chunk_limit(seq)
+        plan = plan_long_context_chunk(
+            seq,
+            max_prefill_num,
+            include_multimodals=False,
+        )
+        return plan.chunk_end
+
+    def _prefill_kv_token_limit(self, seq: SchedulerSequence):
+        """Limit KV allocation for a non-final long-context prefill chunk."""
+        max_prefill_num = self._long_context_chunk_limit(seq)
+        if seq.num_token_ids <= max_prefill_num:
+            return None
+        return self._next_long_context_chunk_end(seq, max_prefill_num)
+
+    def _prefill_admission_token_count(self, seq: SchedulerSequence):
+        """Return token budget cost for the next prefill or chunk."""
+        kv_token_limit = self._prefill_kv_token_limit(seq)
+        if kv_token_limit is None:
+            return seq.num_token_ids
+        return max(0, kv_token_limit - seq.num_history_ids)
+
+    def _prepare_prefill_allocation(
+        self,
+        seq: SchedulerSequence,
+        prealloc_size: int,
+    ):
+        """Apply chunk KV limit and return the effective prealloc size."""
+        kv_token_limit = self._prefill_kv_token_limit(seq)
+        if kv_token_limit is None:
+            seq.kv_token_limit = None
+            return prealloc_size
+
+        seq.kv_token_limit = kv_token_limit
+        return 0
+
+    def has_waiting_long_prefill(self, waiting: SeqList):
+        """Whether a waiting request needs a non-final prefill chunk."""
+        return any(
+            self._prefill_kv_token_limit(seq) is not None
+            for seq in waiting
+        )
+
+    def reserve_long_context_chunk(
+        self,
+        seq: SchedulerSequence,
+        *,
+        hanging: SeqList,
+        waiting: SeqList,
+        chunk_size: int,
+        prealloc_size: int = 0,
+        is_last_chunk: bool = False,
+    ):
+        """Reserve KV blocks for the next chunk of a running long prefill."""
+        old_kv_token_limit = seq.kv_token_limit
+        if is_last_chunk:
+            seq.kv_token_limit = None
+        else:
+            seq.kv_token_limit = seq.num_history_ids + chunk_size
+            prealloc_size = 0
+
+        evictable = hanging + waiting
+        if not self.eviction_helper.evict_for_seq(
+            seq,
+            evictable,
+            prealloc_size,
+        ):
+            seq.kv_token_limit = old_kv_token_limit
+            return False
+
+        self.block_manager.allocate(seq, prealloc_size)
+        self.block_trie.allocate(seq)
+        return True
+
+    @record_function('schedule_prefill')
+    def schedule(
+        self,
+        *,
+        waiting: SeqList,
+        hanging: SeqList,
+        num_ready: int,
+        num_running: int,
+        prealloc_size: int = 0,
+        allow_long_prefill: bool = True,
+        prefer_long_prefill: bool = False,
+    ):
+        """Select and activate one prefill batch."""
+        self.last_schedule_had_pending_lookup = False
+        max_batches = self.scheduler_config.max_batches - num_ready - num_running
+        running: SeqList = []
+        token_count = 0
+
+        def _to_running(
+            seq: SchedulerSequence,
+            prefill_token_count: int,
+        ):
+            """Activate an admitted sequence and count its prefill tokens."""
+            seq.state.activate()
+            running.append(seq)
+            nonlocal token_count
+            token_count += prefill_token_count
+
+        if len(running) >= max_batches or len(waiting) == 0:
+            return running
+
+        waiting = _PrefillReorderer(self).reorder(
+            waiting,
+            allow_long_prefill=allow_long_prefill,
+            prefer_long_prefill=prefer_long_prefill,
+        )
+        skipped_waiting: SeqList = []
+        while len(waiting) > 0 and len(running) < max_batches:
+            seq = waiting.pop(0)
+            evictable_waiting = skipped_waiting + waiting
+            admission = _PrefillAdmissionAttempt(
+                self,
+                seq,
+                hanging=hanging,
+                evictable_waiting=evictable_waiting,
+                prealloc_size=prealloc_size,
+                token_count=token_count,
+                has_admitted=len(running) > 0,
+                allow_long_prefill=allow_long_prefill,
+            ).run()
+
+            if admission.action is _PrefillAdmissionAction.LOAD_STARTED:
+                # The request left WAITING for asynchronous load without using
+                # a model-batch slot or prefill token budget.
+                continue
+            if admission.action is _PrefillAdmissionAction.SKIP:
+                skipped_waiting.append(seq)
+                continue
+            if admission.action is _PrefillAdmissionAction.STOP:
+                break
+
+            assert admission.action is _PrefillAdmissionAction.ADMIT
+            _to_running(seq, admission.prefill_token_count)
+            seq.record_event(EventType.SCHEDULED)
+
+            if seq.kv_token_limit is not None:
+                break
+
+        return running
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index d7e33d7b01..ffe2bfa08c 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -1,92 +1,29 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 # modify from: https://github.com/vllm-project/vllm
-"""Request scheduling and prefix-cache side-effect boundaries.
-
-The scheduler is the first owner of prefix-cache side effects.  In prefill,
-``BlockTrie.match()`` is intentionally called before eviction and allocation so
-the scheduler can account for reused KV/state.  That match is tentative:
-rollback is required if checkpoint pinning, KV eviction, or runtime state
-allocation means the request cannot safely run now.  Long-context suffixes can
-continue chunking from the accepted prefix hit.
-
-Successful prefill scheduling keeps this order:
-
-1. ``block_trie.match(seq)`` mutates sequence state to skip a cached prefix.
-2. eviction and SSM runtime-state availability are checked.
-3. ``block_manager.allocate(seq)`` allocates missing KV blocks.
-4. ``block_trie.allocate(seq)`` publishes newly allocated full blocks.
-5. For SSM, downstream input/model/engine code restores and saves checkpoint
-   states; the scheduler only owns resource decisions and rollback.
-
-SSM scheduling detail:
-
-* ``block_trie.match(seq)`` may find a published checkpoint and record
-  ``seq.prefix_cache.restore`` before the request owns a runtime state.
-  The scheduler must treat that as tentative until KV blocks and one runtime
-  state slot are guaranteed.
-* A matched restore checkpoint can be pinned before eviction so checkpoint LRU
-  cannot free the source slot.  If that pin prevents eviction from finding
-  enough resources, the scheduler rolls the match back, releases the pin, and
-  retries eviction once without the tentative hit.
-* Runtime state availability is checked after KV eviction because old unpinned
-  checkpoints may be dropped to free state-cache slots.  If no runtime slot can
-  be recovered, the tentative prefix hit is rolled back and the request waits.
-* ``state_manager.allocate(seq)`` assigns the request runtime state only after
-  ``block_manager.allocate(seq)`` and ``block_trie.allocate(seq)`` succeed.
-  Later, ``InputsMaker`` may reserve checkpoint saves for the exact produced
-  step; scheduler code does not perform state-cache tensor copies or publish
-  checkpoint readiness.
-
-External KV scheduling detail:
-
-* External lookup is enabled only for a KV consumer with a connector, and is
-  kept separate from the SSM checkpoint path. Local ``BlockTrie.match()`` runs
-  first so the connector searches only beyond KV already resident on this node.
-* Lookup is asynchronous. A pending result must leave the request schedulable
-  for a later tick without retaining a tentative local match, so multi-turn
-  sequence state is snapshotted and restored exactly.
-* A positive hit is block-aligned, allocated, and handed to
-  ``KVLoadCoordinator``. While workers may write those blocks, the sequence is
-  in ``WAITING_FOR_REMOTE_KVS`` and paging cleanup is deferred.
-* A successful load is published into the local trie and prioritized for its
-  remaining prefill. A failed or cancelled load returns to the last safe
-  block-aligned prefix because partially written destinations are untrusted.
-* Prefill saves take a physical block snapshot for workers and a logical block
-  lease for paging. ``KVSaveCoordinator`` keeps those blocks alive until every
-  TP rank reports terminal progress or worker queues are drained.
-"""
-
-import enum
-import time
-from collections import Counter, OrderedDict
+"""Public sequence lifecycle and prefill/decode scheduling facade."""
+
+from collections import OrderedDict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from torch.profiler import record_function
 
-from lmdeploy.messages import EventType, ScheduleMetrics
-from lmdeploy.pytorch import envs as _envs
-from lmdeploy.pytorch.long_context import get_long_context_chunk_limit, plan_long_context_chunk
-from lmdeploy.utils import get_logger
+from lmdeploy.messages import ScheduleMetrics
 
 from ..config import CacheConfig, SchedulerConfig
 from ..messages import MessageStatus, SchedulerSequence, SchedulerSession, SequenceManager, SequenceMeta
 from .block_manager import build_block_manager
 from .block_trie import BlockTrie
 from .eviction_helper import build_eviction_helper
-from .kv_load_coordinator import KVLoadAdmission, KVLoadCoordinator
+from .kv_load_coordinator import KVLoadCoordinator
 from .kv_save_coordinator import KVSaveCoordinator
-from .state_manager import StateManager, build_state_manager
+from .prefill_scheduler import _PrefillScheduler
+from .state_manager import build_state_manager
 
 if TYPE_CHECKING:
     from lmdeploy.pytorch.kv_connector.base import KVConnectorBase
 
-    from .block_manager.base_block_manager import BaseBlockManager
-    from .eviction_helper.base_eviction_helper import BaseEvictionHelper
-
-logger = get_logger('lmdeploy')
-
 MapType = dict[int, int]
 SeqList = list[SchedulerSequence]
 
@@ -111,902 +48,6 @@ class SchedulerOutput:
     connector_logical_block_ids: tuple[tuple[int, ...], ...] = ()
 
 
-@dataclass(frozen=True)
-class _PrefillReorderInfo:
-    """Immutable pre-admission metadata used only for waiting-list ordering."""
-
-    prefill_token_count: int
-    is_nonfinal_long_prefill: bool
-    estimated_long_chunks: int
-
-
-class _PrefillReorderer:
-    """Order waiting prefills without applying scheduler side effects."""
-
-    def __init__(self, prefill_scheduler: '_PrefillScheduler'):
-        self.prefill_scheduler = prefill_scheduler
-        self._info_cache: dict[int, _PrefillReorderInfo] = {}
-
-    def reorder(self,
-                waiting: SeqList,
-                allow_long_prefill: bool,
-                prefer_long_prefill: bool):
-        """Return waiting requests in the order the prefill loop should try."""
-        waiting = sorted(waiting, key=lambda seq: seq.arrive_time)
-        # A completed load already owns destination blocks and a soft prefill
-        # reservation. Admit it first to shorten that ownership window.
-        remote_ready = [
-            seq for seq in waiting
-            if self.prefill_scheduler.load_coordinator.is_remote_ready(seq)
-        ]
-        waiting = [
-            seq for seq in waiting
-            if not self.prefill_scheduler.load_coordinator.is_remote_ready(seq)
-        ]
-        if prefer_long_prefill:
-            # Long-work turns choose one long waiter first. The size policy only
-            # reorders this long lane; it is not global shortest-prefill-first
-            # admission.
-            long_turn_order = self._reorder_for_long_turn(waiting)
-            if long_turn_order is not None:
-                return remote_ready + self._warn_if_not_permutation(waiting, long_turn_order)
-
-        if allow_long_prefill:
-            return remote_ready + self._warn_if_not_permutation(waiting, waiting)
-
-        reordered = self._reorder_for_short_turn(waiting)
-        return remote_ready + self._warn_if_not_permutation(waiting, reordered)
-
-    def _warn_if_not_permutation(self, original: SeqList, reordered: SeqList):
-        """Warn if reorder drops, duplicates, or substitutes waiting
-        sequences."""
-        original_ids = [id(seq) for seq in original]
-        reordered_ids = [id(seq) for seq in reordered]
-        if len(original_ids) == len(reordered_ids) and Counter(original_ids) == Counter(reordered_ids):
-            return reordered
-
-        logger.warning('Unexpected prefill reorder result: original_len=%s reordered_len=%s '
-                       'original_sample=%s reordered_sample=%s',
-                       len(original), len(reordered), self._seq_id_sample(original), self._seq_id_sample(reordered))
-        return reordered
-
-    @staticmethod
-    def _seq_id_sample(seqs: SeqList):
-        return [(seq.session_id, seq.seq_id) for seq in seqs[:5]]
-
-    def _get_reorder_info(self, seq: SchedulerSequence):
-        """Return reorder-only info before prefix-cache side effects.
-
-        Prefix-cache match/rollback mutates the remaining prompt. Keep this cache confined to waiting-list ordering and
-        recompute fresh values in the admission path.
-        """
-        seq_key = id(seq)
-        info = self._info_cache.get(seq_key)
-        if info is not None:
-            return info
-
-        prefill = self.prefill_scheduler
-        chunk_limit = prefill._long_context_chunk_limit(seq)
-        if seq.num_token_ids <= chunk_limit:
-            info = _PrefillReorderInfo(prefill_token_count=seq.num_token_ids,
-                                       is_nonfinal_long_prefill=False,
-                                       estimated_long_chunks=1)
-        else:
-            kv_token_limit = prefill._next_long_context_chunk_end(seq, chunk_limit)
-            safe_chunk_limit = max(1, chunk_limit)
-            info = _PrefillReorderInfo(
-                prefill_token_count=max(0, kv_token_limit - seq.num_history_ids),
-                is_nonfinal_long_prefill=True,
-                estimated_long_chunks=max(1, (seq.num_token_ids + safe_chunk_limit - 1) // safe_chunk_limit),
-            )
-        self._info_cache[seq_key] = info
-        return info
-
-    def _long_priority_key(self, seq: SchedulerSequence, now: float):
-        """Prefer smaller long prompts, with age credit to avoid starvation."""
-        prefill = self.prefill_scheduler
-        info = self._get_reorder_info(seq)
-        wait_age = max(0.0, now - seq.arrive_time)
-        age_credit = int(wait_age // prefill._long_prefill_aging_seconds_per_chunk)
-        age_adjusted_chunks = info.estimated_long_chunks - age_credit
-        return age_adjusted_chunks, info.estimated_long_chunks, seq.arrive_time
-
-    def _split_by_prefill_kind(self, waiting: SeqList):
-        """Split waiting requests into normal/final and non-final long
-        prefill."""
-        normal_waiting: SeqList = []
-        long_waiting: SeqList = []
-        for seq in waiting:
-            if self._get_reorder_info(seq).is_nonfinal_long_prefill:
-                long_waiting.append(seq)
-            else:
-                normal_waiting.append(seq)
-        return normal_waiting, long_waiting
-
-    def _sort_normal_prefills(self, waiting: SeqList):
-        return sorted(waiting,
-                      key=lambda seq: (self._get_reorder_info(seq).prefill_token_count, seq.arrive_time))
-
-    def _sort_long_prefills(self, waiting: SeqList):
-        prefill = self.prefill_scheduler
-        if prefill._long_prefill_policy != 'size':
-            return waiting
-        now = time.perf_counter()
-        return sorted(waiting, key=lambda seq: self._long_priority_key(seq, now))
-
-    def _reorder_for_long_turn(self, waiting: SeqList):
-        """Choose one long waiter, then fill the turn with normal prefills."""
-        normal_waiting, long_waiting = self._split_by_prefill_kind(waiting)
-        if len(long_waiting) == 0:
-            return None
-
-        long_waiting = self._sort_long_prefills(long_waiting)
-        normal_waiting = self._sort_normal_prefills(normal_waiting)
-        return [long_waiting[0]] + normal_waiting + long_waiting[1:]
-
-    def _reorder_for_short_turn(self, waiting: SeqList):
-        """Prioritize normal/final prefills while preserving long waiters."""
-        normal_waiting, long_waiting = self._split_by_prefill_kind(waiting)
-        return self._sort_normal_prefills(normal_waiting) + long_waiting
-
-
-class _PrefillAdmissionAction(enum.Enum):
-    ADMIT = enum.auto()
-    SKIP = enum.auto()
-    STOP = enum.auto()
-    LOAD_STARTED = enum.auto()
-
-
-@dataclass(frozen=True)
-class _PrefillAdmissionResult:
-    """Outcome from trying to admit one waiting prefill request.
-
-    The outer loop distinguishes four outcomes:
-
-    * ``ADMIT``: include the request in this tick's model batch.
-    * ``SKIP``: leave it waiting but continue trying later candidates.
-    * ``STOP``: resource pressure ends this prefill admission turn.
-    * ``LOAD_STARTED``: no model work was selected, but the request left the
-      waiting queue for asynchronous KV load.
-    """
-
-    action: _PrefillAdmissionAction
-    prefill_token_count: int = 0
-
-    @classmethod
-    def admit(cls, prefill_token_count: int):
-        return cls(action=_PrefillAdmissionAction.ADMIT,
-                   prefill_token_count=prefill_token_count)
-
-    @classmethod
-    def skip(cls):
-        return cls(action=_PrefillAdmissionAction.SKIP)
-
-    @classmethod
-    def stop(cls):
-        return cls(action=_PrefillAdmissionAction.STOP)
-
-    @classmethod
-    def load(cls):
-        return cls(action=_PrefillAdmissionAction.LOAD_STARTED)
-
-
-@dataclass(frozen=True, slots=True)
-class _PrefixMatchStateSnapshot:
-    """Exact sequence state captured before a tentative local trie match.
-
-    External lookup itself does not mutate sequence paging state. The scheduler
-    may, however, run ``block_trie.match`` first so the connector queries only
-    beyond the locally resident prefix. If that non-blocking lookup returns
-    pending, or a positive hit cannot be admitted before worker writes start,
-    the request will not run this tick and the tentative local match must be
-    undone.
-
-    A multi-turn request may already own valid history, blocks, and model
-    metadata before this attempt. Restoring this baseline preserves that exact
-    committed state; the legacy new-request rollback to step zero would discard
-    it. This snapshot is not used after an asynchronous load starts--load
-    failure then rolls back to its block-aligned ``fallback_step`` through
-    ``KVLoadCoordinator`` because workers may have partially written KV.
-    """
-
-    # Committed sequence progress and block ownership before tentative match.
-    num_history_ids: int
-    num_blocks: int
-    # Prefix-cache cursor, public hit accounting, and temporary overlap state.
-    trie_cursor: Any
-    match_start_step: int
-    cached_tokens: int
-    # Request-local allocation limit that a multi-turn attempt may carry.
-    kv_token_limit: int | None
-    # Temporary recompute-overlap identities created by local trie matching.
-    fresh_block_range: range | None
-    trie_block_map: dict[int, int]
-    # Model state that must remain aligned with the committed history step.
-    model_meta: Any
-
-    @classmethod
-    def capture(cls, seq: SchedulerSequence):
-        overlap = seq.prefix_cache.recompute_overlap
-        return cls(
-            num_history_ids=int(seq.num_history_ids),
-            num_blocks=int(seq.num_blocks),
-            trie_cursor=seq.prefix_cache.trie_cursor,
-            match_start_step=int(seq.prefix_cache.match_start_step),
-            cached_tokens=int(seq.cached_tokens),
-            kv_token_limit=seq.kv_token_limit,
-            fresh_block_range=overlap.fresh_block_range,
-            trie_block_map=dict(overlap.trie_block_map),
-            model_meta=seq.model_meta,
-        )
-
-
-class _TentativePrefixMatch:
-    """Request-local transaction around ``BlockTrie.match`` side effects.
-
-    Ordinary and SSM admission preserve the historical fallback to an unmatched request. External lookup instead needs
-    an exact pre-match snapshot because a multi-turn request may already own committed progress. Both contracts share
-    one stats snapshot, restore-pin boundary, and explicit commit/rollback lifecycle without changing their rollback
-    semantics.
-    """
-
-    __slots__ = (
-        'seq',
-        'block_trie',
-        'block_manager',
-        'is_ssm',
-        '_preserve_existing_state',
-        '_stats_snapshot',
-        '_state_snapshot',
-        '_rejection_on_rollback',
-        '_started',
-        'matched',
-    )
-
-    def __init__(self,
-                 seq: SchedulerSequence,
-                 block_trie: BlockTrie,
-                 block_manager,
-                 *,
-                 is_ssm: bool,
-                 preserve_existing_state: bool):
-        self.seq = seq
-        self.block_trie = block_trie
-        self.block_manager = block_manager
-        self.is_ssm = is_ssm
-        self._preserve_existing_state = preserve_existing_state
-        self._stats_snapshot = None
-        self._state_snapshot: _PrefixMatchStateSnapshot | None = None
-        self._rejection_on_rollback: _PrefillAdmissionResult | None = None
-        self._started = False
-        self.matched = False
-
-    def begin(self) -> None:
-        """Start the transaction before gates can mutate exact external state.
-
-        Ordinary admission starts lazily from ``match``. External admission
-        starts before gates so rollback can restore existing request state even
-        when a private partial block prevents another trie match.
-        """
-        if self._started or not self.block_trie.enabled:
-            return
-        self._stats_snapshot = self.block_trie.stats.snapshot()
-        if self._preserve_existing_state:
-            self._state_snapshot = _PrefixMatchStateSnapshot.capture(self.seq)
-        self._started = True
-
-    def match(self) -> None:
-        """Apply one tentative match after capturing its rollback boundary."""
-        assert not self.matched
-        self.begin()
-        self.block_trie.match(self.seq)
-        self.matched = True
-
-    def retain_for_admission(self, rejection_on_rollback: _PrefillAdmissionResult) -> None:
-        """Keep a gate-enabling match and remember its original rejection."""
-        assert self.matched
-        self._rejection_on_rollback = rejection_on_rollback
-
-    def pin_restore(self) -> bool:
-        """Pin an SSM restore selected by this tentative match."""
-        restore = self.seq.prefix_cache.restore
-        if not self.is_ssm or not restore.is_selected:
-            return True
-        return self.block_trie.state_checkpoints.pin_restore(self.seq)
-
-    def commit(self) -> None:
-        """Accept the match and discard request-local rollback state."""
-        self._clear()
-
-    def rollback(self, reason: str):
-        """Undo the transaction and return any gate-defined rejection."""
-        rejection = self._rejection_on_rollback
-        if not self._started:
-            return rejection
-
-        seq = self.seq
-        logger.debug('Rollback tentative prefix-cache match: session_id=%s seq_id=%s reason=%s '
-                     'num_history_ids=%s restore_state=%s', seq.session_id, seq.seq_id, reason, seq.num_history_ids,
-                     seq.prefix_cache.restore.slot)
-        self.block_trie.stats.restore(self._stats_snapshot)
-        snapshot = self._state_snapshot
-        if snapshot is None:
-            self._reset_to_unmatched()
-        else:
-            self._restore_snapshot(snapshot)
-        self._clear()
-        return rejection
-
-    def _restore_snapshot(self, snapshot: _PrefixMatchStateSnapshot) -> None:
-        seq = self.seq
-        if seq.num_blocks < snapshot.num_blocks:
-            raise RuntimeError(
-                'tentative prefix match removed sequence-owned baseline blocks')
-        if seq.num_blocks > snapshot.num_blocks:
-            self.block_manager.truncate(seq, snapshot.num_blocks)
-        seq.set_step(snapshot.num_history_ids)
-        seq.model_meta = snapshot.model_meta
-        seq.kv_token_limit = snapshot.kv_token_limit
-        prefix_cache = seq.prefix_cache
-        prefix_cache.trie_cursor = snapshot.trie_cursor
-        prefix_cache.match_start_step = snapshot.match_start_step
-        overlap = prefix_cache.recompute_overlap
-        overlap.fresh_block_range = snapshot.fresh_block_range
-        overlap.trie_block_map.clear()
-        overlap.trie_block_map.update(snapshot.trie_block_map)
-        seq.cached_tokens = snapshot.cached_tokens
-
-    def _reset_to_unmatched(self) -> None:
-        seq = self.seq
-        if self.is_ssm:
-            self.block_trie.state_checkpoints.unpin_restore(seq)
-        if seq.num_blocks > 0 or seq.logical_state >= 0:
-            seq.state.free()
-        elif seq.num_history_ids > 0:
-            seq.set_step(0)
-        seq.kv_token_limit = None
-        prefix_cache = seq.prefix_cache
-        prefix_cache.trie_cursor = None
-        prefix_cache.restore.clear()
-        prefix_cache.match_start_step = -1
-        prefix_cache.recompute_overlap.clear_tracking()
-        seq.cached_tokens = 0
-
-    def _clear(self) -> None:
-        self._stats_snapshot = None
-        self._state_snapshot = None
-        self._rejection_on_rollback = None
-        self._started = False
-        self.matched = False
-
-
-class _PrefillAdmissionAttempt:
-    """Try to admit one waiting prefill sequence.
-
-    The attempt owns all tentative prefix-cache side effects for the sequence:
-    match, SSM restore pinning, eviction, runtime-state checks, allocation, and
-    rollback. The outer prefill loop still owns queue traversal and decides
-    whether a rejected candidate is skipped or ends the current prefill turn.
-    """
-
-    def __init__(self,
-                 prefill_scheduler: '_PrefillScheduler',
-                 seq: SchedulerSequence,
-                 hanging: SeqList,
-                 evictable_waiting: SeqList,
-                 prealloc_size: int,
-                 token_count: int,
-                 has_admitted: bool,
-                 allow_long_prefill: bool):
-        self.prefill_scheduler = prefill_scheduler
-        self.seq = seq
-        self.hanging = hanging
-        self.evictable_waiting = evictable_waiting
-        self.prealloc_size = prealloc_size
-        self.token_count = token_count
-        self.has_admitted = has_admitted
-        self.load_coordinator = prefill_scheduler.load_coordinator
-        self._remote_ready = self.load_coordinator.is_remote_ready(seq)
-        self.allow_long_prefill = allow_long_prefill
-        self._alloc_size = prealloc_size
-        self._prefix_match = _TentativePrefixMatch(
-            seq,
-            prefill_scheduler.block_trie,
-            prefill_scheduler.block_manager,
-            is_ssm=prefill_scheduler.is_ssm,
-            preserve_existing_state=(
-                self.load_coordinator.lookup_enabled and not self._remote_ready),
-        )
-
-    def run(self):
-        """Run the admission route for one waiting prefill.
-
-        1. If a previous external lookup is pending, skip without applying new
-           local prefix-cache side effects.
-        2. Snapshot multi-turn state before a local match may become tentative.
-        3. Apply long-prefill and token-budget gates.
-        4. Prefer a local trie hit, then query/load only its remote extension.
-        5. Admit KV/state resources or roll the tentative match back precisely.
-        6. On success, allocate blocks/states and publish the accepted hit.
-        """
-        if self.load_coordinator.lookup_enabled and not self._remote_ready:
-            if self._lookup_is_pending():
-                return _PrefillAdmissionResult.skip()
-            self._prefix_match.begin()
-
-        gate_result = self._check_prefill_admission_gates()
-        if gate_result is not None:
-            return gate_result
-
-        resource_result = self._admit_resources()
-        if resource_result is not None:
-            return resource_result
-
-        return self._finish_admission()
-
-    def _lookup_is_pending(self) -> bool:
-        """Skip without touching local prefix state while lookup is running.
-
-        The connector owns the Future and deduplicates polls. Marking this turn lets EngineLoop use a short I/O poll
-        delay instead of diagnosing an empty batch as GPU-cache pressure.
-        """
-        prefill = self.prefill_scheduler
-        if not self.load_coordinator.is_lookup_pending(self.seq):
-            return False
-        prefill.last_schedule_had_pending_lookup = True
-        return True
-
-    def _admit_resources(self):
-        if self.prefill_scheduler.block_trie.enabled:
-            return self._admit_prefix_cache_resources()
-        if self.load_coordinator.lookup_enabled:
-            lookup_result = self._query_external_prefix()
-            if lookup_result is not None:
-                return lookup_result
-        if not self._prepare_and_evict():
-            return _PrefillAdmissionResult.stop()
-        return None
-
-    def _admit_prefix_cache_resources(self):
-        """Admit resources for prefix-cache scheduling.
-
-        Route map:
-        1. Use or create the tentative prefix-cache match.
-        2. For external consumers, query only beyond that local match.
-        3. Pin any SSM restore state required by the match.
-        4. Prepare allocation limits and evict KV/state resources.
-        5. For SSM, verify a runtime state slot is still available.
-
-        Any failure rolls the tentative match back. A match created only to pass
-        a prefill gate returns that gate's skip/stop result after rollback;
-        normal resource failures keep their local retry/stop behavior here.
-        """
-        prefill = self.prefill_scheduler
-        seq = self.seq
-        if not self._prefix_match.matched:
-            # A completed external load has already published the accepted
-            # prefix interval. Matching again would restart accounting at the
-            # remote step and drop the restored tokens from request metrics.
-            if not self._remote_ready and not self._has_private_local_tail():
-                self._prefix_match.match()
-
-        if self.load_coordinator.lookup_enabled:
-            lookup_result = self._query_external_prefix()
-            if lookup_result is not None:
-                return lookup_result
-
-        had_ssm_restore = prefill.is_ssm and seq.prefix_cache.restore.is_selected
-        if not self._prefix_match.pin_restore():
-            result = self._prefix_match.rollback(
-                'failed to pin SSM restore checkpoint')
-            if result is not None:
-                return result
-
-        if not self._prepare_and_evict():
-            if not had_ssm_restore:
-                result = self._prefix_match.rollback('eviction failed')
-                if result is not None:
-                    return result
-                return _PrefillAdmissionResult.stop()
-
-            # A matched SSM restore may be pinning the only checkpoint state
-            # that eviction would otherwise free. Roll it back once and retry
-            # eviction before declaring the sequence unschedulable.
-            result = self._prefix_match.rollback(
-                'eviction failed with pinned SSM restore')
-            if result is not None:
-                return result
-            if not self._prepare_and_evict():
-                return _PrefillAdmissionResult.stop()
-
-        if prefill.is_ssm and not prefill._ensure_runtime_state_available():
-            result = self._prefix_match.rollback(
-                'no runtime SSM state available')
-            if result is not None:
-                return result
-            if not self._prepare_and_evict():
-                return _PrefillAdmissionResult.stop()
-            if not prefill._ensure_runtime_state_available():
-                seq.kv_token_limit = None
-                return _PrefillAdmissionResult.stop()
-
-        return None
-
-    def _query_external_prefix(self):
-        """Map connector/paging admission to prefill queue policy."""
-        prefill = self.prefill_scheduler
-        if self._remote_ready:
-            return None
-        admission = self.load_coordinator.try_load(
-            self.seq,
-            prealloc_size=self.prealloc_size,
-            evictable_seqs=self._evictable_sequences(),
-        )
-        if admission is KVLoadAdmission.NO_LOAD:
-            return None
-        if admission is KVLoadAdmission.PENDING:
-            self._prefix_match.rollback('external lookup pending')
-            prefill.last_schedule_had_pending_lookup = True
-            return _PrefillAdmissionResult.skip()
-        if admission is KVLoadAdmission.STARTED:
-            self._prefix_match.commit()
-            return _PrefillAdmissionResult.load()
-        if admission is KVLoadAdmission.FULL_PREFILL_UNAVAILABLE:
-            reason = 'full prefill capacity unavailable'
-        else:
-            assert admission is KVLoadAdmission.SOFT_BUDGET_UNAVAILABLE
-            reason = 'soft prefill budget unavailable'
-        # No worker has seen a destination on rejected admission, so the
-        # request-local prefix transaction remains exactly reversible.
-        self._prefix_match.rollback(reason)
-        return _PrefillAdmissionResult.stop()
-
-    def _evictable_sequences(self):
-        """Iterate queue-owned eviction candidates in historical order."""
-        yield from reversed(self.hanging)
-        yield from reversed(self.evictable_waiting)
-
-    def _match_prefix_for_prefill_gate(self):
-        """Tentatively match once so a request can be rechecked by a gate."""
-        prefill = self.prefill_scheduler
-        if (self._remote_ready or not prefill.block_trie.enabled
-                or self._has_private_local_tail()):
-            return None
-        self._prefix_match.match()
-        return True
-
-    def _has_private_local_tail(self) -> bool:
-        """Whether blocks exist beyond the full-block part of local history.
-
-        ``num_history_ids // block_size`` counts the completely computed
-        blocks before the current step.  A larger ``num_blocks`` means that the
-        sequence also owns the block containing a non-aligned current step, or
-        blocks preallocated after it.  Those blocks are private to this
-        sequence because their KV is partial or not computed yet, so they
-        cannot be published as complete reusable trie blocks.
-
-        For example, with block size 4, a chunked prefill may stop at step 5
-        with block table ``[P0, P1]``::
-
-            P0 -> tokens [0, 4), complete
-            P1 -> tokens [4, 8), only the KV at token 4 is valid
-
-        The trie cursor is at step 4 while ``P1`` already occupies logical
-        block index 1.  If another ``block_trie.match`` finds a shared block
-        ``S1`` for tokens [4, 8), matching appends it after ``P1`` instead of
-        filling ``P1``.  The resulting table ``[P0, P1, S1]`` is misaligned:
-        ``S1`` describes logical block 1 but resides at block-table index 2.
-
-        External lookup starts at the exact local step 5, but a block-granular
-        transfer rounds its start down to step 4.  It must therefore reuse
-        ``P1`` at index 1 as the first destination and overwrite the incomplete
-        KV there.  Skipping trie rematch keeps that destination stable until
-        the load is bound.  Without this guard, lookup may start from an
-        incorrectly advanced step or the load/model may address a block table
-        whose logical token ranges no longer match its indices.
-
-        This state means that a sequence retains local progress across
-        scheduling attempts.  It can result from chunked prefill, repeated
-        model forwards, preemption/resume, or a continued chat session; it is
-        not specific to multi-turn conversation.
-        """
-        if not self.load_coordinator.lookup_enabled:
-            return False
-        seq = self.seq
-        return seq.num_blocks > int(seq.num_history_ids) // seq.block_size
-
-    def _token_budget_rejection(self):
-        if self.allow_long_prefill:
-            return _PrefillAdmissionResult.stop()
-        return _PrefillAdmissionResult.skip()
-
-    def _check_prefill_admission_gates(self):
-        """Apply prefill gates, tentatively matching only when it may help."""
-        prefill = self.prefill_scheduler
-        seq = self.seq
-        token_budget = prefill.cache_config.max_prefill_token_num
-        prefill_token_count = prefill._prefill_admission_token_count(seq)
-        is_nonfinal_long_prefill = prefill._prefill_kv_token_limit(seq) is not None
-
-        if is_nonfinal_long_prefill and not self.allow_long_prefill:
-            matched = self._match_prefix_for_prefill_gate()
-            if matched is None:
-                return _PrefillAdmissionResult.skip()
-            if prefill._prefill_kv_token_limit(seq) is not None:
-                self._prefix_match.rollback('still non-final long prefill on short turn')
-                return _PrefillAdmissionResult.skip()
-            self._prefix_match.retain_for_admission(
-                _PrefillAdmissionResult.skip())
-            prefill_token_count = prefill._prefill_admission_token_count(seq)
-
-        exceeds_token_budget = self.has_admitted and self.token_count + prefill_token_count > token_budget
-        if not exceeds_token_budget:
-            return None
-
-        if not self._prefix_match.matched:
-            matched = self._match_prefix_for_prefill_gate()
-            if matched is not None:
-                prefill_token_count = prefill._prefill_admission_token_count(seq)
-                if self.token_count + prefill_token_count <= token_budget:
-                    self._prefix_match.retain_for_admission(
-                        self._token_budget_rejection())
-                    return None
-                self._prefix_match.rollback('still exceeds prefill token budget')
-        else:
-            self._prefix_match.rollback('still exceeds prefill token budget')
-        return self._token_budget_rejection()
-
-    def _prepare_and_evict(self):
-        """Apply chunk allocation limits and evict for this prefill."""
-        prefill = self.prefill_scheduler
-        seq = self.seq
-        alloc_size = prefill._prepare_prefill_allocation(seq, self.prealloc_size)
-        self._alloc_size = alloc_size
-        if self._evict_for_seq(alloc_size):
-            return True
-        seq.kv_token_limit = None
-        return False
-
-    def _evict_for_seq(self, alloc_size: int):
-        """Evict stopped or skipped waiters until this sequence can run."""
-        prefill = self.prefill_scheduler
-        return prefill.eviction_helper.evict_for_seq(
-            self.seq,
-            list(self._evictable_sequences()),
-            alloc_size,
-        )
-
-    def _finish_admission(self):
-        prefill = self.prefill_scheduler
-        seq = self.seq
-        # Prefix-cache matching can advance the sequence step and shrink the
-        # remaining prefill tail. Charge the admitted batch with the
-        # post-match/post-rollback cost, not the conservative pre-match
-        # estimate used to decide whether this sequence is worth trying.
-        prefill_token_count = prefill._prefill_admission_token_count(seq)
-        prefill.block_manager.allocate(seq, self._alloc_size)
-        if prefill.block_trie.enabled:
-            prefill.block_trie.allocate(seq)
-        if prefill.is_ssm:
-            prefill.state_manager.allocate(seq)
-        if prefill.block_trie.enabled:
-            prefill.block_trie.finalize_match(seq)
-        self.load_coordinator.track_prefill(
-            seq,
-            prealloc_size=self.prealloc_size,
-        )
-        if self._remote_ready:
-            # Preserve the load record through the remaining prefill so its
-            # reservation can be released only after model output advances the
-            # sequence to input_end_pos.
-            self.load_coordinator.mark_scheduled(seq)
-        self._prefix_match.commit()
-        return _PrefillAdmissionResult.admit(prefill_token_count)
-
-
-class _PrefillScheduler:
-    """Own prefill ordering, admission, and long-context reservation.
-
-    Long-lived dependencies are the resource owners used by prefill. Queue
-    contents and active-batch counts remain request-local inputs supplied by
-    the public :class:`Scheduler` facade for each scheduling turn.
-    """
-
-    def __init__(
-        self,
-        scheduler_config: SchedulerConfig,
-        cache_config: CacheConfig,
-        *,
-        is_ssm: bool,
-        block_manager: 'BaseBlockManager',
-        block_trie: BlockTrie,
-        state_manager: StateManager,
-        eviction_helper: 'BaseEvictionHelper',
-        load_coordinator: KVLoadCoordinator,
-    ) -> None:
-        self.scheduler_config = scheduler_config
-        self.cache_config = cache_config
-        self.is_ssm = is_ssm
-        self.block_manager = block_manager
-        self.block_trie = block_trie
-        self.state_manager = state_manager
-        self.eviction_helper = eviction_helper
-        self.load_coordinator = load_coordinator
-        self.last_schedule_had_pending_lookup = False
-        self._long_prefill_policy = _envs.opt_ttft_policy
-        self._long_prefill_aging_seconds_per_chunk = max(
-            0.001,
-            _envs.opt_ttft_aging_sec,
-        )
-
-    def _ensure_runtime_state_available(self):
-        """Make one state-cache slot available for an SSM runtime state."""
-        if not self.is_ssm:
-            return True
-        if self.state_manager.get_num_free_runtime() > 0:
-            return True
-        self.block_trie.state_checkpoints.evict(1)
-        return self.state_manager.get_num_free_runtime() > 0
-
-    def _long_context_chunk_limit(self, seq: SchedulerSequence):
-        """Return the token budget for one long-context chunk."""
-        return get_long_context_chunk_limit(
-            seq,
-            self.cache_config.max_prefill_token_num,
-        )
-
-    def _next_long_context_chunk_end(
-        self,
-        seq: SchedulerSequence,
-        max_prefill_num: int | None = None,
-    ):
-        """Return the exclusive absolute token end for the next chunk."""
-        if max_prefill_num is None:
-            max_prefill_num = self._long_context_chunk_limit(seq)
-        plan = plan_long_context_chunk(
-            seq,
-            max_prefill_num,
-            include_multimodals=False,
-        )
-        return plan.chunk_end
-
-    def _prefill_kv_token_limit(self, seq: SchedulerSequence):
-        """Limit KV allocation for a non-final long-context prefill chunk."""
-        max_prefill_num = self._long_context_chunk_limit(seq)
-        if seq.num_token_ids <= max_prefill_num:
-            return None
-        return self._next_long_context_chunk_end(seq, max_prefill_num)
-
-    def _prefill_admission_token_count(self, seq: SchedulerSequence):
-        """Return token budget cost for the next prefill or chunk."""
-        kv_token_limit = self._prefill_kv_token_limit(seq)
-        if kv_token_limit is None:
-            return seq.num_token_ids
-        return max(0, kv_token_limit - seq.num_history_ids)
-
-    def _prepare_prefill_allocation(
-        self,
-        seq: SchedulerSequence,
-        prealloc_size: int,
-    ):
-        """Apply chunk KV limit and return the effective prealloc size."""
-        kv_token_limit = self._prefill_kv_token_limit(seq)
-        if kv_token_limit is None:
-            seq.kv_token_limit = None
-            return prealloc_size
-
-        seq.kv_token_limit = kv_token_limit
-        return 0
-
-    def has_waiting_long_prefill(self, waiting: SeqList):
-        """Whether a waiting request needs a non-final prefill chunk."""
-        return any(
-            self._prefill_kv_token_limit(seq) is not None
-            for seq in waiting
-        )
-
-    def reserve_long_context_chunk(
-        self,
-        seq: SchedulerSequence,
-        *,
-        hanging: SeqList,
-        waiting: SeqList,
-        chunk_size: int,
-        prealloc_size: int = 0,
-        is_last_chunk: bool = False,
-    ):
-        """Reserve KV blocks for the next chunk of a running long prefill."""
-        old_kv_token_limit = seq.kv_token_limit
-        if is_last_chunk:
-            seq.kv_token_limit = None
-        else:
-            seq.kv_token_limit = seq.num_history_ids + chunk_size
-            prealloc_size = 0
-
-        evictable = hanging + waiting
-        if not self.eviction_helper.evict_for_seq(
-            seq,
-            evictable,
-            prealloc_size,
-        ):
-            seq.kv_token_limit = old_kv_token_limit
-            return False
-
-        self.block_manager.allocate(seq, prealloc_size)
-        self.block_trie.allocate(seq)
-        return True
-
-    @record_function('schedule_prefill')
-    def schedule(
-        self,
-        *,
-        waiting: SeqList,
-        hanging: SeqList,
-        num_ready: int,
-        num_running: int,
-        prealloc_size: int = 0,
-        allow_long_prefill: bool = True,
-        prefer_long_prefill: bool = False,
-    ):
-        """Select and activate one prefill batch."""
-        self.last_schedule_had_pending_lookup = False
-        max_batches = self.scheduler_config.max_batches - num_ready - num_running
-        running: SeqList = []
-        token_count = 0
-
-        def _to_running(
-            seq: SchedulerSequence,
-            prefill_token_count: int,
-        ):
-            """Activate an admitted sequence and count its prefill tokens."""
-            seq.state.activate()
-            running.append(seq)
-            nonlocal token_count
-            token_count += prefill_token_count
-
-        if len(running) >= max_batches or len(waiting) == 0:
-            return running
-
-        waiting = _PrefillReorderer(self).reorder(
-            waiting,
-            allow_long_prefill=allow_long_prefill,
-            prefer_long_prefill=prefer_long_prefill,
-        )
-        skipped_waiting: SeqList = []
-        while len(waiting) > 0 and len(running) < max_batches:
-            seq = waiting.pop(0)
-            evictable_waiting = skipped_waiting + waiting
-            admission = _PrefillAdmissionAttempt(
-                self,
-                seq,
-                hanging=hanging,
-                evictable_waiting=evictable_waiting,
-                prealloc_size=prealloc_size,
-                token_count=token_count,
-                has_admitted=len(running) > 0,
-                allow_long_prefill=allow_long_prefill,
-            ).run()
-
-            if admission.action is _PrefillAdmissionAction.LOAD_STARTED:
-                # The request left WAITING for asynchronous load without using
-                # a model-batch slot or prefill token budget.
-                continue
-            if admission.action is _PrefillAdmissionAction.SKIP:
-                skipped_waiting.append(seq)
-                continue
-            if admission.action is _PrefillAdmissionAction.STOP:
-                break
-
-            assert admission.action is _PrefillAdmissionAction.ADMIT
-            _to_running(seq, admission.prefill_token_count)
-            seq.record_event(EventType.SCHEDULED)
-
-            if seq.kv_token_limit is not None:
-                break
-
-        return running
-
-
 class Scheduler:
     """Tools to schedule next step.
 
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index 96dcabd580..8dd31462ff 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-import lmdeploy.pytorch.paging.scheduler as scheduler_module
+import lmdeploy.pytorch.paging.prefill_scheduler as prefill_scheduler_module
 from lmdeploy.messages import KVTransferConfig
 from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
 from lmdeploy.pytorch.disagg.conn.protocol import MigrationProtocol, MigrationRequest
@@ -1924,8 +1924,10 @@ def test_schedule_prefill_prefer_long_admits_oldest_long_waiter_first():
 
 
 def test_scheduler_reads_opt_ttft_env(monkeypatch):
-    monkeypatch.setattr(scheduler_module._envs, 'opt_ttft_policy', 'fifo')
-    monkeypatch.setattr(scheduler_module._envs, 'opt_ttft_aging_sec', 0.25)
+    monkeypatch.setattr(prefill_scheduler_module._envs, 'opt_ttft_policy',
+                        'fifo')
+    monkeypatch.setattr(prefill_scheduler_module._envs, 'opt_ttft_aging_sec',
+                        0.25)
 
     scheduler, _ = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
 

From baf4987c3a258559e1d8dbfbefd6caff54114802 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 16:11:49 +0800
Subject: [PATCH 07/22] refactor: clarify scheduler flow and test ownership

---
 lmdeploy/pytorch/paging/prefill_scheduler.py  |  311 +--
 lmdeploy/pytorch/paging/scheduler.py          |  163 +-
 .../pytorch/paging/test_prefill_scheduler.py  |  635 ++++++
 tests/pytorch/paging/test_scheduler.py        | 1918 +----------------
 .../paging/test_scheduler_kv_transfer.py      |  787 +++++++
 tests/pytorch/paging/test_scheduler_ssm.py    |  393 ++++
 tests/pytorch/paging/test_state_manager.py    |   39 +-
 7 files changed, 2048 insertions(+), 2198 deletions(-)
 create mode 100644 tests/pytorch/paging/test_prefill_scheduler.py
 create mode 100644 tests/pytorch/paging/test_scheduler_kv_transfer.py
 create mode 100644 tests/pytorch/paging/test_scheduler_ssm.py

diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index 8492566e10..ba50a0982c 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -1,58 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-"""Request scheduling and prefix-cache side-effect boundaries.
-
-The scheduler is the first owner of prefix-cache side effects.  In prefill,
-``BlockTrie.match()`` is intentionally called before eviction and allocation so
-the scheduler can account for reused KV/state.  That match is tentative:
-rollback is required if checkpoint pinning, KV eviction, or runtime state
-allocation means the request cannot safely run now.  Long-context suffixes can
-continue chunking from the accepted prefix hit.
-
-Successful prefill scheduling keeps this order:
-
-1. ``block_trie.match(seq)`` mutates sequence state to skip a cached prefix.
-2. eviction and SSM runtime-state availability are checked.
-3. ``block_manager.allocate(seq)`` allocates missing KV blocks.
-4. ``block_trie.allocate(seq)`` publishes newly allocated full blocks.
-5. For SSM, downstream input/model/engine code restores and saves checkpoint
-   states; the scheduler only owns resource decisions and rollback.
-
-SSM scheduling detail:
-
-* ``block_trie.match(seq)`` may find a published checkpoint and record
-  ``seq.prefix_cache.restore`` before the request owns a runtime state.
-  The scheduler must treat that as tentative until KV blocks and one runtime
-  state slot are guaranteed.
-* A matched restore checkpoint can be pinned before eviction so checkpoint LRU
-  cannot free the source slot.  If that pin prevents eviction from finding
-  enough resources, the scheduler rolls the match back, releases the pin, and
-  retries eviction once without the tentative hit.
-* Runtime state availability is checked after KV eviction because old unpinned
-  checkpoints may be dropped to free state-cache slots.  If no runtime slot can
-  be recovered, the tentative prefix hit is rolled back and the request waits.
-* ``state_manager.allocate(seq)`` assigns the request runtime state only after
-  ``block_manager.allocate(seq)`` and ``block_trie.allocate(seq)`` succeed.
-  Later, ``InputsMaker`` may reserve checkpoint saves for the exact produced
-  step; scheduler code does not perform state-cache tensor copies or publish
-  checkpoint readiness.
-
-External KV scheduling detail:
-
-* External lookup is enabled only for a KV consumer with a connector, and is
-  kept separate from the SSM checkpoint path. Local ``BlockTrie.match()`` runs
-  first so the connector searches only beyond KV already resident on this node.
-* Lookup is asynchronous. A pending result must leave the request schedulable
-  for a later tick without retaining a tentative local match, so multi-turn
-  sequence state is snapshotted and restored exactly.
-* A positive hit is block-aligned, allocated, and handed to
-  ``KVLoadCoordinator``. While workers may write those blocks, the sequence is
-  in ``WAITING_FOR_REMOTE_KVS`` and paging cleanup is deferred.
-* A successful load is published into the local trie and prioritized for its
-  remaining prefill. A failed or cancelled load returns to the last safe
-  block-aligned prefix because partially written destinations are untrusted.
-* Prefill saves take a physical block snapshot for workers and a logical block
-  lease for paging. ``KVSaveCoordinator`` keeps those blocks alive until every
-  TP rank reports terminal progress or worker queues are drained.
+"""Prefill ordering and prefix-cache admission.
+
+Each candidate is checked against long-context and token-budget policy before acquiring KV or runtime-state resources.
+Local trie matches remain tentative until those resources are available; rejected attempts restore the sequence's exact
+pre-match state. External loads use a later rollback boundary once a worker may have written block-aligned destinations.
 """
 
 import enum
@@ -259,7 +210,7 @@ def stop(cls):
         return cls(action=_PrefillAdmissionAction.STOP)
 
     @classmethod
-    def load(cls):
+    def load_started(cls):
         return cls(action=_PrefillAdmissionAction.LOAD_STARTED)
 
 
@@ -464,45 +415,37 @@ class _PrefillAdmissionAttempt:
     def __init__(self,
                  prefill_scheduler: '_PrefillScheduler',
                  seq: SchedulerSequence,
-                 hanging: SeqList,
+                 stopped: SeqList,
                  evictable_waiting: SeqList,
                  prealloc_size: int,
-                 token_count: int,
-                 has_admitted: bool,
+                 batch_prefill_tokens: int,
+                 batch_has_prefill: bool,
                  allow_long_prefill: bool):
         self.prefill_scheduler = prefill_scheduler
         self.seq = seq
-        self.hanging = hanging
+        self.stopped = stopped
         self.evictable_waiting = evictable_waiting
         self.prealloc_size = prealloc_size
-        self.token_count = token_count
-        self.has_admitted = has_admitted
+        self.batch_prefill_tokens = batch_prefill_tokens
+        self.batch_has_prefill = batch_has_prefill
         self.load_coordinator = prefill_scheduler.load_coordinator
-        self._remote_ready = self.load_coordinator.is_remote_ready(seq)
+        self._load_ready = self.load_coordinator.is_remote_ready(seq)
         self.allow_long_prefill = allow_long_prefill
-        self._alloc_size = prealloc_size
+        self._effective_prealloc_size = prealloc_size
         self._prefix_match = _TentativePrefixMatch(
             seq,
             prefill_scheduler.block_trie,
             prefill_scheduler.block_manager,
             is_ssm=prefill_scheduler.is_ssm,
             preserve_existing_state=(
-                self.load_coordinator.lookup_enabled and not self._remote_ready),
+                self.load_coordinator.lookup_enabled and not self._load_ready),
         )
 
     def run(self):
-        """Run the admission route for one waiting prefill.
-
-        1. If a previous external lookup is pending, skip without applying new
-           local prefix-cache side effects.
-        2. Snapshot multi-turn state before a local match may become tentative.
-        3. Apply long-prefill and token-budget gates.
-        4. Prefer a local trie hit, then query/load only its remote extension.
-        5. Admit KV/state resources or roll the tentative match back precisely.
-        6. On success, allocate blocks/states and publish the accepted hit.
-        """
-        if self.load_coordinator.lookup_enabled and not self._remote_ready:
-            if self._lookup_is_pending():
+        """Apply policy, acquire resources, and commit one admission."""
+        if self.load_coordinator.lookup_enabled and not self._load_ready:
+            if self.load_coordinator.is_lookup_pending(self.seq):
+                self.prefill_scheduler.last_schedule_had_pending_lookup = True
                 return _PrefillAdmissionResult.skip()
             self._prefix_match.begin()
 
@@ -514,59 +457,42 @@ def run(self):
         if resource_result is not None:
             return resource_result
 
-        return self._finish_admission()
-
-    def _lookup_is_pending(self) -> bool:
-        """Skip without touching local prefix state while lookup is running.
-
-        The connector owns the Future and deduplicates polls. Marking this turn lets EngineLoop use a short I/O poll
-        delay instead of diagnosing an empty batch as GPU-cache pressure.
-        """
-        prefill = self.prefill_scheduler
-        if not self.load_coordinator.is_lookup_pending(self.seq):
-            return False
-        prefill.last_schedule_had_pending_lookup = True
-        return True
+        return self._allocate_and_commit()
 
     def _admit_resources(self):
         if self.prefill_scheduler.block_trie.enabled:
             return self._admit_prefix_cache_resources()
         if self.load_coordinator.lookup_enabled:
-            lookup_result = self._query_external_prefix()
-            if lookup_result is not None:
-                return lookup_result
+            load_result = self._try_external_load()
+            if load_result is not None:
+                return load_result
         if not self._prepare_and_evict():
             return _PrefillAdmissionResult.stop()
         return None
 
     def _admit_prefix_cache_resources(self):
-        """Admit resources for prefix-cache scheduling.
-
-        Route map:
-        1. Use or create the tentative prefix-cache match.
-        2. For external consumers, query only beyond that local match.
-        3. Pin any SSM restore state required by the match.
-        4. Prepare allocation limits and evict KV/state resources.
-        5. For SSM, verify a runtime state slot is still available.
-
-        Any failure rolls the tentative match back. A match created only to pass
-        a prefill gate returns that gate's skip/stop result after rollback;
-        normal resource failures keep their local retry/stop behavior here.
-        """
-        prefill = self.prefill_scheduler
-        seq = self.seq
+        """Resolve the prefix source, then admit its paging resources."""
+        load_result = self._resolve_prefix_source()
+        if load_result is not None:
+            return load_result
+        return self._admit_matched_resources()
+
+    def _resolve_prefix_source(self):
+        """Match local cache first, then try loading its remote extension."""
         if not self._prefix_match.matched:
             # A completed external load has already published the accepted
-            # prefix interval. Matching again would restart accounting at the
-            # remote step and drop the restored tokens from request metrics.
-            if not self._remote_ready and not self._has_private_local_tail():
+            # prefix. Matching again would lose its cached-token accounting.
+            if not self._load_ready and not self._has_private_local_tail():
                 self._prefix_match.match()
 
         if self.load_coordinator.lookup_enabled:
-            lookup_result = self._query_external_prefix()
-            if lookup_result is not None:
-                return lookup_result
+            return self._try_external_load()
+        return None
 
+    def _admit_matched_resources(self):
+        """Pin matched state, evict for KV, and reserve runtime state."""
+        prefill = self.prefill_scheduler
+        seq = self.seq
         had_ssm_restore = prefill.is_ssm and seq.prefix_cache.restore.is_selected
         if not self._prefix_match.pin_restore():
             result = self._prefix_match.rollback(
@@ -574,40 +500,51 @@ def _admit_prefix_cache_resources(self):
             if result is not None:
                 return result
 
-        if not self._prepare_and_evict():
-            if not had_ssm_restore:
-                result = self._prefix_match.rollback('eviction failed')
-                if result is not None:
-                    return result
-                return _PrefillAdmissionResult.stop()
-
-            # A matched SSM restore may be pinning the only checkpoint state
-            # that eviction would otherwise free. Roll it back once and retry
-            # eviction before declaring the sequence unschedulable.
-            result = self._prefix_match.rollback(
-                'eviction failed with pinned SSM restore')
-            if result is not None:
-                return result
-            if not self._prepare_and_evict():
-                return _PrefillAdmissionResult.stop()
+        kv_result = self._admit_kv_resources(had_ssm_restore)
+        if kv_result is not None:
+            return kv_result
+        return self._admit_runtime_state()
 
-        if prefill.is_ssm and not prefill._ensure_runtime_state_available():
-            result = self._prefix_match.rollback(
-                'no runtime SSM state available')
-            if result is not None:
-                return result
-            if not self._prepare_and_evict():
-                return _PrefillAdmissionResult.stop()
-            if not prefill._ensure_runtime_state_available():
-                seq.kv_token_limit = None
-                return _PrefillAdmissionResult.stop()
+    def _admit_kv_resources(self, had_ssm_restore: bool):
+        """Evict for KV, retrying once without a pinned SSM restore."""
+        if self._prepare_and_evict():
+            return None
+
+        reason = 'eviction failed'
+        if had_ssm_restore:
+            reason = 'eviction failed with pinned SSM restore'
+        result = self._prefix_match.rollback(reason)
+        if result is not None:
+            return result
+
+        # The matched restore may pin the only checkpoint state that eviction
+        # can free. Retrying after rollback preserves the unmatched fallback.
+        if had_ssm_restore and self._prepare_and_evict():
+            return None
+        return _PrefillAdmissionResult.stop()
 
+    def _admit_runtime_state(self):
+        """Ensure an SSM runtime slot, retrying after match rollback."""
+        prefill = self.prefill_scheduler
+        seq = self.seq
+        if not prefill.is_ssm or prefill._ensure_runtime_state_available():
+            return None
+
+        result = self._prefix_match.rollback(
+            'no runtime SSM state available')
+        if result is not None:
+            return result
+        if not self._prepare_and_evict():
+            return _PrefillAdmissionResult.stop()
+        if not prefill._ensure_runtime_state_available():
+            seq.kv_token_limit = None
+            return _PrefillAdmissionResult.stop()
         return None
 
-    def _query_external_prefix(self):
+    def _try_external_load(self):
         """Map connector/paging admission to prefill queue policy."""
         prefill = self.prefill_scheduler
-        if self._remote_ready:
+        if self._load_ready:
             return None
         admission = self.load_coordinator.try_load(
             self.seq,
@@ -622,7 +559,7 @@ def _query_external_prefix(self):
             return _PrefillAdmissionResult.skip()
         if admission is KVLoadAdmission.STARTED:
             self._prefix_match.commit()
-            return _PrefillAdmissionResult.load()
+            return _PrefillAdmissionResult.load_started()
         if admission is KVLoadAdmission.FULL_PREFILL_UNAVAILABLE:
             reason = 'full prefill capacity unavailable'
         else:
@@ -635,52 +572,25 @@ def _query_external_prefix(self):
 
     def _evictable_sequences(self):
         """Iterate queue-owned eviction candidates in historical order."""
-        yield from reversed(self.hanging)
+        yield from reversed(self.stopped)
         yield from reversed(self.evictable_waiting)
 
     def _match_prefix_for_prefill_gate(self):
         """Tentatively match once so a request can be rechecked by a gate."""
         prefill = self.prefill_scheduler
-        if (self._remote_ready or not prefill.block_trie.enabled
+        if (self._load_ready or not prefill.block_trie.enabled
                 or self._has_private_local_tail()):
             return None
         self._prefix_match.match()
         return True
 
     def _has_private_local_tail(self) -> bool:
-        """Whether blocks exist beyond the full-block part of local history.
-
-        ``num_history_ids // block_size`` counts the completely computed
-        blocks before the current step.  A larger ``num_blocks`` means that the
-        sequence also owns the block containing a non-aligned current step, or
-        blocks preallocated after it.  Those blocks are private to this
-        sequence because their KV is partial or not computed yet, so they
-        cannot be published as complete reusable trie blocks.
-
-        For example, with block size 4, a chunked prefill may stop at step 5
-        with block table ``[P0, P1]``::
-
-            P0 -> tokens [0, 4), complete
-            P1 -> tokens [4, 8), only the KV at token 4 is valid
-
-        The trie cursor is at step 4 while ``P1`` already occupies logical
-        block index 1.  If another ``block_trie.match`` finds a shared block
-        ``S1`` for tokens [4, 8), matching appends it after ``P1`` instead of
-        filling ``P1``.  The resulting table ``[P0, P1, S1]`` is misaligned:
-        ``S1`` describes logical block 1 but resides at block-table index 2.
-
-        External lookup starts at the exact local step 5, but a block-granular
-        transfer rounds its start down to step 4.  It must therefore reuse
-        ``P1`` at index 1 as the first destination and overwrite the incomplete
-        KV there.  Skipping trie rematch keeps that destination stable until
-        the load is bound.  Without this guard, lookup may start from an
-        incorrectly advanced step or the load/model may address a block table
-        whose logical token ranges no longer match its indices.
-
-        This state means that a sequence retains local progress across
-        scheduling attempts.  It can result from chunked prefill, repeated
-        model forwards, preemption/resume, or a continued chat session; it is
-        not specific to multi-turn conversation.
+        """Whether a private partial block prevents another trie match.
+
+        A block beyond ``history // block_size`` occupies the logical range
+        that another shared match would append, misaligning table positions.
+        Block-granular external load must instead reuse that private boundary
+        block as its first destination.
         """
         if not self.load_coordinator.lookup_enabled:
             return False
@@ -711,7 +621,10 @@ def _check_prefill_admission_gates(self):
                 _PrefillAdmissionResult.skip())
             prefill_token_count = prefill._prefill_admission_token_count(seq)
 
-        exceeds_token_budget = self.has_admitted and self.token_count + prefill_token_count > token_budget
+        exceeds_token_budget = (
+            self.batch_has_prefill
+            and self.batch_prefill_tokens + prefill_token_count > token_budget
+        )
         if not exceeds_token_budget:
             return None
 
@@ -719,7 +632,7 @@ def _check_prefill_admission_gates(self):
             matched = self._match_prefix_for_prefill_gate()
             if matched is not None:
                 prefill_token_count = prefill._prefill_admission_token_count(seq)
-                if self.token_count + prefill_token_count <= token_budget:
+                if self.batch_prefill_tokens + prefill_token_count <= token_budget:
                     self._prefix_match.retain_for_admission(
                         self._token_budget_rejection())
                     return None
@@ -733,7 +646,7 @@ def _prepare_and_evict(self):
         prefill = self.prefill_scheduler
         seq = self.seq
         alloc_size = prefill._prepare_prefill_allocation(seq, self.prealloc_size)
-        self._alloc_size = alloc_size
+        self._effective_prealloc_size = alloc_size
         if self._evict_for_seq(alloc_size):
             return True
         seq.kv_token_limit = None
@@ -748,7 +661,7 @@ def _evict_for_seq(self, alloc_size: int):
             alloc_size,
         )
 
-    def _finish_admission(self):
+    def _allocate_and_commit(self):
         prefill = self.prefill_scheduler
         seq = self.seq
         # Prefix-cache matching can advance the sequence step and shrink the
@@ -756,7 +669,7 @@ def _finish_admission(self):
         # post-match/post-rollback cost, not the conservative pre-match
         # estimate used to decide whether this sequence is worth trying.
         prefill_token_count = prefill._prefill_admission_token_count(seq)
-        prefill.block_manager.allocate(seq, self._alloc_size)
+        prefill.block_manager.allocate(seq, self._effective_prealloc_size)
         if prefill.block_trie.enabled:
             prefill.block_trie.allocate(seq)
         if prefill.is_ssm:
@@ -767,7 +680,7 @@ def _finish_admission(self):
             seq,
             prealloc_size=self.prealloc_size,
         )
-        if self._remote_ready:
+        if self._load_ready:
             # Preserve the load record through the remaining prefill so its
             # reservation can be released only after model output advances the
             # sequence to input_end_pos.
@@ -881,7 +794,7 @@ def reserve_long_context_chunk(
         self,
         seq: SchedulerSequence,
         *,
-        hanging: SeqList,
+        stopped: SeqList,
         waiting: SeqList,
         chunk_size: int,
         prealloc_size: int = 0,
@@ -895,7 +808,7 @@ def reserve_long_context_chunk(
             seq.kv_token_limit = seq.num_history_ids + chunk_size
             prealloc_size = 0
 
-        evictable = hanging + waiting
+        evictable = stopped + waiting
         if not self.eviction_helper.evict_for_seq(
             seq,
             evictable,
@@ -913,7 +826,7 @@ def schedule(
         self,
         *,
         waiting: SeqList,
-        hanging: SeqList,
+        stopped: SeqList,
         num_ready: int,
         num_running: int,
         prealloc_size: int = 0,
@@ -924,19 +837,9 @@ def schedule(
         self.last_schedule_had_pending_lookup = False
         max_batches = self.scheduler_config.max_batches - num_ready - num_running
         running: SeqList = []
-        token_count = 0
-
-        def _to_running(
-            seq: SchedulerSequence,
-            prefill_token_count: int,
-        ):
-            """Activate an admitted sequence and count its prefill tokens."""
-            seq.state.activate()
-            running.append(seq)
-            nonlocal token_count
-            token_count += prefill_token_count
+        batch_prefill_tokens = 0
 
-        if len(running) >= max_batches or len(waiting) == 0:
+        if max_batches <= 0 or not waiting:
             return running
 
         waiting = _PrefillReorderer(self).reorder(
@@ -945,17 +848,17 @@ def _to_running(
             prefer_long_prefill=prefer_long_prefill,
         )
         skipped_waiting: SeqList = []
-        while len(waiting) > 0 and len(running) < max_batches:
+        while waiting and len(running) < max_batches:
             seq = waiting.pop(0)
             evictable_waiting = skipped_waiting + waiting
             admission = _PrefillAdmissionAttempt(
                 self,
                 seq,
-                hanging=hanging,
+                stopped=stopped,
                 evictable_waiting=evictable_waiting,
                 prealloc_size=prealloc_size,
-                token_count=token_count,
-                has_admitted=len(running) > 0,
+                batch_prefill_tokens=batch_prefill_tokens,
+                batch_has_prefill=bool(running),
                 allow_long_prefill=allow_long_prefill,
             ).run()
 
@@ -970,7 +873,9 @@ def _to_running(
                 break
 
             assert admission.action is _PrefillAdmissionAction.ADMIT
-            _to_running(seq, admission.prefill_token_count)
+            seq.state.activate()
+            running.append(seq)
+            batch_prefill_tokens += admission.prefill_token_count
             seq.record_event(EventType.SCHEDULED)
 
             if seq.kv_token_limit is not None:
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index ffe2bfa08c..2bf2cf73c8 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -1,10 +1,11 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 # modify from: https://github.com/vllm-project/vllm
-"""Public sequence lifecycle and prefill/decode scheduling facade."""
+"""Public paging scheduler and sequence-lifecycle facade."""
 
 from collections import OrderedDict
 from contextlib import contextmanager
 from dataclasses import dataclass
+from itertools import chain
 from typing import TYPE_CHECKING
 
 from torch.profiler import record_function
@@ -30,7 +31,7 @@
 
 @dataclass
 class SchedulerOutput:
-    """Output of schedule."""
+    """Paging selection and connector snapshots for one model step."""
 
     running: SeqList
     swap_in_map: MapType
@@ -49,11 +50,11 @@ class SchedulerOutput:
 
 
 class Scheduler:
-    """Tools to schedule next step.
+    """Coordinate sequence lifecycle and paging-resource admission.
 
     Args:
-        scheduler_config (SchedulerConfig): The config of scheduler.
-        cache_config (CacheConfig): The config of cache info.
+        scheduler_config: Batch and eviction policy.
+        cache_config: KV and state-cache configuration.
     """
 
     def __init__(
@@ -152,7 +153,7 @@ def reserve_long_context_chunk(self,
         """Reserve KV blocks for the next chunk of a running long prefill."""
         return self._prefill_scheduler.reserve_long_context_chunk(
             seq,
-            hanging=self.hanging,
+            stopped=self.hanging,
             waiting=self.waiting,
             chunk_size=chunk_size,
             prealloc_size=prealloc_size,
@@ -227,137 +228,67 @@ def add_session(self, session_id: int):
 
     def _schedule_migration(self):
         migration_ready: SeqList = []
-        migrating_token_count = 0
-
-        def _to_running(seq: SchedulerSequence):
-            """Activate a migrated sequence and count its tokens."""
-            seq.state.activate()
-            migration_ready.append(seq)
-            nonlocal migrating_token_count
-            migrating_token_count += seq.num_token_ids
-
-        def __evict_for_seq(seq: SchedulerSequence, waiting):
-            """Evict until can append."""
-            from itertools import chain
-
-            hanging = reversed(self.hanging)
-            waiting = reversed(waiting)
-            evictable = list(chain(hanging, waiting))
-            return self.eviction_helper.evict_for_seq(seq, evictable, 0)
-
-        def _reorder_migrating():
-            """Reorder waiting."""
-            return sorted(self.migration_waiting, key=lambda seq: seq.arrive_time)
-
-        migration_waiting = _reorder_migrating()
+        migration_waiting = sorted(
+            self.migration_waiting,
+            key=lambda seq: seq.arrive_time,
+        )
 
         max_batches = self.scheduler_config.max_batches - self.num_ready() - self.num_running()
-        while len(migration_waiting) > 0 and len(migration_ready) < max_batches:
+        while migration_waiting and len(migration_ready) < max_batches:
             seq = migration_waiting.pop(0)
             self.block_trie.match(seq)
-            if not __evict_for_seq(seq, migration_waiting):
+            evictable = list(
+                chain(
+                    reversed(self.hanging),
+                    reversed(migration_waiting),
+                ))
+            if not self.eviction_helper.evict_for_seq(seq, evictable, 0):
                 break
 
             # allocate session memory
             self.block_manager.allocate(seq)
             self.block_trie.finalize_match(seq)
-            _to_running(seq)
+            seq.state.activate()
+            migration_ready.append(seq)
 
         return migration_ready
 
-    @record_function('schedule_decoding')
-    def _schedule_decoding(self, prealloc_size: int = 0):
-        """Schedule decoding."""
-
-        def _reorder_running():
-            """Reorder running."""
-            return sorted(self.ready, key=lambda seq: seq.arrive_time)
-
-        running = _reorder_running()
-        assert len(running) != 0
-
-        eviction_helper = self.eviction_helper
-        swap_out_map: MapType = dict()
-        swap_in_map: MapType = dict()
-        copy_map: MapType = dict()
-
-        def __evict_for_seq(seq: SchedulerSequence, num_required_blocks: int):
-            """Evict until can append."""
-            if num_required_blocks == 0:
-                # No need to evict, just return True.
-                return True
-            elif num_required_blocks <= self.block_manager.get_num_free_gpu_blocks():
-                # Enough free blocks, just return True.
-                return True
-
-            from itertools import chain
-            hanging = reversed(self.hanging)
-            waiting = reversed(self.waiting)
-            evictable = list(chain(hanging, waiting))
-            return eviction_helper.evict_for_seq(seq, evictable, prealloc_size)
-
-        # 1. running
-        while len(running) > 0:
-            # token + n
-            seq = running.pop(0)
-            num_required_blocks = self.block_manager.num_required_blocks(seq, prealloc_size)
-            assert seq.num_blocks + num_required_blocks <= self.block_manager.num_gpu_blocks, (
-                'Sequence requires more blocks than total gpu blocks.')
-
-            while not __evict_for_seq(seq, num_required_blocks):
-                if len(running) == 0:
-                    break
-                seq_preempted = running.pop(-1)
-                # Preemption abandons the tracked full-prefill target. Keeping
-                # it would reserve blocks for work no longer admitted.
-                self.kv_load_coordinator.release(seq_preempted)
-                seq_preempted.state.evict()
-
-            if self.block_manager.get_num_free_gpu_blocks() < num_required_blocks:
-                self.kv_load_coordinator.release(seq)
-                seq.state.evict()
-                continue
-
-            self.block_manager.allocate(seq, prealloc_size)
-            self.block_trie.allocate(seq)
-
-        return self.ready[:self.scheduler_config.max_batches], swap_in_map, swap_out_map, copy_map
-
     def schedule(self,
                  is_prefill: bool,
                  prealloc_size: int = 0,
                  allow_long_prefill: bool = True,
                  prefer_long_prefill: bool = False):
-        """Schedule inputs for next steps."""
-        self.last_schedule_had_pending_lookup = False
-        if is_prefill:
-            running = self._prefill_scheduler.schedule(
-                waiting=self.waiting,
-                hanging=self.hanging,
-                num_ready=self.num_ready(),
-                num_running=self.num_running(),
-                prealloc_size=prealloc_size,
-                allow_long_prefill=allow_long_prefill,
-                prefer_long_prefill=prefer_long_prefill,
-            )
-            self.last_schedule_had_pending_lookup = (
-                self._prefill_scheduler.last_schedule_had_pending_lookup)
-            swap_in_map: MapType = {}
-            swap_out_map: MapType = {}
-            copy_map: MapType = {}
-        else:
-            running, swap_in_map, swap_out_map, copy_map = self._schedule_decoding(
-                prealloc_size)
+        """Select the next prefill batch.
 
-        return SchedulerOutput(running=running, swap_in_map=swap_in_map, swap_out_map=swap_out_map, copy_map=copy_map)
+        Decode capacity is admitted by :meth:`schedule_running`.
+        """
+        if not is_prefill:
+            raise ValueError(
+                'schedule only selects prefill work; use schedule_running '
+                'for decode capacity admission')
+
+        self.last_schedule_had_pending_lookup = False
+        running = self._prefill_scheduler.schedule(
+            waiting=self.waiting,
+            stopped=self.hanging,
+            num_ready=self.num_ready(),
+            num_running=self.num_running(),
+            prealloc_size=prealloc_size,
+            allow_long_prefill=allow_long_prefill,
+            prefer_long_prefill=prefer_long_prefill,
+        )
+        self.last_schedule_had_pending_lookup = (
+            self._prefill_scheduler.last_schedule_had_pending_lookup)
+        return SchedulerOutput(
+            running=running,
+            swap_in_map={},
+            swap_out_map={},
+            copy_map={},
+        )
 
     @record_function('schedule_running')
     def schedule_running(self, running: SeqList, num_required_tokens: int = 1, prealloc_size: int = 1):
-        """Schedule running sequences.
-
-        This function is used to add blocks for running sequences request would be marked as invalid if not enough
-        blocks can be allocated.
-        """
+        """Admit KV growth for running sequences and return their validity."""
         assert len(running) > 0
         eviction_helper = self.eviction_helper
 
diff --git a/tests/pytorch/paging/test_prefill_scheduler.py b/tests/pytorch/paging/test_prefill_scheduler.py
new file mode 100644
index 0000000000..2a70cc4a1b
--- /dev/null
+++ b/tests/pytorch/paging/test_prefill_scheduler.py
@@ -0,0 +1,635 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import time
+from unittest.mock import Mock
+
+import torch
+
+import lmdeploy.pytorch.paging.prefill_scheduler as prefill_scheduler_module
+from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
+from lmdeploy.pytorch.messages import MessageStatus, SequenceMeta, UpdateTokenMode
+from lmdeploy.pytorch.paging.scheduler import Scheduler
+
+
+def test_scheduler_publishes_cached_tokens_for_accepted_prefix_hit():
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2] * block_size + [3])
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    seq = scheduler.add_session(1).add_sequence([1] * block_size + [2] * block_size + [4])
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.num_history_ids == block_size * 2
+    assert seq.cached_tokens == block_size * 2
+
+    seq.update_token_ids(torch.tensor([5]))
+
+    assert seq.cached_tokens == 0
+    assert seq.prefix_cache.match_start_step == -1
+
+
+def test_scheduler_ar_spec_prefix_hit_recomputes_overlap_block():
+    from lmdeploy.pytorch.strategies.ar_spec.sequence import ARSpecSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSpecSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    token_ids = [1] * block_size + [2] * block_size + [3] * block_size + [4]
+    cached = scheduler.add_session(0).add_sequence(token_ids)
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached_blocks = cached.logical_blocks.get_real_blocks().copy()
+    cached.state.stop()
+
+    seq = scheduler.add_session(1).add_sequence(token_ids)
+    scheduler.block_trie.stats.reset()
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.prefix_cache.recompute_overlap.recompute_blocks == 1
+    assert seq.num_history_ids == block_size * 2
+    assert seq.cached_tokens == block_size * 2
+    assert seq.logical_blocks[2] != cached_blocks[2]
+    assert seq.prefix_cache.recompute_overlap.fresh_block_range is None
+    assert scheduler.block_trie.stats.num_query_tokens == len(token_ids)
+    assert scheduler.block_trie.stats.num_hit_tokens == block_size * 2
+
+
+def test_scheduler_prefix_match_rollback_clears_recompute_overlap_window(monkeypatch):
+    from lmdeploy.pytorch.strategies.ar_spec.sequence import ARSpecSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSpecSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    token_ids = [1] * block_size + [2] * block_size + [3] * block_size + [4]
+    cached = scheduler.add_session(0).add_sequence(token_ids)
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached.state.stop()
+
+    seq = scheduler.add_session(1).add_sequence(token_ids)
+    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', Mock(return_value=False))
+    scheduler.block_trie.stats.reset()
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.num_history_ids == 0
+    assert seq.num_token_ids == len(token_ids)
+    assert seq.cached_tokens == 0
+    assert seq.prefix_cache.recompute_overlap.fresh_block_range is None
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def test_scheduler_recomputes_prefill_budget_after_prefix_hit():
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=2,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               max_prefill_token_num=block_size,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=2,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2])
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    cache_hit_tail = scheduler.add_session(1).add_sequence([1] * block_size + [3])
+    short = scheduler.add_session(2).add_sequence([4])
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [cache_hit_tail, short]
+    assert cache_hit_tail.num_history_ids == block_size
+    assert cache_hit_tail.num_token_ids == 1
+    assert short.status == MessageStatus.READY
+
+
+def _make_prefix_cache_scheduler(max_batches: int = 2, max_prefill_token_num: int = 16):
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=max_batches,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               max_prefill_token_num=max_prefill_token_num,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=max_batches,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+    return scheduler, block_size
+
+
+def test_scheduler_short_turn_uses_prefix_hit_to_admit_long_looking_sibling():
+    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size)
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    short = scheduler.add_session(1).add_sequence([4])
+    cache_hit_tail = scheduler.add_session(2).add_sequence([1] * block_size + [3])
+
+    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert output.running == [short, cache_hit_tail]
+    assert cache_hit_tail.num_history_ids == block_size
+    assert cache_hit_tail.num_token_ids == 1
+    assert cache_hit_tail.cached_tokens == block_size
+
+
+def test_scheduler_budget_gate_uses_prefix_hit_to_admit_sibling():
+    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size)
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    almost_full = scheduler.add_session(1).add_sequence([4] * (block_size - 1))
+    cache_hit_tail = scheduler.add_session(2).add_sequence([1] * block_size + [3])
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [almost_full, cache_hit_tail]
+    assert cache_hit_tail.num_history_ids == block_size
+    assert cache_hit_tail.num_token_ids == 1
+
+
+def test_scheduler_reorder_cache_stays_order_only_after_prefix_hit():
+    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size)
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    cache_hit_tail = scheduler.add_session(1).add_sequence([1] * block_size + [3])
+    normal = scheduler.add_session(2).add_sequence([4] * (block_size - 1))
+
+    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
+
+    assert output.running == [cache_hit_tail, normal]
+    assert cache_hit_tail.num_history_ids == block_size
+    assert cache_hit_tail.num_token_ids == 1
+    assert cache_hit_tail.cached_tokens == block_size
+    assert normal.status == MessageStatus.READY
+
+
+def test_scheduler_resource_rejection_rolls_back_tentative_prefix_match(monkeypatch):
+    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=1)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2])
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+    cached_block = cached.logical_blocks.get_real_blocks()[:1]
+    ref_count = scheduler.block_manager.allocator.get_ref_count(cached_block).copy()
+    scheduler.block_trie.stats.reset()
+
+    seq = scheduler.add_session(1).add_sequence([1] * block_size + [3])
+    evict_for_seq = Mock(return_value=False)
+    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_history_ids == 0
+    assert seq.num_blocks == 0
+    assert seq.kv_token_limit is None
+    assert seq.cached_tokens == 0
+    assert seq.prefix_cache.trie_cursor is None
+    assert seq.prefix_cache.match_start_step == -1
+    assert evict_for_seq.call_count == 1
+    assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def test_scheduler_rolls_back_prefix_match_for_prefill_gate_when_tail_still_exceeds_budget():
+    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size)
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    full = scheduler.add_session(1).add_sequence([4] * block_size)
+    cache_hit_tail = scheduler.add_session(2).add_sequence([1] * block_size + [3])
+
+    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert output.running == [full]
+    assert cache_hit_tail.status == MessageStatus.WAITING
+    assert cache_hit_tail.num_history_ids == 0
+    assert cache_hit_tail.cached_tokens == 0
+    assert cache_hit_tail.prefix_cache.trie_cursor is None
+    assert cache_hit_tail.prefix_cache.match_start_step == -1
+
+
+def test_scheduler_rolls_back_prefix_match_for_prefill_gate_that_still_needs_long_chunk():
+    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=1, max_prefill_token_num=16)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size)
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+    scheduler.block_trie.stats.reset()
+
+    still_long = scheduler.add_session(1).add_sequence([1] * block_size + [3] * (block_size + 1))
+
+    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert output.running == []
+    assert still_long.status == MessageStatus.WAITING
+    assert still_long.num_history_ids == 0
+    assert still_long.cached_tokens == 0
+    assert still_long.prefix_cache.trie_cursor is None
+    assert still_long.prefix_cache.match_start_step == -1
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def test_scheduler_reports_zero_cached_tokens_for_prefix_miss():
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2])
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    seq = scheduler.add_session(1).add_sequence([3] * block_size + [4])
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.num_history_ids == 0
+    assert seq.cached_tokens == 0
+
+
+def test_scheduler_cached_tokens_only_count_current_prompt_after_session_eviction():
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    session = scheduler.add_session(0)
+    seq = session.add_sequence([1] * block_size + [2] * block_size + [3])
+    scheduler.schedule(is_prefill=True)
+    seq.update_token_ids(torch.tensor([9]), mode=UpdateTokenMode.PREFILL)
+    seq.state.stop()
+    seq.state.free()
+
+    seq.update_token_ids(torch.tensor([4] * 4))
+    assert seq.input_start_pos == block_size * 2 + 2
+    assert seq.input_end_pos == block_size * 2 + 6
+    seq.state.activate()
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.num_history_ids == block_size * 2
+    assert seq.cached_tokens == 0
+
+
+def test_scheduler_excludes_recompute_eviction_prefix_hits_from_stats():
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=4,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    seq = scheduler.add_session(0).add_sequence([1] * block_size + [2] * block_size + [3])
+    output = scheduler.schedule(is_prefill=True)
+    assert output.running == [seq]
+
+    seq.state.evict()
+    pressure = scheduler.add_session(1).add_sequence([9] * block_size * 3)
+    scheduler.block_trie.stats.reset()
+
+    assert scheduler.eviction_helper.evict_for_seq(pressure, [seq], 0)
+    assert seq.prefix_cache.suppress_match_stats
+    pressure.session.remove_sequence(pressure)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.num_history_ids >= block_size
+    assert seq.cached_tokens == 0
+    assert not seq.prefix_cache.suppress_match_stats
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def _make_scheduler_for_long_context_chunks(num_gpu_blocks: int = 6):
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 4
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=2,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=num_gpu_blocks,
+                               max_prefill_token_num=block_size * 2)
+    scheduler_config = SchedulerConfig(max_batches=2,
+                                       max_session_len=64,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+    return scheduler, block_size
+
+
+def test_schedule_prefill_allocates_only_first_long_context_chunk():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=2)
+    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
+
+    output = scheduler.schedule(is_prefill=True, prealloc_size=1)
+
+    assert output.running == [long_seq]
+    assert long_seq.status == MessageStatus.READY
+    assert long_seq.kv_token_limit == block_size * 2
+    assert long_seq.num_blocks == 2
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
+
+
+def test_schedule_prefill_short_only_skips_long_waiter_without_mutation():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
+    head_long = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
+    short_a = scheduler.add_session(101).add_sequence([2] * (block_size // 2))
+    short_b = scheduler.add_session(102).add_sequence([3] * (block_size // 2))
+
+    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert output.running == [short_a, short_b]
+    assert head_long.status == MessageStatus.WAITING
+    assert head_long.num_blocks == 0
+    assert head_long.kv_token_limit is None
+    assert short_a.status == MessageStatus.READY
+    assert short_b.status == MessageStatus.READY
+
+    short_a.session.remove_sequence(short_a)
+    short_b.session.remove_sequence(short_b)
+    next_output = scheduler.schedule(is_prefill=True)
+
+    assert next_output.running == [head_long]
+    assert head_long.status == MessageStatus.READY
+    assert head_long.kv_token_limit == block_size * 2
+    assert head_long.num_blocks == 2
+
+
+def test_schedule_prefill_prefer_long_admits_oldest_long_waiter_first():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
+    short_a = scheduler.add_session(100).add_sequence([1] * (block_size // 2))
+    old_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
+    short_b = scheduler.add_session(102).add_sequence([3] * (block_size // 2))
+    new_long = scheduler.add_session(103).add_sequence([4] * (block_size * 4))
+
+    assert scheduler.has_waiting_long_prefill()
+
+    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
+
+    assert output.running == [old_long]
+    assert old_long.status == MessageStatus.READY
+    assert old_long.kv_token_limit == block_size * 2
+    assert old_long.num_blocks == 2
+    assert short_a.status == MessageStatus.WAITING
+    assert short_a.num_blocks == 0
+    assert short_b.status == MessageStatus.WAITING
+    assert short_b.num_blocks == 0
+    assert new_long.status == MessageStatus.WAITING
+    assert new_long.num_blocks == 0
+    assert new_long.kv_token_limit is None
+
+
+def test_scheduler_reads_opt_ttft_env(monkeypatch):
+    monkeypatch.setattr(prefill_scheduler_module._envs, 'opt_ttft_policy',
+                        'fifo')
+    monkeypatch.setattr(prefill_scheduler_module._envs, 'opt_ttft_aging_sec',
+                        0.25)
+
+    scheduler, _ = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
+
+    assert scheduler._prefill_scheduler._long_prefill_policy == 'fifo'
+    assert scheduler._prefill_scheduler._long_prefill_aging_seconds_per_chunk == 0.25
+
+
+def test_schedule_prefill_prefer_long_fifo_policy_keeps_oldest_huge_waiter_first():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
+    scheduler._prefill_scheduler._long_prefill_policy = 'fifo'
+    now = time.perf_counter()
+    huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
+    huge_long.arrive_time = now - 1.0
+    moderate_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
+    moderate_long.arrive_time = now
+
+    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
+
+    assert output.running == [huge_long]
+    assert huge_long.status == MessageStatus.READY
+    assert huge_long.kv_token_limit == block_size * 2
+    assert huge_long.num_blocks == 2
+    assert moderate_long.status == MessageStatus.WAITING
+    assert moderate_long.num_blocks == 0
+    assert moderate_long.kv_token_limit is None
+
+
+def test_schedule_prefill_prefer_long_admits_smaller_long_waiter_first():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
+    now = time.perf_counter()
+    huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
+    huge_long.arrive_time = now - 1.0
+    moderate_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
+    moderate_long.arrive_time = now
+    short = scheduler.add_session(102).add_sequence([3] * (block_size // 2))
+
+    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
+
+    assert output.running == [moderate_long]
+    assert moderate_long.status == MessageStatus.READY
+    assert moderate_long.kv_token_limit == block_size * 2
+    assert moderate_long.num_blocks == 2
+    assert huge_long.status == MessageStatus.WAITING
+    assert huge_long.num_blocks == 0
+    assert huge_long.kv_token_limit is None
+    assert short.status == MessageStatus.WAITING
+    assert short.num_blocks == 0
+
+
+def test_schedule_prefill_prefer_long_ages_huge_long_waiter():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
+    scheduler._prefill_scheduler._long_prefill_aging_seconds_per_chunk = 0.01
+    now = time.perf_counter()
+    huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
+    huge_long.arrive_time = now - 1.0
+    moderate_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
+    moderate_long.arrive_time = now
+
+    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
+
+    assert output.running == [huge_long]
+    assert huge_long.status == MessageStatus.READY
+    assert huge_long.kv_token_limit == block_size * 2
+    assert huge_long.num_blocks == 2
+    assert moderate_long.status == MessageStatus.WAITING
+    assert moderate_long.num_blocks == 0
+    assert moderate_long.kv_token_limit is None
+
+
+def test_reserve_long_context_chunk_grows_one_chunk_at_a_time():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=6)
+    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 5))
+
+    output = scheduler.schedule(is_prefill=True, prealloc_size=1)
+    assert output.running == [long_seq]
+    assert long_seq.kv_token_limit == block_size * 2
+    assert long_seq.num_blocks == 2
+
+    scheduler.activate_seqs([long_seq])
+    long_seq.set_step(block_size * 2)
+
+    assert scheduler.reserve_long_context_chunk(long_seq, block_size * 2)
+    assert long_seq.status == MessageStatus.RUNNING
+    assert long_seq.kv_token_limit == block_size * 4
+    assert long_seq.num_blocks == 4
+
+    long_seq.set_step(block_size * 4)
+
+    assert scheduler.reserve_long_context_chunk(long_seq, block_size, prealloc_size=1, is_last_chunk=True)
+    assert long_seq.kv_token_limit is None
+    assert long_seq.num_blocks == 6
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
+
+
+def test_reserve_long_context_chunk_failure_preserves_committed_prefix():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=2)
+    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
+
+    output = scheduler.schedule(is_prefill=True)
+    assert output.running == [long_seq]
+    scheduler.activate_seqs([long_seq])
+    long_seq.set_step(block_size * 2)
+
+    assert not scheduler.reserve_long_context_chunk(long_seq, block_size * 2)
+    assert long_seq.status == MessageStatus.RUNNING
+    assert long_seq.kv_token_limit == block_size * 2
+    assert long_seq.num_blocks == 2
+
+
+def test_reserve_last_long_context_chunk_failure_restores_chunk_limit():
+    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=3)
+    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
+
+    output = scheduler.schedule(is_prefill=True)
+    assert output.running == [long_seq]
+    scheduler.activate_seqs([long_seq])
+    long_seq.set_step(block_size * 2)
+
+    assert not scheduler.reserve_long_context_chunk(long_seq,
+                                                    block_size * 2,
+                                                    prealloc_size=1,
+                                                    is_last_chunk=True)
+    assert long_seq.status == MessageStatus.RUNNING
+    assert long_seq.kv_token_limit == block_size * 2
+    assert long_seq.num_blocks == 2
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 1
+
+
+def test_scheduler_accepts_prefix_hit_that_starts_middle_long_context_chunk():
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=8,
+                               max_prefill_token_num=block_size * 2,
+                               enable_prefix_caching=True)
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2] * block_size)
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached.state.stop()
+
+    token_ids = [1] * block_size + [2] * block_size + [3] * block_size
+    token_ids += [4] * block_size + [5] * block_size
+    seq = scheduler.add_session(1).add_sequence(token_ids)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.num_history_ids == block_size * 2
+    assert seq.num_token_ids == len(token_ids) - block_size * 2
+    assert seq.cached_tokens == block_size * 2
+    assert scheduler.block_trie.stats.num_query_tokens == len(token_ids)
+    assert scheduler.block_trie.stats.num_hit_tokens == block_size * 2
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index 8dd31462ff..ca4261f92a 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -1,116 +1,12 @@
-import time
-from unittest.mock import Mock
+# Copyright (c) OpenMMLab. All rights reserved.
 
 import pytest
 import torch
 
-import lmdeploy.pytorch.paging.prefill_scheduler as prefill_scheduler_module
-from lmdeploy.messages import KVTransferConfig
 from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
 from lmdeploy.pytorch.disagg.conn.protocol import MigrationProtocol, MigrationRequest
-from lmdeploy.pytorch.engine.inputs_maker import _make_state_prefix_cache_save_plan
-from lmdeploy.pytorch.kv_connector import (
-    KVConnectorMetadata,
-    KVConnectorOutput,
-    KVConnectorResult,
-    KVLoadResult,
-    KVSaveBlockLease,
-)
-from lmdeploy.pytorch.kv_connector.mooncake.store.scheduler import MooncakeStoreScheduler
-from lmdeploy.pytorch.messages import MessageStatus, SequenceMeta, UpdateTokenMode
+from lmdeploy.pytorch.messages import MessageStatus, SequenceMeta
 from lmdeploy.pytorch.paging.scheduler import Scheduler
-from lmdeploy.pytorch.paging.state_manager import StateManager
-
-
-class _AsyncLookupConnector:
-
-    def __init__(self, results, failed_ids=()):
-        self.results = iter(results)
-        self.failed_ids = set(failed_ids)
-        self.pending_ids = set()
-        self.lookup_calls = []
-        self.cancelled = []
-        self.finished = []
-        self.allocations = []
-
-    def on_new_request(self, request):
-        pass
-
-    def is_lookup_pending(self, request_id):
-        return request_id in self.pending_ids
-
-    def get_num_new_matched_tokens(self, request, num_computed_tokens):
-        self.lookup_calls.append((request.seq_id, num_computed_tokens))
-        result = next(self.results)
-        if result[0] is None:
-            self.pending_ids.add(request.seq_id)
-        else:
-            self.pending_ids.discard(request.seq_id)
-        return result
-
-    def cancel_lookup(self, request_id):
-        self.pending_ids.discard(request_id)
-        self.cancelled.append(request_id)
-
-    def update_state_after_alloc(self, request, block_ids, num_external_tokens):
-        self.allocations.append((request.seq_id, tuple(block_ids), num_external_tokens))
-
-    def build_connector_meta(self, scheduler_output):
-        return None
-
-    def update_connector_output(self, connector_output):
-        return KVConnectorResult(
-            load_results=tuple(
-                KVLoadResult(
-                    request_id=request_id,
-                    success=request_id not in self.failed_ids,
-                )
-                for request_id in (connector_output.finished_receiving or set())
-            )
-        )
-
-    def request_finished(self, request):
-        self.finished.append(request.seq_id)
-
-    def finish_transfers_after_worker_drain(self):
-        pass
-
-    def shutdown(self):
-        pass
-
-
-def _make_async_lookup_scheduler(
-    connector,
-    *,
-    enable_prefix_caching=True,
-    max_batches=1,
-    num_gpu_blocks=16,
-    max_prefill_token_num=8192,
-):
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 4
-    return Scheduler(
-        scheduler_config=SchedulerConfig(
-            max_batches=max_batches,
-            max_session_len=64,
-            max_request_output_len=16,
-            eviction_type='recompute',
-        ),
-        cache_config=CacheConfig(
-            max_batches=max_batches,
-            block_size=block_size,
-            num_cpu_blocks=0,
-            num_gpu_blocks=num_gpu_blocks,
-            max_prefill_token_num=max_prefill_token_num,
-            enable_prefix_caching=enable_prefix_caching,
-            kv_transfer_config=KVTransferConfig(
-                kv_connector='MooncakeStoreConnector',
-                kv_role='kv_both',
-            ),
-        ),
-        seq_meta=SequenceMeta(block_size, strategy=ARSequenceStrategy()),
-        kv_connector=connector,
-    )
 
 
 class TestScheduler:
@@ -284,307 +180,9 @@ def test_evict(self, scheduler, block_size, num_gpu_blocks, num_cpu_blocks):
         assert seq2.status == MessageStatus.READY
         assert block_manager.get_num_free_gpu_blocks() == 0
 
-        # test running append
-        seq1.update_token_ids(torch.tensor([1] * block_size))
-        seq2.update_token_ids(torch.tensor([1] * block_size))
-        assert len(scheduler.ready) == 2
-        scheduler.schedule(is_prefill=False)
-        # seq1: 2 running gpu
-        # seq2: 4 waiting cpu
-        # seq3: 3 nan
-        assert seq1.status == MessageStatus.READY
-        assert seq2.status == MessageStatus.WAITING
-        assert block_manager.get_num_free_gpu_blocks() == 2
-
-
-def test_state_manager_reserves_system_state_slot():
-    manager = StateManager(num_states=3, num_reserved=1)
-
-    assert manager.allocate_state() == 1
-    assert manager.allocate_state() == 2
-    with pytest.raises(RuntimeError, match='No free states'):
-        manager.allocate_state()
-
-
-def test_state_manager_checkpoint_can_borrow_idle_runtime_slots():
-    manager = StateManager(num_states=5, num_reserved=1, num_runtime_states=2)
-
-    checkpoints = [manager.allocate_checkpoint_state() for _ in range(4)]
-    assert checkpoints == [1, 2, 3, 4]
-    with pytest.raises(RuntimeError, match='No free states'):
-        manager.allocate_checkpoint_state()
-
-    manager.free_checkpoint_state(checkpoints[0])
-    manager.free_checkpoint_state(checkpoints[1])
-    assert manager.allocate_state() == checkpoints[1]
-    assert manager.allocate_state() == checkpoints[0]
-    with pytest.raises(RuntimeError, match='No free states'):
-        manager.allocate_state()
-
-
-def test_state_manager_caps_runtime_count_even_with_extra_free_slots():
-    manager = StateManager(num_states=6, num_reserved=1, num_runtime_states=2)
-
-    assert manager.num_runtime_states == 2
-    assert manager.allocate_state() == 1
-    assert manager.allocate_state() == 2
-    assert manager.get_num_free() == 3
-    assert manager.get_num_free_runtime() == 0
-    with pytest.raises(RuntimeError, match='No free states'):
-        manager.allocate_state()
-
-
-def _make_ssm_scheduler(max_batch_size: int = 1, prefix_cache_state_budget: int = 0, num_gpu_blocks: int = 16):
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
-    cache_config = CacheConfig(max_batches=max_batch_size,
-                               block_size=block_size,
-                               num_cpu_blocks=4,
-                               num_gpu_blocks=num_gpu_blocks,
-                               enable_prefix_caching=True,
-                               num_state_caches=max_batch_size + 1 + prefix_cache_state_budget,
-                               prefix_cache_state_budget=prefix_cache_state_budget,
-                               states_shapes=[((1, ), torch.float32)])
-    scheduler_config = SchedulerConfig(max_batches=max_batch_size,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    return Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-
-def _add_published_ssm_checkpoint(scheduler: Scheduler, token_ids: list[int]):
-    session = scheduler.add_session(len(scheduler.sessions))
-    seq = session.add_sequence(token_ids)
-    scheduler.block_manager.allocate(seq)
-    scheduler.block_trie.allocate(seq)
-    state_idx = scheduler.block_trie.state_checkpoints.reserve_save(seq)
-    assert state_idx >= 0
-    assert scheduler.block_trie.state_checkpoints.publish_save(seq)
-    node = seq.prefix_cache.trie_cursor
-    session.remove_sequence(seq)
-    return node, state_idx
-
-
-def test_ssm_runtime_state_reclaims_borrowed_checkpoint_slot():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
-    block_size = scheduler.seq_meta.block_size
-    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    seq = scheduler.add_session(100).add_sequence([2] * block_size * 2)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.logical_state == state_idx
-    assert node.state_checkpoint is None
-    assert scheduler.state_manager.get_num_runtime_states() == 1
-    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0
-
-
-def test_ssm_long_chunked_request_schedules_with_only_runtime_state_slot():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
-    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size * 2
-    block_size = scheduler.seq_meta.block_size
-    token_ids = [1] * block_size + [2] * block_size + [3] * block_size
-    seq = scheduler.add_session(100).add_sequence(token_ids)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.logical_state >= 0
-    assert scheduler.state_manager.get_num_runtime_states() == 1
-    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0
-    assert scheduler.block_trie.state_checkpoints.reserve_save(seq, step=block_size * 2) == -1
-
-
-def test_ssm_running_request_reuses_own_runtime_state_without_spare_slot():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
-    block_size = scheduler.seq_meta.block_size
-    seq = scheduler.add_session(100).add_sequence([1] * block_size)
-
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [seq]
-    assert scheduler.state_manager.get_num_free_runtime() == 0
-    seq.state.activate()
-
-    seq.update_token_ids([2] * block_size, mode=UpdateTokenMode.DECODE)
-    valid_mask = scheduler.schedule_running([seq], num_required_tokens=0, prealloc_size=0)
-
-    assert valid_mask == [True]
-    assert seq.status == MessageStatus.RUNNING
-    assert seq.logical_state >= 0
-    assert seq.num_blocks == 2
-    assert scheduler.state_manager.get_num_runtime_states() == 1
-    assert scheduler.state_manager.get_num_free_runtime() == 0
-
-
-def test_ssm_runtime_state_waits_when_only_checkpoint_slot_is_pinned():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
-    block_size = scheduler.seq_meta.block_size
-    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    node.state_checkpoint.pin_count = 1
-    seq = scheduler.add_session(100).add_sequence([2] * block_size * 2)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert seq.status == MessageStatus.WAITING
-    assert seq.logical_state == -1
-    assert node.state_checkpoint.slot == state_idx
-    assert node.state_checkpoint.published
-
-
-def test_ssm_same_batch_duplicate_checkpoint_save_has_unique_dst_offsets():
-    scheduler = _make_ssm_scheduler(max_batch_size=2, prefix_cache_state_budget=2)
-    block_size = scheduler.seq_meta.block_size
-    token_ids = [1] * block_size * 2
-
-    seq_a = scheduler.add_session(100).add_sequence(token_ids)
-    seq_b = scheduler.add_session(101).add_sequence(token_ids)
-
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [seq_a, seq_b]
-    assert seq_a.logical_state >= 0
-    assert seq_b.logical_state >= 0
-    assert seq_a.logical_state != seq_b.logical_state
-    assert seq_a.prefix_cache.trie_cursor is seq_b.prefix_cache.trie_cursor
-
-    save_state_offsets = [
-        scheduler.block_trie.state_checkpoints.reserve_save(seq) for seq in output.running
-    ]
-    save_plan = _make_state_prefix_cache_save_plan(output.running, save_state_offsets)
-    assert save_plan is not None
-    save_src_offsets, save_dst_offsets = save_plan
-
-    assert save_src_offsets == (seq_a.logical_state, )
-    assert save_dst_offsets == (save_state_offsets[0], )
-    assert save_state_offsets[0] >= 0
-    assert save_state_offsets[1] == -1
-    assert len(save_dst_offsets) == len(set(save_dst_offsets))
-
-
-def test_ssm_end_session_discards_pending_checkpoint_reservation():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
-    block_size = scheduler.seq_meta.block_size
-    session = scheduler.add_session(100)
-    seq = session.add_sequence([1] * block_size * 2)
-    scheduler.block_manager.allocate(seq)
-    scheduler.block_trie.allocate(seq)
-    scheduler.state_manager.allocate(seq)
-
-    state_idx = scheduler.block_trie.state_checkpoints.reserve_save(seq)
-    node = seq.prefix_cache.pending_save.node
-    assert state_idx >= 0
-    assert node is not None
-    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 1
-
-    scheduler.end_session(100)
-
-    assert 100 not in scheduler.sessions
-    assert node.state_checkpoint is None
-    assert scheduler.state_manager.get_num_runtime_states() == 0
-    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0
-
-
-def test_ssm_end_session_unpins_restore_checkpoint():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
-    block_size = scheduler.seq_meta.block_size
-    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    seq = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [2])
-
-    scheduler.block_trie.match(seq)
-    assert seq.prefix_cache.restore.slot == state_idx
-    assert scheduler.block_trie.state_checkpoints.pin_restore(seq)
-    assert node.state_checkpoint.pin_count == 1
-
-    scheduler.end_session(100)
-
-    assert 100 not in scheduler.sessions
-    assert node.state_checkpoint.slot == state_idx
-    assert node.state_checkpoint.published
-    assert node.state_checkpoint.pin_count == 0
-
-
-def test_ssm_failed_restore_schedule_rolls_back_match():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
-    block_size = scheduler.seq_meta.block_size
-    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    node.state_checkpoint.pin_count = 1
-    seq = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [2])
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert seq.status == MessageStatus.WAITING
-    assert seq.num_history_ids == 0
-    assert len(seq.logical_blocks) == 0
-    assert seq.cached_tokens == 0
-    assert seq.prefix_cache.trie_cursor is None
-    assert seq.prefix_cache.restore.slot == -1
-    assert seq.prefix_cache.restore.node is None
-    assert node.state_checkpoint.slot == state_idx
-    assert node.state_checkpoint.published
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-    node.state_checkpoint.pin_count = 0
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.status == MessageStatus.READY
-    assert seq.num_history_ids == 0
-    assert seq.prefix_cache.restore.slot == -1
-    assert seq.logical_state == state_idx
-    assert node.state_checkpoint is None
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def test_ssm_scheduler_preserves_matched_checkpoint_when_evicting_for_runtime_state():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
-    block_size = scheduler.seq_meta.block_size
-    node_a, state_idx_a = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    node_b, state_idx_b = _add_published_ssm_checkpoint(scheduler, [2] * block_size * 2)
-    seq = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.num_history_ids == block_size * 2
-    assert seq.cached_tokens == block_size * 2
-    assert seq.prefix_cache.restore.slot == state_idx_a
-    assert seq.prefix_cache.restore.node is node_a
-    assert seq.prefix_cache.restore.pinned
-    assert seq.logical_state == state_idx_b
-    assert node_a.state_checkpoint.slot == state_idx_a
-    assert node_a.state_checkpoint.published
-    assert node_a.state_checkpoint.pin_count == 1
-    assert node_b.state_checkpoint is None
-    assert scheduler.block_trie.stats.num_hit_tokens == block_size * 2
-
-    assert scheduler.block_trie.state_checkpoints.unpin_restore(seq)
-
-
-def test_ssm_scheduler_evicts_stopped_runtime_state_with_free_checkpoint_slot():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
-    block_size = scheduler.seq_meta.block_size
-    seq_a = scheduler.add_session(100).add_sequence([1] * block_size)
-
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [seq_a]
-    assert seq_a.logical_state >= 0
-    assert scheduler.state_manager.get_num_free() == 1
-    assert scheduler.state_manager.get_num_free_runtime() == 0
-
-    seq_a.state.stop()
-    seq_b = scheduler.add_session(101).add_sequence([2] * block_size)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq_b]
-    assert seq_b.logical_state >= 0
-    assert seq_a.logical_state == -1
-    assert seq_a.status == MessageStatus.STOPPED
+    def test_decode_requires_schedule_running(self, scheduler):
+        with pytest.raises(ValueError, match='schedule_running'):
+            scheduler.schedule(is_prefill=False)
 
 
 def test_schedule_migration_matches_current_sequence():
@@ -614,1498 +212,62 @@ def test_schedule_migration_matches_current_sequence():
     assert seq.status == MessageStatus.MIGRATION_READY
 
 
-def test_scheduler_publishes_cached_tokens_for_accepted_prefix_hit():
+def _make_scheduler_for_decode_growth(num_gpu_blocks: int = 2):
     from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
+    block_size = 4
     seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2] * block_size + [3])
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-
-    seq = scheduler.add_session(1).add_sequence([1] * block_size + [2] * block_size + [4])
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.num_history_ids == block_size * 2
-    assert seq.cached_tokens == block_size * 2
-
-    seq.update_token_ids(torch.tensor([5]))
-
-    assert seq.cached_tokens == 0
-    assert seq.prefix_cache.match_start_step == -1
-
-
-def test_scheduler_ar_spec_prefix_hit_recomputes_overlap_block():
-    from lmdeploy.pytorch.strategies.ar_spec.sequence import ARSpecSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSpecSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    token_ids = [1] * block_size + [2] * block_size + [3] * block_size + [4]
-    cached = scheduler.add_session(0).add_sequence(token_ids)
-    scheduler.block_manager.allocate(cached)
-    scheduler.block_trie.allocate(cached)
-    cached_blocks = cached.logical_blocks.get_real_blocks().copy()
-    cached.state.stop()
-
-    seq = scheduler.add_session(1).add_sequence(token_ids)
-    scheduler.block_trie.stats.reset()
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.prefix_cache.recompute_overlap.recompute_blocks == 1
-    assert seq.num_history_ids == block_size * 2
-    assert seq.cached_tokens == block_size * 2
-    assert seq.logical_blocks[2] != cached_blocks[2]
-    assert seq.prefix_cache.recompute_overlap.fresh_block_range is None
-    assert scheduler.block_trie.stats.num_query_tokens == len(token_ids)
-    assert scheduler.block_trie.stats.num_hit_tokens == block_size * 2
-
-
-def test_scheduler_prefix_match_rollback_clears_recompute_overlap_window(monkeypatch):
-    from lmdeploy.pytorch.strategies.ar_spec.sequence import ARSpecSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSpecSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
+    cache_config = CacheConfig(max_batches=2,
                                block_size=block_size,
                                num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=128,
+                               num_gpu_blocks=num_gpu_blocks,
+                               max_prefill_token_num=block_size * 4)
+    scheduler_config = SchedulerConfig(max_batches=2,
+                                       max_session_len=64,
                                        max_request_output_len=64,
                                        eviction_type='recompute')
     scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    token_ids = [1] * block_size + [2] * block_size + [3] * block_size + [4]
-    cached = scheduler.add_session(0).add_sequence(token_ids)
-    scheduler.block_manager.allocate(cached)
-    scheduler.block_trie.allocate(cached)
-    cached.state.stop()
-
-    seq = scheduler.add_session(1).add_sequence(token_ids)
-    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', Mock(return_value=False))
-    scheduler.block_trie.stats.reset()
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert seq.num_history_ids == 0
-    assert seq.num_token_ids == len(token_ids)
-    assert seq.cached_tokens == 0
-    assert seq.prefix_cache.recompute_overlap.fresh_block_range is None
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def test_async_lookup_pending_rolls_back_a_new_request_once():
-    connector = _AsyncLookupConnector([(None, False), (0, False)])
-    scheduler = _make_async_lookup_scheduler(connector)
-    seq = scheduler.add_session(70).add_sequence(torch.arange(9))
-
-    first = scheduler.schedule(is_prefill=True)
-    assert first.running == []
-    assert scheduler.last_schedule_had_pending_lookup
-    assert seq.num_history_ids == 0
-    assert seq.num_blocks == 0
-    assert seq.prefix_cache.trie_cursor is None
-    assert seq.prefix_cache.match_start_step == -1
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-
-    second = scheduler.schedule(is_prefill=True)
-    assert second.running == []
-    assert connector.lookup_calls == [(seq.seq_id, 0)]
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-
-    connector.pending_ids.clear()
-    third = scheduler.schedule(is_prefill=True)
-    assert third.running == [seq]
-    assert connector.lookup_calls == [(seq.seq_id, 0), (seq.seq_id, 0)]
-
-
-def test_async_lookup_rebases_remote_hit_after_local_trie_grows(monkeypatch):
-    connector = MooncakeStoreScheduler(
-        CacheConfig(
-            max_batches=1,
-            block_size=4,
-            num_cpu_blocks=0,
-            num_gpu_blocks=16,
-            enable_prefix_caching=True,
-            kv_transfer_config=KVTransferConfig(
-                kv_connector='MooncakeStoreConnector',
-                kv_role='kv_both',
-            ),
-        ))
-    assert connector.client is not None
-    # The same asynchronous lookup first reports pending, then returns its
-    # absolute remote prefix boundary from the original snapshot.
-    monkeypatch.setattr(connector.client, 'lookup', Mock(side_effect=(None, 16)))
-    observed_results = []
-    get_matched_tokens = connector.get_num_new_matched_tokens
-
-    def _record_result(request, num_computed_tokens):
-        result = get_matched_tokens(request, num_computed_tokens)
-        observed_results.append((num_computed_tokens, result))
-        return result
-
-    monkeypatch.setattr(connector, 'get_num_new_matched_tokens', _record_result)
-    scheduler = _make_async_lookup_scheduler(connector)
-    tokens = torch.arange(17)
-
-    cached_to_8 = scheduler.add_session(80).add_sequence(tokens[:9])
-    scheduler.block_manager.allocate(cached_to_8)
-    scheduler.block_trie.allocate(cached_to_8)
-    cached_to_8.state.stop()
-
-    seq = scheduler.add_session(81).add_sequence(tokens)
-    first = scheduler.schedule(is_prefill=True)
-
-    assert first.running == []
-    assert seq.num_history_ids == 0
-    assert observed_results == [(8, (None, False))]
-
-    # While the remote lookup is pending, another sequence publishes the next
-    # complete local block. The retried match must now advance from 8 to 12.
-    cached_to_12 = scheduler.add_session(82).add_sequence(tokens[:13])
-    scheduler.block_manager.allocate(cached_to_12)
-    scheduler.block_trie.allocate(cached_to_12)
-    cached_to_12.state.stop()
-
-    second = scheduler.schedule(is_prefill=True)
-
-    assert second.running == []
-    assert seq.num_history_ids == 12
-    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert observed_results == [
-        (8, (None, False)),
-        (12, (4, True)),
-    ]
-
-    metadata = connector.build_connector_meta(second)
-    assert metadata is not None
-    assert len(metadata.load_requests) == 1
-    load_request = metadata.load_requests[0]
-    block_table = scheduler.block_manager.get_block_table(seq)
-    assert load_request.block_ids == (int(block_table[3]), )
-    assert load_request.remote_block_count == 4
-    scheduler.shutdown()
-
-
-def test_async_lookup_pending_request_does_not_block_later_waiter():
-    connector = _AsyncLookupConnector([(None, False), (0, False)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-    )
-    pending = scheduler.add_session(74).add_sequence(torch.arange(9))
-    schedulable = scheduler.add_session(75).add_sequence(torch.arange(9, 18))
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [schedulable]
-    assert pending.status == MessageStatus.WAITING
-    assert schedulable.status == MessageStatus.READY
-    assert pending.seq_id in connector.pending_ids
-    assert connector.lookup_calls == [
-        (pending.seq_id, 0),
-        (schedulable.seq_id, 0),
-    ]
-    assert scheduler.last_schedule_had_pending_lookup
-
-
-def test_async_lookup_precisely_restores_a_multiturn_local_prefix():
-    connector = _AsyncLookupConnector([(None, False), (4, True)])
-    scheduler = _make_async_lookup_scheduler(connector)
-    tokens = torch.arange(13)
-    seq = scheduler.add_session(71).add_sequence(tokens)
-
-    seq.kv_token_limit = 4
-    scheduler.block_manager.allocate(seq)
-    scheduler.block_trie.allocate(seq)
-    seq.set_step(4)
-    seq.kv_token_limit = 5
-    seq.cached_tokens = 3
-    seq.model_meta = {'state': 'keep'}
-    seq.prefix_cache.recompute_overlap.fresh_block_range = range(0, 1)
-    seq.prefix_cache.recompute_overlap.trie_block_map[0] = seq.logical_blocks[0]
-    baseline_blocks = seq.logical_blocks.get_real_blocks().copy()
-    baseline_cursor = seq.prefix_cache.trie_cursor
-
-    cached = scheduler.add_session(72).add_sequence(tokens[:9])
-    scheduler.block_manager.allocate(cached)
-    scheduler.block_trie.allocate(cached)
-    cached.state.stop()
-    matched_block = cached.logical_blocks[1]
-    matched_ref_count = scheduler.block_manager.allocator.get_ref_count(
-        cached.logical_blocks.get_real_blocks()[1:2]).copy()
-    scheduler.block_trie.stats.reset()
-
-    first = scheduler.schedule(is_prefill=True)
-    assert first.running == []
-    assert scheduler.last_schedule_had_pending_lookup
-    assert seq.num_history_ids == 4
-    assert seq.num_blocks == 1
-    assert torch.equal(torch.from_numpy(seq.logical_blocks.get_real_blocks()),
-                       torch.from_numpy(baseline_blocks))
-    assert seq.prefix_cache.trie_cursor is baseline_cursor
-    assert seq.prefix_cache.match_start_step == -1
-    assert seq.prefix_cache.recompute_overlap.fresh_block_range == range(0, 1)
-    assert seq.prefix_cache.recompute_overlap.trie_block_map == {0: baseline_blocks[0]}
-    assert seq.cached_tokens == 3
-    assert seq.kv_token_limit == 5
-    assert seq.model_meta == {'state': 'keep'}
-    assert scheduler.block_manager.allocator.get_ref_count(
-        cached.logical_blocks.get_real_blocks()[1:2]).tolist() == matched_ref_count.tolist()
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-
-    second = scheduler.schedule(is_prefill=True)
-    assert second.running == []
-    assert connector.lookup_calls == [(seq.seq_id, 8)]
-    assert seq.num_history_ids == 4
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-
-    connector.pending_ids.clear()
-    seq.kv_token_limit = None
-    seq.prefix_cache.recompute_overlap.clear_tracking()
-    third = scheduler.schedule(is_prefill=True)
-    assert third.running == []
-    assert connector.lookup_calls == [(seq.seq_id, 8), (seq.seq_id, 8)]
-    assert seq.num_history_ids == 8
-    assert matched_block in seq.logical_blocks.get_real_blocks()
-    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert connector.allocations[0][2] == 4
-
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={seq.seq_id}))
-    assert seq.num_history_ids == 12
-    assert seq.status == MessageStatus.WAITING
-
-    fourth = scheduler.schedule(is_prefill=True)
-    assert fourth.running == [seq]
-
-
-def test_async_lookup_pending_preserves_private_partial_prefix():
-    connector = _AsyncLookupConnector([(None, False)])
-    scheduler = _make_async_lookup_scheduler(connector)
-    tokens = torch.arange(13)
-    seq = scheduler.add_session(72).add_sequence(tokens)
-
-    seq.kv_token_limit = 5
-    scheduler.block_manager.allocate(seq)
-    scheduler.block_trie.allocate(seq)
-    seq.set_step(5)
-    seq.kv_token_limit = 7
-    seq.cached_tokens = 3
-    seq.model_meta = {'state': 'keep'}
-    baseline_blocks = seq.logical_blocks.get_real_blocks().copy()
-    baseline_cursor = seq.prefix_cache.trie_cursor
-    scheduler.block_trie.stats.reset()
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert scheduler.last_schedule_had_pending_lookup
-    assert connector.lookup_calls == [(seq.seq_id, 5)]
-    assert seq.num_history_ids == 5
-    assert torch.equal(torch.from_numpy(seq.logical_blocks.get_real_blocks()),
-                       torch.from_numpy(baseline_blocks))
-    assert seq.prefix_cache.trie_cursor is baseline_cursor
-    assert seq.prefix_cache.match_start_step == -1
-    assert seq.cached_tokens == 3
-    assert seq.kv_token_limit == 7
-    assert seq.model_meta == {'state': 'keep'}
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-
-
-def test_external_cached_tokens_survive_remote_ready_admission():
-    connector = _AsyncLookupConnector([(8, True)])
-    scheduler = _make_async_lookup_scheduler(connector)
-    seq = scheduler.add_session(73).add_sequence(torch.arange(13))
-
-    started = scheduler.schedule(is_prefill=True)
-    assert started.running == []
-
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={seq.seq_id}))
-    assert seq.num_history_ids == 8
-    assert seq.cached_tokens == 8
-    assert seq.prefix_cache.match_start_step == 0
-
-    admitted = scheduler.schedule(is_prefill=True)
-
-    assert admitted.running == [seq]
-    assert seq.cached_tokens == 8
-    assert seq.prefix_cache.match_start_step == 0
-
-
-def test_remote_ready_long_prefill_respects_short_only_turn():
-    connector = _AsyncLookupConnector([(8, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        max_prefill_token_num=4,
-    )
-    seq = scheduler.add_session(74).add_sequence(torch.arange(17))
-
-    started = scheduler.schedule(is_prefill=True)
-    assert started.running == []
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={seq.seq_id}))
-    assert seq.num_history_ids == 8
-    assert scheduler.kv_load_coordinator.is_remote_ready(seq)
-
-    short_turn = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert short_turn.running == []
-    assert seq.status == MessageStatus.WAITING
-    assert seq.num_history_ids == 8
-    assert scheduler.kv_load_coordinator.is_remote_ready(seq)
-
-    long_turn = scheduler.schedule(is_prefill=True)
-    assert long_turn.running == [seq]
-    assert seq.status == MessageStatus.READY
-
-
-def test_external_cached_tokens_survive_prefill_budget_rejection():
-    connector = _AsyncLookupConnector([(8, True), (8, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        max_batches=2,
-        max_prefill_token_num=8,
-    )
-    admitted_seq = scheduler.add_session(74).add_sequence(torch.arange(13))
-    waiting_seq = scheduler.add_session(75).add_sequence(torch.arange(20, 33))
-
-    started = scheduler.schedule(is_prefill=True)
-    assert started.running == []
-    scheduler.update_connector_output(
-        KVConnectorOutput(
-            finished_receiving={admitted_seq.seq_id, waiting_seq.seq_id}))
-
-    admitted = scheduler.schedule(is_prefill=True)
-
-    assert admitted.running == [admitted_seq]
-    assert waiting_seq.status == MessageStatus.WAITING
-    assert waiting_seq.num_history_ids == 8
-    assert waiting_seq.num_blocks == 2
-    assert waiting_seq.cached_tokens == 8
-    assert waiting_seq.prefix_cache.match_start_step == 0
-    assert scheduler.kv_load_coordinator.is_remote_ready(waiting_seq)
-
-
-def test_async_load_keeps_a_private_partial_block_at_the_suffix_start():
-    connector = _AsyncLookupConnector([(7, True)])
-    scheduler = _make_async_lookup_scheduler(connector)
-    tokens = torch.arange(13)
-    cached = scheduler.add_session(76).add_sequence(tokens)
-    scheduler.block_manager.allocate(cached)
-    scheduler.block_trie.allocate(cached)
-    cached.state.stop()
-
-    seq = scheduler.add_session(77).add_sequence(tokens)
-    seq.kv_token_limit = 5
-    scheduler.block_manager.allocate(seq)
-    scheduler.block_trie.allocate(seq)
-    seq.set_step(5)
-    seq.kv_token_limit = None
-    private_block = int(scheduler.block_manager.get_block_table(seq)[1])
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert seq.num_blocks == 3
-    assert connector.lookup_calls == [(seq.seq_id, 5)]
-    load_blocks = connector.allocations[0][1]
-    assert load_blocks == tuple(
-        int(block_id)
-        for block_id in scheduler.block_manager.get_block_table(seq)[1:3]
-    )
-    assert load_blocks[0] == private_block
-
-
-def test_async_load_requires_capacity_for_the_complete_prefill():
-    connector = _AsyncLookupConnector([(8, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        num_gpu_blocks=3,
-    )
-    seq = scheduler.add_session(76).add_sequence(torch.arange(13))
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert seq.status == MessageStatus.WAITING
-    assert seq.num_blocks == 0
-    assert connector.allocations == []
-
-
-def test_async_load_capacity_failure_restores_tentative_local_prefix(monkeypatch):
-    connector = _AsyncLookupConnector([(4, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        num_gpu_blocks=3,
-    )
-    tokens = torch.arange(13)
-    cached = scheduler.add_session(76).add_sequence(tokens[:5])
-    scheduler.block_manager.allocate(cached)
-    scheduler.block_trie.allocate(cached)
-    cached.state.stop()
-    cached_block = cached.logical_blocks.get_real_blocks()[:1]
-    ref_count = scheduler.block_manager.allocator.get_ref_count(cached_block).copy()
-    scheduler.block_trie.stats.reset()
-
-    seq = scheduler.add_session(77).add_sequence(tokens)
-    evict_for_seq = Mock(return_value=False)
-    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert seq.status == MessageStatus.WAITING
-    assert seq.num_history_ids == 0
-    assert seq.num_blocks == 0
-    assert seq.kv_token_limit is None
-    assert seq.cached_tokens == 0
-    assert seq.prefix_cache.trie_cursor is None
-    assert seq.prefix_cache.match_start_step == -1
-    assert connector.lookup_calls == [(seq.seq_id, 4)]
-    assert connector.allocations == []
-    assert evict_for_seq.call_count == 1
-    assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
+    return scheduler, block_size
 
 
-def test_async_load_does_not_consume_model_batch_slot():
-    connector = _AsyncLookupConnector([(8, True), (0, False)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        num_gpu_blocks=5,
-    )
-    loading = scheduler.add_session(77).add_sequence(torch.arange(13))
-    later = scheduler.add_session(78).add_sequence(torch.arange(8))
+def test_schedule_running_reclaims_waiting_blocks_for_decode_growth():
+    scheduler, block_size = _make_scheduler_for_decode_growth(num_gpu_blocks=2)
+    decode = scheduler.add_session(100).add_sequence([1] * block_size)
+    waiting = scheduler.add_session(101).add_sequence([2] * block_size)
 
     output = scheduler.schedule(is_prefill=True)
+    assert output.running == [decode, waiting]
+    scheduler.activate_seqs([decode])
+    waiting.state.evict()
+    assert decode.status == MessageStatus.RUNNING
+    assert waiting.status == MessageStatus.WAITING
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
 
-    assert output.running == [later]
-    assert loading.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert loading.num_blocks == 2
-    assert later.status == MessageStatus.READY
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 2
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 1
+    valid_mask = scheduler.schedule_running([decode], num_required_tokens=1, prealloc_size=1)
 
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={loading.seq_id}))
-    assert loading.num_history_ids == 8
-    assert loading.cached_tokens == 8
+    assert valid_mask == [True]
+    assert decode.status == MessageStatus.RUNNING
+    assert decode.num_blocks == 2
+    assert waiting.status == MessageStatus.WAITING
+    assert waiting.num_blocks == 0
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
 
 
-def test_multiple_async_loads_start_in_one_prefill_turn():
-    connector = _AsyncLookupConnector([(8, True), (8, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        max_batches=1,
-        num_gpu_blocks=16,
-    )
-    first = scheduler.add_session(91).add_sequence(torch.arange(13))
-    second = scheduler.add_session(92).add_sequence(torch.arange(20, 33))
+def test_schedule_running_keeps_other_running_sequence_when_decode_growth_fails():
+    scheduler, block_size = _make_scheduler_for_decode_growth(num_gpu_blocks=2)
+    decode = scheduler.add_session(100).add_sequence([1] * block_size)
+    long_chunk = scheduler.add_session(101).add_sequence([2] * block_size)
 
     output = scheduler.schedule(is_prefill=True)
+    assert output.running == [decode, long_chunk]
+    scheduler.activate_seqs([decode, long_chunk])
+    assert decode.status == MessageStatus.RUNNING
+    assert long_chunk.status == MessageStatus.RUNNING
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
 
-    assert output.running == []
-    assert first.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert second.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert [allocation[0] for allocation in connector.allocations] == [
-        first.seq_id,
-        second.seq_id,
-    ]
-
-
-def test_soft_reservation_blocks_new_load_until_capacity_is_released():
-    connector = _AsyncLookupConnector([
-        (4, True),
-        (12, True),
-        (12, True),
-    ])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        num_gpu_blocks=6,
-    )
-    first = scheduler.add_session(86).add_sequence(torch.arange(13))
-
-    scheduler.schedule(is_prefill=True)
-    assert first.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert first.num_blocks == 1
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 3
-
-    second = scheduler.add_session(87).add_sequence(torch.arange(17))
-    blocked = scheduler.schedule(is_prefill=True)
-    assert blocked.running == []
-    assert second.status == MessageStatus.WAITING
-    assert second.num_blocks == 0
-    assert [allocation[0] for allocation in connector.allocations] == [first.seq_id]
-
-    scheduler.end_session(86)
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={first.seq_id}))
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
-
-    retried = scheduler.schedule(is_prefill=True)
-    assert retried.running == []
-    assert second.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert second.num_blocks == 3
-    assert [allocation[0] for allocation in connector.allocations] == [
-        first.seq_id,
-        second.seq_id,
-    ]
-
-
-def test_failed_async_load_preserves_local_prefix_and_releases_remote_tail():
-    connector = _AsyncLookupConnector([(4, True)])
-    scheduler = _make_async_lookup_scheduler(connector)
-    tokens = torch.arange(13)
-    cached = scheduler.add_session(79).add_sequence(tokens[:9])
-    scheduler.block_manager.allocate(cached)
-    scheduler.block_trie.allocate(cached)
-    cached.state.stop()
-    seq = scheduler.add_session(80).add_sequence(tokens)
-    scheduler.schedule(is_prefill=True)
-    connector.failed_ids.add(seq.seq_id)
-
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={seq.seq_id}))
-
-    assert seq.status == MessageStatus.WAITING
-    assert seq.num_history_ids == 8
-    assert seq.num_blocks == 2
-    assert seq.cached_tokens == 8
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
-
-
-def test_async_load_soft_reservation_shrinks_across_chunks():
-    connector = _AsyncLookupConnector([(4, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        max_prefill_token_num=4,
-    )
-    seq = scheduler.add_session(81).add_sequence(torch.arange(13))
-
-    scheduler.schedule(is_prefill=True)
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 3
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={seq.seq_id}))
-
-    assert scheduler.schedule(is_prefill=True).running == [seq]
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 2
-    seq.set_step(8)
-    scheduler.release_completed_prefill_reservations([seq])
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 2
-
-    assert scheduler.reserve_long_context_chunk(seq, chunk_size=4)
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 1
-    seq.set_step(12)
-    assert scheduler.reserve_long_context_chunk(seq, chunk_size=1, is_last_chunk=True)
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
-
-
-def test_end_session_waits_for_active_async_load_before_freeing_blocks():
-    connector = _AsyncLookupConnector([(8, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-    )
-    seq = scheduler.add_session(82).add_sequence(torch.arange(13))
-    scheduler.schedule(is_prefill=True)
-
-    scheduler.end_session(82)
-    assert 82 in scheduler.sessions
-    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
-    assert seq.num_blocks == 2
-    assert connector.finished == []
-
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={seq.seq_id}))
-    assert 82 not in scheduler.sessions
-    assert connector.finished == [seq.seq_id]
-
-
-def test_worker_drain_finishes_an_ended_session_with_a_dropped_load_output():
-    connector = _AsyncLookupConnector([(8, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-    )
-    seq = scheduler.add_session(85).add_sequence(torch.arange(13))
-    scheduler.schedule(is_prefill=True)
-    scheduler.end_session(85)
-
-    scheduler.finish_deferred_kv_transfers_after_worker_drain()
-
-    assert 85 not in scheduler.sessions
-    assert connector.finished == [seq.seq_id]
-    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
-
-
-def test_completed_async_load_is_admitted_before_an_older_waiter():
-    connector = _AsyncLookupConnector([(8, True)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        num_gpu_blocks=4,
-    )
-    loaded = scheduler.add_session(83).add_sequence(torch.arange(13))
-    scheduler.schedule(is_prefill=True)
-    scheduler.update_connector_output(
-        KVConnectorOutput(finished_receiving={loaded.seq_id}))
-    newcomer = scheduler.add_session(84).add_sequence(torch.arange(4))
-    newcomer.arrive_time = loaded.arrive_time - 1
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [loaded]
-    assert newcomer.status == MessageStatus.WAITING
-
-
-def test_stop_and_end_session_cancel_lookup_before_request_cleanup():
-    connector = _AsyncLookupConnector([(None, False)])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-    )
-    seq = scheduler.add_session(73).add_sequence(torch.arange(9))
-    scheduler.schedule(is_prefill=True)
-
-    scheduler.stop_session(73)
-    assert connector.cancelled == [seq.seq_id]
-
-    scheduler.end_session(73)
-    assert connector.finished == [seq.seq_id]
-
-
-def test_async_save_lease_keeps_exact_blocks_alive_until_all_tp_complete():
-
-    class _SaveMetadata(KVConnectorMetadata):
-
-        def __init__(self, logical_block_ids):
-            self.logical_block_ids = tuple(logical_block_ids)
-
-        def get_save_block_leases(self):
-            return (KVSaveBlockLease(7, self.logical_block_ids), )
-
-    class _SaveConnector(_AsyncLookupConnector):
-
-        def __init__(self):
-            super().__init__([])
-            self.metadata = None
-
-        def build_connector_meta(self, scheduler_output):
-            metadata, self.metadata = self.metadata, None
-            return metadata
-
-        def update_connector_output(self, connector_output):
-            return KVConnectorResult(
-                completed_save_ids=frozenset(
-                    connector_output.completed_save_ids or ()),
-            )
-
-    connector = _SaveConnector()
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        num_gpu_blocks=4,
-    )
-    seq = scheduler.add_session(88).add_sequence(torch.arange(8))
-    scheduler.block_manager.allocate(seq)
-    logical_blocks = seq.logical_blocks.get_real_blocks().copy()
-    allocator = scheduler.block_manager.allocator
-    connector.metadata = _SaveMetadata(logical_blocks)
-
-    metadata = scheduler.build_connector_meta(
-        [seq],
-        connector_token_lens=(8, ),
-    )
-
-    assert metadata is not None
-    assert allocator.get_ref_count(logical_blocks).tolist() == [2, 2]
-    assert scheduler.has_unfinished()
-
-    # Sequence ownership may disappear before remote I/O completes. The save
-    # lease remains as the only reference and prevents physical reuse.
-    scheduler.block_manager.free(seq)
-    assert allocator.get_ref_count(logical_blocks).tolist() == [1, 1]
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 2
-
-    scheduler.update_connector_output(
-        KVConnectorOutput(completed_save_ids={7}))
-    assert allocator.get_ref_count(logical_blocks).tolist() == [0, 0]
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 4
-    assert not scheduler.kv_save_coordinator.has_pending()
-
-
-def test_worker_drain_releases_save_leases_when_outputs_are_discarded():
-
-    class _SaveMetadata(KVConnectorMetadata):
-
-        def __init__(self, logical_block_ids):
-            self.logical_block_ids = tuple(logical_block_ids)
-
-        def get_save_block_leases(self):
-            return (KVSaveBlockLease(9, self.logical_block_ids), )
-
-    connector = _AsyncLookupConnector([])
-    scheduler = _make_async_lookup_scheduler(
-        connector,
-        enable_prefix_caching=False,
-        num_gpu_blocks=2,
-    )
-    seq = scheduler.add_session(89).add_sequence(torch.arange(8))
-    scheduler.block_manager.allocate(seq)
-    logical_blocks = seq.logical_blocks.get_real_blocks().copy()
-    metadata = _SaveMetadata(logical_blocks)
-    connector.build_connector_meta = lambda scheduler_output: metadata
-
-    scheduler.build_connector_meta([seq], connector_token_lens=(8, ))
-    scheduler.block_manager.free(seq)
-    scheduler.finish_deferred_kv_transfers_after_worker_drain()
-
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 2
-    assert not scheduler.kv_save_coordinator.has_pending()
-
-
-def test_scheduler_recomputes_prefill_budget_after_prefix_hit():
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=2,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               max_prefill_token_num=block_size,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=2,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2])
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-
-    cache_hit_tail = scheduler.add_session(1).add_sequence([1] * block_size + [3])
-    short = scheduler.add_session(2).add_sequence([4])
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [cache_hit_tail, short]
-    assert cache_hit_tail.num_history_ids == block_size
-    assert cache_hit_tail.num_token_ids == 1
-    assert short.status == MessageStatus.READY
-
-
-def _make_prefix_cache_scheduler(max_batches: int = 2, max_prefill_token_num: int = 16):
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=max_batches,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               max_prefill_token_num=max_prefill_token_num,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=max_batches,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-    return scheduler, block_size
-
-
-def test_scheduler_short_turn_uses_prefix_hit_to_admit_long_looking_sibling():
-    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size)
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-
-    short = scheduler.add_session(1).add_sequence([4])
-    cache_hit_tail = scheduler.add_session(2).add_sequence([1] * block_size + [3])
-
-    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert output.running == [short, cache_hit_tail]
-    assert cache_hit_tail.num_history_ids == block_size
-    assert cache_hit_tail.num_token_ids == 1
-    assert cache_hit_tail.cached_tokens == block_size
-
-
-def test_scheduler_budget_gate_uses_prefix_hit_to_admit_sibling():
-    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size)
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-
-    almost_full = scheduler.add_session(1).add_sequence([4] * (block_size - 1))
-    cache_hit_tail = scheduler.add_session(2).add_sequence([1] * block_size + [3])
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [almost_full, cache_hit_tail]
-    assert cache_hit_tail.num_history_ids == block_size
-    assert cache_hit_tail.num_token_ids == 1
-
-
-def test_scheduler_reorder_cache_stays_order_only_after_prefix_hit():
-    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size)
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-
-    cache_hit_tail = scheduler.add_session(1).add_sequence([1] * block_size + [3])
-    normal = scheduler.add_session(2).add_sequence([4] * (block_size - 1))
-
-    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
-
-    assert output.running == [cache_hit_tail, normal]
-    assert cache_hit_tail.num_history_ids == block_size
-    assert cache_hit_tail.num_token_ids == 1
-    assert cache_hit_tail.cached_tokens == block_size
-    assert normal.status == MessageStatus.READY
-
-
-def test_scheduler_resource_rejection_rolls_back_tentative_prefix_match(monkeypatch):
-    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=1)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2])
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-    cached_block = cached.logical_blocks.get_real_blocks()[:1]
-    ref_count = scheduler.block_manager.allocator.get_ref_count(cached_block).copy()
-    scheduler.block_trie.stats.reset()
-
-    seq = scheduler.add_session(1).add_sequence([1] * block_size + [3])
-    evict_for_seq = Mock(return_value=False)
-    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == []
-    assert seq.status == MessageStatus.WAITING
-    assert seq.num_history_ids == 0
-    assert seq.num_blocks == 0
-    assert seq.kv_token_limit is None
-    assert seq.cached_tokens == 0
-    assert seq.prefix_cache.trie_cursor is None
-    assert seq.prefix_cache.match_start_step == -1
-    assert evict_for_seq.call_count == 1
-    assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def test_scheduler_rolls_back_prefix_match_for_prefill_gate_when_tail_still_exceeds_budget():
-    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size)
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-
-    full = scheduler.add_session(1).add_sequence([4] * block_size)
-    cache_hit_tail = scheduler.add_session(2).add_sequence([1] * block_size + [3])
-
-    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert output.running == [full]
-    assert cache_hit_tail.status == MessageStatus.WAITING
-    assert cache_hit_tail.num_history_ids == 0
-    assert cache_hit_tail.cached_tokens == 0
-    assert cache_hit_tail.prefix_cache.trie_cursor is None
-    assert cache_hit_tail.prefix_cache.match_start_step == -1
-
-
-def test_scheduler_rolls_back_prefix_match_for_prefill_gate_that_still_needs_long_chunk():
-    scheduler, block_size = _make_prefix_cache_scheduler(max_batches=1, max_prefill_token_num=16)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size)
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-    scheduler.block_trie.stats.reset()
-
-    still_long = scheduler.add_session(1).add_sequence([1] * block_size + [3] * (block_size + 1))
-
-    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert output.running == []
-    assert still_long.status == MessageStatus.WAITING
-    assert still_long.num_history_ids == 0
-    assert still_long.cached_tokens == 0
-    assert still_long.prefix_cache.trie_cursor is None
-    assert still_long.prefix_cache.match_start_step == -1
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def test_ssm_scheduler_rolls_back_prefix_match_for_prefill_gate_without_pinning_restore_state():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
-    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size
-    block_size = scheduler.seq_meta.block_size
-    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    scheduler.block_trie.stats.reset()
-
-    still_long = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3] * (block_size + 1))
-
-    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert output.running == []
-    assert still_long.status == MessageStatus.WAITING
-    assert still_long.num_history_ids == 0
-    assert still_long.cached_tokens == 0
-    assert still_long.prefix_cache.trie_cursor is None
-    assert still_long.prefix_cache.restore.slot == -1
-    assert still_long.prefix_cache.restore.node is None
-    assert not still_long.prefix_cache.restore.pinned
-    assert node.state_checkpoint.slot == state_idx
-    assert node.state_checkpoint.pin_count == 0
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def test_ssm_scheduler_rejects_prefix_match_for_prefill_gate_after_pinned_restore_rollback():
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1, num_gpu_blocks=2)
-    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size
-    block_size = scheduler.seq_meta.block_size
-    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    scheduler.block_trie.stats.reset()
-
-    cache_hit_tail = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
-
-    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert output.running == []
-    assert cache_hit_tail.status == MessageStatus.WAITING
-    assert cache_hit_tail.num_history_ids == 0
-    assert cache_hit_tail.num_token_ids == block_size * 2 + 1
-    assert cache_hit_tail.num_blocks == 0
-    assert cache_hit_tail.kv_token_limit is None
-    assert cache_hit_tail.logical_state == -1
-    assert cache_hit_tail.cached_tokens == 0
-    assert cache_hit_tail.prefix_cache.trie_cursor is None
-    assert cache_hit_tail.prefix_cache.restore.slot == -1
-    assert cache_hit_tail.prefix_cache.restore.node is None
-    assert not cache_hit_tail.prefix_cache.restore.pinned
-    assert node.state_checkpoint.slot == state_idx
-    assert node.state_checkpoint.published
-    assert node.state_checkpoint.pin_count == 0
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def test_ssm_scheduler_rejects_prefix_match_for_prefill_gate_after_runtime_state_rollback(monkeypatch):
-    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1, num_gpu_blocks=4)
-    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size
-    block_size = scheduler.seq_meta.block_size
-    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
-    ensure_results = iter([False, True])
-
-    def _ensure_runtime_state_available_once_then_succeed():
-        return next(ensure_results)
-
-    monkeypatch.setattr(scheduler._prefill_scheduler,
-                        '_ensure_runtime_state_available',
-                        _ensure_runtime_state_available_once_then_succeed)
-    scheduler.block_trie.stats.reset()
-
-    cache_hit_tail = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
-
-    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert output.running == []
-    assert cache_hit_tail.status == MessageStatus.WAITING
-    assert cache_hit_tail.num_history_ids == 0
-    assert cache_hit_tail.num_token_ids == block_size * 2 + 1
-    assert cache_hit_tail.num_blocks == 0
-    assert cache_hit_tail.kv_token_limit is None
-    assert cache_hit_tail.logical_state == -1
-    assert cache_hit_tail.cached_tokens == 0
-    assert cache_hit_tail.prefix_cache.trie_cursor is None
-    assert cache_hit_tail.prefix_cache.restore.slot == -1
-    assert cache_hit_tail.prefix_cache.restore.node is None
-    assert not cache_hit_tail.prefix_cache.restore.pinned
-    assert node.state_checkpoint.slot == state_idx
-    assert node.state_checkpoint.published
-    assert node.state_checkpoint.pin_count == 0
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def test_scheduler_reports_zero_cached_tokens_for_prefix_miss():
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2])
-    scheduler.schedule(is_prefill=True)
-    cached.state.stop()
-
-    seq = scheduler.add_session(1).add_sequence([3] * block_size + [4])
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.num_history_ids == 0
-    assert seq.cached_tokens == 0
-
-
-def test_scheduler_cached_tokens_only_count_current_prompt_after_session_eviction():
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    session = scheduler.add_session(0)
-    seq = session.add_sequence([1] * block_size + [2] * block_size + [3])
-    scheduler.schedule(is_prefill=True)
-    seq.update_token_ids(torch.tensor([9]), mode=UpdateTokenMode.PREFILL)
-    seq.state.stop()
-    seq.state.free()
-
-    seq.update_token_ids(torch.tensor([4] * 4))
-    assert seq.input_start_pos == block_size * 2 + 2
-    assert seq.input_end_pos == block_size * 2 + 6
-    seq.state.activate()
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.num_history_ids == block_size * 2
-    assert seq.cached_tokens == 0
-
-
-def test_scheduler_excludes_recompute_eviction_prefix_hits_from_stats():
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=4,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    seq = scheduler.add_session(0).add_sequence([1] * block_size + [2] * block_size + [3])
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [seq]
-
-    seq.state.evict()
-    pressure = scheduler.add_session(1).add_sequence([9] * block_size * 3)
-    scheduler.block_trie.stats.reset()
-
-    assert scheduler.eviction_helper.evict_for_seq(pressure, [seq], 0)
-    assert seq.prefix_cache.suppress_match_stats
-    pressure.session.remove_sequence(pressure)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.num_history_ids >= block_size
-    assert seq.cached_tokens == 0
-    assert not seq.prefix_cache.suppress_match_stats
-    assert scheduler.block_trie.stats.num_query_tokens == 0
-    assert scheduler.block_trie.stats.num_hit_tokens == 0
-
-
-def _make_scheduler_for_decode_growth(num_gpu_blocks: int = 2):
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 4
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=2,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=num_gpu_blocks,
-                               max_prefill_token_num=block_size * 4)
-    scheduler_config = SchedulerConfig(max_batches=2,
-                                       max_session_len=64,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-    return scheduler, block_size
-
-
-def _make_scheduler_for_long_context_chunks(num_gpu_blocks: int = 6):
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 4
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=2,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=num_gpu_blocks,
-                               max_prefill_token_num=block_size * 2)
-    scheduler_config = SchedulerConfig(max_batches=2,
-                                       max_session_len=64,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-    return scheduler, block_size
-
-
-def _make_ssm_scheduler_for_long_context_chunks(num_gpu_blocks: int = 2):
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 4
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=num_gpu_blocks,
-                               max_prefill_token_num=block_size * 2,
-                               num_state_caches=2,
-                               states_shapes=[((1, ), torch.float32)])
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=64,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-    return scheduler, block_size
-
-
-def test_schedule_running_reclaims_waiting_blocks_for_decode_growth():
-    scheduler, block_size = _make_scheduler_for_decode_growth(num_gpu_blocks=2)
-    decode = scheduler.add_session(100).add_sequence([1] * block_size)
-    waiting = scheduler.add_session(101).add_sequence([2] * block_size)
-
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [decode, waiting]
-    scheduler.activate_seqs([decode])
-    waiting.state.evict()
-    assert decode.status == MessageStatus.RUNNING
-    assert waiting.status == MessageStatus.WAITING
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
-
-    valid_mask = scheduler.schedule_running([decode], num_required_tokens=1, prealloc_size=1)
-
-    assert valid_mask == [True]
-    assert decode.status == MessageStatus.RUNNING
-    assert decode.num_blocks == 2
-    assert waiting.status == MessageStatus.WAITING
-    assert waiting.num_blocks == 0
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
-
-
-def test_schedule_running_keeps_other_running_sequence_when_decode_growth_fails():
-    scheduler, block_size = _make_scheduler_for_decode_growth(num_gpu_blocks=2)
-    decode = scheduler.add_session(100).add_sequence([1] * block_size)
-    long_chunk = scheduler.add_session(101).add_sequence([2] * block_size)
-
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [decode, long_chunk]
-    scheduler.activate_seqs([decode, long_chunk])
-    assert decode.status == MessageStatus.RUNNING
-    assert long_chunk.status == MessageStatus.RUNNING
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
-
-    valid_mask = scheduler.schedule_running([decode], num_required_tokens=1, prealloc_size=1)
+    valid_mask = scheduler.schedule_running([decode], num_required_tokens=1, prealloc_size=1)
 
     assert valid_mask == [False]
     assert decode.status == MessageStatus.WAITING
     assert long_chunk.status == MessageStatus.RUNNING
     assert long_chunk.num_blocks == 1
     assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
-
-
-def test_schedule_prefill_allocates_only_first_long_context_chunk():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=2)
-    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
-
-    output = scheduler.schedule(is_prefill=True, prealloc_size=1)
-
-    assert output.running == [long_seq]
-    assert long_seq.status == MessageStatus.READY
-    assert long_seq.kv_token_limit == block_size * 2
-    assert long_seq.num_blocks == 2
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
-
-
-def test_schedule_prefill_short_only_skips_long_waiter_without_mutation():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-    head_long = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
-    short_a = scheduler.add_session(101).add_sequence([2] * (block_size // 2))
-    short_b = scheduler.add_session(102).add_sequence([3] * (block_size // 2))
-
-    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
-
-    assert output.running == [short_a, short_b]
-    assert head_long.status == MessageStatus.WAITING
-    assert head_long.num_blocks == 0
-    assert head_long.kv_token_limit is None
-    assert short_a.status == MessageStatus.READY
-    assert short_b.status == MessageStatus.READY
-
-    short_a.session.remove_sequence(short_a)
-    short_b.session.remove_sequence(short_b)
-    next_output = scheduler.schedule(is_prefill=True)
-
-    assert next_output.running == [head_long]
-    assert head_long.status == MessageStatus.READY
-    assert head_long.kv_token_limit == block_size * 2
-    assert head_long.num_blocks == 2
-
-
-def test_schedule_prefill_prefer_long_admits_oldest_long_waiter_first():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-    short_a = scheduler.add_session(100).add_sequence([1] * (block_size // 2))
-    old_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
-    short_b = scheduler.add_session(102).add_sequence([3] * (block_size // 2))
-    new_long = scheduler.add_session(103).add_sequence([4] * (block_size * 4))
-
-    assert scheduler.has_waiting_long_prefill()
-
-    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
-
-    assert output.running == [old_long]
-    assert old_long.status == MessageStatus.READY
-    assert old_long.kv_token_limit == block_size * 2
-    assert old_long.num_blocks == 2
-    assert short_a.status == MessageStatus.WAITING
-    assert short_a.num_blocks == 0
-    assert short_b.status == MessageStatus.WAITING
-    assert short_b.num_blocks == 0
-    assert new_long.status == MessageStatus.WAITING
-    assert new_long.num_blocks == 0
-    assert new_long.kv_token_limit is None
-
-
-def test_scheduler_reads_opt_ttft_env(monkeypatch):
-    monkeypatch.setattr(prefill_scheduler_module._envs, 'opt_ttft_policy',
-                        'fifo')
-    monkeypatch.setattr(prefill_scheduler_module._envs, 'opt_ttft_aging_sec',
-                        0.25)
-
-    scheduler, _ = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-
-    assert scheduler._prefill_scheduler._long_prefill_policy == 'fifo'
-    assert scheduler._prefill_scheduler._long_prefill_aging_seconds_per_chunk == 0.25
-
-
-def test_schedule_prefill_prefer_long_fifo_policy_keeps_oldest_huge_waiter_first():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-    scheduler._prefill_scheduler._long_prefill_policy = 'fifo'
-    now = time.perf_counter()
-    huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
-    huge_long.arrive_time = now - 1.0
-    moderate_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
-    moderate_long.arrive_time = now
-
-    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
-
-    assert output.running == [huge_long]
-    assert huge_long.status == MessageStatus.READY
-    assert huge_long.kv_token_limit == block_size * 2
-    assert huge_long.num_blocks == 2
-    assert moderate_long.status == MessageStatus.WAITING
-    assert moderate_long.num_blocks == 0
-    assert moderate_long.kv_token_limit is None
-
-
-def test_schedule_prefill_prefer_long_admits_smaller_long_waiter_first():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-    now = time.perf_counter()
-    huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
-    huge_long.arrive_time = now - 1.0
-    moderate_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
-    moderate_long.arrive_time = now
-    short = scheduler.add_session(102).add_sequence([3] * (block_size // 2))
-
-    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
-
-    assert output.running == [moderate_long]
-    assert moderate_long.status == MessageStatus.READY
-    assert moderate_long.kv_token_limit == block_size * 2
-    assert moderate_long.num_blocks == 2
-    assert huge_long.status == MessageStatus.WAITING
-    assert huge_long.num_blocks == 0
-    assert huge_long.kv_token_limit is None
-    assert short.status == MessageStatus.WAITING
-    assert short.num_blocks == 0
-
-
-def test_schedule_prefill_prefer_long_ages_huge_long_waiter():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=8)
-    scheduler._prefill_scheduler._long_prefill_aging_seconds_per_chunk = 0.01
-    now = time.perf_counter()
-    huge_long = scheduler.add_session(100).add_sequence([1] * (block_size * 16))
-    huge_long.arrive_time = now - 1.0
-    moderate_long = scheduler.add_session(101).add_sequence([2] * (block_size * 4))
-    moderate_long.arrive_time = now
-
-    output = scheduler.schedule(is_prefill=True, prefer_long_prefill=True)
-
-    assert output.running == [huge_long]
-    assert huge_long.status == MessageStatus.READY
-    assert huge_long.kv_token_limit == block_size * 2
-    assert huge_long.num_blocks == 2
-    assert moderate_long.status == MessageStatus.WAITING
-    assert moderate_long.num_blocks == 0
-    assert moderate_long.kv_token_limit is None
-
-
-def test_schedule_prefill_reapplies_chunk_limit_after_ssm_state_rollback():
-    scheduler, block_size = _make_ssm_scheduler_for_long_context_chunks(num_gpu_blocks=2)
-    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
-
-    ensure_results = iter([False, True])
-
-    def _ensure_runtime_state_available_once_then_succeed():
-        return next(ensure_results)
-
-    scheduler._prefill_scheduler._ensure_runtime_state_available = (
-        _ensure_runtime_state_available_once_then_succeed)
-
-    output = scheduler.schedule(is_prefill=True, prealloc_size=1)
-
-    assert output.running == [long_seq]
-    assert long_seq.status == MessageStatus.READY
-    assert long_seq.kv_token_limit == block_size * 2
-    assert long_seq.num_blocks == 2
-
-
-def test_reserve_long_context_chunk_grows_one_chunk_at_a_time():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=6)
-    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 5))
-
-    output = scheduler.schedule(is_prefill=True, prealloc_size=1)
-    assert output.running == [long_seq]
-    assert long_seq.kv_token_limit == block_size * 2
-    assert long_seq.num_blocks == 2
-
-    scheduler.activate_seqs([long_seq])
-    long_seq.set_step(block_size * 2)
-
-    assert scheduler.reserve_long_context_chunk(long_seq, block_size * 2)
-    assert long_seq.status == MessageStatus.RUNNING
-    assert long_seq.kv_token_limit == block_size * 4
-    assert long_seq.num_blocks == 4
-
-    long_seq.set_step(block_size * 4)
-
-    assert scheduler.reserve_long_context_chunk(long_seq, block_size, prealloc_size=1, is_last_chunk=True)
-    assert long_seq.kv_token_limit is None
-    assert long_seq.num_blocks == 6
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 0
-
-
-def test_reserve_long_context_chunk_failure_preserves_committed_prefix():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=2)
-    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
-
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [long_seq]
-    scheduler.activate_seqs([long_seq])
-    long_seq.set_step(block_size * 2)
-
-    assert not scheduler.reserve_long_context_chunk(long_seq, block_size * 2)
-    assert long_seq.status == MessageStatus.RUNNING
-    assert long_seq.kv_token_limit == block_size * 2
-    assert long_seq.num_blocks == 2
-
-
-def test_reserve_last_long_context_chunk_failure_restores_chunk_limit():
-    scheduler, block_size = _make_scheduler_for_long_context_chunks(num_gpu_blocks=3)
-    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
-
-    output = scheduler.schedule(is_prefill=True)
-    assert output.running == [long_seq]
-    scheduler.activate_seqs([long_seq])
-    long_seq.set_step(block_size * 2)
-
-    assert not scheduler.reserve_long_context_chunk(long_seq,
-                                                    block_size * 2,
-                                                    prealloc_size=1,
-                                                    is_last_chunk=True)
-    assert long_seq.status == MessageStatus.RUNNING
-    assert long_seq.kv_token_limit == block_size * 2
-    assert long_seq.num_blocks == 2
-    assert scheduler.block_manager.get_num_free_gpu_blocks() == 1
-
-
-def test_scheduler_accepts_prefix_hit_that_starts_middle_long_context_chunk():
-    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
-    block_size = 16
-    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
-    cache_config = CacheConfig(max_batches=1,
-                               block_size=block_size,
-                               num_cpu_blocks=0,
-                               num_gpu_blocks=8,
-                               max_prefill_token_num=block_size * 2,
-                               enable_prefix_caching=True)
-    scheduler_config = SchedulerConfig(max_batches=1,
-                                       max_session_len=128,
-                                       max_request_output_len=64,
-                                       eviction_type='recompute')
-    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
-
-    cached = scheduler.add_session(0).add_sequence([1] * block_size + [2] * block_size)
-    scheduler.block_manager.allocate(cached)
-    scheduler.block_trie.allocate(cached)
-    cached.state.stop()
-
-    token_ids = [1] * block_size + [2] * block_size + [3] * block_size
-    token_ids += [4] * block_size + [5] * block_size
-    seq = scheduler.add_session(1).add_sequence(token_ids)
-
-    output = scheduler.schedule(is_prefill=True)
-
-    assert output.running == [seq]
-    assert seq.num_history_ids == block_size * 2
-    assert seq.num_token_ids == len(token_ids) - block_size * 2
-    assert seq.cached_tokens == block_size * 2
-    assert scheduler.block_trie.stats.num_query_tokens == len(token_ids)
-    assert scheduler.block_trie.stats.num_hit_tokens == block_size * 2
diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py
new file mode 100644
index 0000000000..8abc12e35c
--- /dev/null
+++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py
@@ -0,0 +1,787 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from unittest.mock import Mock
+
+import torch
+
+from lmdeploy.messages import KVTransferConfig
+from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
+from lmdeploy.pytorch.kv_connector import (
+    KVConnectorMetadata,
+    KVConnectorOutput,
+    KVConnectorResult,
+    KVLoadResult,
+    KVSaveBlockLease,
+)
+from lmdeploy.pytorch.kv_connector.mooncake.store.scheduler import MooncakeStoreScheduler
+from lmdeploy.pytorch.messages import MessageStatus, SequenceMeta
+from lmdeploy.pytorch.paging.scheduler import Scheduler
+
+
+class _AsyncLookupConnector:
+
+    def __init__(self, results, failed_ids=()):
+        self.results = iter(results)
+        self.failed_ids = set(failed_ids)
+        self.pending_ids = set()
+        self.lookup_calls = []
+        self.cancelled = []
+        self.finished = []
+        self.allocations = []
+
+    def on_new_request(self, request):
+        pass
+
+    def is_lookup_pending(self, request_id):
+        return request_id in self.pending_ids
+
+    def get_num_new_matched_tokens(self, request, num_computed_tokens):
+        self.lookup_calls.append((request.seq_id, num_computed_tokens))
+        result = next(self.results)
+        if result[0] is None:
+            self.pending_ids.add(request.seq_id)
+        else:
+            self.pending_ids.discard(request.seq_id)
+        return result
+
+    def cancel_lookup(self, request_id):
+        self.pending_ids.discard(request_id)
+        self.cancelled.append(request_id)
+
+    def update_state_after_alloc(self, request, block_ids, num_external_tokens):
+        self.allocations.append((request.seq_id, tuple(block_ids), num_external_tokens))
+
+    def build_connector_meta(self, scheduler_output):
+        return None
+
+    def update_connector_output(self, connector_output):
+        return KVConnectorResult(
+            load_results=tuple(
+                KVLoadResult(
+                    request_id=request_id,
+                    success=request_id not in self.failed_ids,
+                )
+                for request_id in (connector_output.finished_receiving or set())
+            )
+        )
+
+    def request_finished(self, request):
+        self.finished.append(request.seq_id)
+
+    def finish_transfers_after_worker_drain(self):
+        pass
+
+    def shutdown(self):
+        pass
+
+
+def _make_async_lookup_scheduler(
+    connector,
+    *,
+    enable_prefix_caching=True,
+    max_batches=1,
+    num_gpu_blocks=16,
+    max_prefill_token_num=8192,
+):
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 4
+    return Scheduler(
+        scheduler_config=SchedulerConfig(
+            max_batches=max_batches,
+            max_session_len=64,
+            max_request_output_len=16,
+            eviction_type='recompute',
+        ),
+        cache_config=CacheConfig(
+            max_batches=max_batches,
+            block_size=block_size,
+            num_cpu_blocks=0,
+            num_gpu_blocks=num_gpu_blocks,
+            max_prefill_token_num=max_prefill_token_num,
+            enable_prefix_caching=enable_prefix_caching,
+            kv_transfer_config=KVTransferConfig(
+                kv_connector='MooncakeStoreConnector',
+                kv_role='kv_both',
+            ),
+        ),
+        seq_meta=SequenceMeta(block_size, strategy=ARSequenceStrategy()),
+        kv_connector=connector,
+    )
+
+
+def test_async_lookup_pending_rolls_back_a_new_request_once():
+    connector = _AsyncLookupConnector([(None, False), (0, False)])
+    scheduler = _make_async_lookup_scheduler(connector)
+    seq = scheduler.add_session(70).add_sequence(torch.arange(9))
+
+    first = scheduler.schedule(is_prefill=True)
+    assert first.running == []
+    assert scheduler.last_schedule_had_pending_lookup
+    assert seq.num_history_ids == 0
+    assert seq.num_blocks == 0
+    assert seq.prefix_cache.trie_cursor is None
+    assert seq.prefix_cache.match_start_step == -1
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+
+    second = scheduler.schedule(is_prefill=True)
+    assert second.running == []
+    assert connector.lookup_calls == [(seq.seq_id, 0)]
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+
+    connector.pending_ids.clear()
+    third = scheduler.schedule(is_prefill=True)
+    assert third.running == [seq]
+    assert connector.lookup_calls == [(seq.seq_id, 0), (seq.seq_id, 0)]
+
+
+def test_async_lookup_rebases_remote_hit_after_local_trie_grows(monkeypatch):
+    connector = MooncakeStoreScheduler(
+        CacheConfig(
+            max_batches=1,
+            block_size=4,
+            num_cpu_blocks=0,
+            num_gpu_blocks=16,
+            enable_prefix_caching=True,
+            kv_transfer_config=KVTransferConfig(
+                kv_connector='MooncakeStoreConnector',
+                kv_role='kv_both',
+            ),
+        ))
+    assert connector.client is not None
+    # The same asynchronous lookup first reports pending, then returns its
+    # absolute remote prefix boundary from the original snapshot.
+    monkeypatch.setattr(connector.client, 'lookup', Mock(side_effect=(None, 16)))
+    observed_results = []
+    get_matched_tokens = connector.get_num_new_matched_tokens
+
+    def _record_result(request, num_computed_tokens):
+        result = get_matched_tokens(request, num_computed_tokens)
+        observed_results.append((num_computed_tokens, result))
+        return result
+
+    monkeypatch.setattr(connector, 'get_num_new_matched_tokens', _record_result)
+    scheduler = _make_async_lookup_scheduler(connector)
+    tokens = torch.arange(17)
+
+    cached_to_8 = scheduler.add_session(80).add_sequence(tokens[:9])
+    scheduler.block_manager.allocate(cached_to_8)
+    scheduler.block_trie.allocate(cached_to_8)
+    cached_to_8.state.stop()
+
+    seq = scheduler.add_session(81).add_sequence(tokens)
+    first = scheduler.schedule(is_prefill=True)
+
+    assert first.running == []
+    assert seq.num_history_ids == 0
+    assert observed_results == [(8, (None, False))]
+
+    # While the remote lookup is pending, another sequence publishes the next
+    # complete local block. The retried match must now advance from 8 to 12.
+    cached_to_12 = scheduler.add_session(82).add_sequence(tokens[:13])
+    scheduler.block_manager.allocate(cached_to_12)
+    scheduler.block_trie.allocate(cached_to_12)
+    cached_to_12.state.stop()
+
+    second = scheduler.schedule(is_prefill=True)
+
+    assert second.running == []
+    assert seq.num_history_ids == 12
+    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert observed_results == [
+        (8, (None, False)),
+        (12, (4, True)),
+    ]
+
+    metadata = connector.build_connector_meta(second)
+    assert metadata is not None
+    assert len(metadata.load_requests) == 1
+    load_request = metadata.load_requests[0]
+    block_table = scheduler.block_manager.get_block_table(seq)
+    assert load_request.block_ids == (int(block_table[3]), )
+    assert load_request.remote_block_count == 4
+    scheduler.shutdown()
+
+
+def test_async_lookup_pending_request_does_not_block_later_waiter():
+    connector = _AsyncLookupConnector([(None, False), (0, False)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+    )
+    pending = scheduler.add_session(74).add_sequence(torch.arange(9))
+    schedulable = scheduler.add_session(75).add_sequence(torch.arange(9, 18))
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [schedulable]
+    assert pending.status == MessageStatus.WAITING
+    assert schedulable.status == MessageStatus.READY
+    assert pending.seq_id in connector.pending_ids
+    assert connector.lookup_calls == [
+        (pending.seq_id, 0),
+        (schedulable.seq_id, 0),
+    ]
+    assert scheduler.last_schedule_had_pending_lookup
+
+
+def test_async_lookup_precisely_restores_a_multiturn_local_prefix():
+    connector = _AsyncLookupConnector([(None, False), (4, True)])
+    scheduler = _make_async_lookup_scheduler(connector)
+    tokens = torch.arange(13)
+    seq = scheduler.add_session(71).add_sequence(tokens)
+
+    seq.kv_token_limit = 4
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    seq.set_step(4)
+    seq.kv_token_limit = 5
+    seq.cached_tokens = 3
+    seq.model_meta = {'state': 'keep'}
+    seq.prefix_cache.recompute_overlap.fresh_block_range = range(0, 1)
+    seq.prefix_cache.recompute_overlap.trie_block_map[0] = seq.logical_blocks[0]
+    baseline_blocks = seq.logical_blocks.get_real_blocks().copy()
+    baseline_cursor = seq.prefix_cache.trie_cursor
+
+    cached = scheduler.add_session(72).add_sequence(tokens[:9])
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached.state.stop()
+    matched_block = cached.logical_blocks[1]
+    matched_ref_count = scheduler.block_manager.allocator.get_ref_count(
+        cached.logical_blocks.get_real_blocks()[1:2]).copy()
+    scheduler.block_trie.stats.reset()
+
+    first = scheduler.schedule(is_prefill=True)
+    assert first.running == []
+    assert scheduler.last_schedule_had_pending_lookup
+    assert seq.num_history_ids == 4
+    assert seq.num_blocks == 1
+    assert torch.equal(torch.from_numpy(seq.logical_blocks.get_real_blocks()),
+                       torch.from_numpy(baseline_blocks))
+    assert seq.prefix_cache.trie_cursor is baseline_cursor
+    assert seq.prefix_cache.match_start_step == -1
+    assert seq.prefix_cache.recompute_overlap.fresh_block_range == range(0, 1)
+    assert seq.prefix_cache.recompute_overlap.trie_block_map == {0: baseline_blocks[0]}
+    assert seq.cached_tokens == 3
+    assert seq.kv_token_limit == 5
+    assert seq.model_meta == {'state': 'keep'}
+    assert scheduler.block_manager.allocator.get_ref_count(
+        cached.logical_blocks.get_real_blocks()[1:2]).tolist() == matched_ref_count.tolist()
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+
+    second = scheduler.schedule(is_prefill=True)
+    assert second.running == []
+    assert connector.lookup_calls == [(seq.seq_id, 8)]
+    assert seq.num_history_ids == 4
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+
+    connector.pending_ids.clear()
+    seq.kv_token_limit = None
+    seq.prefix_cache.recompute_overlap.clear_tracking()
+    third = scheduler.schedule(is_prefill=True)
+    assert third.running == []
+    assert connector.lookup_calls == [(seq.seq_id, 8), (seq.seq_id, 8)]
+    assert seq.num_history_ids == 8
+    assert matched_block in seq.logical_blocks.get_real_blocks()
+    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert connector.allocations[0][2] == 4
+
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={seq.seq_id}))
+    assert seq.num_history_ids == 12
+    assert seq.status == MessageStatus.WAITING
+
+    fourth = scheduler.schedule(is_prefill=True)
+    assert fourth.running == [seq]
+
+
+def test_async_lookup_pending_preserves_private_partial_prefix():
+    connector = _AsyncLookupConnector([(None, False)])
+    scheduler = _make_async_lookup_scheduler(connector)
+    tokens = torch.arange(13)
+    seq = scheduler.add_session(72).add_sequence(tokens)
+
+    seq.kv_token_limit = 5
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    seq.set_step(5)
+    seq.kv_token_limit = 7
+    seq.cached_tokens = 3
+    seq.model_meta = {'state': 'keep'}
+    baseline_blocks = seq.logical_blocks.get_real_blocks().copy()
+    baseline_cursor = seq.prefix_cache.trie_cursor
+    scheduler.block_trie.stats.reset()
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert scheduler.last_schedule_had_pending_lookup
+    assert connector.lookup_calls == [(seq.seq_id, 5)]
+    assert seq.num_history_ids == 5
+    assert torch.equal(torch.from_numpy(seq.logical_blocks.get_real_blocks()),
+                       torch.from_numpy(baseline_blocks))
+    assert seq.prefix_cache.trie_cursor is baseline_cursor
+    assert seq.prefix_cache.match_start_step == -1
+    assert seq.cached_tokens == 3
+    assert seq.kv_token_limit == 7
+    assert seq.model_meta == {'state': 'keep'}
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+
+
+def test_external_cached_tokens_survive_remote_ready_admission():
+    connector = _AsyncLookupConnector([(8, True)])
+    scheduler = _make_async_lookup_scheduler(connector)
+    seq = scheduler.add_session(73).add_sequence(torch.arange(13))
+
+    started = scheduler.schedule(is_prefill=True)
+    assert started.running == []
+
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={seq.seq_id}))
+    assert seq.num_history_ids == 8
+    assert seq.cached_tokens == 8
+    assert seq.prefix_cache.match_start_step == 0
+
+    admitted = scheduler.schedule(is_prefill=True)
+
+    assert admitted.running == [seq]
+    assert seq.cached_tokens == 8
+    assert seq.prefix_cache.match_start_step == 0
+
+
+def test_remote_ready_long_prefill_respects_short_only_turn():
+    connector = _AsyncLookupConnector([(8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        max_prefill_token_num=4,
+    )
+    seq = scheduler.add_session(74).add_sequence(torch.arange(17))
+
+    started = scheduler.schedule(is_prefill=True)
+    assert started.running == []
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={seq.seq_id}))
+    assert seq.num_history_ids == 8
+    assert scheduler.kv_load_coordinator.is_remote_ready(seq)
+
+    short_turn = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert short_turn.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_history_ids == 8
+    assert scheduler.kv_load_coordinator.is_remote_ready(seq)
+
+    long_turn = scheduler.schedule(is_prefill=True)
+    assert long_turn.running == [seq]
+    assert seq.status == MessageStatus.READY
+
+
+def test_external_cached_tokens_survive_prefill_budget_rejection():
+    connector = _AsyncLookupConnector([(8, True), (8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        max_batches=2,
+        max_prefill_token_num=8,
+    )
+    admitted_seq = scheduler.add_session(74).add_sequence(torch.arange(13))
+    waiting_seq = scheduler.add_session(75).add_sequence(torch.arange(20, 33))
+
+    started = scheduler.schedule(is_prefill=True)
+    assert started.running == []
+    scheduler.update_connector_output(
+        KVConnectorOutput(
+            finished_receiving={admitted_seq.seq_id, waiting_seq.seq_id}))
+
+    admitted = scheduler.schedule(is_prefill=True)
+
+    assert admitted.running == [admitted_seq]
+    assert waiting_seq.status == MessageStatus.WAITING
+    assert waiting_seq.num_history_ids == 8
+    assert waiting_seq.num_blocks == 2
+    assert waiting_seq.cached_tokens == 8
+    assert waiting_seq.prefix_cache.match_start_step == 0
+    assert scheduler.kv_load_coordinator.is_remote_ready(waiting_seq)
+
+
+def test_async_load_keeps_a_private_partial_block_at_the_suffix_start():
+    connector = _AsyncLookupConnector([(7, True)])
+    scheduler = _make_async_lookup_scheduler(connector)
+    tokens = torch.arange(13)
+    cached = scheduler.add_session(76).add_sequence(tokens)
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached.state.stop()
+
+    seq = scheduler.add_session(77).add_sequence(tokens)
+    seq.kv_token_limit = 5
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    seq.set_step(5)
+    seq.kv_token_limit = None
+    private_block = int(scheduler.block_manager.get_block_table(seq)[1])
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert seq.num_blocks == 3
+    assert connector.lookup_calls == [(seq.seq_id, 5)]
+    load_blocks = connector.allocations[0][1]
+    assert load_blocks == tuple(
+        int(block_id)
+        for block_id in scheduler.block_manager.get_block_table(seq)[1:3]
+    )
+    assert load_blocks[0] == private_block
+
+
+def test_async_load_requires_capacity_for_the_complete_prefill():
+    connector = _AsyncLookupConnector([(8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        num_gpu_blocks=3,
+    )
+    seq = scheduler.add_session(76).add_sequence(torch.arange(13))
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_blocks == 0
+    assert connector.allocations == []
+
+
+def test_async_load_capacity_failure_restores_tentative_local_prefix(monkeypatch):
+    connector = _AsyncLookupConnector([(4, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        num_gpu_blocks=3,
+    )
+    tokens = torch.arange(13)
+    cached = scheduler.add_session(76).add_sequence(tokens[:5])
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached.state.stop()
+    cached_block = cached.logical_blocks.get_real_blocks()[:1]
+    ref_count = scheduler.block_manager.allocator.get_ref_count(cached_block).copy()
+    scheduler.block_trie.stats.reset()
+
+    seq = scheduler.add_session(77).add_sequence(tokens)
+    evict_for_seq = Mock(return_value=False)
+    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_history_ids == 0
+    assert seq.num_blocks == 0
+    assert seq.kv_token_limit is None
+    assert seq.cached_tokens == 0
+    assert seq.prefix_cache.trie_cursor is None
+    assert seq.prefix_cache.match_start_step == -1
+    assert connector.lookup_calls == [(seq.seq_id, 4)]
+    assert connector.allocations == []
+    assert evict_for_seq.call_count == 1
+    assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def test_async_load_does_not_consume_model_batch_slot():
+    connector = _AsyncLookupConnector([(8, True), (0, False)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        num_gpu_blocks=5,
+    )
+    loading = scheduler.add_session(77).add_sequence(torch.arange(13))
+    later = scheduler.add_session(78).add_sequence(torch.arange(8))
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [later]
+    assert loading.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert loading.num_blocks == 2
+    assert later.status == MessageStatus.READY
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 2
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 1
+
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={loading.seq_id}))
+    assert loading.num_history_ids == 8
+    assert loading.cached_tokens == 8
+
+
+def test_multiple_async_loads_start_in_one_prefill_turn():
+    connector = _AsyncLookupConnector([(8, True), (8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        max_batches=1,
+        num_gpu_blocks=16,
+    )
+    first = scheduler.add_session(91).add_sequence(torch.arange(13))
+    second = scheduler.add_session(92).add_sequence(torch.arange(20, 33))
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert first.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert second.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert [allocation[0] for allocation in connector.allocations] == [
+        first.seq_id,
+        second.seq_id,
+    ]
+
+
+def test_soft_reservation_blocks_new_load_until_capacity_is_released():
+    connector = _AsyncLookupConnector([
+        (4, True),
+        (12, True),
+        (12, True),
+    ])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        num_gpu_blocks=6,
+    )
+    first = scheduler.add_session(86).add_sequence(torch.arange(13))
+
+    scheduler.schedule(is_prefill=True)
+    assert first.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert first.num_blocks == 1
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 3
+
+    second = scheduler.add_session(87).add_sequence(torch.arange(17))
+    blocked = scheduler.schedule(is_prefill=True)
+    assert blocked.running == []
+    assert second.status == MessageStatus.WAITING
+    assert second.num_blocks == 0
+    assert [allocation[0] for allocation in connector.allocations] == [first.seq_id]
+
+    scheduler.end_session(86)
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={first.seq_id}))
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
+
+    retried = scheduler.schedule(is_prefill=True)
+    assert retried.running == []
+    assert second.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert second.num_blocks == 3
+    assert [allocation[0] for allocation in connector.allocations] == [
+        first.seq_id,
+        second.seq_id,
+    ]
+
+
+def test_failed_async_load_preserves_local_prefix_and_releases_remote_tail():
+    connector = _AsyncLookupConnector([(4, True)])
+    scheduler = _make_async_lookup_scheduler(connector)
+    tokens = torch.arange(13)
+    cached = scheduler.add_session(79).add_sequence(tokens[:9])
+    scheduler.block_manager.allocate(cached)
+    scheduler.block_trie.allocate(cached)
+    cached.state.stop()
+    seq = scheduler.add_session(80).add_sequence(tokens)
+    scheduler.schedule(is_prefill=True)
+    connector.failed_ids.add(seq.seq_id)
+
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={seq.seq_id}))
+
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_history_ids == 8
+    assert seq.num_blocks == 2
+    assert seq.cached_tokens == 8
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
+
+
+def test_async_load_soft_reservation_shrinks_across_chunks():
+    connector = _AsyncLookupConnector([(4, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        max_prefill_token_num=4,
+    )
+    seq = scheduler.add_session(81).add_sequence(torch.arange(13))
+
+    scheduler.schedule(is_prefill=True)
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 3
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={seq.seq_id}))
+
+    assert scheduler.schedule(is_prefill=True).running == [seq]
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 2
+    seq.set_step(8)
+    scheduler.release_completed_prefill_reservations([seq])
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 2
+
+    assert scheduler.reserve_long_context_chunk(seq, chunk_size=4)
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 1
+    seq.set_step(12)
+    assert scheduler.reserve_long_context_chunk(seq, chunk_size=1, is_last_chunk=True)
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
+
+
+def test_end_session_waits_for_active_async_load_before_freeing_blocks():
+    connector = _AsyncLookupConnector([(8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+    )
+    seq = scheduler.add_session(82).add_sequence(torch.arange(13))
+    scheduler.schedule(is_prefill=True)
+
+    scheduler.end_session(82)
+    assert 82 in scheduler.sessions
+    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert seq.num_blocks == 2
+    assert connector.finished == []
+
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={seq.seq_id}))
+    assert 82 not in scheduler.sessions
+    assert connector.finished == [seq.seq_id]
+
+
+def test_worker_drain_finishes_an_ended_session_with_a_dropped_load_output():
+    connector = _AsyncLookupConnector([(8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+    )
+    seq = scheduler.add_session(85).add_sequence(torch.arange(13))
+    scheduler.schedule(is_prefill=True)
+    scheduler.end_session(85)
+
+    scheduler.finish_deferred_kv_transfers_after_worker_drain()
+
+    assert 85 not in scheduler.sessions
+    assert connector.finished == [seq.seq_id]
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
+
+
+def test_completed_async_load_is_admitted_before_an_older_waiter():
+    connector = _AsyncLookupConnector([(8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        num_gpu_blocks=4,
+    )
+    loaded = scheduler.add_session(83).add_sequence(torch.arange(13))
+    scheduler.schedule(is_prefill=True)
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={loaded.seq_id}))
+    newcomer = scheduler.add_session(84).add_sequence(torch.arange(4))
+    newcomer.arrive_time = loaded.arrive_time - 1
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [loaded]
+    assert newcomer.status == MessageStatus.WAITING
+
+
+def test_stop_and_end_session_cancel_lookup_before_request_cleanup():
+    connector = _AsyncLookupConnector([(None, False)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+    )
+    seq = scheduler.add_session(73).add_sequence(torch.arange(9))
+    scheduler.schedule(is_prefill=True)
+
+    scheduler.stop_session(73)
+    assert connector.cancelled == [seq.seq_id]
+
+    scheduler.end_session(73)
+    assert connector.finished == [seq.seq_id]
+
+
+def test_async_save_lease_keeps_exact_blocks_alive_until_all_tp_complete():
+
+    class _SaveMetadata(KVConnectorMetadata):
+
+        def __init__(self, logical_block_ids):
+            self.logical_block_ids = tuple(logical_block_ids)
+
+        def get_save_block_leases(self):
+            return (KVSaveBlockLease(7, self.logical_block_ids), )
+
+    class _SaveConnector(_AsyncLookupConnector):
+
+        def __init__(self):
+            super().__init__([])
+            self.metadata = None
+
+        def build_connector_meta(self, scheduler_output):
+            metadata, self.metadata = self.metadata, None
+            return metadata
+
+        def update_connector_output(self, connector_output):
+            return KVConnectorResult(
+                completed_save_ids=frozenset(
+                    connector_output.completed_save_ids or ()),
+            )
+
+    connector = _SaveConnector()
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        num_gpu_blocks=4,
+    )
+    seq = scheduler.add_session(88).add_sequence(torch.arange(8))
+    scheduler.block_manager.allocate(seq)
+    logical_blocks = seq.logical_blocks.get_real_blocks().copy()
+    allocator = scheduler.block_manager.allocator
+    connector.metadata = _SaveMetadata(logical_blocks)
+
+    metadata = scheduler.build_connector_meta(
+        [seq],
+        connector_token_lens=(8, ),
+    )
+
+    assert metadata is not None
+    assert allocator.get_ref_count(logical_blocks).tolist() == [2, 2]
+    assert scheduler.has_unfinished()
+
+    # Sequence ownership may disappear before remote I/O completes. The save
+    # lease remains as the only reference and prevents physical reuse.
+    scheduler.block_manager.free(seq)
+    assert allocator.get_ref_count(logical_blocks).tolist() == [1, 1]
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 2
+
+    scheduler.update_connector_output(
+        KVConnectorOutput(completed_save_ids={7}))
+    assert allocator.get_ref_count(logical_blocks).tolist() == [0, 0]
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 4
+    assert not scheduler.kv_save_coordinator.has_pending()
+
+
+def test_worker_drain_releases_save_leases_when_outputs_are_discarded():
+
+    class _SaveMetadata(KVConnectorMetadata):
+
+        def __init__(self, logical_block_ids):
+            self.logical_block_ids = tuple(logical_block_ids)
+
+        def get_save_block_leases(self):
+            return (KVSaveBlockLease(9, self.logical_block_ids), )
+
+    connector = _AsyncLookupConnector([])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        num_gpu_blocks=2,
+    )
+    seq = scheduler.add_session(89).add_sequence(torch.arange(8))
+    scheduler.block_manager.allocate(seq)
+    logical_blocks = seq.logical_blocks.get_real_blocks().copy()
+    metadata = _SaveMetadata(logical_blocks)
+    connector.build_connector_meta = lambda scheduler_output: metadata
+
+    scheduler.build_connector_meta([seq], connector_token_lens=(8, ))
+    scheduler.block_manager.free(seq)
+    scheduler.finish_deferred_kv_transfers_after_worker_drain()
+
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 2
+    assert not scheduler.kv_save_coordinator.has_pending()
diff --git a/tests/pytorch/paging/test_scheduler_ssm.py b/tests/pytorch/paging/test_scheduler_ssm.py
new file mode 100644
index 0000000000..399a24812f
--- /dev/null
+++ b/tests/pytorch/paging/test_scheduler_ssm.py
@@ -0,0 +1,393 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+import torch
+
+from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
+from lmdeploy.pytorch.engine.inputs_maker import _make_state_prefix_cache_save_plan
+from lmdeploy.pytorch.messages import MessageStatus, SequenceMeta, UpdateTokenMode
+from lmdeploy.pytorch.paging.scheduler import Scheduler
+
+
+def _make_ssm_scheduler(max_batch_size: int = 1, prefix_cache_state_budget: int = 0, num_gpu_blocks: int = 16):
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 16
+    cache_config = CacheConfig(max_batches=max_batch_size,
+                               block_size=block_size,
+                               num_cpu_blocks=4,
+                               num_gpu_blocks=num_gpu_blocks,
+                               enable_prefix_caching=True,
+                               num_state_caches=max_batch_size + 1 + prefix_cache_state_budget,
+                               prefix_cache_state_budget=prefix_cache_state_budget,
+                               states_shapes=[((1, ), torch.float32)])
+    scheduler_config = SchedulerConfig(max_batches=max_batch_size,
+                                       max_session_len=128,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    return Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+
+
+def _add_published_ssm_checkpoint(scheduler: Scheduler, token_ids: list[int]):
+    session = scheduler.add_session(len(scheduler.sessions))
+    seq = session.add_sequence(token_ids)
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    state_idx = scheduler.block_trie.state_checkpoints.reserve_save(seq)
+    assert state_idx >= 0
+    assert scheduler.block_trie.state_checkpoints.publish_save(seq)
+    node = seq.prefix_cache.trie_cursor
+    session.remove_sequence(seq)
+    return node, state_idx
+
+
+def test_ssm_runtime_state_reclaims_borrowed_checkpoint_slot():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
+    block_size = scheduler.seq_meta.block_size
+    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    seq = scheduler.add_session(100).add_sequence([2] * block_size * 2)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.logical_state == state_idx
+    assert node.state_checkpoint is None
+    assert scheduler.state_manager.get_num_runtime_states() == 1
+    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0
+
+
+def test_ssm_long_chunked_request_schedules_with_only_runtime_state_slot():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
+    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size * 2
+    block_size = scheduler.seq_meta.block_size
+    token_ids = [1] * block_size + [2] * block_size + [3] * block_size
+    seq = scheduler.add_session(100).add_sequence(token_ids)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.logical_state >= 0
+    assert scheduler.state_manager.get_num_runtime_states() == 1
+    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0
+    assert scheduler.block_trie.state_checkpoints.reserve_save(seq, step=block_size * 2) == -1
+
+
+def test_ssm_running_request_reuses_own_runtime_state_without_spare_slot():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
+    block_size = scheduler.seq_meta.block_size
+    seq = scheduler.add_session(100).add_sequence([1] * block_size)
+
+    output = scheduler.schedule(is_prefill=True)
+    assert output.running == [seq]
+    assert scheduler.state_manager.get_num_free_runtime() == 0
+    seq.state.activate()
+
+    seq.update_token_ids([2] * block_size, mode=UpdateTokenMode.DECODE)
+    valid_mask = scheduler.schedule_running([seq], num_required_tokens=0, prealloc_size=0)
+
+    assert valid_mask == [True]
+    assert seq.status == MessageStatus.RUNNING
+    assert seq.logical_state >= 0
+    assert seq.num_blocks == 2
+    assert scheduler.state_manager.get_num_runtime_states() == 1
+    assert scheduler.state_manager.get_num_free_runtime() == 0
+
+
+def test_ssm_runtime_state_waits_when_only_checkpoint_slot_is_pinned():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
+    block_size = scheduler.seq_meta.block_size
+    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    node.state_checkpoint.pin_count = 1
+    seq = scheduler.add_session(100).add_sequence([2] * block_size * 2)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.logical_state == -1
+    assert node.state_checkpoint.slot == state_idx
+    assert node.state_checkpoint.published
+
+
+def test_ssm_same_batch_duplicate_checkpoint_save_has_unique_dst_offsets():
+    scheduler = _make_ssm_scheduler(max_batch_size=2, prefix_cache_state_budget=2)
+    block_size = scheduler.seq_meta.block_size
+    token_ids = [1] * block_size * 2
+
+    seq_a = scheduler.add_session(100).add_sequence(token_ids)
+    seq_b = scheduler.add_session(101).add_sequence(token_ids)
+
+    output = scheduler.schedule(is_prefill=True)
+    assert output.running == [seq_a, seq_b]
+    assert seq_a.logical_state >= 0
+    assert seq_b.logical_state >= 0
+    assert seq_a.logical_state != seq_b.logical_state
+    assert seq_a.prefix_cache.trie_cursor is seq_b.prefix_cache.trie_cursor
+
+    save_state_offsets = [
+        scheduler.block_trie.state_checkpoints.reserve_save(seq) for seq in output.running
+    ]
+    save_plan = _make_state_prefix_cache_save_plan(output.running, save_state_offsets)
+    assert save_plan is not None
+    save_src_offsets, save_dst_offsets = save_plan
+
+    assert save_src_offsets == (seq_a.logical_state, )
+    assert save_dst_offsets == (save_state_offsets[0], )
+    assert save_state_offsets[0] >= 0
+    assert save_state_offsets[1] == -1
+    assert len(save_dst_offsets) == len(set(save_dst_offsets))
+
+
+def test_ssm_end_session_discards_pending_checkpoint_reservation():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
+    block_size = scheduler.seq_meta.block_size
+    session = scheduler.add_session(100)
+    seq = session.add_sequence([1] * block_size * 2)
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    scheduler.state_manager.allocate(seq)
+
+    state_idx = scheduler.block_trie.state_checkpoints.reserve_save(seq)
+    node = seq.prefix_cache.pending_save.node
+    assert state_idx >= 0
+    assert node is not None
+    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 1
+
+    scheduler.end_session(100)
+
+    assert 100 not in scheduler.sessions
+    assert node.state_checkpoint is None
+    assert scheduler.state_manager.get_num_runtime_states() == 0
+    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0
+
+
+def test_ssm_end_session_unpins_restore_checkpoint():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
+    block_size = scheduler.seq_meta.block_size
+    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    seq = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [2])
+
+    scheduler.block_trie.match(seq)
+    assert seq.prefix_cache.restore.slot == state_idx
+    assert scheduler.block_trie.state_checkpoints.pin_restore(seq)
+    assert node.state_checkpoint.pin_count == 1
+
+    scheduler.end_session(100)
+
+    assert 100 not in scheduler.sessions
+    assert node.state_checkpoint.slot == state_idx
+    assert node.state_checkpoint.published
+    assert node.state_checkpoint.pin_count == 0
+
+
+def test_ssm_failed_restore_schedule_rolls_back_match():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=0)
+    block_size = scheduler.seq_meta.block_size
+    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    node.state_checkpoint.pin_count = 1
+    seq = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [2])
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == []
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_history_ids == 0
+    assert len(seq.logical_blocks) == 0
+    assert seq.cached_tokens == 0
+    assert seq.prefix_cache.trie_cursor is None
+    assert seq.prefix_cache.restore.slot == -1
+    assert seq.prefix_cache.restore.node is None
+    assert node.state_checkpoint.slot == state_idx
+    assert node.state_checkpoint.published
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+    node.state_checkpoint.pin_count = 0
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.status == MessageStatus.READY
+    assert seq.num_history_ids == 0
+    assert seq.prefix_cache.restore.slot == -1
+    assert seq.logical_state == state_idx
+    assert node.state_checkpoint is None
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def test_ssm_scheduler_preserves_matched_checkpoint_when_evicting_for_runtime_state():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
+    block_size = scheduler.seq_meta.block_size
+    node_a, state_idx_a = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    node_b, state_idx_b = _add_published_ssm_checkpoint(scheduler, [2] * block_size * 2)
+    seq = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq]
+    assert seq.num_history_ids == block_size * 2
+    assert seq.cached_tokens == block_size * 2
+    assert seq.prefix_cache.restore.slot == state_idx_a
+    assert seq.prefix_cache.restore.node is node_a
+    assert seq.prefix_cache.restore.pinned
+    assert seq.logical_state == state_idx_b
+    assert node_a.state_checkpoint.slot == state_idx_a
+    assert node_a.state_checkpoint.published
+    assert node_a.state_checkpoint.pin_count == 1
+    assert node_b.state_checkpoint is None
+    assert scheduler.block_trie.stats.num_hit_tokens == block_size * 2
+
+    assert scheduler.block_trie.state_checkpoints.unpin_restore(seq)
+
+
+def test_ssm_scheduler_evicts_stopped_runtime_state_with_free_checkpoint_slot():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
+    block_size = scheduler.seq_meta.block_size
+    seq_a = scheduler.add_session(100).add_sequence([1] * block_size)
+
+    output = scheduler.schedule(is_prefill=True)
+    assert output.running == [seq_a]
+    assert seq_a.logical_state >= 0
+    assert scheduler.state_manager.get_num_free() == 1
+    assert scheduler.state_manager.get_num_free_runtime() == 0
+
+    seq_a.state.stop()
+    seq_b = scheduler.add_session(101).add_sequence([2] * block_size)
+
+    output = scheduler.schedule(is_prefill=True)
+
+    assert output.running == [seq_b]
+    assert seq_b.logical_state >= 0
+    assert seq_a.logical_state == -1
+    assert seq_a.status == MessageStatus.STOPPED
+
+
+def test_ssm_scheduler_rolls_back_prefix_match_for_prefill_gate_without_pinning_restore_state():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1)
+    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size
+    block_size = scheduler.seq_meta.block_size
+    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    scheduler.block_trie.stats.reset()
+
+    still_long = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3] * (block_size + 1))
+
+    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert output.running == []
+    assert still_long.status == MessageStatus.WAITING
+    assert still_long.num_history_ids == 0
+    assert still_long.cached_tokens == 0
+    assert still_long.prefix_cache.trie_cursor is None
+    assert still_long.prefix_cache.restore.slot == -1
+    assert still_long.prefix_cache.restore.node is None
+    assert not still_long.prefix_cache.restore.pinned
+    assert node.state_checkpoint.slot == state_idx
+    assert node.state_checkpoint.pin_count == 0
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def test_ssm_scheduler_rejects_prefix_match_for_prefill_gate_after_pinned_restore_rollback():
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1, num_gpu_blocks=2)
+    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size
+    block_size = scheduler.seq_meta.block_size
+    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    scheduler.block_trie.stats.reset()
+
+    cache_hit_tail = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
+
+    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert output.running == []
+    assert cache_hit_tail.status == MessageStatus.WAITING
+    assert cache_hit_tail.num_history_ids == 0
+    assert cache_hit_tail.num_token_ids == block_size * 2 + 1
+    assert cache_hit_tail.num_blocks == 0
+    assert cache_hit_tail.kv_token_limit is None
+    assert cache_hit_tail.logical_state == -1
+    assert cache_hit_tail.cached_tokens == 0
+    assert cache_hit_tail.prefix_cache.trie_cursor is None
+    assert cache_hit_tail.prefix_cache.restore.slot == -1
+    assert cache_hit_tail.prefix_cache.restore.node is None
+    assert not cache_hit_tail.prefix_cache.restore.pinned
+    assert node.state_checkpoint.slot == state_idx
+    assert node.state_checkpoint.published
+    assert node.state_checkpoint.pin_count == 0
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def test_ssm_scheduler_rejects_prefix_match_for_prefill_gate_after_runtime_state_rollback(monkeypatch):
+    scheduler = _make_ssm_scheduler(max_batch_size=1, prefix_cache_state_budget=1, num_gpu_blocks=4)
+    scheduler.cache_config.max_prefill_token_num = scheduler.seq_meta.block_size
+    block_size = scheduler.seq_meta.block_size
+    node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
+    ensure_results = iter([False, True])
+
+    def _ensure_runtime_state_available_once_then_succeed():
+        return next(ensure_results)
+
+    monkeypatch.setattr(scheduler._prefill_scheduler,
+                        '_ensure_runtime_state_available',
+                        _ensure_runtime_state_available_once_then_succeed)
+    scheduler.block_trie.stats.reset()
+
+    cache_hit_tail = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
+
+    output = scheduler.schedule(is_prefill=True, allow_long_prefill=False)
+
+    assert output.running == []
+    assert cache_hit_tail.status == MessageStatus.WAITING
+    assert cache_hit_tail.num_history_ids == 0
+    assert cache_hit_tail.num_token_ids == block_size * 2 + 1
+    assert cache_hit_tail.num_blocks == 0
+    assert cache_hit_tail.kv_token_limit is None
+    assert cache_hit_tail.logical_state == -1
+    assert cache_hit_tail.cached_tokens == 0
+    assert cache_hit_tail.prefix_cache.trie_cursor is None
+    assert cache_hit_tail.prefix_cache.restore.slot == -1
+    assert cache_hit_tail.prefix_cache.restore.node is None
+    assert not cache_hit_tail.prefix_cache.restore.pinned
+    assert node.state_checkpoint.slot == state_idx
+    assert node.state_checkpoint.published
+    assert node.state_checkpoint.pin_count == 0
+    assert scheduler.block_trie.stats.num_query_tokens == 0
+    assert scheduler.block_trie.stats.num_hit_tokens == 0
+
+
+def _make_ssm_scheduler_for_long_context_chunks(num_gpu_blocks: int = 2):
+    from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+    block_size = 4
+    seq_meta = SequenceMeta(block_size, strategy=ARSequenceStrategy())
+    cache_config = CacheConfig(max_batches=1,
+                               block_size=block_size,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=num_gpu_blocks,
+                               max_prefill_token_num=block_size * 2,
+                               num_state_caches=2,
+                               states_shapes=[((1, ), torch.float32)])
+    scheduler_config = SchedulerConfig(max_batches=1,
+                                       max_session_len=64,
+                                       max_request_output_len=64,
+                                       eviction_type='recompute')
+    scheduler = Scheduler(scheduler_config=scheduler_config, cache_config=cache_config, seq_meta=seq_meta)
+    return scheduler, block_size
+
+
+def test_schedule_prefill_reapplies_chunk_limit_after_ssm_state_rollback():
+    scheduler, block_size = _make_ssm_scheduler_for_long_context_chunks(num_gpu_blocks=2)
+    long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
+
+    ensure_results = iter([False, True])
+
+    def _ensure_runtime_state_available_once_then_succeed():
+        return next(ensure_results)
+
+    scheduler._prefill_scheduler._ensure_runtime_state_available = (
+        _ensure_runtime_state_available_once_then_succeed)
+
+    output = scheduler.schedule(is_prefill=True, prealloc_size=1)
+
+    assert output.running == [long_seq]
+    assert long_seq.status == MessageStatus.READY
+    assert long_seq.kv_token_limit == block_size * 2
+    assert long_seq.num_blocks == 2
diff --git a/tests/pytorch/paging/test_state_manager.py b/tests/pytorch/paging/test_state_manager.py
index f0643d3513..2a7e2389c5 100644
--- a/tests/pytorch/paging/test_state_manager.py
+++ b/tests/pytorch/paging/test_state_manager.py
@@ -5,7 +5,7 @@
 import torch
 
 from lmdeploy.pytorch.config import CacheConfig
-from lmdeploy.pytorch.paging.state_manager import build_state_manager
+from lmdeploy.pytorch.paging.state_manager import StateManager, build_state_manager
 
 
 def test_reserved_state_cache_is_not_allocatable():
@@ -82,3 +82,40 @@ def test_non_ssm_state_manager_without_state_caches():
     assert state_manager.get_num_free() == 0
     with pytest.raises(RuntimeError, match='No free states.'):
         state_manager.allocate(SimpleNamespace(logical_state=-1))
+
+
+def test_state_manager_reserves_system_state_slot():
+    manager = StateManager(num_states=3, num_reserved=1)
+
+    assert manager.allocate_state() == 1
+    assert manager.allocate_state() == 2
+    with pytest.raises(RuntimeError, match='No free states'):
+        manager.allocate_state()
+
+
+def test_state_manager_checkpoint_can_borrow_idle_runtime_slots():
+    manager = StateManager(num_states=5, num_reserved=1, num_runtime_states=2)
+
+    checkpoints = [manager.allocate_checkpoint_state() for _ in range(4)]
+    assert checkpoints == [1, 2, 3, 4]
+    with pytest.raises(RuntimeError, match='No free states'):
+        manager.allocate_checkpoint_state()
+
+    manager.free_checkpoint_state(checkpoints[0])
+    manager.free_checkpoint_state(checkpoints[1])
+    assert manager.allocate_state() == checkpoints[1]
+    assert manager.allocate_state() == checkpoints[0]
+    with pytest.raises(RuntimeError, match='No free states'):
+        manager.allocate_state()
+
+
+def test_state_manager_caps_runtime_count_even_with_extra_free_slots():
+    manager = StateManager(num_states=6, num_reserved=1, num_runtime_states=2)
+
+    assert manager.num_runtime_states == 2
+    assert manager.allocate_state() == 1
+    assert manager.allocate_state() == 2
+    assert manager.get_num_free() == 3
+    assert manager.get_num_free_runtime() == 0
+    with pytest.raises(RuntimeError, match='No free states'):
+        manager.allocate_state()

From 5197bd700a0df08f685d19a889a5d97b8f993d94 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 16:21:40 +0800
Subject: [PATCH 08/22] refactor: make scheduler status facade explicit

---
 lmdeploy/pytorch/paging/scheduler.py   | 103 ++++++++++++++-----------
 tests/pytorch/paging/test_scheduler.py |  53 +++++++++++++
 2 files changed, 111 insertions(+), 45 deletions(-)

diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 2bf2cf73c8..a704fc8515 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -160,60 +160,73 @@ def reserve_long_context_chunk(self,
             is_last_chunk=is_last_chunk,
         )
 
-    @staticmethod
-    def create_status_list_property(status: MessageStatus):
-        """Create status list property."""
+    # Remote-loading sequences are intentionally separate from WAITING: workers
+    # may address their destination blocks, so ordinary scheduling/eviction
+    # must not treat them as candidates until the coordinator publishes them.
 
-        def _get_status_list(self):
-            seq_map = self.seq_manager.get_sequences(status)
-            return list(seq_map.values())
+    # Sequence views.
+    @property
+    def waiting(self) -> SeqList:
+        return list(self.seq_manager.get_sequences(MessageStatus.WAITING).values())
 
-        return property(_get_status_list)
+    @property
+    def remote_loading(self) -> SeqList:
+        return list(self.seq_manager.get_sequences(MessageStatus.WAITING_FOR_REMOTE_KVS).values())
 
-    @staticmethod
-    def create_num_status_method(status: MessageStatus):
-        """Create num status method."""
+    @property
+    def ready(self) -> SeqList:
+        return list(self.seq_manager.get_sequences(MessageStatus.READY).values())
 
-        def _num_status(self):
-            return self.seq_manager.num_sequences(status)
+    @property
+    def hanging(self) -> SeqList:
+        return list(self.seq_manager.get_sequences(MessageStatus.STOPPED).values())
 
-        return _num_status
+    @property
+    def running(self) -> SeqList:
+        return list(self.seq_manager.get_sequences(MessageStatus.RUNNING).values())
 
-    @staticmethod
-    def create_has_status_method(status: MessageStatus):
-        """Create has status method."""
+    @property
+    def migration_waiting(self) -> SeqList:
+        return list(self.seq_manager.get_sequences(MessageStatus.MIGRATION_WAITING).values())
 
-        def _has_status(self):
-            return self.seq_manager.num_sequences(status) > 0
+    @property
+    def migration_done(self) -> SeqList:
+        return list(self.seq_manager.get_sequences(MessageStatus.MIGRATION_DONE).values())
 
-        return _has_status
+    # Sequence counts.
+    def num_waiting(self) -> int:
+        return self.seq_manager.num_sequences(MessageStatus.WAITING)
 
-    # Remote-loading sequences are intentionally separate from WAITING: workers
-    # may address their destination blocks, so ordinary scheduling/eviction
-    # must not treat them as candidates until the coordinator publishes them.
-    # status list properties
-    waiting = create_status_list_property(MessageStatus.WAITING)
-    remote_loading = create_status_list_property(MessageStatus.WAITING_FOR_REMOTE_KVS)
-    ready = create_status_list_property(MessageStatus.READY)
-    hanging = create_status_list_property(MessageStatus.STOPPED)
-    running = create_status_list_property(MessageStatus.RUNNING)
-    migration_waiting = create_status_list_property(MessageStatus.MIGRATION_WAITING)
-    migration_done = create_status_list_property(MessageStatus.MIGRATION_DONE)
-
-    # num status methods
-    num_waiting = create_num_status_method(MessageStatus.WAITING)
-    num_remote_loading = create_num_status_method(MessageStatus.WAITING_FOR_REMOTE_KVS)
-    num_ready = create_num_status_method(MessageStatus.READY)
-    num_running = create_num_status_method(MessageStatus.RUNNING)
-    num_migration_waiting = create_num_status_method(MessageStatus.MIGRATION_WAITING)
-    num_migration_done = create_num_status_method(MessageStatus.MIGRATION_DONE)
-
-    # has status methods
-    has_waiting = create_has_status_method(MessageStatus.WAITING)
-    has_remote_loading = create_has_status_method(MessageStatus.WAITING_FOR_REMOTE_KVS)
-    has_ready = create_has_status_method(MessageStatus.READY)
-    has_migration_waiting = create_has_status_method(MessageStatus.MIGRATION_WAITING)
-    has_migration_done = create_has_status_method(MessageStatus.MIGRATION_DONE)
+    def num_remote_loading(self) -> int:
+        return self.seq_manager.num_sequences(MessageStatus.WAITING_FOR_REMOTE_KVS)
+
+    def num_ready(self) -> int:
+        return self.seq_manager.num_sequences(MessageStatus.READY)
+
+    def num_running(self) -> int:
+        return self.seq_manager.num_sequences(MessageStatus.RUNNING)
+
+    def num_migration_waiting(self) -> int:
+        return self.seq_manager.num_sequences(MessageStatus.MIGRATION_WAITING)
+
+    def num_migration_done(self) -> int:
+        return self.seq_manager.num_sequences(MessageStatus.MIGRATION_DONE)
+
+    # Non-empty status checks used by engine control flow.
+    def has_waiting(self) -> bool:
+        return self.seq_manager.num_sequences(MessageStatus.WAITING) > 0
+
+    def has_remote_loading(self) -> bool:
+        return self.seq_manager.num_sequences(MessageStatus.WAITING_FOR_REMOTE_KVS) > 0
+
+    def has_ready(self) -> bool:
+        return self.seq_manager.num_sequences(MessageStatus.READY) > 0
+
+    def has_migration_waiting(self) -> bool:
+        return self.seq_manager.num_sequences(MessageStatus.MIGRATION_WAITING) > 0
+
+    def has_migration_done(self) -> bool:
+        return self.seq_manager.num_sequences(MessageStatus.MIGRATION_DONE) > 0
 
     def add_session(self, session_id: int):
         """Add new session.
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index ca4261f92a..321ab33072 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -184,6 +184,59 @@ def test_decode_requires_schedule_running(self, scheduler):
         with pytest.raises(ValueError, match='schedule_running'):
             scheduler.schedule(is_prefill=False)
 
+    @pytest.mark.parametrize(
+        ('name', 'status'),
+        [
+            ('waiting', MessageStatus.WAITING),
+            ('remote_loading', MessageStatus.WAITING_FOR_REMOTE_KVS),
+            ('ready', MessageStatus.READY),
+            ('hanging', MessageStatus.STOPPED),
+            ('running', MessageStatus.RUNNING),
+            ('migration_waiting', MessageStatus.MIGRATION_WAITING),
+            ('migration_done', MessageStatus.MIGRATION_DONE),
+        ],
+    )
+    def test_status_sequence_views(self, scheduler, monkeypatch, name, status):
+        seq = object()
+        queried = []
+
+        def get_sequences(actual_status):
+            queried.append(actual_status)
+            return {0: seq}
+
+        monkeypatch.setattr(scheduler.seq_manager, 'get_sequences', get_sequences)
+
+        assert getattr(scheduler, name) == [seq]
+        assert queried == [status]
+
+    @pytest.mark.parametrize(
+        ('name', 'status', 'expected'),
+        [
+            ('num_waiting', MessageStatus.WAITING, 3),
+            ('num_remote_loading', MessageStatus.WAITING_FOR_REMOTE_KVS, 3),
+            ('num_ready', MessageStatus.READY, 3),
+            ('num_running', MessageStatus.RUNNING, 3),
+            ('num_migration_waiting', MessageStatus.MIGRATION_WAITING, 3),
+            ('num_migration_done', MessageStatus.MIGRATION_DONE, 3),
+            ('has_waiting', MessageStatus.WAITING, True),
+            ('has_remote_loading', MessageStatus.WAITING_FOR_REMOTE_KVS, True),
+            ('has_ready', MessageStatus.READY, True),
+            ('has_migration_waiting', MessageStatus.MIGRATION_WAITING, True),
+            ('has_migration_done', MessageStatus.MIGRATION_DONE, True),
+        ],
+    )
+    def test_status_queries(self, scheduler, monkeypatch, name, status, expected):
+        queried = []
+
+        def num_sequences(actual_status):
+            queried.append(actual_status)
+            return 3
+
+        monkeypatch.setattr(scheduler.seq_manager, 'num_sequences', num_sequences)
+
+        assert getattr(scheduler, name)() == expected
+        assert queried == [status]
+
 
 def test_schedule_migration_matches_current_sequence():
     from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy

From 6e15c34986b0d95c905b4970029f7bf50e01982f Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 16:43:52 +0800
Subject: [PATCH 09/22] refactor: centralize scheduler signal ownership

---
 lmdeploy/pytorch/paging/scheduler.py          | 28 +++++++++++--------
 .../engine/test_kv_connector_wiring.py        |  9 +++++-
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index a704fc8515..e2367ad1a0 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -69,9 +69,6 @@ def __init__(
         self.sessions: dict[int, SchedulerSession] = OrderedDict()
         self.kv_connector = kv_connector
 
-        # For Disaggregation
-        self.locked_sessions: dict[int, SchedulerSession] = OrderedDict()
-
         self.state_manager = build_state_manager(self.cache_config)
         self.block_manager = build_block_manager(cache_config)
         self.is_ssm = len(self.cache_config.states_shapes) > 0
@@ -79,7 +76,7 @@ def __init__(
         # A producer-only connector still needs the save path below, but must
         # not issue lookups. SSM restore owns a different state-cache protocol
         # and is deliberately excluded from external KV load admission.
-        self._external_lookup_enabled = (
+        external_lookup_enabled = (
             kv_connector is not None
             and transfer_config is not None
             and transfer_config.is_kv_consumer
@@ -95,7 +92,7 @@ def __init__(
         # Load admission receives only paging owners plus request-local queue
         # candidates from its caller; it does not reach back through Scheduler.
         self.kv_load_coordinator = KVLoadCoordinator(
-            lookup_enabled=self._external_lookup_enabled,
+            lookup_enabled=external_lookup_enabled,
             connector=kv_connector,
             block_manager=self.block_manager,
             block_trie=self.block_trie,
@@ -114,9 +111,6 @@ def __init__(
         )
         # Keep save call sites uniform even when the producer role is disabled.
         self.kv_save_coordinator = KVSaveCoordinator(self)
-        # Per-tick signal consumed by EngineLoop to distinguish asynchronous
-        # lookup latency from actual cache-allocation pressure.
-        self.last_schedule_had_pending_lookup = False
 
         seq_meta = seq_meta or SequenceMeta(self.cache_config.block_size)
         self.seq_meta = seq_meta
@@ -135,12 +129,25 @@ def shutdown(self) -> None:
         """
         connector = self.kv_connector
         self.kv_connector = None
-        self._external_lookup_enabled = False
         self.kv_load_coordinator.disable()
         self.kv_save_coordinator.clear()
         if connector is not None:
             connector.shutdown()
 
+    @property
+    def _external_lookup_enabled(self) -> bool:
+        """Whether external KV lookup admission is currently enabled."""
+        return self.kv_load_coordinator.lookup_enabled
+
+    @property
+    def last_schedule_had_pending_lookup(self) -> bool:
+        """Whether the latest prefill turn encountered a pending lookup."""
+        return self._prefill_scheduler.last_schedule_had_pending_lookup
+
+    @last_schedule_had_pending_lookup.setter
+    def last_schedule_had_pending_lookup(self, value: bool) -> None:
+        self._prefill_scheduler.last_schedule_had_pending_lookup = value
+
     def has_waiting_long_prefill(self):
         """Whether a waiting request would need a non-final prefill chunk."""
         return self._prefill_scheduler.has_waiting_long_prefill(self.waiting)
@@ -280,7 +287,6 @@ def schedule(self,
                 'schedule only selects prefill work; use schedule_running '
                 'for decode capacity admission')
 
-        self.last_schedule_had_pending_lookup = False
         running = self._prefill_scheduler.schedule(
             waiting=self.waiting,
             stopped=self.hanging,
@@ -290,8 +296,6 @@ def schedule(self,
             allow_long_prefill=allow_long_prefill,
             prefer_long_prefill=prefer_long_prefill,
         )
-        self.last_schedule_had_pending_lookup = (
-            self._prefill_scheduler.last_schedule_had_pending_lookup)
         return SchedulerOutput(
             running=running,
             swap_in_map={},
diff --git a/tests/pytorch/engine/test_kv_connector_wiring.py b/tests/pytorch/engine/test_kv_connector_wiring.py
index 26af606bae..8c310ec615 100644
--- a/tests/pytorch/engine/test_kv_connector_wiring.py
+++ b/tests/pytorch/engine/test_kv_connector_wiring.py
@@ -218,9 +218,16 @@ def test_prepare_kv_connector_config_does_not_change_disabled_config(transfer_co
 
 def test_scheduler_shutdown_releases_injected_connector_once():
     connector = Mock()
+    load_coordinator = Mock()
+    load_coordinator.lookup_enabled = True
+
+    def disable_loads():
+        load_coordinator.lookup_enabled = False
+
+    load_coordinator.disable.side_effect = disable_loads
     scheduler = Scheduler.__new__(Scheduler)
     scheduler.kv_connector = connector
-    scheduler.kv_load_coordinator = Mock()
+    scheduler.kv_load_coordinator = load_coordinator
     scheduler.kv_save_coordinator = Mock()
 
     scheduler.shutdown()

From 269cd4b1509de650caad00f5f8c40d917adfea89 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 16:54:40 +0800
Subject: [PATCH 10/22] refactor: make prefill states explicit

---
 lmdeploy/pytorch/paging/prefill_scheduler.py  | 83 +++++++++++++------
 lmdeploy/pytorch/paging/scheduler.py          |  9 +-
 .../pytorch/paging/test_prefill_scheduler.py  | 26 ++++++
 3 files changed, 88 insertions(+), 30 deletions(-)

diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index ba50a0982c..08adb426b7 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -43,6 +43,27 @@ class _PrefillReorderInfo:
     estimated_long_chunks: int
 
 
+class _PrefillTurnPolicy(enum.Enum):
+    """Complete internal policy derived from the public long-prefill flags."""
+
+    STANDARD = (True, False)
+    SHORT_ONLY = (False, False)
+    LONG_FIRST = (True, True)
+    # Preserve the fourth public flag combination: try a long-looking request
+    # first, but admit it only if prefix matching makes this its final prefill.
+    LONG_FIRST_IF_FINAL = (False, True)
+
+    def __init__(self, allows_nonfinal_long_prefill: bool,
+                 prefers_long_prefill: bool):
+        self.allows_nonfinal_long_prefill = allows_nonfinal_long_prefill
+        self.prefers_long_prefill = prefers_long_prefill
+
+    @classmethod
+    def from_flags(cls, allow_long_prefill: bool, prefer_long_prefill: bool):
+        """Normalize the stable public flag pair at the scheduler boundary."""
+        return cls((allow_long_prefill, prefer_long_prefill))
+
+
 class _PrefillReorderer:
     """Order waiting prefills without applying scheduler side effects."""
 
@@ -52,8 +73,7 @@ def __init__(self, prefill_scheduler: '_PrefillScheduler'):
 
     def reorder(self,
                 waiting: SeqList,
-                allow_long_prefill: bool,
-                prefer_long_prefill: bool):
+                turn_policy: _PrefillTurnPolicy):
         """Return waiting requests in the order the prefill loop should try."""
         waiting = sorted(waiting, key=lambda seq: seq.arrive_time)
         # A completed load already owns destination blocks and a soft prefill
@@ -66,7 +86,7 @@ def reorder(self,
             seq for seq in waiting
             if not self.prefill_scheduler.load_coordinator.is_remote_ready(seq)
         ]
-        if prefer_long_prefill:
+        if turn_policy.prefers_long_prefill:
             # Long-work turns choose one long waiter first. The size policy only
             # reorders this long lane; it is not global shortest-prefill-first
             # admission.
@@ -74,7 +94,7 @@ def reorder(self,
             if long_turn_order is not None:
                 return remote_ready + self._warn_if_not_permutation(waiting, long_turn_order)
 
-        if allow_long_prefill:
+        if turn_policy.allows_nonfinal_long_prefill:
             return remote_ready + self._warn_if_not_permutation(waiting, waiting)
 
         reordered = self._reorder_for_short_turn(waiting)
@@ -264,6 +284,14 @@ def capture(cls, seq: SchedulerSequence):
         )
 
 
+class _PrefixMatchPhase(enum.Enum):
+    """Lifecycle of one request-local tentative prefix transaction."""
+
+    IDLE = enum.auto()
+    TRACKING = enum.auto()
+    MATCHED = enum.auto()
+
+
 class _TentativePrefixMatch:
     """Request-local transaction around ``BlockTrie.match`` side effects.
 
@@ -282,8 +310,7 @@ class _TentativePrefixMatch:
         '_stats_snapshot',
         '_state_snapshot',
         '_rejection_on_rollback',
-        '_started',
-        'matched',
+        '_phase',
     )
 
     def __init__(self,
@@ -301,8 +328,11 @@ def __init__(self,
         self._stats_snapshot = None
         self._state_snapshot: _PrefixMatchStateSnapshot | None = None
         self._rejection_on_rollback: _PrefillAdmissionResult | None = None
-        self._started = False
-        self.matched = False
+        self._phase = _PrefixMatchPhase.IDLE
+
+    @property
+    def is_matched(self) -> bool:
+        return self._phase is _PrefixMatchPhase.MATCHED
 
     def begin(self) -> None:
         """Start the transaction before gates can mutate exact external state.
@@ -311,23 +341,24 @@ def begin(self) -> None:
         starts before gates so rollback can restore existing request state even
         when a private partial block prevents another trie match.
         """
-        if self._started or not self.block_trie.enabled:
+        if self._phase is not _PrefixMatchPhase.IDLE or not self.block_trie.enabled:
             return
         self._stats_snapshot = self.block_trie.stats.snapshot()
         if self._preserve_existing_state:
             self._state_snapshot = _PrefixMatchStateSnapshot.capture(self.seq)
-        self._started = True
+        self._phase = _PrefixMatchPhase.TRACKING
 
     def match(self) -> None:
         """Apply one tentative match after capturing its rollback boundary."""
-        assert not self.matched
+        assert not self.is_matched
         self.begin()
+        assert self._phase is _PrefixMatchPhase.TRACKING
         self.block_trie.match(self.seq)
-        self.matched = True
+        self._phase = _PrefixMatchPhase.MATCHED
 
     def retain_for_admission(self, rejection_on_rollback: _PrefillAdmissionResult) -> None:
         """Keep a gate-enabling match and remember its original rejection."""
-        assert self.matched
+        assert self.is_matched
         self._rejection_on_rollback = rejection_on_rollback
 
     def pin_restore(self) -> bool:
@@ -344,7 +375,7 @@ def commit(self) -> None:
     def rollback(self, reason: str):
         """Undo the transaction and return any gate-defined rejection."""
         rejection = self._rejection_on_rollback
-        if not self._started:
+        if self._phase is _PrefixMatchPhase.IDLE:
             return rejection
 
         seq = self.seq
@@ -399,8 +430,7 @@ def _clear(self) -> None:
         self._stats_snapshot = None
         self._state_snapshot = None
         self._rejection_on_rollback = None
-        self._started = False
-        self.matched = False
+        self._phase = _PrefixMatchPhase.IDLE
 
 
 class _PrefillAdmissionAttempt:
@@ -420,7 +450,7 @@ def __init__(self,
                  prealloc_size: int,
                  batch_prefill_tokens: int,
                  batch_has_prefill: bool,
-                 allow_long_prefill: bool):
+                 turn_policy: _PrefillTurnPolicy):
         self.prefill_scheduler = prefill_scheduler
         self.seq = seq
         self.stopped = stopped
@@ -430,7 +460,7 @@ def __init__(self,
         self.batch_has_prefill = batch_has_prefill
         self.load_coordinator = prefill_scheduler.load_coordinator
         self._load_ready = self.load_coordinator.is_remote_ready(seq)
-        self.allow_long_prefill = allow_long_prefill
+        self.turn_policy = turn_policy
         self._effective_prealloc_size = prealloc_size
         self._prefix_match = _TentativePrefixMatch(
             seq,
@@ -479,7 +509,7 @@ def _admit_prefix_cache_resources(self):
 
     def _resolve_prefix_source(self):
         """Match local cache first, then try loading its remote extension."""
-        if not self._prefix_match.matched:
+        if not self._prefix_match.is_matched:
             # A completed external load has already published the accepted
             # prefix. Matching again would lose its cached-token accounting.
             if not self._load_ready and not self._has_private_local_tail():
@@ -598,7 +628,7 @@ def _has_private_local_tail(self) -> bool:
         return seq.num_blocks > int(seq.num_history_ids) // seq.block_size
 
     def _token_budget_rejection(self):
-        if self.allow_long_prefill:
+        if self.turn_policy.allows_nonfinal_long_prefill:
             return _PrefillAdmissionResult.stop()
         return _PrefillAdmissionResult.skip()
 
@@ -610,7 +640,8 @@ def _check_prefill_admission_gates(self):
         prefill_token_count = prefill._prefill_admission_token_count(seq)
         is_nonfinal_long_prefill = prefill._prefill_kv_token_limit(seq) is not None
 
-        if is_nonfinal_long_prefill and not self.allow_long_prefill:
+        if (is_nonfinal_long_prefill
+                and not self.turn_policy.allows_nonfinal_long_prefill):
             matched = self._match_prefix_for_prefill_gate()
             if matched is None:
                 return _PrefillAdmissionResult.skip()
@@ -628,7 +659,7 @@ def _check_prefill_admission_gates(self):
         if not exceeds_token_budget:
             return None
 
-        if not self._prefix_match.matched:
+        if not self._prefix_match.is_matched:
             matched = self._match_prefix_for_prefill_gate()
             if matched is not None:
                 prefill_token_count = prefill._prefill_admission_token_count(seq)
@@ -829,9 +860,8 @@ def schedule(
         stopped: SeqList,
         num_ready: int,
         num_running: int,
+        turn_policy: _PrefillTurnPolicy,
         prealloc_size: int = 0,
-        allow_long_prefill: bool = True,
-        prefer_long_prefill: bool = False,
     ):
         """Select and activate one prefill batch."""
         self.last_schedule_had_pending_lookup = False
@@ -844,8 +874,7 @@ def schedule(
 
         waiting = _PrefillReorderer(self).reorder(
             waiting,
-            allow_long_prefill=allow_long_prefill,
-            prefer_long_prefill=prefer_long_prefill,
+            turn_policy=turn_policy,
         )
         skipped_waiting: SeqList = []
         while waiting and len(running) < max_batches:
@@ -859,7 +888,7 @@ def schedule(
                 prealloc_size=prealloc_size,
                 batch_prefill_tokens=batch_prefill_tokens,
                 batch_has_prefill=bool(running),
-                allow_long_prefill=allow_long_prefill,
+                turn_policy=turn_policy,
             ).run()
 
             if admission.action is _PrefillAdmissionAction.LOAD_STARTED:
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index e2367ad1a0..ec01db0ea2 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -19,7 +19,7 @@
 from .eviction_helper import build_eviction_helper
 from .kv_load_coordinator import KVLoadCoordinator
 from .kv_save_coordinator import KVSaveCoordinator
-from .prefill_scheduler import _PrefillScheduler
+from .prefill_scheduler import _PrefillScheduler, _PrefillTurnPolicy
 from .state_manager import build_state_manager
 
 if TYPE_CHECKING:
@@ -287,14 +287,17 @@ def schedule(self,
                 'schedule only selects prefill work; use schedule_running '
                 'for decode capacity admission')
 
+        turn_policy = _PrefillTurnPolicy.from_flags(
+            allow_long_prefill,
+            prefer_long_prefill,
+        )
         running = self._prefill_scheduler.schedule(
             waiting=self.waiting,
             stopped=self.hanging,
             num_ready=self.num_ready(),
             num_running=self.num_running(),
+            turn_policy=turn_policy,
             prealloc_size=prealloc_size,
-            allow_long_prefill=allow_long_prefill,
-            prefer_long_prefill=prefer_long_prefill,
         )
         return SchedulerOutput(
             running=running,
diff --git a/tests/pytorch/paging/test_prefill_scheduler.py b/tests/pytorch/paging/test_prefill_scheduler.py
index 2a70cc4a1b..9a2ea705e1 100644
--- a/tests/pytorch/paging/test_prefill_scheduler.py
+++ b/tests/pytorch/paging/test_prefill_scheduler.py
@@ -182,6 +182,32 @@ def test_scheduler_short_turn_uses_prefix_hit_to_admit_long_looking_sibling():
     assert cache_hit_tail.cached_tokens == block_size
 
 
+def test_scheduler_long_first_short_turn_admits_only_final_prefix_hit():
+    scheduler, block_size = _make_prefix_cache_scheduler(
+        max_batches=1,
+        max_prefill_token_num=16,
+    )
+
+    cached = scheduler.add_session(0).add_sequence([1] * block_size)
+    scheduler.schedule(is_prefill=True)
+    cached.state.stop()
+
+    short = scheduler.add_session(1).add_sequence([4])
+    cache_hit_tail = scheduler.add_session(2).add_sequence(
+        [1] * block_size + [3])
+
+    output = scheduler.schedule(
+        is_prefill=True,
+        allow_long_prefill=False,
+        prefer_long_prefill=True,
+    )
+
+    assert output.running == [cache_hit_tail]
+    assert short.status == MessageStatus.WAITING
+    assert cache_hit_tail.num_history_ids == block_size
+    assert cache_hit_tail.num_token_ids == 1
+
+
 def test_scheduler_budget_gate_uses_prefix_hit_to_admit_sibling():
     scheduler, block_size = _make_prefix_cache_scheduler(max_batches=2, max_prefill_token_num=16)
 

From cb96529598d51281d69adcd882dd5cc5e3ddee63 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 17:05:40 +0800
Subject: [PATCH 11/22] refactor: clarify external load lifecycle

---
 .../pytorch/paging/kv_load_coordinator.py     | 117 ++++++++++++++----
 .../paging/test_scheduler_kv_transfer.py      |  46 ++++++-
 2 files changed, 135 insertions(+), 28 deletions(-)

diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index 4341ff9087..38103a6e66 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -66,6 +66,24 @@ class _LoadPhase(enum.Enum):
     PREFILLING = enum.auto()
 
 
+class _DeferredLoadCleanup(enum.Enum):
+    """Cleanup to apply after an active device write becomes safe."""
+
+    NONE = enum.auto()
+    STOP = enum.auto()
+    END = enum.auto()
+
+
+@dataclass(frozen=True, slots=True)
+class _LoadPlan:
+    """Block-aligned remote interval and its admission rollback boundary."""
+
+    fallback_step: int
+    remote_step: int
+    target_blocks: int
+    original_kv_token_limit: int | None
+
+
 @dataclass
 class _LoadRecord:
     """Paging state retained for one asynchronous load.
@@ -73,15 +91,15 @@ class _LoadRecord:
     ``fallback_step`` is the block-aligned prefix that remains trustworthy if
     a worker fails or is cancelled after partially writing a destination.
     ``remote_step`` is published only after every TP rank reports success.
-    Stop/end flags defer user-requested cleanup until device writes are safe.
+    ``deferred_cleanup`` records the strongest user-requested action until
+    device writes are safe; ending a request takes precedence over stopping it.
     """
 
     seq: SchedulerSequence
     fallback_step: int
     remote_step: int
     phase: _LoadPhase = _LoadPhase.LOADING
-    stop_requested: bool = False
-    end_requested: bool = False
+    deferred_cleanup: _DeferredLoadCleanup = _DeferredLoadCleanup.NONE
 
 
 class KVLoadCoordinator:
@@ -234,8 +252,29 @@ def _admit_load(
         evictable_seqs: Iterable[SchedulerSequence],
     ) -> KVLoadAdmission:
         """Admit the complete prefill, then allocate the remote interval."""
-        connector = self.connector
-        assert connector is not None
+        plan = self._plan_load(seq, num_external_tokens, prealloc_size)
+        if plan is None:
+            return KVLoadAdmission.NO_LOAD
+
+        failure = self._admit_load_capacity(
+            seq,
+            plan,
+            prealloc_size,
+            evictable_seqs,
+        )
+        if failure is not None:
+            return failure
+
+        self._allocate_and_start_load(seq, plan)
+        return KVLoadAdmission.STARTED
+
+    def _plan_load(
+        self,
+        seq: SchedulerSequence,
+        num_external_tokens: int,
+        prealloc_size: int,
+    ) -> _LoadPlan | None:
+        """Plan the safe block-aligned interval before mutating ownership."""
         block_size = seq.block_size
         local_step = int(seq.num_history_ids)
         # Transfers are block-granular. Reuse a private partial boundary block,
@@ -245,10 +284,23 @@ def _admit_load(
         remote_step = min(remote_step, int(seq.get_prefix_cache_max_match_step()))
         remote_step = remote_step // block_size * block_size
         if remote_step <= fallback_step:
-            return KVLoadAdmission.NO_LOAD
+            return None
 
-        target_blocks = self.prefill_target_blocks(seq, prealloc_size)
-        old_kv_token_limit = seq.kv_token_limit
+        return _LoadPlan(
+            fallback_step=fallback_step,
+            remote_step=remote_step,
+            target_blocks=self.prefill_target_blocks(seq, prealloc_size),
+            original_kv_token_limit=seq.kv_token_limit,
+        )
+
+    def _admit_load_capacity(
+        self,
+        seq: SchedulerSequence,
+        plan: _LoadPlan,
+        prealloc_size: int,
+        evictable_seqs: Iterable[SchedulerSequence],
+    ) -> KVLoadAdmission | None:
+        """Admit the full prefill against physical and soft capacity."""
         # Only the remote hit is allocated now, but admission guarantees the
         # complete prefill can finish beside every existing soft reservation.
         seq.kv_token_limit = None
@@ -258,21 +310,30 @@ def _admit_load(
             prealloc_size,
         )
         if not full_prefill_fits:
-            seq.kv_token_limit = old_kv_token_limit
+            seq.kv_token_limit = plan.original_kv_token_limit
             return KVLoadAdmission.FULL_PREFILL_UNAVAILABLE
-        if not self.can_admit_load(seq, target_blocks):
-            seq.kv_token_limit = old_kv_token_limit
+        if not self.can_admit_load(seq, plan.target_blocks):
+            seq.kv_token_limit = plan.original_kv_token_limit
             return KVLoadAdmission.SOFT_BUDGET_UNAVAILABLE
+        return None
 
+    def _allocate_and_start_load(
+        self,
+        seq: SchedulerSequence,
+        plan: _LoadPlan,
+    ) -> None:
+        """Allocate destinations, bind them, then transfer paging ownership."""
+        connector = self.connector
+        assert connector is not None
         original_num_blocks = seq.num_blocks
         try:
             # Allocate only the checked remote interval. The unallocated local
-            # tail remains represented by the soft target above.
-            seq.kv_token_limit = remote_step
+            # tail remains represented by the plan's soft target.
+            seq.kv_token_limit = plan.remote_step
             self.block_manager.allocate(seq)
             block_table = self.block_manager.get_block_table(seq)
-            fallback_block = fallback_step // block_size
-            remote_block = remote_step // block_size
+            fallback_block = plan.fallback_step // seq.block_size
+            remote_block = plan.remote_step // seq.block_size
             load_block_ids = tuple(
                 int(block_id)
                 for block_id in block_table[fallback_block:remote_block]
@@ -280,23 +341,22 @@ def _admit_load(
             connector.update_state_after_alloc(
                 seq,
                 load_block_ids,
-                remote_step - fallback_step,
+                plan.remote_step - plan.fallback_step,
             )
             # From start_load onward, cleanup must retain destinations until
             # workers report terminal progress or their queues are drained.
             self.start_load(
                 seq,
-                fallback_step=fallback_step,
-                remote_step=remote_step,
-                target_blocks=target_blocks,
+                fallback_step=plan.fallback_step,
+                remote_step=plan.remote_step,
+                target_blocks=plan.target_blocks,
             )
         except Exception:
             if seq.num_blocks > original_num_blocks:
                 self.block_manager.truncate(seq, original_num_blocks)
-            seq.kv_token_limit = old_kv_token_limit
+            seq.kv_token_limit = plan.original_kv_token_limit
             raise
         seq.kv_token_limit = None
-        return KVLoadAdmission.STARTED
 
     def start_load(
         self,
@@ -346,7 +406,8 @@ def update(self, results: tuple[KVLoadResult, ...]) -> None:
             record = self._loads.get(int(result.request_id))
             if record is None or record.phase is not _LoadPhase.LOADING:
                 continue
-            if record.stop_requested or record.end_requested or not result.success:
+            if (record.deferred_cleanup is not _DeferredLoadCleanup.NONE
+                    or not result.success):
                 self._rollback(record)
                 self._finish_cancelled_or_failed(record)
             else:
@@ -404,9 +465,9 @@ def _finish_cancelled_or_failed(self, record: _LoadRecord) -> None:
         self._loads.pop(request_id, None)
         self._prefill_targets.pop(request_id, None)
         seq.state.finish_remote_load()
-        if record.end_requested:
+        if record.deferred_cleanup is _DeferredLoadCleanup.END:
             self._remove_sequence(seq)
-        elif record.stop_requested:
+        elif record.deferred_cleanup is _DeferredLoadCleanup.STOP:
             seq.state.stop()
 
     def request_stop(self, seq: SchedulerSequence) -> bool:
@@ -420,7 +481,8 @@ def request_stop(self, seq: SchedulerSequence) -> bool:
         if record is None or record.phase is not _LoadPhase.LOADING:
             self.release(seq)
             return False
-        record.stop_requested = True
+        if record.deferred_cleanup is _DeferredLoadCleanup.NONE:
+            record.deferred_cleanup = _DeferredLoadCleanup.STOP
         return True
 
     def request_end(self, seq: SchedulerSequence) -> bool:
@@ -433,7 +495,7 @@ def request_end(self, seq: SchedulerSequence) -> bool:
         if record is None or record.phase is not _LoadPhase.LOADING:
             self.release(seq)
             return False
-        record.end_requested = True
+        record.deferred_cleanup = _DeferredLoadCleanup.END
         return True
 
     def release(self, seq: SchedulerSequence) -> None:
@@ -475,7 +537,8 @@ def finish_deferred_loads_after_worker_drain(self) -> None:
         records = [
             record
             for record in self._loads.values()
-            if record.phase is _LoadPhase.LOADING and record.end_requested
+            if (record.phase is _LoadPhase.LOADING
+                and record.deferred_cleanup is _DeferredLoadCleanup.END)
         ]
         for record in records:
             self._rollback(record)
diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py
index 8abc12e35c..f13589439a 100644
--- a/tests/pytorch/paging/test_scheduler_kv_transfer.py
+++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py
@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from unittest.mock import Mock
 
+import pytest
 import torch
 
 from lmdeploy.messages import KVTransferConfig
@@ -487,6 +488,27 @@ def test_async_load_capacity_failure_restores_tentative_local_prefix(monkeypatch
     assert scheduler.block_trie.stats.num_hit_tokens == 0
 
 
+def test_async_load_binding_failure_releases_allocated_destinations():
+    connector = _AsyncLookupConnector([(8, True)])
+    connector.update_state_after_alloc = Mock(
+        side_effect=RuntimeError('binding failed'))
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+        num_gpu_blocks=4,
+    )
+    seq = scheduler.add_session(77).add_sequence(torch.arange(13))
+
+    with pytest.raises(RuntimeError, match='binding failed'):
+        scheduler.schedule(is_prefill=True)
+
+    assert seq.status == MessageStatus.WAITING
+    assert seq.num_blocks == 0
+    assert seq.kv_token_limit is None
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 4
+    assert not scheduler.has_remote_loading()
+
+
 def test_async_load_does_not_consume_model_batch_slot():
     connector = _AsyncLookupConnector([(8, True), (0, False)])
     scheduler = _make_async_lookup_scheduler(
@@ -623,7 +645,28 @@ def test_async_load_soft_reservation_shrinks_across_chunks():
     assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
 
 
-def test_end_session_waits_for_active_async_load_before_freeing_blocks():
+def test_stop_session_waits_for_active_async_load_before_stopping():
+    connector = _AsyncLookupConnector([(8, True)])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+    )
+    seq = scheduler.add_session(82).add_sequence(torch.arange(13))
+    scheduler.schedule(is_prefill=True)
+
+    scheduler.stop_session(82)
+    assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS
+    assert seq.num_blocks == 2
+
+    scheduler.update_connector_output(
+        KVConnectorOutput(finished_receiving={seq.seq_id}))
+    assert seq.status == MessageStatus.STOPPED
+    assert seq.num_blocks == 0
+    assert 82 in scheduler.sessions
+    assert connector.finished == []
+
+
+def test_end_session_overrides_stop_deferred_during_active_async_load():
     connector = _AsyncLookupConnector([(8, True)])
     scheduler = _make_async_lookup_scheduler(
         connector,
@@ -632,6 +675,7 @@ def test_end_session_waits_for_active_async_load_before_freeing_blocks():
     seq = scheduler.add_session(82).add_sequence(torch.arange(13))
     scheduler.schedule(is_prefill=True)
 
+    scheduler.stop_session(82)
     scheduler.end_session(82)
     assert 82 in scheduler.sessions
     assert seq.status == MessageStatus.WAITING_FOR_REMOTE_KVS

From c5a8cfad92a11bd5b137f29715c3e4b6c4f85512 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 17:17:31 +0800
Subject: [PATCH 12/22] refactor: simplify scheduler admission flow

---
 lmdeploy/pytorch/paging/prefill_scheduler.py | 101 +++++++------------
 lmdeploy/pytorch/paging/scheduler.py         |  13 +--
 tests/pytorch/paging/test_scheduler.py       |  13 +++
 tests/pytorch/paging/test_scheduler_ssm.py   |  12 +--
 4 files changed, 62 insertions(+), 77 deletions(-)

diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index 08adb426b7..2a1b8dc674 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -202,16 +202,7 @@ class _PrefillAdmissionAction(enum.Enum):
 
 @dataclass(frozen=True)
 class _PrefillAdmissionResult:
-    """Outcome from trying to admit one waiting prefill request.
-
-    The outer loop distinguishes four outcomes:
-
-    * ``ADMIT``: include the request in this tick's model batch.
-    * ``SKIP``: leave it waiting but continue trying later candidates.
-    * ``STOP``: resource pressure ends this prefill admission turn.
-    * ``LOAD_STARTED``: no model work was selected, but the request left the
-      waiting queue for asynchronous KV load.
-    """
+    """Outcome from trying to admit one waiting prefill request."""
 
     action: _PrefillAdmissionAction
     prefill_token_count: int = 0
@@ -236,21 +227,9 @@ def load_started(cls):
 
 @dataclass(frozen=True, slots=True)
 class _PrefixMatchStateSnapshot:
-    """Exact sequence state captured before a tentative local trie match.
-
-    External lookup itself does not mutate sequence paging state. The scheduler
-    may, however, run ``block_trie.match`` first so the connector queries only
-    beyond the locally resident prefix. If that non-blocking lookup returns
-    pending, or a positive hit cannot be admitted before worker writes start,
-    the request will not run this tick and the tentative local match must be
-    undone.
-
-    A multi-turn request may already own valid history, blocks, and model
-    metadata before this attempt. Restoring this baseline preserves that exact
-    committed state; the legacy new-request rollback to step zero would discard
-    it. This snapshot is not used after an asynchronous load starts--load
-    failure then rolls back to its block-aligned ``fallback_step`` through
-    ``KVLoadCoordinator`` because workers may have partially written KV.
+    """Committed state restored when a tentative local match is rejected.
+
+    Load failure after worker writes uses its block-aligned fallback instead.
     """
 
     # Committed sequence progress and block ownership before tentative match.
@@ -293,13 +272,7 @@ class _PrefixMatchPhase(enum.Enum):
 
 
 class _TentativePrefixMatch:
-    """Request-local transaction around ``BlockTrie.match`` side effects.
-
-    Ordinary and SSM admission preserve the historical fallback to an unmatched request. External lookup instead needs
-    an exact pre-match snapshot because a multi-turn request may already own committed progress. Both contracts share
-    one stats snapshot, restore-pin boundary, and explicit commit/rollback lifecycle without changing their rollback
-    semantics.
-    """
+    """Manage one request's tentative ``BlockTrie.match`` side effects."""
 
     __slots__ = (
         'seq',
@@ -434,13 +407,7 @@ def _clear(self) -> None:
 
 
 class _PrefillAdmissionAttempt:
-    """Try to admit one waiting prefill sequence.
-
-    The attempt owns all tentative prefix-cache side effects for the sequence:
-    match, SSM restore pinning, eviction, runtime-state checks, allocation, and
-    rollback. The outer prefill loop still owns queue traversal and decides
-    whether a rejected candidate is skipped or ends the current prefill turn.
-    """
+    """Own one waiting sequence's tentative admission side effects."""
 
     def __init__(self,
                  prefill_scheduler: '_PrefillScheduler',
@@ -557,7 +524,7 @@ def _admit_runtime_state(self):
         """Ensure an SSM runtime slot, retrying after match rollback."""
         prefill = self.prefill_scheduler
         seq = self.seq
-        if not prefill.is_ssm or prefill._ensure_runtime_state_available():
+        if not prefill.is_ssm or prefill._make_runtime_state_available():
             return None
 
         result = self._prefix_match.rollback(
@@ -566,7 +533,7 @@ def _admit_runtime_state(self):
             return result
         if not self._prepare_and_evict():
             return _PrefillAdmissionResult.stop()
-        if not prefill._ensure_runtime_state_available():
+        if not prefill._make_runtime_state_available():
             seq.kv_token_limit = None
             return _PrefillAdmissionResult.stop()
         return None
@@ -633,30 +600,40 @@ def _token_budget_rejection(self):
         return _PrefillAdmissionResult.skip()
 
     def _check_prefill_admission_gates(self):
-        """Apply prefill gates, tentatively matching only when it may help."""
+        """Apply long-prefill and token-budget admission gates."""
+        result = self._apply_nonfinal_long_prefill_gate()
+        if result is not None:
+            return result
+        return self._apply_prefill_token_budget_gate()
+
+    def _apply_nonfinal_long_prefill_gate(self):
+        """Reject non-final long prefills when this turn excludes them."""
         prefill = self.prefill_scheduler
         seq = self.seq
-        token_budget = prefill.cache_config.max_prefill_token_num
-        prefill_token_count = prefill._prefill_admission_token_count(seq)
-        is_nonfinal_long_prefill = prefill._prefill_kv_token_limit(seq) is not None
+        if (self.turn_policy.allows_nonfinal_long_prefill
+                or prefill._prefill_kv_token_limit(seq) is None):
+            return None
 
-        if (is_nonfinal_long_prefill
-                and not self.turn_policy.allows_nonfinal_long_prefill):
-            matched = self._match_prefix_for_prefill_gate()
-            if matched is None:
-                return _PrefillAdmissionResult.skip()
-            if prefill._prefill_kv_token_limit(seq) is not None:
-                self._prefix_match.rollback('still non-final long prefill on short turn')
-                return _PrefillAdmissionResult.skip()
-            self._prefix_match.retain_for_admission(
-                _PrefillAdmissionResult.skip())
-            prefill_token_count = prefill._prefill_admission_token_count(seq)
+        matched = self._match_prefix_for_prefill_gate()
+        if matched is None:
+            return _PrefillAdmissionResult.skip()
+        if prefill._prefill_kv_token_limit(seq) is not None:
+            self._prefix_match.rollback(
+                'still non-final long prefill on short turn')
+            return _PrefillAdmissionResult.skip()
+        self._prefix_match.retain_for_admission(
+            _PrefillAdmissionResult.skip())
+        return None
 
-        exceeds_token_budget = (
-            self.batch_has_prefill
-            and self.batch_prefill_tokens + prefill_token_count > token_budget
-        )
-        if not exceeds_token_budget:
+    def _apply_prefill_token_budget_gate(self):
+        """Reject a second prefill that exceeds this turn's token budget."""
+        prefill = self.prefill_scheduler
+        seq = self.seq
+        prefill_token_count = prefill._prefill_admission_token_count(seq)
+        token_budget = prefill.cache_config.max_prefill_token_num
+        if (not self.batch_has_prefill
+                or self.batch_prefill_tokens + prefill_token_count
+                <= token_budget):
             return None
 
         if not self._prefix_match.is_matched:
@@ -755,7 +732,7 @@ def __init__(
             _envs.opt_ttft_aging_sec,
         )
 
-    def _ensure_runtime_state_available(self):
+    def _make_runtime_state_available(self):
         """Make one state-cache slot available for an SSM runtime state."""
         if not self.is_ssm:
             return True
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index ec01db0ea2..ed9b8a89ef 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -312,12 +312,10 @@ def schedule_running(self, running: SeqList, num_required_tokens: int = 1, preal
         assert len(running) > 0
         eviction_helper = self.eviction_helper
 
-        valid_mask = [True for _ in running]
-
-        # loop over reverse running
-        rev_running = reversed(running)
-        for idx, seq in enumerate(rev_running):
-            if not seq.status == MessageStatus.RUNNING:
+        valid_mask = [True] * len(running)
+        for idx in reversed(range(len(running))):
+            seq = running[idx]
+            if seq.status != MessageStatus.RUNNING:
                 valid_mask[idx] = False
                 continue
             num_required_blocks = self.block_manager.num_required_blocks(seq, num_required_tokens)
@@ -329,13 +327,10 @@ def schedule_running(self, running: SeqList, num_required_tokens: int = 1, preal
                 self.block_trie.allocate(seq)
                 continue
 
-            # running to ready
             seq.state.deactivate()
-            # ready to waiting
             self.kv_load_coordinator.release(seq)
             seq.state.evict()
             valid_mask[idx] = False
-        valid_mask = list(reversed(valid_mask))
         return valid_mask
 
     def stop_session(self, session_id: int):
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index 321ab33072..ed070f67f8 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -1,5 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -184,6 +186,17 @@ def test_decode_requires_schedule_running(self, scheduler):
         with pytest.raises(ValueError, match='schedule_running'):
             scheduler.schedule(is_prefill=False)
 
+    def test_schedule_running_validity_uses_input_indices(self, scheduler,
+                                                          monkeypatch):
+        waiting = SimpleNamespace(status=MessageStatus.WAITING)
+        running = SimpleNamespace(status=MessageStatus.RUNNING)
+        monkeypatch.setattr(scheduler.block_manager, 'num_required_blocks',
+                            lambda seq, num_tokens: 0)
+
+        valid_mask = scheduler.schedule_running([waiting, running])
+
+        assert valid_mask == [False, True]
+
     @pytest.mark.parametrize(
         ('name', 'status'),
         [
diff --git a/tests/pytorch/paging/test_scheduler_ssm.py b/tests/pytorch/paging/test_scheduler_ssm.py
index 399a24812f..74938d4c51 100644
--- a/tests/pytorch/paging/test_scheduler_ssm.py
+++ b/tests/pytorch/paging/test_scheduler_ssm.py
@@ -323,12 +323,12 @@ def test_ssm_scheduler_rejects_prefix_match_for_prefill_gate_after_runtime_state
     node, state_idx = _add_published_ssm_checkpoint(scheduler, [1] * block_size * 2)
     ensure_results = iter([False, True])
 
-    def _ensure_runtime_state_available_once_then_succeed():
+    def _make_runtime_state_available_once_then_succeed():
         return next(ensure_results)
 
     monkeypatch.setattr(scheduler._prefill_scheduler,
-                        '_ensure_runtime_state_available',
-                        _ensure_runtime_state_available_once_then_succeed)
+                        '_make_runtime_state_available',
+                        _make_runtime_state_available_once_then_succeed)
     scheduler.block_trie.stats.reset()
 
     cache_hit_tail = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
@@ -379,11 +379,11 @@ def test_schedule_prefill_reapplies_chunk_limit_after_ssm_state_rollback():
 
     ensure_results = iter([False, True])
 
-    def _ensure_runtime_state_available_once_then_succeed():
+    def _make_runtime_state_available_once_then_succeed():
         return next(ensure_results)
 
-    scheduler._prefill_scheduler._ensure_runtime_state_available = (
-        _ensure_runtime_state_available_once_then_succeed)
+    scheduler._prefill_scheduler._make_runtime_state_available = (
+        _make_runtime_state_available_once_then_succeed)
 
     output = scheduler.schedule(is_prefill=True, prealloc_size=1)
 

From c8eeb8fc147da27c57b3652960703ee0bdff8d99 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 17:28:14 +0800
Subject: [PATCH 13/22] refactor: flatten prefill admission guards

---
 lmdeploy/pytorch/paging/prefill_scheduler.py | 33 ++++++++++++--------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index 2a1b8dc674..ac425a8d8e 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -575,8 +575,11 @@ def _evictable_sequences(self):
     def _match_prefix_for_prefill_gate(self):
         """Tentatively match once so a request can be rechecked by a gate."""
         prefill = self.prefill_scheduler
-        if (self._load_ready or not prefill.block_trie.enabled
-                or self._has_private_local_tail()):
+        if self._load_ready:
+            return None
+        if not prefill.block_trie.enabled:
+            return None
+        if self._has_private_local_tail():
             return None
         self._prefix_match.match()
         return True
@@ -636,18 +639,22 @@ def _apply_prefill_token_budget_gate(self):
                 <= token_budget):
             return None
 
-        if not self._prefix_match.is_matched:
-            matched = self._match_prefix_for_prefill_gate()
-            if matched is not None:
-                prefill_token_count = prefill._prefill_admission_token_count(seq)
-                if self.batch_prefill_tokens + prefill_token_count <= token_budget:
-                    self._prefix_match.retain_for_admission(
-                        self._token_budget_rejection())
-                    return None
-                self._prefix_match.rollback('still exceeds prefill token budget')
-        else:
+        rejection = self._token_budget_rejection()
+        if self._prefix_match.is_matched:
+            self._prefix_match.rollback('still exceeds prefill token budget')
+            return rejection
+
+        matched = self._match_prefix_for_prefill_gate()
+        if matched is None:
+            return rejection
+
+        prefill_token_count = prefill._prefill_admission_token_count(seq)
+        if self.batch_prefill_tokens + prefill_token_count > token_budget:
             self._prefix_match.rollback('still exceeds prefill token budget')
-        return self._token_budget_rejection()
+            return rejection
+
+        self._prefix_match.retain_for_admission(rejection)
+        return None
 
     def _prepare_and_evict(self):
         """Apply chunk allocation limits and evict for this prefill."""

From 0d4b05056c52b4927d002623d47555fadda8cb32 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 17:55:47 +0800
Subject: [PATCH 14/22] test: characterize scheduler API boundaries

---
 tests/pytorch/engine/test_inputs_maker.py     | 34 +++++++++++++++++++
 .../paging/test_scheduler_kv_transfer.py      | 18 +++++++++-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/tests/pytorch/engine/test_inputs_maker.py b/tests/pytorch/engine/test_inputs_maker.py
index 985231e4b8..044e56fc8a 100644
--- a/tests/pytorch/engine/test_inputs_maker.py
+++ b/tests/pytorch/engine/test_inputs_maker.py
@@ -379,6 +379,40 @@ async def send_next_inputs(self):
     assert events == ['collect_migration_done', 'send_next_inputs']
 
 
+def test_migration_loop_schedules_and_processes_ready_batch():
+    events = []
+    migration_ready = [object()]
+
+    class _Scheduler:
+
+        def _schedule_migration(self):
+            events.append('schedule')
+            return migration_ready
+
+        def has_migration_waiting(self):
+            raise AssertionError('a ready migration batch must be processed')
+
+    class _MigrationEvent:
+
+        def clear(self):
+            events.append('clear')
+
+    async def _process_ready(actual):
+        events.append(('process', actual))
+        loop.stop_event.set()
+
+    loop = EngineLoop.__new__(EngineLoop)
+    loop.stop_event = asyncio.Event()
+    loop._sleep_requested = False
+    loop.scheduler = _Scheduler()
+    loop.migration_event = _MigrationEvent()
+    loop._migration_loop_process_ready = _process_ready
+
+    asyncio.run(loop.migration_loop())
+
+    assert events == ['schedule', 'clear', ('process', migration_ready)]
+
+
 def test_engine_loop_uses_short_yield_only_for_pending_lookup(monkeypatch):
     sleeps = []
 
diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py
index f13589439a..a7156c2fd5 100644
--- a/tests/pytorch/paging/test_scheduler_kv_transfer.py
+++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py
@@ -24,13 +24,14 @@ def __init__(self, results, failed_ids=()):
         self.results = iter(results)
         self.failed_ids = set(failed_ids)
         self.pending_ids = set()
+        self.new_requests = []
         self.lookup_calls = []
         self.cancelled = []
         self.finished = []
         self.allocations = []
 
     def on_new_request(self, request):
-        pass
+        self.new_requests.append(request.seq_id)
 
     def is_lookup_pending(self, request_id):
         return request_id in self.pending_ids
@@ -109,6 +110,21 @@ def _make_async_lookup_scheduler(
     )
 
 
+def test_sequence_lifecycle_notifies_connector_on_add_and_end():
+    connector = _AsyncLookupConnector([])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+    )
+
+    seq = scheduler.add_session(71).add_sequence(torch.arange(4))
+    assert connector.new_requests == [seq.seq_id]
+
+    scheduler.end_session(71)
+    assert connector.finished == [seq.seq_id]
+    assert 71 not in scheduler.sessions
+
+
 def test_async_lookup_pending_rolls_back_a_new_request_once():
     connector = _AsyncLookupConnector([(None, False), (0, False)])
     scheduler = _make_async_lookup_scheduler(connector)

From 9e2908610582b9dae435a4398df19bb4cfe60b85 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 18:09:41 +0800
Subject: [PATCH 15/22] refactor: narrow scheduler engine boundary

---
 lmdeploy/pytorch/engine/engine_loop.py    | 22 +++---
 lmdeploy/pytorch/engine/inputs_maker.py   | 35 ++++++---
 lmdeploy/pytorch/paging/scheduler.py      | 18 ++++-
 tests/pytorch/engine/test_inputs_maker.py | 93 +++++++++++++----------
 tests/pytorch/paging/test_scheduler.py    |  2 +-
 5 files changed, 106 insertions(+), 64 deletions(-)

diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py
index bd240f776e..fd8ffce42a 100644
--- a/lmdeploy/pytorch/engine/engine_loop.py
+++ b/lmdeploy/pytorch/engine/engine_loop.py
@@ -24,6 +24,7 @@
     from lmdeploy.pytorch.engine.model_agent import BatchedOutputs
     from lmdeploy.pytorch.model_inputs import ModelInputs, ModelInputsDelta
     from lmdeploy.pytorch.paging import Scheduler
+    from lmdeploy.pytorch.paging.block_trie.checkpoint_lifecycle import StateCheckpointLifecycle
     from lmdeploy.pytorch.strategies.base.sequence import SequenceStrategy
 
     from .engine import Engine, SeqList
@@ -117,6 +118,7 @@ class EngineLoop:
     def __init__(self,
                  req_manager: 'RequestManager',
                  scheduler: 'Scheduler',
+                 state_checkpoints: 'StateCheckpointLifecycle',
                  executor: 'ExecutorBase',
                  seq_strategy: 'SequenceStrategy',
                  inputs_maker: 'InputsMakerAsync',
@@ -124,6 +126,7 @@ def __init__(self,
                  engine_conn: Optional['EngineP2PConnection'] = None):
         self.req_manager = req_manager
         self.scheduler = scheduler
+        self.state_checkpoints = state_checkpoints
         self.executor = executor
         self.seq_strategy = seq_strategy
         self.inputs_maker = inputs_maker
@@ -323,7 +326,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
             seq.append_routed_experts(all_routed_experts)
             seq.append_logits(logits)
             seq.append_ce_loss(ce_loss, finish=False)
-            self.scheduler.block_trie.cache_routed_experts_for_seq(seq)
+            self.scheduler.cache_routed_experts([seq])
             return dict()
 
         new_token_timestamp = batched_outputs.new_token_timestamp
@@ -337,7 +340,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
                                          batched_outputs=batched_outputs,
                                          model_inputs=model_inputs,
                                          delta=delta)
-        self.scheduler.block_trie.cache_routed_experts(running)
+        self.scheduler.cache_routed_experts(running)
 
         # generate output
         outputs: dict[int, InferOutput] = dict()
@@ -354,7 +357,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
                 continue
             session_id = msg.session_id
             if msg.resp_cache:
-                cache_block_ids = self.scheduler.block_manager.get_block_table(msg).tolist()
+                cache_block_ids = self.scheduler.get_block_tables([msg])[0].tolist()
             else:
                 cache_block_ids = None
 
@@ -422,21 +425,21 @@ async def _wait_for_schedulable_prefill(self):
             # warning or adding the full pressure backoff to TTFT.
             await asyncio.sleep(0.001)
             return
+        cache_usage = scheduler.schedule_metrics.cache_usage
         logger.warning(f'no next prefill running request, Maybe cache is full, '
-                       f'free gpu cache blocks: {scheduler.block_manager.get_num_free_gpu_blocks()}, '
-                       f'total gpu cache blocks: {scheduler.block_manager.num_gpu_blocks}')
+                       f'gpu cache usage: {cache_usage:.1%}')
         await asyncio.sleep(0.1)
 
     def _publish_forward_checkpoints(self, running: 'SeqList', has_state_checkpoint_save: bool):
         """Publish per-forward prefix-cache ownership before prefetching."""
-        state_checkpoints = self.scheduler.block_trie.state_checkpoints
+        state_checkpoints = self.state_checkpoints
         if has_state_checkpoint_save:
             state_checkpoints.publish_saves(running, pin_saves=True)
         state_checkpoints.unpin_restores(running)
 
     def _release_forward_save_pins(self, running: 'SeqList'):
         """Unpin producers after the forward output/event boundary."""
-        self.scheduler.block_trie.state_checkpoints.unpin_saves(running)
+        self.state_checkpoints.unpin_saves(running)
 
     def _finish_forward_output(self,
                                out: 'BatchedOutputs | None',
@@ -558,7 +561,7 @@ async def _migration_loop_migrate(self, migration_ready: 'SeqList'):
             migration_execution_requests: list[tuple[int, list[tuple[int, int]]]] = []
             migration_request = msg.migration_request
             prefill_block_ids = migration_request.remote_block_ids
-            decode_block_ids = list(self.scheduler.block_manager.get_block_table(msg=msg))
+            decode_block_ids = list(self.scheduler.get_block_tables([msg])[0])
 
             assert len(prefill_block_ids) == len(decode_block_ids), (
                 f'#prefill block ids ({len(prefill_block_ids)}) must equal to '
@@ -615,7 +618,7 @@ async def migration_loop(self):
                 await self._sleep_resume_event.wait()
                 continue
 
-            migration_ready = self.scheduler._schedule_migration()
+            migration_ready = self.scheduler.schedule_migration()
             if not migration_ready and not self.scheduler.has_migration_waiting():
                 await self.migration_event.wait()
             elif migration_ready:
@@ -692,6 +695,7 @@ def build_engine_loop(engine: 'Engine'):
     return EngineLoop(
         req_manager=engine.req_manager,
         scheduler=engine.scheduler,
+        state_checkpoints=engine.scheduler.state_checkpoints,
         executor=engine.executor,
         seq_strategy=engine.seq_strategy,
         inputs_maker=inputs_maker,
diff --git a/lmdeploy/pytorch/engine/inputs_maker.py b/lmdeploy/pytorch/engine/inputs_maker.py
index 3da3eab2f1..164eb1446d 100644
--- a/lmdeploy/pytorch/engine/inputs_maker.py
+++ b/lmdeploy/pytorch/engine/inputs_maker.py
@@ -34,6 +34,7 @@
     from lmdeploy.pytorch.messages import SchedulerSequence
     from lmdeploy.pytorch.multimodal.data_type import MultiModalInputs
     from lmdeploy.pytorch.paging import Scheduler
+    from lmdeploy.pytorch.paging.block_trie.checkpoint_lifecycle import StateCheckpointLifecycle
     from lmdeploy.pytorch.strategies.base.engine import EngineStrategy
     from lmdeploy.pytorch.strategies.base.model_agent import ModelAgentStrategy
     from lmdeploy.pytorch.strategies.base.sampling import SamplingStrategy
@@ -95,6 +96,11 @@ class InputsMakerConfig:
     max_batches: int
     max_prefill_token_num: int
     role: EngineRole
+    block_size: int
+    kernel_block_size: int
+    window_size: int
+    enable_prefix_caching: bool
+    prefix_cache_decode_state_interval: int
     is_ssm: bool = False
     dp: int = 1
     spec_decoding: bool = False
@@ -118,6 +124,11 @@ def from_engine(engine: 'Engine'):
             max_batches=cache_config.max_batches,
             max_prefill_token_num=cache_config.max_prefill_token_num,
             role=cache_config.role,
+            block_size=cache_config.block_size,
+            kernel_block_size=cache_config.kernel_block_size,
+            window_size=cache_config.window_size,
+            enable_prefix_caching=cache_config.enable_prefix_caching,
+            prefix_cache_decode_state_interval=cache_config.prefix_cache_decode_state_interval,
             is_ssm=len(cache_config.states_shapes) > 0,
             dp=engine.dist_config.dp,
             enable_chunked_prefill=engine.misc_config.enable_chunked_prefill,
@@ -355,7 +366,7 @@ def run(self):
 
         result = self.result
         connector_token_lens = ()
-        connector_enabled = maker.scheduler.kv_connector is not None
+        connector_enabled = maker.scheduler.has_kv_connector()
         if (connector_enabled and result.inputs is not None
                 and not result.inputs.is_decoding and not result.inputs.is_dummy):
             # A prefill writes KV through the end of its query. The connector
@@ -689,6 +700,7 @@ def __init__(
         self,
         executor: 'ExecutorBase',
         scheduler: 'Scheduler',
+        state_checkpoints: 'StateCheckpointLifecycle',
         adapter_manager: 'AdapterManager',
         engine_strategy: 'EngineStrategy',
         sampling_strategy: 'SamplingStrategy',
@@ -697,11 +709,11 @@ def __init__(
     ):
         self.executor = executor
         self.scheduler = scheduler
+        self.state_checkpoints = state_checkpoints
         self.adapter_manager = adapter_manager
         self.config = config
         self.spec_decoding = config.spec_decoding
-        self.cache_config = scheduler.cache_config
-        self.kernel_blocks_per_kv = self.cache_config.block_size // self.cache_config.kernel_block_size
+        self.kernel_blocks_per_kv = config.block_size // config.kernel_block_size
         self.kernel_block_arange = torch.arange(self.kernel_blocks_per_kv, dtype=self.torch_int_dtype)
 
         # strategies
@@ -893,7 +905,7 @@ def _make_kv_prefix_cache_copy_plan(
 
     def _ssm_prefix_cache_enabled(self):
         """Check whether this input maker emits SSM checkpoint operations."""
-        return self.config.is_ssm and self.cache_config.enable_prefix_caching
+        return self.config.is_ssm and self.config.enable_prefix_caching
 
     def _prepare_prefill_cache_restore(
             self, messages: 'SeqList') -> tuple[torch.LongTensor | None, StateCacheCopyPlan | None]:
@@ -902,7 +914,7 @@ def _prepare_prefill_cache_restore(
         if state_restore_plan is None:
             return None, None
 
-        state_checkpoints = self.scheduler.block_trie.state_checkpoints
+        state_checkpoints = self.state_checkpoints
         # Keep checkpoint sources alive while the prefetched forward waits to
         # copy them into request-owned KV and runtime state.
         state_checkpoints.pin_restores(messages)
@@ -917,7 +929,7 @@ def _prepare_prefill_cache_restore(
             checkpoint = restore.node.state_checkpoint
             if checkpoint.frozen_block_id < 0:
                 continue
-            dst_block_idx = checkpoint.step // self.cache_config.block_size
+            dst_block_idx = checkpoint.step // self.config.block_size
             if dst_block_idx >= len(msg.logical_blocks):
                 raise RuntimeError('SSM prefix-cache restore destination block is missing.')
             logical_pairs.append((checkpoint.frozen_block_id, msg.logical_blocks[dst_block_idx]))
@@ -933,7 +945,7 @@ def _prepare_prefill_cache_save(
         save_steps: tuple[int, ...] | None,
     ) -> tuple[torch.LongTensor | None, StateCacheCopyPlan | None]:
         """Reserve checkpoints and build prefill save plans."""
-        state_checkpoints = self.scheduler.block_trie.state_checkpoints
+        state_checkpoints = self.state_checkpoints
         if save_steps is None:
             save_state_offsets = [state_checkpoints.reserve_save(msg) for msg in messages]
         else:
@@ -949,7 +961,7 @@ def _prepare_prefill_cache_save(
             checkpoint = pending_save.node.state_checkpoint
             if checkpoint.frozen_block_id < 0:
                 continue
-            src_block_idx = pending_save.step // self.cache_config.block_size
+            src_block_idx = pending_save.step // self.config.block_size
             if src_block_idx >= len(msg.logical_blocks):
                 raise RuntimeError('SSM prefix-cache save source block is missing.')
             logical_pairs.append((msg.logical_blocks[src_block_idx], checkpoint.frozen_block_id))
@@ -984,11 +996,11 @@ def _make_decode_cache_inputs(self, valid_seqs: 'SeqList', delta: ModelInputsDel
         if delta is None or len(valid_seqs) == 0 or not self._ssm_prefix_cache_enabled():
             return None
 
-        decode_state_interval = self.cache_config.prefix_cache_decode_state_interval
+        decode_state_interval = self.config.prefix_cache_decode_state_interval
         if (decode_state_interval <= 0 or self.spec_decoding or delta.max_q_seqlen != 1):
             return None
 
-        state_checkpoints = self.scheduler.block_trie.state_checkpoints
+        state_checkpoints = self.state_checkpoints
         save_state_offsets = [state_checkpoints.reserve_decode_save(seq, decode_state_interval)
                               for seq in valid_seqs]
         state_save_plan = _make_state_prefix_cache_save_plan(valid_seqs, save_state_offsets)
@@ -1161,7 +1173,7 @@ def create_model_inputs_delta(self):
         block_offsets = self._map_to_kernel_block_offsets(block_offsets)
 
         # sliding window
-        if self.scheduler.cache_config.window_size > 0:
+        if self.config.window_size > 0:
             num_ignored_history = torch.tensor([msg.num_ignored_history for msg in valid_seqs])
         else:
             num_ignored_history = torch.zeros(len(valid_seqs), dtype=torch.long)
@@ -1342,6 +1354,7 @@ def build_inputs_maker(engine: 'Engine'):
     return InputsMakerAsync(
         executor=engine.executor,
         scheduler=engine.scheduler,
+        state_checkpoints=engine.scheduler.state_checkpoints,
         adapter_manager=engine.adapter_manager,
         engine_strategy=engine.engine_strategy,
         sampling_strategy=engine.sampling_strategy,
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index ed9b8a89ef..ccc5faca28 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -25,6 +25,8 @@
 if TYPE_CHECKING:
     from lmdeploy.pytorch.kv_connector.base import KVConnectorBase
 
+    from .block_trie.checkpoint_lifecycle import StateCheckpointLifecycle
+
 MapType = dict[int, int]
 SeqList = list[SchedulerSequence]
 
@@ -246,7 +248,8 @@ def add_session(self, session_id: int):
         self.sessions[session_id] = session
         return session
 
-    def _schedule_migration(self):
+    def schedule_migration(self):
+        """Admit waiting migration sequences to paging resources."""
         migration_ready: SeqList = []
         migration_waiting = sorted(
             self.migration_waiting,
@@ -439,6 +442,10 @@ def build_connector_meta(
             self.kv_save_coordinator.acquire(metadata)
         return metadata
 
+    def has_kv_connector(self) -> bool:
+        """Whether connector work can be produced or polled."""
+        return self.kv_connector is not None
+
     def update_connector_output(self, connector_output) -> None:
         """Convert all-TP worker progress into paging state transitions.
 
@@ -471,6 +478,15 @@ def get_block_tables(self, seqs: SeqList):
         """Get block tables for the sequences."""
         return [self.block_manager.get_block_table(seq) for seq in seqs]
 
+    @property
+    def state_checkpoints(self) -> 'StateCheckpointLifecycle':
+        """Return the prefix-checkpoint lifecycle owner."""
+        return self.block_trie.state_checkpoints
+
+    def cache_routed_experts(self, seqs: SeqList) -> None:
+        """Publish routed-expert history to reusable prefix nodes."""
+        self.block_trie.cache_routed_experts(seqs)
+
     def resolve_gpu_block_offsets(self, logical_block_ids):
         """Resolve paging-owned logical ids for a forward cache-copy plan."""
         return self.block_manager.resolve_gpu_block_offsets(logical_block_ids)
diff --git a/tests/pytorch/engine/test_inputs_maker.py b/tests/pytorch/engine/test_inputs_maker.py
index 044e56fc8a..777d49f6cb 100644
--- a/tests/pytorch/engine/test_inputs_maker.py
+++ b/tests/pytorch/engine/test_inputs_maker.py
@@ -130,6 +130,9 @@ def __init__(self, running, waiting=None, num_ready=0, num_running=0):
         self.kv_connector = None
         self.connector_meta_calls = []
 
+    def has_kv_connector(self):
+        return self.kv_connector is not None
+
     def schedule(self,
                  is_prefill: bool,
                  prealloc_size: int,
@@ -186,6 +189,17 @@ def make_stopping_criteria(self, running):
         return None
 
 
+def _make_inputs_maker_config():
+    return InputsMakerConfig(max_batches=1,
+                             max_prefill_token_num=512,
+                             role=EngineRole.Decode,
+                             block_size=16,
+                             kernel_block_size=16,
+                             window_size=-1,
+                             enable_prefix_caching=False,
+                             prefix_cache_decode_state_interval=0)
+
+
 def _fake_model_inputs(is_chunk: bool = False):
     return SimpleNamespace(is_decoding=False,
                            is_dummy=False,
@@ -237,9 +251,9 @@ async def get_output_async(self):
             return None
 
     state_checkpoints = _StateCheckpoints()
-    block_trie = SimpleNamespace(enabled=True, state_checkpoints=state_checkpoints)
     loop = EngineLoop.__new__(EngineLoop)
-    loop.scheduler = SimpleNamespace(block_trie=block_trie, collect_migration_done=lambda: None)
+    loop.scheduler = SimpleNamespace(collect_migration_done=lambda: None)
+    loop.state_checkpoints = state_checkpoints
     loop.inputs_maker = _InputsMaker(state_checkpoints)
     loop.executor = _Executor(state_checkpoints)
     loop._sleep_requested = False
@@ -323,9 +337,9 @@ async def get_output_async(self):
             return None
 
     state_checkpoints = _StateCheckpoints()
-    block_trie = SimpleNamespace(enabled=True, state_checkpoints=state_checkpoints)
     loop = EngineLoop.__new__(EngineLoop)
-    loop.scheduler = SimpleNamespace(block_trie=block_trie, collect_migration_done=lambda: None)
+    loop.scheduler = SimpleNamespace(collect_migration_done=lambda: None)
+    loop.state_checkpoints = state_checkpoints
     loop.inputs_maker = _InputsMaker()
     loop.executor = _Executor()
     loop._sleep_requested = True
@@ -385,7 +399,7 @@ def test_migration_loop_schedules_and_processes_ready_batch():
 
     class _Scheduler:
 
-        def _schedule_migration(self):
+        def schedule_migration(self):
             events.append('schedule')
             return migration_ready
 
@@ -416,36 +430,32 @@ async def _process_ready(actual):
 def test_engine_loop_uses_short_yield_only_for_pending_lookup(monkeypatch):
     sleeps = []
 
-    class _BlockManager:
-        num_gpu_blocks = 8
-
+    class _Scheduler:
         def __init__(self):
             self.reads = 0
+            self.last_schedule_had_pending_lookup = True
 
-        def get_num_free_gpu_blocks(self):
+        @property
+        def schedule_metrics(self):
             self.reads += 1
-            return 4
+            return SimpleNamespace(cache_usage=0.5)
 
     async def record_sleep(delay):
         sleeps.append(delay)
 
-    block_manager = _BlockManager()
-    scheduler = SimpleNamespace(
-        last_schedule_had_pending_lookup=True,
-        block_manager=block_manager,
-    )
+    scheduler = _Scheduler()
     loop = EngineLoop.__new__(EngineLoop)
     loop.scheduler = scheduler
     monkeypatch.setattr(engine_loop_module.asyncio, 'sleep', record_sleep)
 
     asyncio.run(loop._wait_for_schedulable_prefill())
     assert sleeps == [0.001]
-    assert block_manager.reads == 0
+    assert scheduler.reads == 0
 
     scheduler.last_schedule_had_pending_lookup = False
     asyncio.run(loop._wait_for_schedulable_prefill())
     assert sleeps == [0.001, 0.1]
-    assert block_manager.reads == 1
+    assert scheduler.reads == 1
 
 
 def test_engine_loop_reset_runtime_state_delegates_to_inputs_maker():
@@ -485,12 +495,13 @@ def _make_policy_maker(long_seq, decode_seq=None):
 
 def test_inputs_maker_reads_opt_ttft_short_turns_env(monkeypatch):
     monkeypatch.setattr(inputs_maker_module._envs, 'opt_ttft_short_turns', 5)
-    scheduler = SimpleNamespace(cache_config=SimpleNamespace(block_size=16, kernel_block_size=16))
-    config = InputsMakerConfig(max_batches=1, max_prefill_token_num=512, role=EngineRole.Decode)
+    scheduler = SimpleNamespace()
+    config = _make_inputs_maker_config()
 
     maker = InputsMakerAsync(
         executor=SimpleNamespace(device_type='cpu'),
         scheduler=scheduler,
+        state_checkpoints=SimpleNamespace(),
         adapter_manager=SimpleNamespace(),
         engine_strategy=_FakeEngineStrategy(),
         sampling_strategy=_FakeSamplingStrategy(),
@@ -503,12 +514,13 @@ def test_inputs_maker_reads_opt_ttft_short_turns_env(monkeypatch):
 
 def test_inputs_maker_clamps_opt_ttft_short_turns_to_one(monkeypatch):
     monkeypatch.setattr(inputs_maker_module._envs, 'opt_ttft_short_turns', 0)
-    scheduler = SimpleNamespace(cache_config=SimpleNamespace(block_size=16, kernel_block_size=16))
-    config = InputsMakerConfig(max_batches=1, max_prefill_token_num=512, role=EngineRole.Decode)
+    scheduler = SimpleNamespace()
+    config = _make_inputs_maker_config()
 
     maker = InputsMakerAsync(
         executor=SimpleNamespace(device_type='cpu'),
         scheduler=scheduler,
+        state_checkpoints=SimpleNamespace(),
         adapter_manager=SimpleNamespace(),
         engine_strategy=_FakeEngineStrategy(),
         sampling_strategy=_FakeSamplingStrategy(),
@@ -520,11 +532,12 @@ def test_inputs_maker_clamps_opt_ttft_short_turns_to_one(monkeypatch):
 
 
 def test_inputs_maker_reset_runtime_state_discards_request_local_state():
-    scheduler = SimpleNamespace(cache_config=SimpleNamespace(block_size=16, kernel_block_size=16))
-    config = InputsMakerConfig(max_batches=1, max_prefill_token_num=512, role=EngineRole.Decode)
+    scheduler = SimpleNamespace()
+    config = _make_inputs_maker_config()
     maker = InputsMakerAsync(
         executor=SimpleNamespace(device_type='cpu'),
         scheduler=scheduler,
+        state_checkpoints=SimpleNamespace(),
         adapter_manager=SimpleNamespace(),
         engine_strategy=_FakeEngineStrategy(),
         sampling_strategy=_FakeSamplingStrategy(),
@@ -1400,9 +1413,8 @@ def reserve_save(self, seq, step=None):
             return state_idx
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
-    maker.config = SimpleNamespace(is_ssm=True)
-    maker.cache_config = SimpleNamespace(enable_prefix_caching=True)
-    maker.scheduler = SimpleNamespace(block_trie=SimpleNamespace(state_checkpoints=_StateCheckpoints()))
+    maker.config = SimpleNamespace(is_ssm=True, enable_prefix_caching=True)
+    maker.state_checkpoints = _StateCheckpoints()
 
     cache_inputs = maker._prepare_prefill_cache_inputs(messages)
 
@@ -1430,9 +1442,8 @@ def reserve_save(self, seq, step=None):
             return 21
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
-    maker.config = SimpleNamespace(is_ssm=True)
-    maker.cache_config = SimpleNamespace(enable_prefix_caching=True)
-    maker.scheduler = SimpleNamespace(block_trie=SimpleNamespace(state_checkpoints=_StateCheckpoints()))
+    maker.config = SimpleNamespace(is_ssm=True, enable_prefix_caching=True)
+    maker.state_checkpoints = _StateCheckpoints()
 
     cache_inputs = maker._prepare_prefill_cache_inputs([seq], save_steps=(160, ))
 
@@ -1463,11 +1474,10 @@ def reserve_save(self, seq, step=None):
             return state_idx
 
     scheduler = _CopyPlanScheduler()
-    scheduler.block_trie = SimpleNamespace(state_checkpoints=_StateCheckpoints())
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
-    maker.config = SimpleNamespace(is_ssm=True)
-    maker.cache_config = SimpleNamespace(enable_prefix_caching=True, block_size=16)
+    maker.config = SimpleNamespace(is_ssm=True, enable_prefix_caching=True, block_size=16)
     maker.scheduler = scheduler
+    maker.state_checkpoints = _StateCheckpoints()
 
     cache_inputs = maker._prepare_prefill_cache_inputs(messages)
 
@@ -1481,8 +1491,7 @@ def reserve_save(self, seq, step=None):
 
 def test_prepare_prefill_cache_inputs_rejects_mismatched_save_steps():
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
-    maker.config = SimpleNamespace(is_ssm=True)
-    maker.cache_config = SimpleNamespace(enable_prefix_caching=True)
+    maker.config = SimpleNamespace(is_ssm=True, enable_prefix_caching=True)
 
     with pytest.raises(ValueError, match='one entry per prefill sequence'):
         maker._prepare_prefill_cache_inputs([_state_seq(4, 11)], save_steps=())
@@ -1498,10 +1507,10 @@ def reserve_decode_save(self, seq, interval):
             return {4: 31, 5: -1}[seq.logical_state]
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
-    maker.config = SimpleNamespace(is_ssm=True)
-    maker.cache_config = SimpleNamespace(enable_prefix_caching=True,
-                                         prefix_cache_decode_state_interval=16)
-    maker.scheduler = SimpleNamespace(block_trie=SimpleNamespace(state_checkpoints=_StateCheckpoints()))
+    maker.config = SimpleNamespace(is_ssm=True,
+                                   enable_prefix_caching=True,
+                                   prefix_cache_decode_state_interval=16)
+    maker.state_checkpoints = _StateCheckpoints()
     maker.spec_decoding = False
     delta = SimpleNamespace(max_q_seqlen=1)
 
@@ -1530,10 +1539,10 @@ def reserve_decode_save(self, seq, interval):
             raise AssertionError('disabled decode checkpoint path must not reserve state')
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
-    maker.config = SimpleNamespace(is_ssm=is_ssm)
-    maker.cache_config = SimpleNamespace(enable_prefix_caching=enabled,
-                                         prefix_cache_decode_state_interval=interval)
-    maker.scheduler = SimpleNamespace(block_trie=SimpleNamespace(state_checkpoints=_StateCheckpoints()))
+    maker.config = SimpleNamespace(is_ssm=is_ssm,
+                                   enable_prefix_caching=enabled,
+                                   prefix_cache_decode_state_interval=interval)
+    maker.state_checkpoints = _StateCheckpoints()
     maker.spec_decoding = spec_decoding
     delta = SimpleNamespace(max_q_seqlen=max_q_seqlen)
 
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index ed070f67f8..e54f7e0b0d 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -272,7 +272,7 @@ def test_schedule_migration_matches_current_sequence():
                                          remote_block_ids=[1])
     seq = scheduler.add_session(100).add_sequence([1] * block_size, migration_request=migration_request)
 
-    output = scheduler._schedule_migration()
+    output = scheduler.schedule_migration()
 
     assert output == [seq]
     assert seq.status == MessageStatus.MIGRATION_READY

From 9cdd103db1eda3284e378693efb571cace779c1c Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 18:28:46 +0800
Subject: [PATCH 16/22] refactor: centralize scheduler sequence lifecycle

---
 lmdeploy/pytorch/disagg/conn/engine_conn.py   |   4 +-
 lmdeploy/pytorch/engine/engine.py             |  25 ++--
 lmdeploy/pytorch/messages.py                  |  34 ++---
 .../pytorch/paging/kv_load_coordinator.py     |   6 +-
 lmdeploy/pytorch/paging/scheduler.py          |  32 ++++-
 .../pytorch/paging/seq_states/__init__.py     |   2 +-
 lmdeploy/pytorch/paging/seq_states/states.py  | 122 ++++++++++++------
 .../pytorch/engine/test_abort_stopped_seqs.py |   3 +
 tests/pytorch/engine/test_engine_sleep.py     |  11 +-
 .../engine/test_kv_connector_wiring.py        |   2 +
 tests/pytorch/paging/test_scheduler.py        |   2 +
 11 files changed, 154 insertions(+), 89 deletions(-)

diff --git a/lmdeploy/pytorch/disagg/conn/engine_conn.py b/lmdeploy/pytorch/disagg/conn/engine_conn.py
index a322291bbd..87b02468fa 100644
--- a/lmdeploy/pytorch/disagg/conn/engine_conn.py
+++ b/lmdeploy/pytorch/disagg/conn/engine_conn.py
@@ -92,9 +92,7 @@ async def handle_zmq_recv(self, remote_engine_id: str):
                 logger.error(f'invalid zmq request from {remote_engine_id}: {e}')
                 continue
             session_id = req.remote_session_id
-            if session_id in self.engine.scheduler.sessions:
-                self.engine.end_session(session_id=session_id)
-            else:
+            if not self.engine.end_session(session_id=session_id):
                 logger.error(f'invalid free, {remote_engine_id}, {session_id}')
 
     async def zmq_disconnect(self, remote_engine_id: str):
diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py
index 398024bea1..2a5681d47f 100644
--- a/lmdeploy/pytorch/engine/engine.py
+++ b/lmdeploy/pytorch/engine/engine.py
@@ -345,7 +345,7 @@ def _on_add_session(self, reqs: list[Request], **kwargs):
             session_id = req.data['session_id']
             resp = req.data.get('response', True)
             resp_type = ResponseType.SESSION_REPEAT
-            if session_id not in self.scheduler.sessions:
+            if self.scheduler.get_session(session_id) is None:
                 self.scheduler.add_session(session_id)
                 resp_type = ResponseType.SUCCESS
             if resp:
@@ -357,8 +357,8 @@ def _on_stop_session(self, reqs: list[Request], **kwargs):
             session_id = req.data['session_id']
             resp = req.data.get('response', True)
             resp_type = ResponseType.SESSION_NOT_EXIST
-            if session_id in self.scheduler.sessions:
-                session = self.scheduler.sessions[session_id]
+            session = self.scheduler.get_session(session_id)
+            if session is not None:
                 stopped_resp_ids = set()
                 for seq in session.sequences.values():
                     if seq.status not in (MessageStatus.STOPPED, MessageStatus.TO_BE_MIGRATED):
@@ -412,8 +412,9 @@ def _on_end_session(self, reqs: list[Request], **kwargs):
             session_id = req.data['session_id']
             resp = req.data.get('response', True)
             resp_type = ResponseType.SESSION_NOT_EXIST
-            if session_id in self.scheduler.sessions:
-                msgs = list(self.scheduler.sessions[session_id].sequences.values())
+            session = self.scheduler.get_session(session_id)
+            if session is not None:
+                msgs = list(session.sequences.values())
                 if len(msgs) > 0 and msgs[0].preserve_cache:
                     msgs[0].state.finish()
                 else:
@@ -428,7 +429,7 @@ def _on_add_message(self, reqs: list[Request], **kwargs):
         for req in reqs:
             req_data = req.data
             session_id = req_data['session_id']
-            if self.scheduler and session_id not in self.scheduler.sessions:
+            if self.scheduler and self.scheduler.get_session(session_id) is None:
                 self._response(req.resp, ResponseType.SESSION_NOT_EXIST)
                 continue
             valid_reqs.append(req)
@@ -481,7 +482,7 @@ def __update_max_new_tokens(msg):
         scheduler = self.scheduler
         for req in reqs:
             session_id = req.data['session_id']
-            sess = scheduler.sessions.get(session_id, None)
+            sess = scheduler.get_session(session_id)
             if sess is None:
                 self._response(req.resp, ResponseType.SESSION_NOT_EXIST)
                 continue
@@ -576,8 +577,9 @@ def _unblock_new_inputs(self):
     def _cancel_and_end_all_sessions(self):
         """Cancel active responses and remove all scheduler sessions."""
         num_cancelled = 0
-        session_ids = list(self.scheduler.sessions.keys())
-        for session in list(self.scheduler.sessions.values()):
+        sessions = self.scheduler.get_sessions()
+        session_ids = [session.session_id for session in sessions]
+        for session in sessions:
             for seq in list(session.sequences.values()):
                 resp: Response = getattr(seq, 'resp', None)
                 if resp is None or resp.is_done:
@@ -714,8 +716,9 @@ def start_loop(self):
 
     def end_session(self, session_id: int):
         """End session."""
-        if session_id in self.scheduler.sessions:
-            has_multimodal = self._has_multimodal_session(self.scheduler.sessions[session_id])
+        session = self.scheduler.get_session(session_id)
+        if session is not None:
+            has_multimodal = self._has_multimodal_session(session)
             self.scheduler.end_session(session_id)
             self._maybe_trim_multimodal_session(has_multimodal)
             return True
diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py
index ab0056307f..06e6971d7b 100644
--- a/lmdeploy/pytorch/messages.py
+++ b/lmdeploy/pytorch/messages.py
@@ -29,8 +29,7 @@
 from .block import LogicalTokenBlocks
 
 if TYPE_CHECKING:
-    from lmdeploy.pytorch.paging.scheduler import Scheduler
-    from lmdeploy.pytorch.paging.seq_states.states import StateBase
+    from lmdeploy.pytorch.paging.seq_states.states import SequenceLifecycle, StateBase
     from lmdeploy.pytorch.strategies.base.sampling import SamplingStrategy
     from lmdeploy.pytorch.strategies.base.sequence import SequenceStrategy
 
@@ -209,6 +208,7 @@ class SequenceMeta:
     strategy: 'SequenceStrategy' = None
     sampling_strategy: 'SamplingStrategy' = None
     use_mrope: bool = False
+    enable_prefix_caching: bool = False
 
 
 class SequenceManager:
@@ -221,7 +221,7 @@ def __init__(self, seq_meta: SequenceMeta) -> None:
         self.seq_meta = seq_meta
         self._seq_count = 0
 
-    def _new_seq_id(self):
+    def new_sequence_id(self):
         seq_id = self._seq_count
         self._seq_count += 1
         return seq_id
@@ -282,12 +282,11 @@ def _to_ndarray(token_ids) -> np.ndarray:
 class SchedulerSession:
     """Scheduler session."""
 
-    def __init__(self, session_id: int, seq_manager: SequenceManager, scheduler: 'Scheduler') -> None:
+    def __init__(self, session_id: int, seq_meta: SequenceMeta, lifecycle: 'SequenceLifecycle') -> None:
         self.session_id = session_id
-        self.seq_meta = seq_manager.seq_meta
+        self.seq_meta = seq_meta
         self.sequences: SeqMap = dict()
-        self.seq_manager = seq_manager
-        self.scheduler = scheduler
+        self.lifecycle = lifecycle
 
     def add_sequence(self,
                      token_ids: Tensor,
@@ -299,12 +298,10 @@ def add_sequence(self,
                      resp_cache: bool = False,
                      preserve_cache: bool = False) -> 'SchedulerSequence':
         """Add a new message."""
-        from lmdeploy.pytorch.paging.seq_states.states import build_seq_state
-
         if sampling_param is None:
             sampling_param = SamplingParam()
 
-        seq_id = self.seq_manager._new_seq_id()
+        seq_id = self.lifecycle.new_sequence_id()
         seq = self.seq_meta.strategy.make_sequence(seq_id=seq_id,
                                                    session=self,
                                                    sampling_param=sampling_param,
@@ -318,16 +315,8 @@ def add_sequence(self,
             embeddings=input_embeddings,
             mode=UpdateTokenMode.INPUTS,
         )
-        self.sequences[seq.seq_id] = seq
-
-        # set status
-        # update seq manager
         status = MessageStatus.WAITING if migration_request is None else MessageStatus.MIGRATION_WAITING
-        seq.set_state(build_seq_state(self.scheduler, seq, status))
-        self.seq_manager.add_sequence(seq)
-        connector = self.scheduler.kv_connector
-        if connector is not None:
-            connector.on_new_request(seq)
+        self.lifecycle.add_sequence(seq, status)
 
         # metrics
         seq.record_event(EventType.QUEUED)
@@ -336,10 +325,7 @@ def add_sequence(self,
 
     def remove_sequence(self, seq: 'SchedulerSequence'):
         """Remove sequence."""
-        assert seq.seq_id in self.sequences
-        seq.state.free()
-        self.sequences.pop(seq.seq_id)
-        self.seq_manager.remove_sequence(seq)
+        self.lifecycle.remove_sequence(seq)
 
 
 def _div_up(x, n):
@@ -966,7 +952,7 @@ def _update_multimodals(self, multimodals: MultiModalInputs):
         if multimodals is None:
             return
         multimodals = HistoryMultiModals.update_multimodals(multimodals, self.num_valid_ids)
-        if self.session.scheduler.cache_config.enable_prefix_caching:
+        if self._seq_meta.enable_prefix_caching:
             self._update_prefix_cache_spans(multimodals)
         self.history_multimodals.add_inputs(multimodals)
 
diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index 38103a6e66..4ba62bd750 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -545,12 +545,8 @@ def finish_deferred_loads_after_worker_drain(self) -> None:
             self._finish_cancelled_or_failed(record)
 
     def _remove_sequence(self, seq: SchedulerSequence) -> None:
-        connector = self.connector
-        if connector is not None:
-            connector.request_finished(seq)
-
         session = seq.session
-        session.remove_sequence(seq)
+        session.lifecycle.finish_sequence(seq)
         if not session.sequences:
             self.sessions.pop(session.session_id, None)
 
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index ccc5faca28..af53de06f8 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -20,6 +20,7 @@
 from .kv_load_coordinator import KVLoadCoordinator
 from .kv_save_coordinator import KVSaveCoordinator
 from .prefill_scheduler import _PrefillScheduler, _PrefillTurnPolicy
+from .seq_states import SequenceLifecycle
 from .state_manager import build_state_manager
 
 if TYPE_CHECKING:
@@ -70,6 +71,10 @@ def __init__(
         self.cache_config = cache_config
         self.sessions: dict[int, SchedulerSession] = OrderedDict()
         self.kv_connector = kv_connector
+        seq_meta = seq_meta or SequenceMeta(self.cache_config.block_size)
+        seq_meta.enable_prefix_caching = self.cache_config.enable_prefix_caching
+        self.seq_meta = seq_meta
+        self.seq_manager = SequenceManager(seq_meta)
 
         self.state_manager = build_state_manager(self.cache_config)
         self.block_manager = build_block_manager(cache_config)
@@ -89,6 +94,15 @@ def __init__(
                                    block_size=self.cache_config.block_size,
                                    enabled=self.cache_config.enable_prefix_caching,
                                    checkpoint_state_manager=checkpoint_state_manager)
+        self.sequence_lifecycle = SequenceLifecycle(
+            seq_manager=self.seq_manager,
+            block_manager=self.block_manager,
+            state_manager=self.state_manager,
+            state_checkpoints=self.block_trie.state_checkpoints,
+            prefix_cache_enabled=self.block_trie.enabled,
+            is_ssm=self.is_ssm,
+            connector=kv_connector,
+        )
 
         self.eviction_helper = build_eviction_helper(self, self.scheduler_config.eviction_type)
         # Load admission receives only paging owners plus request-local queue
@@ -114,9 +128,6 @@ def __init__(
         # Keep save call sites uniform even when the producer role is disabled.
         self.kv_save_coordinator = KVSaveCoordinator(self)
 
-        seq_meta = seq_meta or SequenceMeta(self.cache_config.block_size)
-        self.seq_meta = seq_meta
-        self.seq_manager = SequenceManager(seq_meta)
         self.scheduler_tick = 0
 
     def tick(self):
@@ -131,6 +142,7 @@ def shutdown(self) -> None:
         """
         connector = self.kv_connector
         self.kv_connector = None
+        self.sequence_lifecycle.disable_connector()
         self.kv_load_coordinator.disable()
         self.kv_save_coordinator.clear()
         if connector is not None:
@@ -244,10 +256,18 @@ def add_session(self, session_id: int):
             session_id (int): New session id.
         """
         assert session_id not in self.sessions
-        session = SchedulerSession(session_id, seq_manager=self.seq_manager, scheduler=self)
+        session = SchedulerSession(session_id, seq_meta=self.seq_meta, lifecycle=self.sequence_lifecycle)
         self.sessions[session_id] = session
         return session
 
+    def get_session(self, session_id: int) -> SchedulerSession | None:
+        """Return one session owner, if it exists."""
+        return self.sessions.get(session_id)
+
+    def get_sessions(self) -> list[SchedulerSession]:
+        """Return a snapshot of current session owners."""
+        return list(self.sessions.values())
+
     def schedule_migration(self):
         """Admit waiting migration sequences to paging resources."""
         migration_ready: SeqList = []
@@ -375,9 +395,7 @@ def end_session(self, session_id: int):
                 continue
             # stop session so it won't get scheduled again
             seq.state.stop()
-            if connector is not None:
-                connector.request_finished(seq)
-            session.remove_sequence(seq)
+            self.sequence_lifecycle.finish_sequence(seq)
         if not session.sequences:
             self.sessions.pop(session_id)
 
diff --git a/lmdeploy/pytorch/paging/seq_states/__init__.py b/lmdeploy/pytorch/paging/seq_states/__init__.py
index bba2109f8e..8c27353978 100644
--- a/lmdeploy/pytorch/paging/seq_states/__init__.py
+++ b/lmdeploy/pytorch/paging/seq_states/__init__.py
@@ -1,2 +1,2 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from .states import StateBase, build_seq_state  # noqa: F401
+from .states import SequenceLifecycle, StateBase  # noqa: F401
diff --git a/lmdeploy/pytorch/paging/seq_states/states.py b/lmdeploy/pytorch/paging/seq_states/states.py
index 76a4ace78d..9ec965b007 100644
--- a/lmdeploy/pytorch/paging/seq_states/states.py
+++ b/lmdeploy/pytorch/paging/seq_states/states.py
@@ -1,28 +1,87 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from typing import TYPE_CHECKING
 
-from lmdeploy.pytorch.messages import MessageStatus, SchedulerSequence
+from lmdeploy.pytorch.messages import MessageStatus, SchedulerSequence, SequenceManager
 
 if TYPE_CHECKING:
-    from lmdeploy.pytorch.paging import Scheduler
-
-
-def _free_seq(seq: SchedulerSequence, scheduler: 'Scheduler'):
-    """Free the sequence."""
-    if scheduler.block_trie.enabled:
-        scheduler.block_trie.state_checkpoints.discard_save(seq)
-        scheduler.block_trie.state_checkpoints.unpin_restore(seq)
-        seq.prefix_cache.restore.clear()
-        seq.prefix_cache.trie_cursor = None
-        seq.prefix_cache.match_start_step = -1
-        seq.prefix_cache.recompute_overlap.clear_tracking()
-    seq.cached_tokens = 0
-    seq.kv_token_limit = None
-    if seq.num_blocks > 0:
-        scheduler.block_manager.free(seq)
-    if seq.logical_state >= 0:
-        scheduler.state_manager.free(seq)
-    seq.set_step(0)
+    from lmdeploy.pytorch.kv_connector.base import KVConnectorBase
+    from lmdeploy.pytorch.paging.block_manager import BaseBlockManager
+    from lmdeploy.pytorch.paging.block_trie.checkpoint_lifecycle import StateCheckpointLifecycle
+    from lmdeploy.pytorch.paging.state_manager import StateManager
+
+
+class SequenceLifecycle:
+    """Own sequence registration, transitions, and paging cleanup."""
+
+    def __init__(
+        self,
+        seq_manager: SequenceManager,
+        block_manager: 'BaseBlockManager',
+        state_manager: 'StateManager',
+        state_checkpoints: 'StateCheckpointLifecycle',
+        prefix_cache_enabled: bool,
+        is_ssm: bool,
+        connector: 'KVConnectorBase | None',
+    ) -> None:
+        self._seq_manager = seq_manager
+        self._block_manager = block_manager
+        self._state_manager = state_manager
+        self._state_checkpoints = state_checkpoints
+        self._prefix_cache_enabled = prefix_cache_enabled
+        self._is_ssm = is_ssm
+        self._connector = connector
+
+    def new_sequence_id(self) -> int:
+        return self._seq_manager.new_sequence_id()
+
+    def add_sequence(self, seq: SchedulerSequence, status: MessageStatus) -> None:
+        seq.session.sequences[seq.seq_id] = seq
+        seq.set_state(StateBase.build(self, seq, status))
+        self._seq_manager.add_sequence(seq)
+        if self._connector is not None:
+            self._connector.on_new_request(seq)
+
+    def remove_sequence(self, seq: SchedulerSequence) -> None:
+        """Release local ownership without a terminal connector event."""
+        assert seq.seq_id in seq.session.sequences
+        self.free_sequence(seq)
+        seq.session.sequences.pop(seq.seq_id)
+        self._seq_manager.remove_sequence(seq)
+
+    def finish_sequence(self, seq: SchedulerSequence) -> None:
+        """Notify connector completion, then release local ownership."""
+        if self._connector is not None:
+            self._connector.request_finished(seq)
+        self.remove_sequence(seq)
+
+    def transition(self, seq: SchedulerSequence, new_state: type['StateBase']) -> None:
+        self._seq_manager.update_sequence_status(seq, new_state.status)
+        seq.set_state(new_state(seq, self))
+
+    def assert_allocated(self, seq: SchedulerSequence) -> None:
+        num_required_blocks = self._block_manager.num_required_blocks(seq)
+        assert seq.num_blocks >= num_required_blocks
+        if self._is_ssm:
+            assert seq.logical_state >= 0
+
+    def free_sequence(self, seq: SchedulerSequence) -> None:
+        if self._prefix_cache_enabled:
+            self._state_checkpoints.discard_save(seq)
+            self._state_checkpoints.unpin_restore(seq)
+            seq.prefix_cache.restore.clear()
+            seq.prefix_cache.trie_cursor = None
+            seq.prefix_cache.match_start_step = -1
+            seq.prefix_cache.recompute_overlap.clear_tracking()
+        seq.cached_tokens = 0
+        seq.kv_token_limit = None
+        if seq.num_blocks > 0:
+            self._block_manager.free(seq)
+        if seq.logical_state >= 0:
+            self._state_manager.free(seq)
+        seq.set_step(0)
+
+    def disable_connector(self) -> None:
+        self._connector = None
 
 
 class StateBase:
@@ -35,20 +94,19 @@ def __init_subclass__(cls, **kargs) -> None:
             cls._registry[cls.status] = cls
 
     @classmethod
-    def build(cls, scheduler: 'Scheduler', seq: 'SchedulerSequence', status: MessageStatus) -> 'StateBase':
+    def build(cls, lifecycle: SequenceLifecycle, seq: 'SchedulerSequence', status: MessageStatus) -> 'StateBase':
         """Build sequence state."""
         if status not in cls._registry:
             raise NotImplementedError(f'Unsupported status {status} for building seq state.')
-        return cls._registry[status](seq, scheduler)
+        return cls._registry[status](seq, lifecycle)
 
-    def __init__(self, seq: SchedulerSequence, scheduler: 'Scheduler'):
+    def __init__(self, seq: SchedulerSequence, lifecycle: SequenceLifecycle):
         self.seq = seq
-        self.scheduler = scheduler
+        self.lifecycle = lifecycle
 
     def to_state(self, new_state):
         """Transition to a new state."""
-        self.scheduler.seq_manager.update_sequence_status(self.seq, new_state.status)
-        self.seq.set_state(new_state(self.seq, self.scheduler))
+        self.lifecycle.transition(self.seq, new_state)
 
     def evict(self):
         """Evict the state."""
@@ -72,7 +130,7 @@ def stop(self):
 
     def free(self):
         """Free the state."""
-        _free_seq(self.seq, self.scheduler)
+        self.lifecycle.free_sequence(self.seq)
 
     def begin_remote_load(self):
         raise NotImplementedError(f'begin_remote_load not implemented for state {self.status}')
@@ -87,10 +145,7 @@ class WaitingState(StateBase):
 
     def activate(self):
         """From WAITING to READY."""
-        num_req_blocks = self.scheduler.block_manager.num_required_blocks(self.seq)
-        assert self.seq.num_blocks >= num_req_blocks
-        if self.scheduler.is_ssm:
-            assert self.seq.logical_state >= 0
+        self.lifecycle.assert_allocated(self.seq)
         self.to_state(ReadyState)
 
     def evict(self):
@@ -201,8 +256,3 @@ def deactivate(self):
 
     def finish(self):
         self.to_state(MigrationDoneState)
-
-
-def build_seq_state(scheduler: 'Scheduler', seq: 'SchedulerSequence', status: MessageStatus) -> StateBase:
-    """Build sequence state."""
-    return StateBase.build(scheduler, seq, status)
diff --git a/tests/pytorch/engine/test_abort_stopped_seqs.py b/tests/pytorch/engine/test_abort_stopped_seqs.py
index 502a22e929..e67157f66c 100644
--- a/tests/pytorch/engine/test_abort_stopped_seqs.py
+++ b/tests/pytorch/engine/test_abort_stopped_seqs.py
@@ -50,6 +50,9 @@ class FakeScheduler:
     def __init__(self, seq):
         self.sessions = {1: SimpleNamespace(sequences={0: seq})}
 
+    def get_session(self, session_id):
+        return self.sessions.get(session_id)
+
     def stop_session(self, session_id):
         for seq in self.sessions[session_id].sequences.values():
             seq.state.stop()
diff --git a/tests/pytorch/engine/test_engine_sleep.py b/tests/pytorch/engine/test_engine_sleep.py
index 1df9d023aa..2b49093481 100644
--- a/tests/pytorch/engine/test_engine_sleep.py
+++ b/tests/pytorch/engine/test_engine_sleep.py
@@ -19,7 +19,8 @@ def __init__(self, resp):
 
 class _FakeSession:
 
-    def __init__(self, seq):
+    def __init__(self, session_id, seq):
+        self.session_id = session_id
         self.sequences = {0: seq}
 
 
@@ -33,6 +34,12 @@ def end_session(self, session_id):
         self.ended_sessions.append(session_id)
         self.sessions.pop(session_id)
 
+    def get_session(self, session_id):
+        return self.sessions.get(session_id)
+
+    def get_sessions(self):
+        return list(self.sessions.values())
+
     def finish_deferred_kv_transfers_after_worker_drain(self):
         pass
 
@@ -104,7 +111,7 @@ def _build_sleeping_test_engine(event_loop):
     engine.req_manager = RequestManager()
     resp = Response(type=ResponseType.INTERNAL_ENGINE_ERROR, sender_id=0, event=asyncio.Event())
     seq = _FakeSequence(resp)
-    session = _FakeSession(seq)
+    session = _FakeSession(1, seq)
     engine.scheduler = _FakeScheduler(session)
     engine._sleeping_tags = set()
     engine.events = []
diff --git a/tests/pytorch/engine/test_kv_connector_wiring.py b/tests/pytorch/engine/test_kv_connector_wiring.py
index 8c310ec615..93cc61d474 100644
--- a/tests/pytorch/engine/test_kv_connector_wiring.py
+++ b/tests/pytorch/engine/test_kv_connector_wiring.py
@@ -227,6 +227,7 @@ def disable_loads():
     load_coordinator.disable.side_effect = disable_loads
     scheduler = Scheduler.__new__(Scheduler)
     scheduler.kv_connector = connector
+    scheduler.sequence_lifecycle = Mock()
     scheduler.kv_load_coordinator = load_coordinator
     scheduler.kv_save_coordinator = Mock()
 
@@ -235,6 +236,7 @@ def disable_loads():
 
     connector.shutdown.assert_called_once_with()
     assert scheduler.kv_load_coordinator.disable.call_count == 2
+    assert scheduler.sequence_lifecycle.disable_connector.call_count == 2
     assert scheduler.kv_save_coordinator.clear.call_count == 2
     assert scheduler.kv_connector is None
     assert not scheduler._external_lookup_enabled
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index e54f7e0b0d..7e5c2a5a42 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -61,6 +61,8 @@ def test_schedule_base(self, scheduler, block_size, num_gpu_blocks):
         session = scheduler.add_session(session_id)
         assert session_id in scheduler.sessions
         assert scheduler.sessions[session_id] == session
+        assert scheduler.get_session(session_id) is session
+        assert scheduler.get_sessions() == [session]
 
         num_blocks = 2
         token_ids = torch.tensor([0] * block_size * num_blocks)

From a74b2570a44fdcb32c5baafe7a3d673c313b403a Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 18:40:31 +0800
Subject: [PATCH 17/22] refactor: clarify scheduler paging ownership

---
 lmdeploy/pytorch/kv_connector/__init__.py     |  2 +
 lmdeploy/pytorch/kv_connector/base.py         | 18 +++++--
 .../kv_connector/mooncake/store/connector.py  |  6 +--
 .../kv_connector/mooncake/store/scheduler.py  | 16 +++---
 .../paging/eviction_helper/__init__.py        | 28 +++++++++-
 .../eviction_helper/base_eviction_helper.py   | 28 +++++++---
 .../recompute_eviction_helper.py              | 43 ++++++++++++----
 .../pytorch/paging/kv_load_coordinator.py     |  9 ++--
 .../pytorch/paging/kv_save_coordinator.py     | 12 ++---
 lmdeploy/pytorch/paging/prefill_scheduler.py  |  1 +
 lmdeploy/pytorch/paging/scheduler.py          | 17 +++++--
 .../test_mooncake_store_connector.py          | 18 +++----
 .../test_mooncake_store_scheduler.py          | 51 ++++++++++---------
 .../paging/test_scheduler_kv_transfer.py      | 11 ++--
 14 files changed, 177 insertions(+), 83 deletions(-)

diff --git a/lmdeploy/pytorch/kv_connector/__init__.py b/lmdeploy/pytorch/kv_connector/__init__.py
index c7ee293b1d..e4b60e7f8b 100644
--- a/lmdeploy/pytorch/kv_connector/__init__.py
+++ b/lmdeploy/pytorch/kv_connector/__init__.py
@@ -6,6 +6,7 @@
     KVConnectorOutputAggregator,
     KVConnectorResult,
     KVConnectorRole,
+    KVConnectorStepInput,
     KVLoadResult,
     KVOperationId,
     KVSaveBlockLease,
@@ -19,6 +20,7 @@
     'KVConnectorOutputAggregator',
     'KVConnectorResult',
     'KVConnectorRole',
+    'KVConnectorStepInput',
     'KVLoadResult',
     'KVOperationId',
     'KVSaveBlockLease',
diff --git a/lmdeploy/pytorch/kv_connector/base.py b/lmdeploy/pytorch/kv_connector/base.py
index d30506f4c8..cfe3520fb8 100644
--- a/lmdeploy/pytorch/kv_connector/base.py
+++ b/lmdeploy/pytorch/kv_connector/base.py
@@ -18,7 +18,6 @@
 
 if TYPE_CHECKING:
     from lmdeploy.pytorch.messages import SchedulerSequence
-    from lmdeploy.pytorch.paging.scheduler import SchedulerOutput
 
 RequestId = int
 KVOperationId = int
@@ -32,6 +31,19 @@ class KVConnectorRole(enum.Enum):
     WORKER = enum.auto()
 
 
+@dataclass
+class KVConnectorStepInput:
+    """Paging snapshot offered to a connector for one engine step."""
+
+    running: list['SchedulerSequence'] = field(default_factory=list)
+    swap_in_map: dict[int, int] = field(default_factory=dict)
+    swap_out_map: dict[int, int] = field(default_factory=dict)
+    copy_map: dict[int, int] = field(default_factory=dict)
+    connector_token_lens: tuple[int, ...] = ()
+    connector_block_ids: tuple[tuple[int, ...], ...] = ()
+    connector_logical_block_ids: tuple[tuple[int, ...], ...] = ()
+
+
 class KVConnectorMetadata(ABC):
     """Scheduler-to-worker metadata for one engine step.
 
@@ -282,10 +294,10 @@ def update_state_after_alloc(
         raise NotImplementedError
 
     @abstractmethod
-    def build_connector_meta(self, scheduler_output: 'SchedulerOutput') -> KVConnectorMetadata | None:
+    def build_connector_meta(self, step_input: KVConnectorStepInput) -> KVConnectorMetadata | None:
         """Build serializable worker metadata for the current scheduler step.
 
-        Implementations must not mutate ``scheduler_output``. They may consume
+        Implementations must not mutate ``step_input``. They may consume
         and reset connector-owned per-step bookkeeping while building the
         returned metadata.
         """
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py b/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
index e2a9299e96..786c633ae0 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
@@ -12,6 +12,7 @@
     KVConnectorOutput,
     KVConnectorResult,
     KVConnectorRole,
+    KVConnectorStepInput,
     RequestId,
 )
 
@@ -22,7 +23,6 @@
 if TYPE_CHECKING:
     from lmdeploy.pytorch.config import CacheConfig
     from lmdeploy.pytorch.messages import SchedulerSequence
-    from lmdeploy.pytorch.paging.scheduler import SchedulerOutput
 
 
 class MooncakeStoreConnector(KVConnectorBase):
@@ -94,9 +94,9 @@ def update_state_after_alloc(
 
     def build_connector_meta(
         self,
-        scheduler_output: SchedulerOutput,
+        step_input: KVConnectorStepInput,
     ) -> MooncakeStoreConnectorMetadata | None:
-        return self._require_scheduler().build_connector_meta(scheduler_output)
+        return self._require_scheduler().build_connector_meta(step_input)
 
     def on_new_request(self, request: SchedulerSequence) -> None:
         return self._require_scheduler().on_new_request(request)
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py b/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py
index 84872cf4e6..19cfb542d0 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/scheduler.py
@@ -10,6 +10,7 @@
 from lmdeploy.pytorch.kv_connector.base import (
     KVConnectorOutput,
     KVConnectorResult,
+    KVConnectorStepInput,
     KVLoadResult,
     RequestId,
 )
@@ -25,7 +26,6 @@
 if TYPE_CHECKING:
     from lmdeploy.pytorch.config import CacheConfig
     from lmdeploy.pytorch.messages import SchedulerSequence
-    from lmdeploy.pytorch.paging.scheduler import SchedulerOutput
 
 
 @dataclass
@@ -229,16 +229,16 @@ def update_state_after_alloc(
 
     def _build_save_requests(
         self,
-        scheduler_output: SchedulerOutput,
+        step_input: KVConnectorStepInput,
     ) -> tuple[MooncakeStoreSaveRequest, ...]:
         """Build newly completed full-block suffixes for prefill work."""
-        token_lens = scheduler_output.connector_token_lens
+        token_lens = step_input.connector_token_lens
         if not self._kv_transfer_config.is_kv_producer or not token_lens:
             return ()
 
-        running = scheduler_output.running
-        block_ids = scheduler_output.connector_block_ids
-        logical_block_ids = scheduler_output.connector_logical_block_ids
+        running = step_input.running
+        block_ids = step_input.connector_block_ids
+        logical_block_ids = step_input.connector_logical_block_ids
         if not (len(running) == len(token_lens) == len(block_ids) == len(logical_block_ids)):
             raise ValueError('connector save fields must contain one value per running request')
 
@@ -283,10 +283,10 @@ def _build_save_requests(
 
     def build_connector_meta(
         self,
-        scheduler_output: SchedulerOutput,
+        step_input: KVConnectorStepInput,
     ) -> MooncakeStoreConnectorMetadata | None:
         """Dispatch new work and keep emitting polling steps while I/O runs."""
-        save_requests = self._build_save_requests(scheduler_output)
+        save_requests = self._build_save_requests(step_input)
         if (not save_requests and not self._pending_loads
                 and not self._inflight_loads and not self._inflight_save_ids):
             return None
diff --git a/lmdeploy/pytorch/paging/eviction_helper/__init__.py b/lmdeploy/pytorch/paging/eviction_helper/__init__.py
index 6b5c44ff97..9a52b580f2 100644
--- a/lmdeploy/pytorch/paging/eviction_helper/__init__.py
+++ b/lmdeploy/pytorch/paging/eviction_helper/__init__.py
@@ -1,10 +1,28 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from lmdeploy.utils import get_logger
 
+if TYPE_CHECKING:
+    from ..block_manager import BaseBlockManager
+    from ..block_trie import BlockTrie
+    from ..kv_load_coordinator import KVLoadCoordinator
+    from ..state_manager import StateManager
+
 logger = get_logger('lmdeploy')
 
 
-def build_eviction_helper(scheduler, eviction_type: str):
+def build_eviction_helper(
+    eviction_type: str,
+    *,
+    block_manager: BaseBlockManager,
+    block_trie: BlockTrie,
+    state_manager: StateManager,
+    load_coordinator: KVLoadCoordinator,
+    is_ssm: bool,
+):
     """Build eviction helper."""
     if eviction_type == 'copy':
         logger.warning('`copy` eviction has been deprecated, '
@@ -12,6 +30,12 @@ def build_eviction_helper(scheduler, eviction_type: str):
         eviction_type = 'recompute'
     if eviction_type == 'recompute':
         from .recompute_eviction_helper import RecomputeEvictionHelper
-        return RecomputeEvictionHelper(scheduler)
+        return RecomputeEvictionHelper(
+            block_manager=block_manager,
+            block_trie=block_trie,
+            state_manager=state_manager,
+            load_coordinator=load_coordinator,
+            is_ssm=is_ssm,
+        )
     else:
         raise TypeError(f'Unknown eviction type: {eviction_type}')
diff --git a/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py b/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py
index f075748f70..ae519aa5b8 100644
--- a/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py
+++ b/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py
@@ -1,7 +1,15 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from ...messages import SchedulerSequence
-from ..scheduler import Scheduler
+
+if TYPE_CHECKING:
+    from ..block_manager import BaseBlockManager
+    from ..block_trie import BlockTrie
+    from ..kv_load_coordinator import KVLoadCoordinator
+    from ..state_manager import StateManager
 
 SeqList = list[SchedulerSequence]
 
@@ -9,12 +17,18 @@
 class BaseEvictionHelper:
     """Base eviction helper."""
 
-    def __init__(self, scheduler: Scheduler):
-        self.scheduler = scheduler
-        self.block_manager = scheduler.block_manager
-        self.block_trie = scheduler.block_trie
-        self.state_manager = scheduler.state_manager
-        self.cache_config = scheduler.cache_config
+    def __init__(
+        self,
+        *,
+        block_manager: BaseBlockManager,
+        block_trie: BlockTrie,
+        state_manager: StateManager,
+        load_coordinator: KVLoadCoordinator,
+    ) -> None:
+        self.block_manager = block_manager
+        self.block_trie = block_trie
+        self.state_manager = state_manager
+        self.load_coordinator = load_coordinator
 
     def need_swap_in(self, seq: SchedulerSequence):
         """Sequence need swap in."""
diff --git a/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py b/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
index bfa39c8919..d0233d6476 100644
--- a/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
+++ b/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
@@ -1,20 +1,41 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from ...messages import SchedulerSequence
-from ..scheduler import Scheduler
 from .base_eviction_helper import BaseEvictionHelper
 
+if TYPE_CHECKING:
+    from ..block_manager import BaseBlockManager
+    from ..block_trie import BlockTrie
+    from ..kv_load_coordinator import KVLoadCoordinator
+    from ..state_manager import StateManager
+
 
 class RecomputeEvictionHelper(BaseEvictionHelper):
     """Recompute eviction."""
 
-    def __init__(self, scheduler: Scheduler):
-        super().__init__(scheduler)
-
-        if len(self.cache_config.states_shapes) == 0:
-            self.evict_for_seq = self._evict_for_seq_default
-        else:
+    def __init__(
+        self,
+        *,
+        block_manager: BaseBlockManager,
+        block_trie: BlockTrie,
+        state_manager: StateManager,
+        load_coordinator: KVLoadCoordinator,
+        is_ssm: bool,
+    ) -> None:
+        super().__init__(
+            block_manager=block_manager,
+            block_trie=block_trie,
+            state_manager=state_manager,
+            load_coordinator=load_coordinator,
+        )
+
+        if is_ssm:
             self.evict_for_seq = self._evict_for_ssm
+        else:
+            self.evict_for_seq = self._evict_for_seq_default
 
     def _evict_for_seq_default(self, seq: SchedulerSequence, evictable_seqs: list[SchedulerSequence],
                                prealloc_size: int):
@@ -33,7 +54,7 @@ def _evict_for_seq_default(self, seq: SchedulerSequence, evictable_seqs: list[Sc
             # A completed remote load has published fresh KV into these blocks.
             # Preserve it until prefill consumes the result instead of paying
             # for the transfer again on a later scheduling turn.
-            if self.scheduler.kv_load_coordinator.is_remote_ready(evict_seq):
+            if self.load_coordinator.is_remote_ready(evict_seq):
                 continue
 
             # skip sequence with no blocks
@@ -44,7 +65,7 @@ def _evict_for_seq_default(self, seq: SchedulerSequence, evictable_seqs: list[Sc
                 evict_seq.prefix_cache.suppress_match_stats = True
             # Eviction also ends the tracked prefill; otherwise its soft block
             # reservation would outlive the local KV blocks freed below.
-            self.scheduler.kv_load_coordinator.release(evict_seq)
+            self.load_coordinator.release(evict_seq)
             evict_seq.state.free()
             num_req = (num_required_blocks - block_manager.get_num_free_gpu_blocks())
             if num_req <= 0:
@@ -92,7 +113,7 @@ def _evict_for_ssm(self, seq: SchedulerSequence, evictable_seqs: list[SchedulerS
 
             # READY remote KV is already transferred and awaiting consumption;
             # do not discard that result merely to admit another prefill.
-            if self.scheduler.kv_load_coordinator.is_remote_ready(evict_seq):
+            if self.load_coordinator.is_remote_ready(evict_seq):
                 continue
 
             # skip sequence with no blocks
@@ -104,7 +125,7 @@ def _evict_for_ssm(self, seq: SchedulerSequence, evictable_seqs: list[SchedulerS
                 evict_seq.prefix_cache.suppress_match_stats = True
             # Keep coordinator ownership and its soft admission budget in sync
             # with the KV blocks and SSM runtime state released by free().
-            self.scheduler.kv_load_coordinator.release(evict_seq)
+            self.load_coordinator.release(evict_seq)
             evict_seq.state.free()
             has_free_state = has_runtime_state or state_manager.get_num_free_runtime() > 0
             if not has_free_state:
diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index 4ba62bd750..ee45af8006 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -125,14 +125,12 @@ def __init__(
         connector: KVConnectorBase | None,
         block_manager: BaseBlockManager,
         block_trie: BlockTrie,
-        eviction_helper: BaseEvictionHelper,
         sessions: dict[int, SchedulerSession],
     ) -> None:
         self.lookup_enabled = lookup_enabled
         self.connector = connector
         self.block_manager = block_manager
         self.block_trie = block_trie
-        self.eviction_helper = eviction_helper
         self.sessions = sessions
         # Active load lifecycle. Records survive the LOADING -> READY ->
         # PREFILLING transitions so stop/end and preemption can find the owner.
@@ -218,6 +216,7 @@ def try_load(
         *,
         prealloc_size: int,
         evictable_seqs: Iterable[SchedulerSequence],
+        eviction_helper: BaseEvictionHelper,
     ) -> KVLoadAdmission:
         """Poll and admit one external prefix without choosing queue policy.
 
@@ -241,6 +240,7 @@ def try_load(
             num_external_tokens=int(num_external_tokens),
             prealloc_size=prealloc_size,
             evictable_seqs=evictable_seqs,
+            eviction_helper=eviction_helper,
         )
 
     def _admit_load(
@@ -250,6 +250,7 @@ def _admit_load(
         num_external_tokens: int,
         prealloc_size: int,
         evictable_seqs: Iterable[SchedulerSequence],
+        eviction_helper: BaseEvictionHelper,
     ) -> KVLoadAdmission:
         """Admit the complete prefill, then allocate the remote interval."""
         plan = self._plan_load(seq, num_external_tokens, prealloc_size)
@@ -261,6 +262,7 @@ def _admit_load(
             plan,
             prealloc_size,
             evictable_seqs,
+            eviction_helper,
         )
         if failure is not None:
             return failure
@@ -299,12 +301,13 @@ def _admit_load_capacity(
         plan: _LoadPlan,
         prealloc_size: int,
         evictable_seqs: Iterable[SchedulerSequence],
+        eviction_helper: BaseEvictionHelper,
     ) -> KVLoadAdmission | None:
         """Admit the full prefill against physical and soft capacity."""
         # Only the remote hit is allocated now, but admission guarantees the
         # complete prefill can finish beside every existing soft reservation.
         seq.kv_token_limit = None
-        full_prefill_fits = self.eviction_helper.evict_for_seq(
+        full_prefill_fits = eviction_helper.evict_for_seq(
             seq,
             list(evictable_seqs),
             prealloc_size,
diff --git a/lmdeploy/pytorch/paging/kv_save_coordinator.py b/lmdeploy/pytorch/paging/kv_save_coordinator.py
index a1e7130060..1009534c4e 100644
--- a/lmdeploy/pytorch/paging/kv_save_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_save_coordinator.py
@@ -27,7 +27,7 @@
 from lmdeploy.pytorch.kv_connector import KVConnectorMetadata, KVOperationId
 
 if TYPE_CHECKING:
-    from .scheduler import Scheduler
+    from .block_manager import BaseBlockManager
 
 
 class KVSaveCoordinator:
@@ -40,8 +40,8 @@ class KVSaveCoordinator:
     saves over its lifetime.
     """
 
-    def __init__(self, scheduler: Scheduler) -> None:
-        self.scheduler = scheduler
+    def __init__(self, block_manager: BaseBlockManager) -> None:
+        self.block_manager = block_manager
         self._leases: dict[KVOperationId, np.ndarray] = {}
 
     def acquire(self, metadata: KVConnectorMetadata) -> None:
@@ -51,7 +51,7 @@ def acquire(self, metadata: KVConnectorMetadata) -> None:
         a worker reads the physical cache slot. Logical IDs let the block manager keep ownership stable even if ordinary
         sequence ownership disappears.
         """
-        block_manager = self.scheduler.block_manager
+        block_manager = self.block_manager
         for lease in metadata.get_save_block_leases():
             if lease.operation_id in self._leases:
                 raise RuntimeError(f'save operation {lease.operation_id} already owns a block lease')
@@ -62,7 +62,7 @@ def acquire(self, metadata: KVConnectorMetadata) -> None:
     def update(self, completed_save_ids: frozenset[KVOperationId]) -> None:
         """Release operations that reached a terminal state on every TP
         rank."""
-        block_manager = self.scheduler.block_manager
+        block_manager = self.block_manager
         for operation_id in completed_save_ids:
             logical_block_ids = self._leases.pop(operation_id, None)
             if logical_block_ids is not None:
@@ -77,7 +77,7 @@ def clear(self) -> None:
         Engine sleep/shutdown can intentionally discard prefetched completion outputs. Once workers are drained they can
         no longer read these blocks, so explicit per-operation completion is no longer required.
         """
-        block_manager = self.scheduler.block_manager
+        block_manager = self.block_manager
         for logical_block_ids in self._leases.values():
             block_manager.release_logical_blocks(logical_block_ids)
         self._leases.clear()
diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index ac425a8d8e..aeb088d57e 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -547,6 +547,7 @@ def _try_external_load(self):
             self.seq,
             prealloc_size=self.prealloc_size,
             evictable_seqs=self._evictable_sequences(),
+            eviction_helper=prefill.eviction_helper,
         )
         if admission is KVLoadAdmission.NO_LOAD:
             return None
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index af53de06f8..c9db079af9 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -13,6 +13,7 @@
 from lmdeploy.messages import ScheduleMetrics
 
 from ..config import CacheConfig, SchedulerConfig
+from ..kv_connector import KVConnectorStepInput
 from ..messages import MessageStatus, SchedulerSequence, SchedulerSession, SequenceManager, SequenceMeta
 from .block_manager import build_block_manager
 from .block_trie import BlockTrie
@@ -104,7 +105,6 @@ def __init__(
             connector=kv_connector,
         )
 
-        self.eviction_helper = build_eviction_helper(self, self.scheduler_config.eviction_type)
         # Load admission receives only paging owners plus request-local queue
         # candidates from its caller; it does not reach back through Scheduler.
         self.kv_load_coordinator = KVLoadCoordinator(
@@ -112,9 +112,16 @@ def __init__(
             connector=kv_connector,
             block_manager=self.block_manager,
             block_trie=self.block_trie,
-            eviction_helper=self.eviction_helper,
             sessions=self.sessions,
         )
+        self.eviction_helper = build_eviction_helper(
+            self.scheduler_config.eviction_type,
+            block_manager=self.block_manager,
+            block_trie=self.block_trie,
+            state_manager=self.state_manager,
+            load_coordinator=self.kv_load_coordinator,
+            is_ssm=self.is_ssm,
+        )
         self._prefill_scheduler = _PrefillScheduler(
             scheduler_config=self.scheduler_config,
             cache_config=self.cache_config,
@@ -126,7 +133,7 @@ def __init__(
             load_coordinator=self.kv_load_coordinator,
         )
         # Keep save call sites uniform even when the producer role is disabled.
-        self.kv_save_coordinator = KVSaveCoordinator(self)
+        self.kv_save_coordinator = KVSaveCoordinator(self.block_manager)
 
         self.scheduler_tick = 0
 
@@ -444,7 +451,7 @@ def build_connector_meta(
         else:
             logical_block_ids = ()
             block_ids = ()
-        scheduler_output = SchedulerOutput(
+        step_input = KVConnectorStepInput(
             running=running,
             swap_in_map=swap_in_map or {},
             swap_out_map=swap_out_map or {},
@@ -453,7 +460,7 @@ def build_connector_meta(
             connector_block_ids=block_ids,
             connector_logical_block_ids=logical_block_ids,
         )
-        metadata = connector.build_connector_meta(scheduler_output)
+        metadata = connector.build_connector_meta(step_input)
         if metadata is not None:
             # Acquire before the caller queues metadata. Sequence cleanup may
             # otherwise release the last block reference before save starts.
diff --git a/tests/pytorch/kv_connector/test_mooncake_store_connector.py b/tests/pytorch/kv_connector/test_mooncake_store_connector.py
index d05595c67a..3850d433ce 100644
--- a/tests/pytorch/kv_connector/test_mooncake_store_connector.py
+++ b/tests/pytorch/kv_connector/test_mooncake_store_connector.py
@@ -1,7 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import json
 import pickle
-from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -12,6 +11,7 @@
     KVConnectorOutput,
     KVConnectorResult,
     KVConnectorRole,
+    KVConnectorStepInput,
 )
 from lmdeploy.pytorch.kv_connector.mooncake.store import worker as worker_module
 from lmdeploy.pytorch.kv_connector.mooncake.store.connector import MooncakeStoreConnector
@@ -121,8 +121,8 @@ def test_scheduler_without_transfer_work_is_fail_closed(cache_config):
 
     assert connector.get_num_new_matched_tokens(request, 0) == (0, False)
     assert connector.update_state_after_alloc(request, [1, 2], 0) is None
-    scheduler_output = SimpleNamespace(connector_token_lens=())
-    assert connector.build_connector_meta(scheduler_output) is None
+    step_input = KVConnectorStepInput()
+    assert connector.build_connector_meta(step_input) is None
     assert connector.on_new_request(request) is None
     assert connector.update_connector_output(KVConnectorOutput()) == KVConnectorResult()
     assert connector.request_finished(request) is None
@@ -157,7 +157,7 @@ def test_scheduler_methods_delegate_arguments_and_results(cache_config):
     assert scheduler is not None
 
     request = object()
-    scheduler_output = object()
+    step_input = object()
     metadata = MooncakeStoreConnectorMetadata()
     scheduler.get_num_new_matched_tokens = MagicMock(return_value=(17, True))
     scheduler.update_state_after_alloc = MagicMock(return_value=None)
@@ -177,8 +177,8 @@ def test_scheduler_methods_delegate_arguments_and_results(cache_config):
     assert connector.update_state_after_alloc(request, [4, 5], 14) is None
     scheduler.update_state_after_alloc.assert_called_once_with(request, [4, 5], 14)
 
-    assert connector.build_connector_meta(scheduler_output) is metadata
-    scheduler.build_connector_meta.assert_called_once_with(scheduler_output)
+    assert connector.build_connector_meta(step_input) is metadata
+    scheduler.build_connector_meta.assert_called_once_with(step_input)
 
     assert connector.on_new_request(request) is None
     scheduler.on_new_request.assert_called_once_with(request)
@@ -234,9 +234,9 @@ def test_worker_methods_delegate_arguments_and_results(cache_config):
 def test_empty_scheduler_has_no_metadata(cache_config):
     connector = MooncakeStoreConnector(KVConnectorRole.SCHEDULER, cache_config)
 
-    scheduler_output = SimpleNamespace(connector_token_lens=())
-    first = connector.build_connector_meta(scheduler_output)
-    second = connector.build_connector_meta(scheduler_output)
+    step_input = KVConnectorStepInput()
+    first = connector.build_connector_meta(step_input)
+    second = connector.build_connector_meta(step_input)
 
     assert first is None
     assert second is None
diff --git a/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py b/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py
index 1dd32f6cb5..4b6eae3422 100644
--- a/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py
+++ b/tests/pytorch/kv_connector/test_mooncake_store_scheduler.py
@@ -7,7 +7,12 @@
 
 from lmdeploy.messages import KVTransferConfig
 from lmdeploy.pytorch.config import CacheConfig
-from lmdeploy.pytorch.kv_connector import KVConnectorOutput, KVConnectorResult, KVLoadResult
+from lmdeploy.pytorch.kv_connector import (
+    KVConnectorOutput,
+    KVConnectorResult,
+    KVConnectorStepInput,
+    KVLoadResult,
+)
 from lmdeploy.pytorch.kv_connector.mooncake.store import scheduler as scheduler_module
 from lmdeploy.pytorch.kv_connector.mooncake.store.data import build_prefix_block_hashes
 from lmdeploy.pytorch.kv_connector.mooncake.store.scheduler import MooncakeStoreScheduler
@@ -45,13 +50,13 @@ def _request(
     )
 
 
-def _scheduler_output(
+def _connector_step(
     running=(),
     token_lens=(),
     block_ids=(),
     logical_block_ids=(),
 ):
-    return SimpleNamespace(
+    return KVConnectorStepInput(
         running=list(running),
         connector_token_lens=tuple(token_lens),
         connector_block_ids=tuple(block_ids),
@@ -169,13 +174,13 @@ def test_scheduler_load_failure_falls_back_until_next_request():
     assert scheduler.get_num_new_matched_tokens(request, 4) == (8, True)
     scheduler.update_state_after_alloc(request, (31, 32), 8)
 
-    metadata = scheduler.build_connector_meta(_scheduler_output())
+    metadata = scheduler.build_connector_meta(_connector_step())
     assert metadata is not None
     assert len(metadata.load_requests) == 1
     load_request = metadata.load_requests[0]
     assert load_request.request_id == request.seq_id
     assert load_request.block_ids == (31, 32)
-    assert scheduler.build_connector_meta(_scheduler_output()).load_requests == ()
+    assert scheduler.build_connector_meta(_connector_step()).load_requests == ()
 
     assert scheduler.update_connector_output(
         KVConnectorOutput(invalid_block_ids={32})) == KVConnectorResult()
@@ -187,7 +192,7 @@ def test_scheduler_load_failure_falls_back_until_next_request():
     )
 
     retry_save = scheduler.build_connector_meta(
-        _scheduler_output(
+        _connector_step(
             running=(request, ),
             token_lens=(12, ),
             block_ids=((30, 31, 32), ),
@@ -214,7 +219,7 @@ def test_scheduler_builds_incremental_save_operations_and_poll_metadata():
     request = _request(range(17), adapter_name='adapter-a')
 
     first = scheduler.build_connector_meta(
-        _scheduler_output(
+        _connector_step(
             running=(request, ),
             token_lens=(10, ),
             block_ids=((31, 32, 33, 34, 35), ),
@@ -234,12 +239,12 @@ def test_scheduler_builds_incremental_save_operations_and_poll_metadata():
 
     # The connector keeps the engine issuing no-forward polling steps while a
     # previous save is still running.
-    poll = scheduler.build_connector_meta(_scheduler_output())
+    poll = scheduler.build_connector_meta(_connector_step())
     assert poll is not None
     assert poll.save_requests == ()
 
     second = scheduler.build_connector_meta(
-        _scheduler_output(
+        _connector_step(
             running=(request, ),
             token_lens=(14, ),
             block_ids=((31, 32, 33, 34, 35), ),
@@ -254,19 +259,19 @@ def test_scheduler_builds_incremental_save_operations_and_poll_metadata():
     result = scheduler.update_connector_output(
         KVConnectorOutput(completed_save_ids={0, 999}))
     assert result.completed_save_ids == frozenset({0})
-    assert scheduler.build_connector_meta(_scheduler_output()) is not None
+    assert scheduler.build_connector_meta(_connector_step()) is not None
 
     result = scheduler.update_connector_output(
         KVConnectorOutput(completed_save_ids={1}))
     assert result.completed_save_ids == frozenset({1})
-    assert scheduler.build_connector_meta(_scheduler_output()) is None
+    assert scheduler.build_connector_meta(_connector_step()) is None
     scheduler.shutdown()
 
 
 def test_new_request_restarts_save_planning_from_first_block():
     scheduler = MooncakeStoreScheduler(_cache_config('kv_producer'))
     request = _request(range(17))
-    output = _scheduler_output(
+    output = _connector_step(
         running=(request, ),
         token_lens=(10, ),
         block_ids=((31, 32, 33, 34), ),
@@ -296,7 +301,7 @@ def test_finished_request_keeps_immutable_save_operation_until_completion():
     scheduler = MooncakeStoreScheduler(_cache_config('kv_producer'))
     request = _request(range(9))
     metadata = scheduler.build_connector_meta(
-        _scheduler_output(
+        _connector_step(
             running=(request, ),
             token_lens=(8, ),
             block_ids=((1, 2, 3), ),
@@ -305,18 +310,18 @@ def test_finished_request_keeps_immutable_save_operation_until_completion():
     save_id = metadata.save_requests[0].save_id
 
     scheduler.request_finished(request)
-    assert scheduler.build_connector_meta(_scheduler_output()) is not None
+    assert scheduler.build_connector_meta(_connector_step()) is not None
     result = scheduler.update_connector_output(
         KVConnectorOutput(completed_save_ids={save_id}))
     assert result.completed_save_ids == frozenset({save_id})
-    assert scheduler.build_connector_meta(_scheduler_output()) is None
+    assert scheduler.build_connector_meta(_connector_step()) is None
     scheduler.shutdown()
 
 
 def test_worker_drain_discards_save_ids_whose_outputs_were_dropped():
     scheduler = MooncakeStoreScheduler(_cache_config('kv_producer'))
     request = _request(range(9))
-    output = _scheduler_output(
+    output = _connector_step(
         running=(request, ),
         token_lens=(8, ),
         block_ids=((1, 2), ),
@@ -324,11 +329,11 @@ def test_worker_drain_discards_save_ids_whose_outputs_were_dropped():
     )
     metadata = scheduler.build_connector_meta(output)
     assert metadata.save_requests
-    assert scheduler.build_connector_meta(_scheduler_output()) is not None
+    assert scheduler.build_connector_meta(_connector_step()) is not None
 
     scheduler.finish_transfers_after_worker_drain()
 
-    assert scheduler.build_connector_meta(_scheduler_output()) is None
+    assert scheduler.build_connector_meta(_connector_step()) is None
     assert scheduler.build_connector_meta(output) is None
 
     next_request = _request(range(9), seq_id=18)
@@ -345,13 +350,13 @@ def test_successful_remote_load_is_not_saved_back_and_save_filters_non_text():
     scheduler.client.lookup = Mock(return_value=12)
     assert scheduler.get_num_new_matched_tokens(request, 0) == (12, True)
     scheduler.update_state_after_alloc(request, (21, 22, 23), 12)
-    scheduler.build_connector_meta(_scheduler_output())
+    scheduler.build_connector_meta(_connector_step())
     result = scheduler.update_connector_output(
         KVConnectorOutput(finished_receiving={request.seq_id}))
     assert result.load_results == (KVLoadResult(request.seq_id, True), )
 
     no_resave = scheduler.build_connector_meta(
-        _scheduler_output(
+        _connector_step(
             running=(request, ),
             token_lens=(13, ),
             block_ids=((21, 22, 23, 24, 25), ),
@@ -361,7 +366,7 @@ def test_successful_remote_load_is_not_saved_back_and_save_filters_non_text():
 
     multimodal = _request(range(17), multimodal=True)
     assert scheduler.build_connector_meta(
-        _scheduler_output(
+        _connector_step(
             running=(multimodal, ),
             token_lens=(16, ),
             block_ids=((1, 2, 3, 4), ),
@@ -377,13 +382,13 @@ def test_shorter_remote_prefix_rewinds_future_save_boundary():
 
     assert scheduler.get_num_new_matched_tokens(request, 0) == (12, True)
     scheduler.update_state_after_alloc(request, (21, 22, 23), 12)
-    scheduler.build_connector_meta(_scheduler_output())
+    scheduler.build_connector_meta(_connector_step())
     scheduler.update_connector_output(
         KVConnectorOutput(finished_receiving={request.seq_id}))
 
     assert scheduler.get_num_new_matched_tokens(request, 4) == (0, False)
     metadata = scheduler.build_connector_meta(
-        _scheduler_output(
+        _connector_step(
             running=(request, ),
             token_lens=(13, ),
             block_ids=((21, 22, 23, 24), ),
diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py
index a7156c2fd5..94a8d9eb91 100644
--- a/tests/pytorch/paging/test_scheduler_kv_transfer.py
+++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py
@@ -10,6 +10,7 @@
     KVConnectorMetadata,
     KVConnectorOutput,
     KVConnectorResult,
+    KVConnectorStepInput,
     KVLoadResult,
     KVSaveBlockLease,
 )
@@ -52,7 +53,7 @@ def cancel_lookup(self, request_id):
     def update_state_after_alloc(self, request, block_ids, num_external_tokens):
         self.allocations.append((request.seq_id, tuple(block_ids), num_external_tokens))
 
-    def build_connector_meta(self, scheduler_output):
+    def build_connector_meta(self, step_input):
         return None
 
     def update_connector_output(self, connector_output):
@@ -772,8 +773,10 @@ class _SaveConnector(_AsyncLookupConnector):
         def __init__(self):
             super().__init__([])
             self.metadata = None
+            self.step_input = None
 
-        def build_connector_meta(self, scheduler_output):
+        def build_connector_meta(self, step_input):
+            self.step_input = step_input
             metadata, self.metadata = self.metadata, None
             return metadata
 
@@ -801,6 +804,8 @@ def update_connector_output(self, connector_output):
     )
 
     assert metadata is not None
+    assert isinstance(connector.step_input, KVConnectorStepInput)
+    assert connector.step_input.running == [seq]
     assert allocator.get_ref_count(logical_blocks).tolist() == [2, 2]
     assert scheduler.has_unfinished()
 
@@ -837,7 +842,7 @@ def get_save_block_leases(self):
     scheduler.block_manager.allocate(seq)
     logical_blocks = seq.logical_blocks.get_real_blocks().copy()
     metadata = _SaveMetadata(logical_blocks)
-    connector.build_connector_meta = lambda scheduler_output: metadata
+    connector.build_connector_meta = lambda step_input: metadata
 
     scheduler.build_connector_meta([seq], connector_token_lens=(8, ))
     scheduler.block_manager.free(seq)

From 2bc1f5edb019d938d5507ed38ccd74e52aaf3d7c Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 18:56:11 +0800
Subject: [PATCH 18/22] refactor: clarify prefill rollback ownership

---
 lmdeploy/pytorch/paging/prefill_scheduler.py | 66 ++++++++++----------
 1 file changed, 34 insertions(+), 32 deletions(-)

diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index aeb088d57e..b18bb2dd46 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -282,7 +282,6 @@ class _TentativePrefixMatch:
         '_preserve_existing_state',
         '_stats_snapshot',
         '_state_snapshot',
-        '_rejection_on_rollback',
         '_phase',
     )
 
@@ -300,7 +299,6 @@ def __init__(self,
         self._preserve_existing_state = preserve_existing_state
         self._stats_snapshot = None
         self._state_snapshot: _PrefixMatchStateSnapshot | None = None
-        self._rejection_on_rollback: _PrefillAdmissionResult | None = None
         self._phase = _PrefixMatchPhase.IDLE
 
     @property
@@ -329,11 +327,6 @@ def match(self) -> None:
         self.block_trie.match(self.seq)
         self._phase = _PrefixMatchPhase.MATCHED
 
-    def retain_for_admission(self, rejection_on_rollback: _PrefillAdmissionResult) -> None:
-        """Keep a gate-enabling match and remember its original rejection."""
-        assert self.is_matched
-        self._rejection_on_rollback = rejection_on_rollback
-
     def pin_restore(self) -> bool:
         """Pin an SSM restore selected by this tentative match."""
         restore = self.seq.prefix_cache.restore
@@ -345,11 +338,10 @@ def commit(self) -> None:
         """Accept the match and discard request-local rollback state."""
         self._clear()
 
-    def rollback(self, reason: str):
-        """Undo the transaction and return any gate-defined rejection."""
-        rejection = self._rejection_on_rollback
+    def rollback(self, reason: str) -> None:
+        """Undo the tentative prefix-cache transaction."""
         if self._phase is _PrefixMatchPhase.IDLE:
-            return rejection
+            return
 
         seq = self.seq
         logger.debug('Rollback tentative prefix-cache match: session_id=%s seq_id=%s reason=%s '
@@ -362,7 +354,6 @@ def rollback(self, reason: str):
         else:
             self._restore_snapshot(snapshot)
         self._clear()
-        return rejection
 
     def _restore_snapshot(self, snapshot: _PrefixMatchStateSnapshot) -> None:
         seq = self.seq
@@ -402,7 +393,6 @@ def _reset_to_unmatched(self) -> None:
     def _clear(self) -> None:
         self._stats_snapshot = None
         self._state_snapshot = None
-        self._rejection_on_rollback = None
         self._phase = _PrefixMatchPhase.IDLE
 
 
@@ -437,6 +427,7 @@ def __init__(self,
             preserve_existing_state=(
                 self.load_coordinator.lookup_enabled and not self._load_ready),
         )
+        self._resource_rollback_rejection: _PrefillAdmissionResult | None = None
 
     def run(self):
         """Apply policy, acquire resources, and commit one admission."""
@@ -492,10 +483,10 @@ def _admit_matched_resources(self):
         seq = self.seq
         had_ssm_restore = prefill.is_ssm and seq.prefix_cache.restore.is_selected
         if not self._prefix_match.pin_restore():
-            result = self._prefix_match.rollback(
+            gate_rejection = self._rollback_match_after_resource_failure(
                 'failed to pin SSM restore checkpoint')
-            if result is not None:
-                return result
+            if gate_rejection is not None:
+                return gate_rejection
 
         kv_result = self._admit_kv_resources(had_ssm_restore)
         if kv_result is not None:
@@ -510,9 +501,9 @@ def _admit_kv_resources(self, had_ssm_restore: bool):
         reason = 'eviction failed'
         if had_ssm_restore:
             reason = 'eviction failed with pinned SSM restore'
-        result = self._prefix_match.rollback(reason)
-        if result is not None:
-            return result
+        gate_rejection = self._rollback_match_after_resource_failure(reason)
+        if gate_rejection is not None:
+            return gate_rejection
 
         # The matched restore may pin the only checkpoint state that eviction
         # can free. Retrying after rollback preserves the unmatched fallback.
@@ -527,10 +518,10 @@ def _admit_runtime_state(self):
         if not prefill.is_ssm or prefill._make_runtime_state_available():
             return None
 
-        result = self._prefix_match.rollback(
+        gate_rejection = self._rollback_match_after_resource_failure(
             'no runtime SSM state available')
-        if result is not None:
-            return result
+        if gate_rejection is not None:
+            return gate_rejection
         if not self._prepare_and_evict():
             return _PrefillAdmissionResult.stop()
         if not prefill._make_runtime_state_available():
@@ -573,18 +564,31 @@ def _evictable_sequences(self):
         yield from reversed(self.stopped)
         yield from reversed(self.evictable_waiting)
 
-    def _match_prefix_for_prefill_gate(self):
+    def _try_match_prefix_for_prefill_gate(self) -> bool:
         """Tentatively match once so a request can be rechecked by a gate."""
         prefill = self.prefill_scheduler
         if self._load_ready:
-            return None
+            return False
         if not prefill.block_trie.enabled:
-            return None
+            return False
         if self._has_private_local_tail():
-            return None
+            return False
         self._prefix_match.match()
         return True
 
+    def _accept_gate_enabling_match(self, rejection: _PrefillAdmissionResult) -> None:
+        """Preserve a gate's outcome if resource admission loses its match."""
+        assert self._prefix_match.is_matched
+        self._resource_rollback_rejection = rejection
+
+    def _rollback_match_after_resource_failure(
+        self,
+        reason: str,
+    ) -> _PrefillAdmissionResult | None:
+        """Undo the match and recover the rejection it allowed us to defer."""
+        self._prefix_match.rollback(reason)
+        return self._resource_rollback_rejection
+
     def _has_private_local_tail(self) -> bool:
         """Whether a private partial block prevents another trie match.
 
@@ -618,14 +622,13 @@ def _apply_nonfinal_long_prefill_gate(self):
                 or prefill._prefill_kv_token_limit(seq) is None):
             return None
 
-        matched = self._match_prefix_for_prefill_gate()
-        if matched is None:
+        if not self._try_match_prefix_for_prefill_gate():
             return _PrefillAdmissionResult.skip()
         if prefill._prefill_kv_token_limit(seq) is not None:
             self._prefix_match.rollback(
                 'still non-final long prefill on short turn')
             return _PrefillAdmissionResult.skip()
-        self._prefix_match.retain_for_admission(
+        self._accept_gate_enabling_match(
             _PrefillAdmissionResult.skip())
         return None
 
@@ -645,8 +648,7 @@ def _apply_prefill_token_budget_gate(self):
             self._prefix_match.rollback('still exceeds prefill token budget')
             return rejection
 
-        matched = self._match_prefix_for_prefill_gate()
-        if matched is None:
+        if not self._try_match_prefix_for_prefill_gate():
             return rejection
 
         prefill_token_count = prefill._prefill_admission_token_count(seq)
@@ -654,7 +656,7 @@ def _apply_prefill_token_budget_gate(self):
             self._prefix_match.rollback('still exceeds prefill token budget')
             return rejection
 
-        self._prefix_match.retain_for_admission(rejection)
+        self._accept_gate_enabling_match(rejection)
         return None
 
     def _prepare_and_evict(self):

From 4601082e4a1c004610c2c764ce49bad585db1913 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 19:39:02 +0800
Subject: [PATCH 19/22] refactor: tighten scheduler lifecycle API

---
 lmdeploy/pytorch/engine/engine.py             |  2 +-
 lmdeploy/pytorch/engine/engine_loop.py        | 10 ++---
 .../pytorch/paging/kv_load_coordinator.py     |  2 +-
 lmdeploy/pytorch/paging/scheduler.py          | 40 +------------------
 tests/pytorch/engine/test_engine_sleep.py     |  2 +-
 .../engine/test_kv_connector_wiring.py        | 27 -------------
 tests/pytorch/paging/test_scheduler.py        | 14 ++-----
 .../paging/test_scheduler_kv_transfer.py      | 27 +++++++++++--
 8 files changed, 37 insertions(+), 87 deletions(-)

diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py
index 2a5681d47f..96c7a2aa0b 100644
--- a/lmdeploy/pytorch/engine/engine.py
+++ b/lmdeploy/pytorch/engine/engine.py
@@ -607,7 +607,7 @@ async def sleep(self, level: int = 1):
         # cancel all remain sessions
         self._cancel_and_end_all_sessions()
         await self.executor.sleep(level)
-        self.scheduler.finish_deferred_kv_transfers_after_worker_drain()
+        self.scheduler.finish_kv_transfers_after_worker_drain()
         if self._engine_loop is not None:
             self._engine_loop.reset_runtime_state()
         logger.info('PyTorch engine entered sleep: level=%s, sleeping_tags=%s.', level, sorted(self._sleeping_tags))
diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py
index fd8ffce42a..9ed392711e 100644
--- a/lmdeploy/pytorch/engine/engine_loop.py
+++ b/lmdeploy/pytorch/engine/engine_loop.py
@@ -449,16 +449,14 @@ def _finish_forward_output(self,
         """Apply connector progress and publish model outputs."""
         if out is None:
             return
-        # A connector polling step intentionally has no token output. Consume
-        # its transfer completions first so newly loaded requests become
-        # schedulable even when no model forward ran in this executor step.
+        # Connector-only polls have no token output; apply completions before
+        # returning.
         self.scheduler.update_connector_output(out.kv_connector_output)
         if out.next_token_ids is None:
             return
         step_outputs = self._make_infer_outputs(out, running=running, model_inputs=model_inputs, delta=delta)
-        # Sequence history is advanced by _make_infer_outputs. Only now can the
-        # scheduler prove that a prefill reached its reserved target and release
-        # the soft block reservation used while admitting external KV loads.
+        # _make_infer_outputs advances history; only then can soft reservations
+        # be released.
         self.scheduler.release_completed_prefill_reservations(running)
         self.resp_queue.put_nowait(step_outputs)
 
diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index ee45af8006..9daccc074b 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -553,7 +553,7 @@ def _remove_sequence(self, seq: SchedulerSequence) -> None:
         if not session.sequences:
             self.sessions.pop(session.session_id, None)
 
-    def disable(self) -> None:
+    def shutdown(self) -> None:
         """Stop new lookup admission and discard scheduler-side ownership."""
         self.lookup_enabled = False
         self.clear()
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index c9db079af9..908cb04138 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -150,25 +150,16 @@ def shutdown(self) -> None:
         connector = self.kv_connector
         self.kv_connector = None
         self.sequence_lifecycle.disable_connector()
-        self.kv_load_coordinator.disable()
+        self.kv_load_coordinator.shutdown()
         self.kv_save_coordinator.clear()
         if connector is not None:
             connector.shutdown()
 
-    @property
-    def _external_lookup_enabled(self) -> bool:
-        """Whether external KV lookup admission is currently enabled."""
-        return self.kv_load_coordinator.lookup_enabled
-
     @property
     def last_schedule_had_pending_lookup(self) -> bool:
         """Whether the latest prefill turn encountered a pending lookup."""
         return self._prefill_scheduler.last_schedule_had_pending_lookup
 
-    @last_schedule_had_pending_lookup.setter
-    def last_schedule_had_pending_lookup(self, value: bool) -> None:
-        self._prefill_scheduler.last_schedule_had_pending_lookup = value
-
     def has_waiting_long_prefill(self):
         """Whether a waiting request would need a non-final prefill chunk."""
         return self._prefill_scheduler.has_waiting_long_prefill(self.waiting)
@@ -197,22 +188,10 @@ def reserve_long_context_chunk(self,
     def waiting(self) -> SeqList:
         return list(self.seq_manager.get_sequences(MessageStatus.WAITING).values())
 
-    @property
-    def remote_loading(self) -> SeqList:
-        return list(self.seq_manager.get_sequences(MessageStatus.WAITING_FOR_REMOTE_KVS).values())
-
-    @property
-    def ready(self) -> SeqList:
-        return list(self.seq_manager.get_sequences(MessageStatus.READY).values())
-
     @property
     def hanging(self) -> SeqList:
         return list(self.seq_manager.get_sequences(MessageStatus.STOPPED).values())
 
-    @property
-    def running(self) -> SeqList:
-        return list(self.seq_manager.get_sequences(MessageStatus.RUNNING).values())
-
     @property
     def migration_waiting(self) -> SeqList:
         return list(self.seq_manager.get_sequences(MessageStatus.MIGRATION_WAITING).values())
@@ -234,12 +213,6 @@ def num_ready(self) -> int:
     def num_running(self) -> int:
         return self.seq_manager.num_sequences(MessageStatus.RUNNING)
 
-    def num_migration_waiting(self) -> int:
-        return self.seq_manager.num_sequences(MessageStatus.MIGRATION_WAITING)
-
-    def num_migration_done(self) -> int:
-        return self.seq_manager.num_sequences(MessageStatus.MIGRATION_DONE)
-
     # Non-empty status checks used by engine control flow.
     def has_waiting(self) -> bool:
         return self.seq_manager.num_sequences(MessageStatus.WAITING) > 0
@@ -487,7 +460,7 @@ def release_completed_prefill_reservations(self, seqs: SeqList) -> None:
         """Release soft targets only after forward output advanced history."""
         self.kv_load_coordinator.release_completed_prefills(seqs)
 
-    def finish_deferred_kv_transfers_after_worker_drain(self) -> None:
+    def finish_kv_transfers_after_worker_drain(self) -> None:
         """Release paging ownership after worker transfer queues have drained.
 
         Engine sleep may discard prefetched completion outputs. Worker drain is
@@ -533,15 +506,6 @@ def deactivate_seqs(self, running: SeqList, filter_status: MessageStatus = Messa
             if seq.status == filter_status:
                 seq.state.deactivate()
 
-    @contextmanager
-    def seqs_activation(self, running: SeqList):
-        """Context manager to activate and deactivate sequences."""
-        self.activate_seqs(running, MessageStatus.READY)
-        try:
-            yield running
-        finally:
-            self.deactivate_seqs(running, MessageStatus.RUNNING)
-
     def activate_migration_seqs(self, running: SeqList):
         """Lock running sequence."""
         return self.activate_seqs(running, filter_status=MessageStatus.MIGRATION_READY)
diff --git a/tests/pytorch/engine/test_engine_sleep.py b/tests/pytorch/engine/test_engine_sleep.py
index 2b49093481..70e6e19a01 100644
--- a/tests/pytorch/engine/test_engine_sleep.py
+++ b/tests/pytorch/engine/test_engine_sleep.py
@@ -40,7 +40,7 @@ def get_session(self, session_id):
     def get_sessions(self):
         return list(self.sessions.values())
 
-    def finish_deferred_kv_transfers_after_worker_drain(self):
+    def finish_kv_transfers_after_worker_drain(self):
         pass
 
 
diff --git a/tests/pytorch/engine/test_kv_connector_wiring.py b/tests/pytorch/engine/test_kv_connector_wiring.py
index 93cc61d474..634a100591 100644
--- a/tests/pytorch/engine/test_kv_connector_wiring.py
+++ b/tests/pytorch/engine/test_kv_connector_wiring.py
@@ -8,7 +8,6 @@
 from lmdeploy.pytorch.engine.config_builder import ConfigBuilder
 from lmdeploy.pytorch.engine.engine import Engine
 from lmdeploy.pytorch.kv_connector import prepare_kv_connector_config
-from lmdeploy.pytorch.paging.scheduler import Scheduler
 
 
 def _make_cache_config(kv_transfer_config=None):
@@ -216,32 +215,6 @@ def test_prepare_kv_connector_config_does_not_change_disabled_config(transfer_co
         assert transfer_config.kv_connector_extra_config == original_extra_config
 
 
-def test_scheduler_shutdown_releases_injected_connector_once():
-    connector = Mock()
-    load_coordinator = Mock()
-    load_coordinator.lookup_enabled = True
-
-    def disable_loads():
-        load_coordinator.lookup_enabled = False
-
-    load_coordinator.disable.side_effect = disable_loads
-    scheduler = Scheduler.__new__(Scheduler)
-    scheduler.kv_connector = connector
-    scheduler.sequence_lifecycle = Mock()
-    scheduler.kv_load_coordinator = load_coordinator
-    scheduler.kv_save_coordinator = Mock()
-
-    scheduler.shutdown()
-    scheduler.shutdown()
-
-    connector.shutdown.assert_called_once_with()
-    assert scheduler.kv_load_coordinator.disable.call_count == 2
-    assert scheduler.sequence_lifecycle.disable_connector.call_count == 2
-    assert scheduler.kv_save_coordinator.clear.call_count == 2
-    assert scheduler.kv_connector is None
-    assert not scheduler._external_lookup_enabled
-
-
 def test_engine_loop_finally_shuts_down_scheduler_before_executor():
     calls = []
     engine = Engine.__new__(Engine)
diff --git a/tests/pytorch/paging/test_scheduler.py b/tests/pytorch/paging/test_scheduler.py
index 7e5c2a5a42..58d6b9b8ae 100644
--- a/tests/pytorch/paging/test_scheduler.py
+++ b/tests/pytorch/paging/test_scheduler.py
@@ -110,19 +110,18 @@ def test_update(self, scheduler, block_size, num_gpu_blocks):
 
         # stop seq
         seq1.state.stop()
-        assert len(scheduler.ready) == 1
+        assert scheduler.num_ready() == 1
         assert seq1 in scheduler.hanging
 
         # end seq
         seq1.session.remove_sequence(seq1)
         assert session_id1 in scheduler.sessions
-        assert seq1 not in scheduler.ready
         assert seq1 not in scheduler.hanging
         assert block_manager.get_num_free_gpu_blocks() == num_gpu_blocks - 2
 
         # stop session
         scheduler.stop_session(session_id2)
-        assert len(scheduler.ready) == 0
+        assert scheduler.num_ready() == 0
         assert len(scheduler.waiting) == 0
         assert len(scheduler.hanging) == 2
 
@@ -155,7 +154,7 @@ def test_evict(self, scheduler, block_size, num_gpu_blocks, num_cpu_blocks):
 
         # test: waiting alloc
         seq2.state.stop()
-        assert len(scheduler.ready) == 1
+        assert scheduler.num_ready() == 1
         assert len(scheduler.waiting) == 1
         assert len(scheduler.hanging) == 1
 
@@ -172,7 +171,7 @@ def test_evict(self, scheduler, block_size, num_gpu_blocks, num_cpu_blocks):
         seq2.state.activate()
         seq3.session.remove_sequence(seq3)
         seq2.update_token_ids(torch.tensor([1] * block_size))
-        assert len(scheduler.ready) == 1
+        assert scheduler.num_ready() == 1
         assert len(scheduler.waiting) == 1
         assert len(scheduler.hanging) == 0
 
@@ -203,10 +202,7 @@ def test_schedule_running_validity_uses_input_indices(self, scheduler,
         ('name', 'status'),
         [
             ('waiting', MessageStatus.WAITING),
-            ('remote_loading', MessageStatus.WAITING_FOR_REMOTE_KVS),
-            ('ready', MessageStatus.READY),
             ('hanging', MessageStatus.STOPPED),
-            ('running', MessageStatus.RUNNING),
             ('migration_waiting', MessageStatus.MIGRATION_WAITING),
             ('migration_done', MessageStatus.MIGRATION_DONE),
         ],
@@ -231,8 +227,6 @@ def get_sequences(actual_status):
             ('num_remote_loading', MessageStatus.WAITING_FOR_REMOTE_KVS, 3),
             ('num_ready', MessageStatus.READY, 3),
             ('num_running', MessageStatus.RUNNING, 3),
-            ('num_migration_waiting', MessageStatus.MIGRATION_WAITING, 3),
-            ('num_migration_done', MessageStatus.MIGRATION_DONE, 3),
             ('has_waiting', MessageStatus.WAITING, True),
             ('has_remote_loading', MessageStatus.WAITING_FOR_REMOTE_KVS, True),
             ('has_ready', MessageStatus.READY, True),
diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py
index 94a8d9eb91..9a30a4c6e6 100644
--- a/tests/pytorch/paging/test_scheduler_kv_transfer.py
+++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py
@@ -30,6 +30,7 @@ def __init__(self, results, failed_ids=()):
         self.cancelled = []
         self.finished = []
         self.allocations = []
+        self.shutdown_calls = 0
 
     def on_new_request(self, request):
         self.new_requests.append(request.seq_id)
@@ -74,7 +75,7 @@ def finish_transfers_after_worker_drain(self):
         pass
 
     def shutdown(self):
-        pass
+        self.shutdown_calls += 1
 
 
 def _make_async_lookup_scheduler(
@@ -126,6 +127,26 @@ def test_sequence_lifecycle_notifies_connector_on_add_and_end():
     assert 71 not in scheduler.sessions
 
 
+def test_scheduler_shutdown_disables_external_loads_and_connector_once():
+    connector = _AsyncLookupConnector([])
+    scheduler = _make_async_lookup_scheduler(
+        connector,
+        enable_prefix_caching=False,
+    )
+    seq = scheduler.add_session(72).add_sequence(torch.arange(9))
+    scheduler.kv_load_coordinator.track_prefill(seq)
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 3
+
+    scheduler.shutdown()
+    scheduler.shutdown()
+
+    assert connector.shutdown_calls == 1
+    assert not scheduler.kv_load_coordinator.lookup_enabled
+    assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
+    assert scheduler.schedule(is_prefill=True).running == [seq]
+    assert connector.lookup_calls == []
+
+
 def test_async_lookup_pending_rolls_back_a_new_request_once():
     connector = _AsyncLookupConnector([(None, False), (0, False)])
     scheduler = _make_async_lookup_scheduler(connector)
@@ -715,7 +736,7 @@ def test_worker_drain_finishes_an_ended_session_with_a_dropped_load_output():
     scheduler.schedule(is_prefill=True)
     scheduler.end_session(85)
 
-    scheduler.finish_deferred_kv_transfers_after_worker_drain()
+    scheduler.finish_kv_transfers_after_worker_drain()
 
     assert 85 not in scheduler.sessions
     assert connector.finished == [seq.seq_id]
@@ -846,7 +867,7 @@ def get_save_block_leases(self):
 
     scheduler.build_connector_meta([seq], connector_token_lens=(8, ))
     scheduler.block_manager.free(seq)
-    scheduler.finish_deferred_kv_transfers_after_worker_drain()
+    scheduler.finish_kv_transfers_after_worker_drain()
 
     assert scheduler.block_manager.get_num_free_gpu_blocks() == 2
     assert not scheduler.kv_save_coordinator.has_pending()

From ed676bc66909248ec38aea1f1fb003ad8fc690e0 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 20:22:33 +0800
Subject: [PATCH 20/22] refactor: align scheduler lifecycle naming

---
 lmdeploy/pytorch/engine/engine_loop.py        |  6 ++--
 lmdeploy/pytorch/messages.py                  |  8 +++---
 .../recompute_eviction_helper.py              |  4 +--
 .../pytorch/paging/kv_load_coordinator.py     | 28 +++++++++----------
 .../pytorch/paging/kv_save_coordinator.py     |  4 +--
 lmdeploy/pytorch/paging/prefill_scheduler.py  |  6 ++--
 lmdeploy/pytorch/paging/scheduler.py          | 20 +++++++------
 lmdeploy/pytorch/paging/seq_states/states.py  |  8 +++---
 tests/pytorch/engine/test_inputs_maker.py     | 10 +++----
 9 files changed, 48 insertions(+), 46 deletions(-)

diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py
index 9ed392711e..55ca40ab98 100644
--- a/lmdeploy/pytorch/engine/engine_loop.py
+++ b/lmdeploy/pytorch/engine/engine_loop.py
@@ -406,14 +406,14 @@ async def _main_loop_try_send_next_inputs(self):
         if self._sleep_requested:
             return None, None
 
-        self.scheduler.collect_migration_done()
+        self.scheduler.resume_completed_migrations()
         return await self.inputs_maker.send_next_inputs()
 
     async def _prefetch_next_inputs(self):
-        """Collect migration completions before prefetching the next batch."""
+        """Resume completed migrations before prefetching the next batch."""
         if self._sleep_requested:
             return None, None
-        self.scheduler.collect_migration_done()
+        self.scheduler.resume_completed_migrations()
         return await self.inputs_maker.prefetch_next_inputs()
 
     async def _wait_for_schedulable_prefill(self):
diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py
index 06e6971d7b..d7b03f7927 100644
--- a/lmdeploy/pytorch/messages.py
+++ b/lmdeploy/pytorch/messages.py
@@ -238,16 +238,16 @@ def num_sequences(self, status: MessageStatus):
         """Num sequences."""
         return len(self.get_sequences(status))
 
-    def add_sequence(self, seq: 'SchedulerSequence'):
-        """Add sequence."""
+    def register_sequence(self, seq: 'SchedulerSequence'):
+        """Register a sequence in the global and status indexes."""
         seq_id = seq.seq_id
         status = seq.status
         status_map = self._status_seq_map[status]
         self._seq_map[seq_id] = seq
         status_map[seq_id] = seq
 
-    def remove_sequence(self, seq: 'SchedulerSequence'):
-        """Remove sequence."""
+    def unregister_sequence(self, seq: 'SchedulerSequence'):
+        """Remove a sequence from the global and status indexes."""
         seq_id = seq.seq_id
         status = seq.status
         status_map = self._status_seq_map[status]
diff --git a/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py b/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
index d0233d6476..05a46867c2 100644
--- a/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
+++ b/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
@@ -65,7 +65,7 @@ def _evict_for_seq_default(self, seq: SchedulerSequence, evictable_seqs: list[Sc
                 evict_seq.prefix_cache.suppress_match_stats = True
             # Eviction also ends the tracked prefill; otherwise its soft block
             # reservation would outlive the local KV blocks freed below.
-            self.load_coordinator.release(evict_seq)
+            self.load_coordinator.release_tracking(evict_seq)
             evict_seq.state.free()
             num_req = (num_required_blocks - block_manager.get_num_free_gpu_blocks())
             if num_req <= 0:
@@ -125,7 +125,7 @@ def _evict_for_ssm(self, seq: SchedulerSequence, evictable_seqs: list[SchedulerS
                 evict_seq.prefix_cache.suppress_match_stats = True
             # Keep coordinator ownership and its soft admission budget in sync
             # with the KV blocks and SSM runtime state released by free().
-            self.load_coordinator.release(evict_seq)
+            self.load_coordinator.release_tracking(evict_seq)
             evict_seq.state.free()
             has_free_state = has_runtime_state or state_manager.get_num_free_runtime() > 0
             if not has_free_state:
diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index 9daccc074b..826fec021a 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -9,8 +9,8 @@
    destinations, binds connector metadata, and calls :meth:`start_load`.
 2. While workers may write those blocks, the sequence stays in
    ``WAITING_FOR_REMOTE_KVS`` and cannot be evicted or removed.
-3. :meth:`update` publishes a successful load or rolls a failed/cancelled load
-   back to the last block-aligned safe step.
+3. :meth:`apply_load_results` publishes a successful load or rolls a
+   failed/cancelled load back to the last block-aligned safe step.
 4. The completed request is admitted for its remaining prefill, after which
    its load record and soft reservation can be released.
 
@@ -392,13 +392,13 @@ def is_remote_ready(self, seq: SchedulerSequence) -> bool:
         record = self._loads.get(int(seq.seq_id))
         return record is not None and record.phase is _LoadPhase.READY
 
-    def mark_scheduled(self, seq: SchedulerSequence) -> None:
+    def mark_prefill_scheduled(self, seq: SchedulerSequence) -> None:
         """Record that the remote-ready request entered remaining prefill."""
         record = self._loads.get(int(seq.seq_id))
         if record is not None and record.phase is _LoadPhase.READY:
             record.phase = _LoadPhase.PREFILLING
 
-    def update(self, results: tuple[KVLoadResult, ...]) -> None:
+    def apply_load_results(self, results: tuple[KVLoadResult, ...]) -> None:
         """Apply terminal load results aggregated across all TP ranks.
 
         Missing or non-``LOADING`` records are stale/duplicate progress and are
@@ -469,11 +469,11 @@ def _finish_cancelled_or_failed(self, record: _LoadRecord) -> None:
         self._prefill_targets.pop(request_id, None)
         seq.state.finish_remote_load()
         if record.deferred_cleanup is _DeferredLoadCleanup.END:
-            self._remove_sequence(seq)
+            self._finish_deferred_end(seq)
         elif record.deferred_cleanup is _DeferredLoadCleanup.STOP:
             seq.state.stop()
 
-    def request_stop(self, seq: SchedulerSequence) -> bool:
+    def defer_stop_if_loading(self, seq: SchedulerSequence) -> bool:
         """Return True when stop must wait for an active device write.
 
         Dropping the record or freeing sequence blocks now could let paging
@@ -482,13 +482,13 @@ def request_stop(self, seq: SchedulerSequence) -> bool:
         """
         record = self._loads.get(int(seq.seq_id))
         if record is None or record.phase is not _LoadPhase.LOADING:
-            self.release(seq)
+            self.release_tracking(seq)
             return False
         if record.deferred_cleanup is _DeferredLoadCleanup.NONE:
             record.deferred_cleanup = _DeferredLoadCleanup.STOP
         return True
 
-    def request_end(self, seq: SchedulerSequence) -> bool:
+    def defer_end_if_loading(self, seq: SchedulerSequence) -> bool:
         """Return True when removal must wait for an active device write.
 
         End differs from stop only in final cleanup: after the write terminates,
@@ -496,12 +496,12 @@ def request_end(self, seq: SchedulerSequence) -> bool:
         """
         record = self._loads.get(int(seq.seq_id))
         if record is None or record.phase is not _LoadPhase.LOADING:
-            self.release(seq)
+            self.release_tracking(seq)
             return False
         record.deferred_cleanup = _DeferredLoadCleanup.END
         return True
 
-    def release(self, seq: SchedulerSequence) -> None:
+    def release_tracking(self, seq: SchedulerSequence) -> None:
         """Drop tracking after prefill completion, preemption, or removal.
 
         ``LOADING`` is intentionally a no-op because only terminal worker
@@ -515,7 +515,7 @@ def release(self, seq: SchedulerSequence) -> None:
         if record is not None:
             self._loads.pop(request_id, None)
 
-    def release_completed_prefills(self, seqs: list[SchedulerSequence]) -> None:
+    def release_completed_prefill_reservations(self, seqs: list[SchedulerSequence]) -> None:
         """Release reservations after model output advances sequence history.
 
         Dispatch alone is insufficient proof: only after EngineLoop applies the
@@ -529,7 +529,7 @@ def release_completed_prefills(self, seqs: list[SchedulerSequence]) -> None:
             if self.is_remote_ready(seq):
                 continue
             if int(seq.num_history_ids) >= int(seq.input_end_pos):
-                self.release(seq)
+                self.release_tracking(seq)
 
     def finish_deferred_loads_after_worker_drain(self) -> None:
         """Remove ended requests after workers can no longer write KV blocks.
@@ -547,9 +547,9 @@ def finish_deferred_loads_after_worker_drain(self) -> None:
             self._rollback(record)
             self._finish_cancelled_or_failed(record)
 
-    def _remove_sequence(self, seq: SchedulerSequence) -> None:
+    def _finish_deferred_end(self, seq: SchedulerSequence) -> None:
         session = seq.session
-        session.lifecycle.finish_sequence(seq)
+        session.lifecycle.end_sequence(seq)
         if not session.sequences:
             self.sessions.pop(session.session_id, None)
 
diff --git a/lmdeploy/pytorch/paging/kv_save_coordinator.py b/lmdeploy/pytorch/paging/kv_save_coordinator.py
index 1009534c4e..257e3ab00a 100644
--- a/lmdeploy/pytorch/paging/kv_save_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_save_coordinator.py
@@ -13,7 +13,7 @@
 1. ``Scheduler.build_connector_meta`` asks the connector for save metadata.
 2. :meth:`acquire` pins all logical blocks before metadata reaches workers.
 3. Worker output is aggregated by the connector into completed operation IDs.
-4. :meth:`update` releases the matching logical references.
+4. :meth:`release_completed_leases` releases the matching logical references.
 5. Engine shutdown drains transfer queues before :meth:`clear` releases any
    leases whose final outputs were intentionally discarded.
 """
@@ -59,7 +59,7 @@ def acquire(self, metadata: KVConnectorMetadata) -> None:
             block_manager.pin_logical_blocks(logical_block_ids)
             self._leases[lease.operation_id] = logical_block_ids
 
-    def update(self, completed_save_ids: frozenset[KVOperationId]) -> None:
+    def release_completed_leases(self, completed_save_ids: frozenset[KVOperationId]) -> None:
         """Release operations that reached a terminal state on every TP
         rank."""
         block_manager = self.block_manager
diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index b18bb2dd46..f847a56b02 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -437,7 +437,7 @@ def run(self):
                 return _PrefillAdmissionResult.skip()
             self._prefix_match.begin()
 
-        gate_result = self._check_prefill_admission_gates()
+        gate_result = self._apply_prefill_admission_gates()
         if gate_result is not None:
             return gate_result
 
@@ -607,7 +607,7 @@ def _token_budget_rejection(self):
             return _PrefillAdmissionResult.stop()
         return _PrefillAdmissionResult.skip()
 
-    def _check_prefill_admission_gates(self):
+    def _apply_prefill_admission_gates(self):
         """Apply long-prefill and token-budget admission gates."""
         result = self._apply_nonfinal_long_prefill_gate()
         if result is not None:
@@ -702,7 +702,7 @@ def _allocate_and_commit(self):
             # Preserve the load record through the remaining prefill so its
             # reservation can be released only after model output advances the
             # sequence to input_end_pos.
-            self.load_coordinator.mark_scheduled(seq)
+            self.load_coordinator.mark_prefill_scheduled(seq)
         self._prefix_match.commit()
         return _PrefillAdmissionResult.admit(prefill_token_count)
 
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 908cb04138..7bb2aaca5c 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -331,7 +331,7 @@ def schedule_running(self, running: SeqList, num_required_tokens: int = 1, preal
                 continue
 
             seq.state.deactivate()
-            self.kv_load_coordinator.release(seq)
+            self.kv_load_coordinator.release_tracking(seq)
             seq.state.evict()
             valid_mask[idx] = False
         return valid_mask
@@ -351,7 +351,7 @@ def stop_session(self, session_id: int):
                 connector.cancel_lookup(seq.seq_id)
             # An active load may still write GPU memory. Defer the state change
             # until all ranks terminate instead of making its blocks evictable.
-            if self.kv_load_coordinator.request_stop(seq):
+            if self.kv_load_coordinator.defer_stop_if_loading(seq):
                 continue
             seq.state.stop()
 
@@ -371,11 +371,11 @@ def end_session(self, session_id: int):
                 connector.cancel_lookup(seq.seq_id)
             # Session removal also frees sequence blocks, so it must be deferred
             # while a worker may still address an in-flight load destination.
-            if self.kv_load_coordinator.request_end(seq):
+            if self.kv_load_coordinator.defer_end_if_loading(seq):
                 continue
             # stop session so it won't get scheduled again
             seq.state.stop()
-            self.sequence_lifecycle.finish_sequence(seq)
+            self.sequence_lifecycle.end_sequence(seq)
         if not session.sequences:
             self.sessions.pop(session_id)
 
@@ -453,12 +453,13 @@ def update_connector_output(self, connector_output) -> None:
         if connector_output is None or self.kv_connector is None:
             return
         result = self.kv_connector.update_connector_output(connector_output)
-        self.kv_load_coordinator.update(result.load_results)
-        self.kv_save_coordinator.update(result.completed_save_ids)
+        self.kv_load_coordinator.apply_load_results(result.load_results)
+        self.kv_save_coordinator.release_completed_leases(
+            result.completed_save_ids)
 
     def release_completed_prefill_reservations(self, seqs: SeqList) -> None:
         """Release soft targets only after forward output advanced history."""
-        self.kv_load_coordinator.release_completed_prefills(seqs)
+        self.kv_load_coordinator.release_completed_prefill_reservations(seqs)
 
     def finish_kv_transfers_after_worker_drain(self) -> None:
         """Release paging ownership after worker transfer queues have drained.
@@ -492,7 +493,7 @@ def resolve_gpu_block_offsets(self, logical_block_ids):
     def evict_seqs(self, running: SeqList):
         """Evict running sequences."""
         for seq in running:
-            self.kv_load_coordinator.release(seq)
+            self.kv_load_coordinator.release_tracking(seq)
             seq.state.evict()
 
     def activate_seqs(self, running: SeqList, filter_status: MessageStatus = MessageStatus.READY):
@@ -523,7 +524,8 @@ def seqs_migration_activation(self, running: SeqList):
         finally:
             self.deactivate_migration_seqs(running)
 
-    def collect_migration_done(self):
+    def resume_completed_migrations(self):
+        """Move completed migration sequences back to the waiting queue."""
         for seq in self.migration_done:
             seq.state.activate()
 
diff --git a/lmdeploy/pytorch/paging/seq_states/states.py b/lmdeploy/pytorch/paging/seq_states/states.py
index 9ec965b007..657e4177bd 100644
--- a/lmdeploy/pytorch/paging/seq_states/states.py
+++ b/lmdeploy/pytorch/paging/seq_states/states.py
@@ -37,7 +37,7 @@ def new_sequence_id(self) -> int:
     def add_sequence(self, seq: SchedulerSequence, status: MessageStatus) -> None:
         seq.session.sequences[seq.seq_id] = seq
         seq.set_state(StateBase.build(self, seq, status))
-        self._seq_manager.add_sequence(seq)
+        self._seq_manager.register_sequence(seq)
         if self._connector is not None:
             self._connector.on_new_request(seq)
 
@@ -46,10 +46,10 @@ def remove_sequence(self, seq: SchedulerSequence) -> None:
         assert seq.seq_id in seq.session.sequences
         self.free_sequence(seq)
         seq.session.sequences.pop(seq.seq_id)
-        self._seq_manager.remove_sequence(seq)
+        self._seq_manager.unregister_sequence(seq)
 
-    def finish_sequence(self, seq: SchedulerSequence) -> None:
-        """Notify connector completion, then release local ownership."""
+    def end_sequence(self, seq: SchedulerSequence) -> None:
+        """Notify the connector that the request ended, then remove it."""
         if self._connector is not None:
             self._connector.request_finished(seq)
         self.remove_sequence(seq)
diff --git a/tests/pytorch/engine/test_inputs_maker.py b/tests/pytorch/engine/test_inputs_maker.py
index 777d49f6cb..588c1c8e36 100644
--- a/tests/pytorch/engine/test_inputs_maker.py
+++ b/tests/pytorch/engine/test_inputs_maker.py
@@ -252,7 +252,7 @@ async def get_output_async(self):
 
     state_checkpoints = _StateCheckpoints()
     loop = EngineLoop.__new__(EngineLoop)
-    loop.scheduler = SimpleNamespace(collect_migration_done=lambda: None)
+    loop.scheduler = SimpleNamespace(resume_completed_migrations=lambda: None)
     loop.state_checkpoints = state_checkpoints
     loop.inputs_maker = _InputsMaker(state_checkpoints)
     loop.executor = _Executor(state_checkpoints)
@@ -338,7 +338,7 @@ async def get_output_async(self):
 
     state_checkpoints = _StateCheckpoints()
     loop = EngineLoop.__new__(EngineLoop)
-    loop.scheduler = SimpleNamespace(collect_migration_done=lambda: None)
+    loop.scheduler = SimpleNamespace(resume_completed_migrations=lambda: None)
     loop.state_checkpoints = state_checkpoints
     loop.inputs_maker = _InputsMaker()
     loop.executor = _Executor()
@@ -369,8 +369,8 @@ class _Scheduler:
         def has_unfinished(self):
             return False
 
-        def collect_migration_done(self):
-            events.append('collect_migration_done')
+        def resume_completed_migrations(self):
+            events.append('resume_completed_migrations')
 
     class _InputsMaker:
 
@@ -390,7 +390,7 @@ async def send_next_inputs(self):
     result = asyncio.run(asyncio.wait_for(loop._main_loop_try_send_next_inputs(), timeout=1.0))
 
     assert result == ('forward_inputs', ['long-seq'])
-    assert events == ['collect_migration_done', 'send_next_inputs']
+    assert events == ['resume_completed_migrations', 'send_next_inputs']
 
 
 def test_migration_loop_schedules_and_processes_ready_batch():

From 27a2578b6869115c181edd8b76b3c2ff0eaccee9 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 20:59:44 +0800
Subject: [PATCH 21/22] refactor: simplify scheduler eviction lifecycle

---
 .../paging/block_trie/checkpoint_lifecycle.py |  10 +
 .../paging/eviction_helper/__init__.py        |   3 +-
 .../eviction_helper/base_eviction_helper.py   |  39 ----
 .../recompute_eviction_helper.py              | 198 ++++++++----------
 .../pytorch/paging/kv_load_coordinator.py     |  10 +-
 lmdeploy/pytorch/paging/prefill_scheduler.py  |  44 ++--
 lmdeploy/pytorch/paging/scheduler.py          |  12 +-
 lmdeploy/pytorch/paging/seq_states/states.py  |  11 +-
 .../test_checkpoint_lifecycle.py              |   4 +-
 .../paging/test_block_trie/test_trie.py       |   2 +-
 .../pytorch/paging/test_prefill_scheduler.py  |  12 +-
 .../paging/test_scheduler_kv_transfer.py      |   6 +-
 tests/pytorch/paging/test_scheduler_ssm.py    |  29 ++-
 13 files changed, 172 insertions(+), 208 deletions(-)
 delete mode 100644 lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py

diff --git a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
index c5e09f19d3..ffac74c1d6 100644
--- a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
+++ b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
@@ -281,6 +281,16 @@ def release_checkpoint(self, node: Node):
             self._state_manager.free_checkpoint_state(checkpoint.slot)
         node.state_checkpoint = None
 
+    def make_runtime_state_available(self) -> bool:
+        """Release one checkpoint when runtime-state capacity is exhausted."""
+        state_manager = self._state_manager
+        if state_manager is None:
+            return False
+        if state_manager.get_num_free_runtime() > 0:
+            return True
+        self.evict(1)
+        return state_manager.get_num_free_runtime() > 0
+
     def evict(self, max_num_checkpoints: int):
         """Evict published checkpoints without removing trie nodes."""
         return self._evict_checkpoints(self._index.unique_nodes(), max_num_checkpoints)
diff --git a/lmdeploy/pytorch/paging/eviction_helper/__init__.py b/lmdeploy/pytorch/paging/eviction_helper/__init__.py
index 9a52b580f2..22a1834fa6 100644
--- a/lmdeploy/pytorch/paging/eviction_helper/__init__.py
+++ b/lmdeploy/pytorch/paging/eviction_helper/__init__.py
@@ -10,6 +10,7 @@
     from ..block_trie import BlockTrie
     from ..kv_load_coordinator import KVLoadCoordinator
     from ..state_manager import StateManager
+    from .recompute_eviction_helper import RecomputeEvictionHelper
 
 logger = get_logger('lmdeploy')
 
@@ -22,7 +23,7 @@ def build_eviction_helper(
     state_manager: StateManager,
     load_coordinator: KVLoadCoordinator,
     is_ssm: bool,
-):
+) -> RecomputeEvictionHelper:
     """Build eviction helper."""
     if eviction_type == 'copy':
         logger.warning('`copy` eviction has been deprecated, '
diff --git a/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py b/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py
deleted file mode 100644
index ae519aa5b8..0000000000
--- a/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-from ...messages import SchedulerSequence
-
-if TYPE_CHECKING:
-    from ..block_manager import BaseBlockManager
-    from ..block_trie import BlockTrie
-    from ..kv_load_coordinator import KVLoadCoordinator
-    from ..state_manager import StateManager
-
-SeqList = list[SchedulerSequence]
-
-
-class BaseEvictionHelper:
-    """Base eviction helper."""
-
-    def __init__(
-        self,
-        *,
-        block_manager: BaseBlockManager,
-        block_trie: BlockTrie,
-        state_manager: StateManager,
-        load_coordinator: KVLoadCoordinator,
-    ) -> None:
-        self.block_manager = block_manager
-        self.block_trie = block_trie
-        self.state_manager = state_manager
-        self.load_coordinator = load_coordinator
-
-    def need_swap_in(self, seq: SchedulerSequence):
-        """Sequence need swap in."""
-        raise NotImplementedError('Not implemented.')
-
-    def evict_for_seq(self, seq: SchedulerSequence, evictable_seqs: list[SchedulerSequence], prealloc_size: int):
-        """Evict seqs."""
-        raise NotImplementedError('Not implemented.')
diff --git a/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py b/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
index 05a46867c2..fbc2a0290f 100644
--- a/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
+++ b/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py
@@ -1,10 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from __future__ import annotations
 
+from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
 from ...messages import SchedulerSequence
-from .base_eviction_helper import BaseEvictionHelper
 
 if TYPE_CHECKING:
     from ..block_manager import BaseBlockManager
@@ -13,8 +13,8 @@
     from ..state_manager import StateManager
 
 
-class RecomputeEvictionHelper(BaseEvictionHelper):
-    """Recompute eviction."""
+class RecomputeEvictionHelper:
+    """Reclaim paging resources so a sequence can be recomputed."""
 
     def __init__(
         self,
@@ -25,132 +25,106 @@ def __init__(
         load_coordinator: KVLoadCoordinator,
         is_ssm: bool,
     ) -> None:
-        super().__init__(
-            block_manager=block_manager,
-            block_trie=block_trie,
-            state_manager=state_manager,
-            load_coordinator=load_coordinator,
-        )
+        self.block_manager = block_manager
+        self.block_trie = block_trie
+        self.state_manager = state_manager
+        self.load_coordinator = load_coordinator
+        self._is_ssm = is_ssm
 
-        if is_ssm:
-            self.evict_for_seq = self._evict_for_ssm
-        else:
-            self.evict_for_seq = self._evict_for_seq_default
+    def try_make_capacity_for(
+        self,
+        seq: SchedulerSequence,
+        evictable_seqs: Iterable[SchedulerSequence],
+        prealloc_size: int,
+    ) -> bool:
+        """Try to reclaim enough paging capacity for one sequence."""
+        if self._is_ssm:
+            return self._try_make_ssm_capacity(
+                seq,
+                evictable_seqs,
+                prealloc_size,
+            )
 
-    def _evict_for_seq_default(self, seq: SchedulerSequence, evictable_seqs: list[SchedulerSequence],
-                               prealloc_size: int):
-        """Evict seqs."""
         block_manager = self.block_manager
-        block_trie = self.block_trie
-        num_required_blocks = block_manager.num_required_blocks(seq, prealloc_size)
-
+        num_required_blocks = block_manager.num_required_blocks(
+            seq,
+            prealloc_size,
+        )
         if block_manager.get_num_free_gpu_blocks() >= num_required_blocks:
             return True
 
-        success = False
-        while len(evictable_seqs) > 0:
-            evict_seq = evictable_seqs.pop(0)
-
-            # A completed remote load has published fresh KV into these blocks.
-            # Preserve it until prefill consumes the result instead of paying
-            # for the transfer again on a later scheduling turn.
-            if self.load_coordinator.is_remote_ready(evict_seq):
+        for evict_seq in evictable_seqs:
+            if not self._reclaim_candidate(evict_seq):
                 continue
+            if self._try_make_block_capacity(num_required_blocks):
+                return True
 
-            # skip sequence with no blocks
-            if evict_seq.num_blocks == 0:
-                continue
+        return self._try_make_block_capacity(num_required_blocks)
 
-            if block_trie.enabled:
-                evict_seq.prefix_cache.suppress_match_stats = True
-            # Eviction also ends the tracked prefill; otherwise its soft block
-            # reservation would outlive the local KV blocks freed below.
-            self.load_coordinator.release_tracking(evict_seq)
-            evict_seq.state.free()
-            num_req = (num_required_blocks - block_manager.get_num_free_gpu_blocks())
-            if num_req <= 0:
-                success = True
-                break
-
-            block_trie.evict(num_req)
-            num_req = (num_required_blocks - block_manager.get_num_free_gpu_blocks())
-            if num_req <= 0:
-                success = True
-                break
-
-        # for empty evictable_seqs case
-        num_req = num_required_blocks - block_manager.get_num_free_gpu_blocks()
-        if num_req > 0:
-            block_trie.evict(num_req)
-            if num_required_blocks <= block_manager.get_num_free_gpu_blocks():
-                success = True
-
-        return success
-
-    def _evict_for_ssm(self, seq: SchedulerSequence, evictable_seqs: list[SchedulerSequence], prealloc_size: int):
-        """Evict blocks and checkpoint states for an SSM sequence.
-
-        SSM scheduling needs both KV blocks and a runtime state slot.  Before evicting live sequences, try dropping old
-        unpinned checkpoints because they are cheaper to recompute than an active request.
-        """
+    def _try_make_ssm_capacity(
+        self,
+        seq: SchedulerSequence,
+        evictable_seqs: Iterable[SchedulerSequence],
+        prealloc_size: int,
+    ) -> bool:
+        """Try to make both KV-block and runtime-state capacity available."""
         block_manager = self.block_manager
         state_manager = self.state_manager
-        block_trie = self.block_trie
-        num_required_blocks = block_manager.num_required_blocks(seq, prealloc_size)
-        # avoid requiring free state when already allocated.
-        has_runtime_state = state_manager.is_allocated(seq)
-        has_free_state = has_runtime_state or state_manager.get_num_free_runtime() > 0
-        if not has_free_state:
-            block_trie.state_checkpoints.evict(1)
-            has_free_state = state_manager.get_num_free_runtime() > 0
+        state_checkpoints = self.block_trie.state_checkpoints
+        num_required_blocks = block_manager.num_required_blocks(
+            seq,
+            prealloc_size,
+        )
 
-        if has_free_state and block_manager.get_num_free_gpu_blocks() >= num_required_blocks:
+        # A running long prefill can reuse its already allocated state.
+        has_runtime_state = state_manager.is_allocated(seq)
+        has_free_state = (
+            has_runtime_state
+            or state_checkpoints.make_runtime_state_available()
+        )
+        if (has_free_state
+                and block_manager.get_num_free_gpu_blocks()
+                >= num_required_blocks):
             return True
 
-        success = False
-        while len(evictable_seqs) > 0:
-            evict_seq = evictable_seqs.pop(0)
-
-            # READY remote KV is already transferred and awaiting consumption;
-            # do not discard that result merely to admit another prefill.
-            if self.load_coordinator.is_remote_ready(evict_seq):
-                continue
-
-            # skip sequence with no blocks
-            if evict_seq.num_blocks == 0 and evict_seq.logical_state < 0:
+        for evict_seq in evictable_seqs:
+            if not self._reclaim_candidate(evict_seq):
                 continue
-
-            # free sequence
-            if block_trie.enabled:
-                evict_seq.prefix_cache.suppress_match_stats = True
-            # Keep coordinator ownership and its soft admission budget in sync
-            # with the KV blocks and SSM runtime state released by free().
-            self.load_coordinator.release_tracking(evict_seq)
-            evict_seq.state.free()
-            has_free_state = has_runtime_state or state_manager.get_num_free_runtime() > 0
-            if not has_free_state:
-                block_trie.state_checkpoints.evict(1)
-                has_free_state = state_manager.get_num_free_runtime() > 0
-            num_req = (num_required_blocks - block_manager.get_num_free_gpu_blocks())
-            if num_req <= 0:
-                success = True
-                break
-
-            # clear cached prefix
-            block_trie.evict(num_req)
-            num_req = (num_required_blocks - block_manager.get_num_free_gpu_blocks())
-            if num_req <= 0:
-                success = True
-                break
+            has_free_state = (
+                has_runtime_state
+                or state_checkpoints.make_runtime_state_available()
+            )
+            if self._try_make_block_capacity(num_required_blocks):
+                return has_free_state
 
         if not has_free_state:
             return False
+        return self._try_make_block_capacity(num_required_blocks)
 
-        # for empty evictable_seqs case
-        num_req = num_required_blocks - block_manager.get_num_free_gpu_blocks()
-        if num_req > 0:
-            block_trie.evict(num_req)
-            if num_required_blocks <= block_manager.get_num_free_gpu_blocks():
-                success = True
+    def _reclaim_candidate(self, seq: SchedulerSequence) -> bool:
+        """Release one eligible candidate's paging ownership."""
+        # Completed remote KV is already transferred and awaits consumption.
+        if self.load_coordinator.is_remote_ready(seq):
+            return False
+        if (seq.num_blocks == 0
+                and (not self._is_ssm or seq.logical_state < 0)):
+            return False
+
+        if self.block_trie.enabled:
+            seq.prefix_cache.suppress_match_stats = True
+        # Keep soft admission accounting aligned with released resources.
+        self.load_coordinator.release_tracking(seq)
+        seq.state.release_paging_resources()
+        return True
 
-        return success
+    def _try_make_block_capacity(self, num_required_blocks: int) -> bool:
+        """Evict cached trie blocks until the required capacity is free."""
+        block_manager = self.block_manager
+        num_missing_blocks = (
+            num_required_blocks - block_manager.get_num_free_gpu_blocks()
+        )
+        if num_missing_blocks > 0:
+            self.block_trie.evict(num_missing_blocks)
+        return (
+            num_required_blocks <= block_manager.get_num_free_gpu_blocks()
+        )
diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index 826fec021a..d9df94a525 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -36,7 +36,7 @@
 
     from .block_manager.base_block_manager import BaseBlockManager
     from .block_trie import BlockTrie
-    from .eviction_helper.base_eviction_helper import BaseEvictionHelper
+    from .eviction_helper.recompute_eviction_helper import RecomputeEvictionHelper
 
 
 class KVLoadAdmission(enum.Enum):
@@ -216,7 +216,7 @@ def try_load(
         *,
         prealloc_size: int,
         evictable_seqs: Iterable[SchedulerSequence],
-        eviction_helper: BaseEvictionHelper,
+        eviction_helper: RecomputeEvictionHelper,
     ) -> KVLoadAdmission:
         """Poll and admit one external prefix without choosing queue policy.
 
@@ -250,7 +250,7 @@ def _admit_load(
         num_external_tokens: int,
         prealloc_size: int,
         evictable_seqs: Iterable[SchedulerSequence],
-        eviction_helper: BaseEvictionHelper,
+        eviction_helper: RecomputeEvictionHelper,
     ) -> KVLoadAdmission:
         """Admit the complete prefill, then allocate the remote interval."""
         plan = self._plan_load(seq, num_external_tokens, prealloc_size)
@@ -301,13 +301,13 @@ def _admit_load_capacity(
         plan: _LoadPlan,
         prealloc_size: int,
         evictable_seqs: Iterable[SchedulerSequence],
-        eviction_helper: BaseEvictionHelper,
+        eviction_helper: RecomputeEvictionHelper,
     ) -> KVLoadAdmission | None:
         """Admit the full prefill against physical and soft capacity."""
         # Only the remote hit is allocated now, but admission guarantees the
         # complete prefill can finish beside every existing soft reservation.
         seq.kv_token_limit = None
-        full_prefill_fits = eviction_helper.evict_for_seq(
+        full_prefill_fits = eviction_helper.try_make_capacity_for(
             seq,
             list(evictable_seqs),
             prealloc_size,
diff --git a/lmdeploy/pytorch/paging/prefill_scheduler.py b/lmdeploy/pytorch/paging/prefill_scheduler.py
index f847a56b02..97b08f111a 100644
--- a/lmdeploy/pytorch/paging/prefill_scheduler.py
+++ b/lmdeploy/pytorch/paging/prefill_scheduler.py
@@ -27,7 +27,7 @@
 
 if TYPE_CHECKING:
     from .block_manager.base_block_manager import BaseBlockManager
-    from .eviction_helper.base_eviction_helper import BaseEvictionHelper
+    from .eviction_helper.recompute_eviction_helper import RecomputeEvictionHelper
 
 logger = get_logger('lmdeploy')
 
@@ -379,7 +379,7 @@ def _reset_to_unmatched(self) -> None:
         if self.is_ssm:
             self.block_trie.state_checkpoints.unpin_restore(seq)
         if seq.num_blocks > 0 or seq.logical_state >= 0:
-            seq.state.free()
+            seq.state.release_paging_resources()
         elif seq.num_history_ids > 0:
             seq.set_step(0)
         seq.kv_token_limit = None
@@ -454,7 +454,7 @@ def _admit_resources(self):
             load_result = self._try_external_load()
             if load_result is not None:
                 return load_result
-        if not self._prepare_and_evict():
+        if not self._prepare_and_make_capacity():
             return _PrefillAdmissionResult.stop()
         return None
 
@@ -495,7 +495,7 @@ def _admit_matched_resources(self):
 
     def _admit_kv_resources(self, had_ssm_restore: bool):
         """Evict for KV, retrying once without a pinned SSM restore."""
-        if self._prepare_and_evict():
+        if self._prepare_and_make_capacity():
             return None
 
         reason = 'eviction failed'
@@ -507,24 +507,25 @@ def _admit_kv_resources(self, had_ssm_restore: bool):
 
         # The matched restore may pin the only checkpoint state that eviction
         # can free. Retrying after rollback preserves the unmatched fallback.
-        if had_ssm_restore and self._prepare_and_evict():
+        if had_ssm_restore and self._prepare_and_make_capacity():
             return None
         return _PrefillAdmissionResult.stop()
 
     def _admit_runtime_state(self):
-        """Ensure an SSM runtime slot, retrying after match rollback."""
+        """Admit an SSM runtime slot, retrying after match rollback."""
         prefill = self.prefill_scheduler
         seq = self.seq
-        if not prefill.is_ssm or prefill._make_runtime_state_available():
+        state_checkpoints = prefill.block_trie.state_checkpoints
+        if not prefill.is_ssm or state_checkpoints.make_runtime_state_available():
             return None
 
         gate_rejection = self._rollback_match_after_resource_failure(
             'no runtime SSM state available')
         if gate_rejection is not None:
             return gate_rejection
-        if not self._prepare_and_evict():
+        if not self._prepare_and_make_capacity():
             return _PrefillAdmissionResult.stop()
-        if not prefill._make_runtime_state_available():
+        if not state_checkpoints.make_runtime_state_available():
             seq.kv_token_limit = None
             return _PrefillAdmissionResult.stop()
         return None
@@ -659,21 +660,21 @@ def _apply_prefill_token_budget_gate(self):
         self._accept_gate_enabling_match(rejection)
         return None
 
-    def _prepare_and_evict(self):
-        """Apply chunk allocation limits and evict for this prefill."""
+    def _prepare_and_make_capacity(self):
+        """Apply chunk allocation limits and reclaim prefill capacity."""
         prefill = self.prefill_scheduler
         seq = self.seq
         alloc_size = prefill._prepare_prefill_allocation(seq, self.prealloc_size)
         self._effective_prealloc_size = alloc_size
-        if self._evict_for_seq(alloc_size):
+        if self._try_make_capacity(alloc_size):
             return True
         seq.kv_token_limit = None
         return False
 
-    def _evict_for_seq(self, alloc_size: int):
-        """Evict stopped or skipped waiters until this sequence can run."""
+    def _try_make_capacity(self, alloc_size: int):
+        """Reclaim stopped or skipped waiters until this sequence can run."""
         prefill = self.prefill_scheduler
-        return prefill.eviction_helper.evict_for_seq(
+        return prefill.eviction_helper.try_make_capacity_for(
             self.seq,
             list(self._evictable_sequences()),
             alloc_size,
@@ -724,7 +725,7 @@ def __init__(
         block_manager: 'BaseBlockManager',
         block_trie: BlockTrie,
         state_manager: StateManager,
-        eviction_helper: 'BaseEvictionHelper',
+        eviction_helper: 'RecomputeEvictionHelper',
         load_coordinator: KVLoadCoordinator,
     ) -> None:
         self.scheduler_config = scheduler_config
@@ -742,15 +743,6 @@ def __init__(
             _envs.opt_ttft_aging_sec,
         )
 
-    def _make_runtime_state_available(self):
-        """Make one state-cache slot available for an SSM runtime state."""
-        if not self.is_ssm:
-            return True
-        if self.state_manager.get_num_free_runtime() > 0:
-            return True
-        self.block_trie.state_checkpoints.evict(1)
-        return self.state_manager.get_num_free_runtime() > 0
-
     def _long_context_chunk_limit(self, seq: SchedulerSequence):
         """Return the token budget for one long-context chunk."""
         return get_long_context_chunk_limit(
@@ -827,7 +819,7 @@ def reserve_long_context_chunk(
             prealloc_size = 0
 
         evictable = stopped + waiting
-        if not self.eviction_helper.evict_for_seq(
+        if not self.eviction_helper.try_make_capacity_for(
             seq,
             evictable,
             prealloc_size,
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 7bb2aaca5c..b5bae8b72b 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -265,7 +265,11 @@ def schedule_migration(self):
                     reversed(self.hanging),
                     reversed(migration_waiting),
                 ))
-            if not self.eviction_helper.evict_for_seq(seq, evictable, 0):
+            if not self.eviction_helper.try_make_capacity_for(
+                seq,
+                evictable,
+                0,
+            ):
                 break
 
             # allocate session memory
@@ -325,7 +329,11 @@ def schedule_running(self, running: SeqList, num_required_tokens: int = 1, preal
             if num_required_blocks == 0:
                 continue
 
-            if eviction_helper.evict_for_seq(seq, self.hanging + self.waiting, prealloc_size):
+            if eviction_helper.try_make_capacity_for(
+                seq,
+                self.hanging + self.waiting,
+                prealloc_size,
+            ):
                 self.block_manager.allocate(seq, prealloc_size)
                 self.block_trie.allocate(seq)
                 continue
diff --git a/lmdeploy/pytorch/paging/seq_states/states.py b/lmdeploy/pytorch/paging/seq_states/states.py
index 657e4177bd..f93e90689b 100644
--- a/lmdeploy/pytorch/paging/seq_states/states.py
+++ b/lmdeploy/pytorch/paging/seq_states/states.py
@@ -44,7 +44,7 @@ def add_sequence(self, seq: SchedulerSequence, status: MessageStatus) -> None:
     def remove_sequence(self, seq: SchedulerSequence) -> None:
         """Release local ownership without a terminal connector event."""
         assert seq.seq_id in seq.session.sequences
-        self.free_sequence(seq)
+        self.release_paging_resources(seq)
         seq.session.sequences.pop(seq.seq_id)
         self._seq_manager.unregister_sequence(seq)
 
@@ -64,7 +64,8 @@ def assert_allocated(self, seq: SchedulerSequence) -> None:
         if self._is_ssm:
             assert seq.logical_state >= 0
 
-    def free_sequence(self, seq: SchedulerSequence) -> None:
+    def release_paging_resources(self, seq: SchedulerSequence) -> None:
+        """Release blocks and state without changing sequence status."""
         if self._prefix_cache_enabled:
             self._state_checkpoints.discard_save(seq)
             self._state_checkpoints.unpin_restore(seq)
@@ -128,9 +129,9 @@ def stop(self):
         """Stop the state."""
         self.to_state(StoppedState)
 
-    def free(self):
-        """Free the state."""
-        self.lifecycle.free_sequence(self.seq)
+    def release_paging_resources(self):
+        """Release blocks and state without changing sequence status."""
+        self.lifecycle.release_paging_resources(self.seq)
 
     def begin_remote_load(self):
         raise NotImplementedError(f'begin_remote_load not implemented for state {self.status}')
diff --git a/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py b/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
index 13eae3d708..47920af6a2 100644
--- a/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
+++ b/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
@@ -72,7 +72,7 @@ def test_ssm_restore_pin_survives_tail_allocation(self, ssm_scheduler):
         assert checkpoint_node.state_checkpoint.pin_count == 0
         assert seq.prefix_cache.restore.node is None
 
-    def test_free_clears_unpinned_ssm_restore(self, ssm_scheduler):
+    def test_release_paging_resources_clears_unpinned_ssm_restore(self, ssm_scheduler):
         block_trie = ssm_scheduler.block_trie
         block_size = ssm_scheduler.seq_meta.block_size
         checkpoint_tokens = [1] * block_size * 2
@@ -85,7 +85,7 @@ def test_free_clears_unpinned_ssm_restore(self, ssm_scheduler):
         assert seq.prefix_cache.restore.node is checkpoint_node
         assert not seq.prefix_cache.restore.pinned
 
-        seq.state.free()
+        seq.state.release_paging_resources()
 
         assert not seq.prefix_cache.restore.is_selected
         assert seq.prefix_cache.restore.node is None
diff --git a/tests/pytorch/paging/test_block_trie/test_trie.py b/tests/pytorch/paging/test_block_trie/test_trie.py
index 7fc130b8a6..331dcda351 100644
--- a/tests/pytorch/paging/test_block_trie/test_trie.py
+++ b/tests/pytorch/paging/test_block_trie/test_trie.py
@@ -387,7 +387,7 @@ def test_match_after_sequence_blocks_are_freed(self, block_trie, block_mgr, sche
 
         block_mgr.allocate(seq)
         block_trie.allocate(seq)
-        seq.state.free()
+        seq.state.release_paging_resources()
 
         assert seq.num_history_ids == 0
         assert len(seq.logical_blocks) == 0
diff --git a/tests/pytorch/paging/test_prefill_scheduler.py b/tests/pytorch/paging/test_prefill_scheduler.py
index 9a2ea705e1..c7bd6d1060 100644
--- a/tests/pytorch/paging/test_prefill_scheduler.py
+++ b/tests/pytorch/paging/test_prefill_scheduler.py
@@ -101,7 +101,7 @@ def test_scheduler_prefix_match_rollback_clears_recompute_overlap_window(monkeyp
     cached.state.stop()
 
     seq = scheduler.add_session(1).add_sequence(token_ids)
-    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', Mock(return_value=False))
+    monkeypatch.setattr(scheduler.eviction_helper, 'try_make_capacity_for', Mock(return_value=False))
     scheduler.block_trie.stats.reset()
 
     output = scheduler.schedule(is_prefill=True)
@@ -255,8 +255,8 @@ def test_scheduler_resource_rejection_rolls_back_tentative_prefix_match(monkeypa
     scheduler.block_trie.stats.reset()
 
     seq = scheduler.add_session(1).add_sequence([1] * block_size + [3])
-    evict_for_seq = Mock(return_value=False)
-    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
+    try_make_capacity = Mock(return_value=False)
+    monkeypatch.setattr(scheduler.eviction_helper, 'try_make_capacity_for', try_make_capacity)
 
     output = scheduler.schedule(is_prefill=True)
 
@@ -268,7 +268,7 @@ def test_scheduler_resource_rejection_rolls_back_tentative_prefix_match(monkeypa
     assert seq.cached_tokens == 0
     assert seq.prefix_cache.trie_cursor is None
     assert seq.prefix_cache.match_start_step == -1
-    assert evict_for_seq.call_count == 1
+    assert try_make_capacity.call_count == 1
     assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
     assert scheduler.block_trie.stats.num_query_tokens == 0
     assert scheduler.block_trie.stats.num_hit_tokens == 0
@@ -363,7 +363,7 @@ def test_scheduler_cached_tokens_only_count_current_prompt_after_session_evictio
     scheduler.schedule(is_prefill=True)
     seq.update_token_ids(torch.tensor([9]), mode=UpdateTokenMode.PREFILL)
     seq.state.stop()
-    seq.state.free()
+    seq.state.release_paging_resources()
 
     seq.update_token_ids(torch.tensor([4] * 4))
     assert seq.input_start_pos == block_size * 2 + 2
@@ -400,7 +400,7 @@ def test_scheduler_excludes_recompute_eviction_prefix_hits_from_stats():
     pressure = scheduler.add_session(1).add_sequence([9] * block_size * 3)
     scheduler.block_trie.stats.reset()
 
-    assert scheduler.eviction_helper.evict_for_seq(pressure, [seq], 0)
+    assert scheduler.eviction_helper.try_make_capacity_for(pressure, [seq], 0)
     assert seq.prefix_cache.suppress_match_stats
     pressure.session.remove_sequence(pressure)
 
diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py
index 9a30a4c6e6..fd7423f050 100644
--- a/tests/pytorch/paging/test_scheduler_kv_transfer.py
+++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py
@@ -505,8 +505,8 @@ def test_async_load_capacity_failure_restores_tentative_local_prefix(monkeypatch
     scheduler.block_trie.stats.reset()
 
     seq = scheduler.add_session(77).add_sequence(tokens)
-    evict_for_seq = Mock(return_value=False)
-    monkeypatch.setattr(scheduler.eviction_helper, 'evict_for_seq', evict_for_seq)
+    try_make_capacity = Mock(return_value=False)
+    monkeypatch.setattr(scheduler.eviction_helper, 'try_make_capacity_for', try_make_capacity)
 
     output = scheduler.schedule(is_prefill=True)
 
@@ -520,7 +520,7 @@ def test_async_load_capacity_failure_restores_tentative_local_prefix(monkeypatch
     assert seq.prefix_cache.match_start_step == -1
     assert connector.lookup_calls == [(seq.seq_id, 4)]
     assert connector.allocations == []
-    assert evict_for_seq.call_count == 1
+    assert try_make_capacity.call_count == 1
     assert scheduler.block_manager.allocator.get_ref_count(cached_block).tolist() == ref_count.tolist()
     assert scheduler.block_trie.stats.num_query_tokens == 0
     assert scheduler.block_trie.stats.num_hit_tokens == 0
diff --git a/tests/pytorch/paging/test_scheduler_ssm.py b/tests/pytorch/paging/test_scheduler_ssm.py
index 74938d4c51..de01ded958 100644
--- a/tests/pytorch/paging/test_scheduler_ssm.py
+++ b/tests/pytorch/paging/test_scheduler_ssm.py
@@ -326,9 +326,16 @@ def test_ssm_scheduler_rejects_prefix_match_for_prefill_gate_after_runtime_state
     def _make_runtime_state_available_once_then_succeed():
         return next(ensure_results)
 
-    monkeypatch.setattr(scheduler._prefill_scheduler,
-                        '_make_runtime_state_available',
-                        _make_runtime_state_available_once_then_succeed)
+    monkeypatch.setattr(
+        scheduler.eviction_helper,
+        'try_make_capacity_for',
+        lambda *args: True,
+    )
+    monkeypatch.setattr(
+        scheduler.block_trie.state_checkpoints,
+        'make_runtime_state_available',
+        _make_runtime_state_available_once_then_succeed,
+    )
     scheduler.block_trie.stats.reset()
 
     cache_hit_tail = scheduler.add_session(100).add_sequence([1] * block_size * 2 + [3])
@@ -373,7 +380,9 @@ def _make_ssm_scheduler_for_long_context_chunks(num_gpu_blocks: int = 2):
     return scheduler, block_size
 
 
-def test_schedule_prefill_reapplies_chunk_limit_after_ssm_state_rollback():
+def test_schedule_prefill_reapplies_chunk_limit_after_ssm_state_rollback(
+    monkeypatch,
+):
     scheduler, block_size = _make_ssm_scheduler_for_long_context_chunks(num_gpu_blocks=2)
     long_seq = scheduler.add_session(100).add_sequence([1] * (block_size * 4))
 
@@ -382,8 +391,16 @@ def test_schedule_prefill_reapplies_chunk_limit_after_ssm_state_rollback():
     def _make_runtime_state_available_once_then_succeed():
         return next(ensure_results)
 
-    scheduler._prefill_scheduler._make_runtime_state_available = (
-        _make_runtime_state_available_once_then_succeed)
+    monkeypatch.setattr(
+        scheduler.eviction_helper,
+        'try_make_capacity_for',
+        lambda *args: True,
+    )
+    monkeypatch.setattr(
+        scheduler.block_trie.state_checkpoints,
+        'make_runtime_state_available',
+        _make_runtime_state_available_once_then_succeed,
+    )
 
     output = scheduler.schedule(is_prefill=True, prealloc_size=1)
 

From 0c7b785557fd9961f6e7ce2af1a92d0e198f3484 Mon Sep 17 00:00:00 2001
From: grimoire <yaoqian@pjlab.org.cn>
Date: Sun, 30 Aug 2026 21:44:06 +0800
Subject: [PATCH 22/22] refactor: consolidate scheduler call boundaries

---
 lmdeploy/pytorch/engine/engine.py             |   4 +-
 lmdeploy/pytorch/engine/engine_loop.py        |  26 +--
 lmdeploy/pytorch/engine/inputs_maker.py       | 113 +++----------
 lmdeploy/pytorch/messages.py                  |  20 ++-
 .../paging/block_trie/checkpoint_lifecycle.py |  96 +++++++++++
 .../pytorch/paging/kv_load_coordinator.py     |   2 +-
 lmdeploy/pytorch/paging/scheduler.py          |  46 +++--
 lmdeploy/pytorch/strategies/ar/sequence.py    |   2 +-
 .../pytorch/strategies/ar_spec/sequence.py    |   2 +-
 lmdeploy/pytorch/strategies/dllm/sequence.py  |   2 +-
 tests/pytorch/engine/test_inputs_maker.py     | 158 ++++++++++--------
 .../test_checkpoint_lifecycle.py              |  48 ++++++
 tests/pytorch/paging/test_scheduler_ssm.py    |  15 +-
 13 files changed, 311 insertions(+), 223 deletions(-)

diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py
index 96c7a2aa0b..c74d2d04d6 100644
--- a/lmdeploy/pytorch/engine/engine.py
+++ b/lmdeploy/pytorch/engine/engine.py
@@ -416,7 +416,7 @@ def _on_end_session(self, reqs: list[Request], **kwargs):
             if session is not None:
                 msgs = list(session.sequences.values())
                 if len(msgs) > 0 and msgs[0].preserve_cache:
-                    msgs[0].state.finish()
+                    msgs[0].finish()
                 else:
                     self.end_session(session_id)
                 resp_type = ResponseType.SUCCESS
@@ -511,7 +511,7 @@ def __update_max_new_tokens(msg):
                     mode=UpdateTokenMode.INPUTS,
                 )
                 msg.sampling_param = sampling_param
-                msg.state.activate()
+                msg.activate()
 
             __update_max_new_tokens(msg)
             msg.resp = req.resp
diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py
index 55ca40ab98..343c873fb1 100644
--- a/lmdeploy/pytorch/engine/engine_loop.py
+++ b/lmdeploy/pytorch/engine/engine_loop.py
@@ -406,14 +406,12 @@ async def _main_loop_try_send_next_inputs(self):
         if self._sleep_requested:
             return None, None
 
-        self.scheduler.resume_completed_migrations()
         return await self.inputs_maker.send_next_inputs()
 
     async def _prefetch_next_inputs(self):
-        """Resume completed migrations before prefetching the next batch."""
+        """Prefetch the next batch unless sleep has started."""
         if self._sleep_requested:
             return None, None
-        self.scheduler.resume_completed_migrations()
         return await self.inputs_maker.prefetch_next_inputs()
 
     async def _wait_for_schedulable_prefill(self):
@@ -430,17 +428,6 @@ async def _wait_for_schedulable_prefill(self):
                        f'gpu cache usage: {cache_usage:.1%}')
         await asyncio.sleep(0.1)
 
-    def _publish_forward_checkpoints(self, running: 'SeqList', has_state_checkpoint_save: bool):
-        """Publish per-forward prefix-cache ownership before prefetching."""
-        state_checkpoints = self.state_checkpoints
-        if has_state_checkpoint_save:
-            state_checkpoints.publish_saves(running, pin_saves=True)
-        state_checkpoints.unpin_restores(running)
-
-    def _release_forward_save_pins(self, running: 'SeqList'):
-        """Unpin producers after the forward output/event boundary."""
-        self.state_checkpoints.unpin_saves(running)
-
     def _finish_forward_output(self,
                                out: 'BatchedOutputs | None',
                                running: 'SeqList',
@@ -482,11 +469,14 @@ async def _main_loop_get_outputs(
         # for GPU output; save checkpoints keep a producer pin until the output
         # event boundary so prefetch cannot evict/reuse their destination slots.
         if has_model_work:
-            self._publish_forward_checkpoints(running, has_state_checkpoint_save)
+            self.state_checkpoints.finish_forward_dispatch(
+                running,
+                has_save_plan=has_state_checkpoint_save,
+            )
         forward_inputs, next_running = await self._prefetch_next_inputs()
         out = await self.executor.get_output_async()
         if has_model_work:
-            self._release_forward_save_pins(running)
+            self.state_checkpoints.unpin_saves(running)
         self._finish_forward_output(out, running, model_inputs, delta)
         # out might come from shared memory, need to explicitly delete to release memory in time
         del out
@@ -529,7 +519,7 @@ async def main_loop(self):
                 running=next_running,
                 forward_inputs=forward_inputs,
             )
-            self.inputs_maker.deactivate_evict_seqs()
+            self.inputs_maker.preempt_invalid_decode_seqs()
             has_runable_event.set()
 
     def update_running_migration(self, running: 'SeqList', next_token_ids: np.ndarray, stopped: torch.Tensor,
@@ -547,7 +537,7 @@ def update_running_migration(self, running: 'SeqList', next_token_ids: np.ndarra
             if stop:
                 update_token = _EMPTY_TOKEN
                 msg.update_token_ids(update_token, model_meta=model_meta, mode=UpdateTokenMode.PREFILL)
-                msg.state.finish()
+                msg.finish()
 
     async def _migration_loop_migrate(self, migration_ready: 'SeqList'):
         """Migration loop migrate."""
diff --git a/lmdeploy/pytorch/engine/inputs_maker.py b/lmdeploy/pytorch/engine/inputs_maker.py
index 164eb1446d..5d1868fe4e 100644
--- a/lmdeploy/pytorch/engine/inputs_maker.py
+++ b/lmdeploy/pytorch/engine/inputs_maker.py
@@ -8,6 +8,7 @@
 """
 import logging
 from collections import defaultdict
+from collections.abc import Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
@@ -58,32 +59,12 @@ def _tensorlize_block_offsets(block_offsets, dtype=torch.int32):
     return torch.as_tensor(out, dtype=dtype)
 
 
-def _make_state_prefix_cache_restore_plan(
-        messages: list['SchedulerSequence']) -> StateCacheCopyPlan | None:
-    """Build a compact host SSM state-restore plan."""
-    src_offsets = []
-    dst_offsets = []
-    for msg in messages:
-        restore = msg.prefix_cache.restore
-        if restore.is_selected:
-            src_offsets.append(restore.slot)
-            dst_offsets.append(msg.logical_state)
-    if len(src_offsets) == 0:
-        return None
-    return tuple(src_offsets), tuple(dst_offsets)
-
-
-def _make_state_prefix_cache_save_plan(messages: list['SchedulerSequence'],
-                                       save_state_offsets: list[int]) -> StateCacheCopyPlan | None:
-    """Build a compact host SSM state-save plan."""
-    src_offsets = []
-    dst_offsets = []
-    for msg, state_idx in zip(messages, save_state_offsets):
-        if state_idx >= 0:
-            src_offsets.append(msg.logical_state)
-            dst_offsets.append(state_idx)
-    if len(src_offsets) == 0:
+def _make_state_checkpoint_copy_plan(
+        pairs: Sequence[tuple[int, int]]) -> StateCacheCopyPlan | None:
+    """Transpose logical owner pairs into the compact engine carrier."""
+    if len(pairs) == 0:
         return None
+    src_offsets, dst_offsets = zip(*pairs)
     return tuple(src_offsets), tuple(dst_offsets)
 
 
@@ -890,7 +871,7 @@ def _map_to_kernel_block_offsets(self, block_offsets: torch.Tensor):
         return block_offsets
 
     def _make_kv_prefix_cache_copy_plan(
-            self, logical_pairs: list[tuple[int, int]]) -> torch.LongTensor:
+            self, logical_pairs: Sequence[tuple[int, int]]) -> torch.LongTensor:
         """Resolve paging ids and build one host KV-block copy plan."""
         logical_ids = np.asarray(logical_pairs, dtype=np.int64).reshape(-1, 2)
         block_offsets = self.scheduler.resolve_gpu_block_offsets(logical_ids.reshape(-1))
@@ -910,33 +891,14 @@ def _ssm_prefix_cache_enabled(self):
     def _prepare_prefill_cache_restore(
             self, messages: 'SeqList') -> tuple[torch.LongTensor | None, StateCacheCopyPlan | None]:
         """Acquire checkpoints and build prefill restore plans."""
-        state_restore_plan = _make_state_prefix_cache_restore_plan(messages)
+        copy_plan = self.state_checkpoints.prepare_restore_batch(messages)
+        state_restore_plan = _make_state_checkpoint_copy_plan(copy_plan.state_pairs)
         if state_restore_plan is None:
             return None, None
 
-        state_checkpoints = self.state_checkpoints
-        # Keep checkpoint sources alive while the prefetched forward waits to
-        # copy them into request-owned KV and runtime state.
-        state_checkpoints.pin_restores(messages)
-        if any(msg.prefix_cache.restore.is_selected and not msg.prefix_cache.restore.pinned for msg in messages):
-            raise RuntimeError('Failed to acquire SSM prefix-cache restore checkpoint.')
-
-        logical_pairs = []
-        for msg in messages:
-            restore = msg.prefix_cache.restore
-            if not restore.is_selected:
-                continue
-            checkpoint = restore.node.state_checkpoint
-            if checkpoint.frozen_block_id < 0:
-                continue
-            dst_block_idx = checkpoint.step // self.config.block_size
-            if dst_block_idx >= len(msg.logical_blocks):
-                raise RuntimeError('SSM prefix-cache restore destination block is missing.')
-            logical_pairs.append((checkpoint.frozen_block_id, msg.logical_blocks[dst_block_idx]))
-
         kv_restore_plan = None
-        if logical_pairs:
-            kv_restore_plan = self._make_kv_prefix_cache_copy_plan(logical_pairs)
+        if copy_plan.kv_block_pairs:
+            kv_restore_plan = self._make_kv_prefix_cache_copy_plan(copy_plan.kv_block_pairs)
         return kv_restore_plan, state_restore_plan
 
     def _prepare_prefill_cache_save(
@@ -945,30 +907,11 @@ def _prepare_prefill_cache_save(
         save_steps: tuple[int, ...] | None,
     ) -> tuple[torch.LongTensor | None, StateCacheCopyPlan | None]:
         """Reserve checkpoints and build prefill save plans."""
-        state_checkpoints = self.state_checkpoints
-        if save_steps is None:
-            save_state_offsets = [state_checkpoints.reserve_save(msg) for msg in messages]
-        else:
-            save_state_offsets = [state_checkpoints.reserve_save(msg, step=step)
-                                  for msg, step in zip(messages, save_steps)]
-        state_save_plan = _make_state_prefix_cache_save_plan(messages, save_state_offsets)
-
-        logical_pairs = []
-        for msg, state_idx in zip(messages, save_state_offsets):
-            if state_idx < 0:
-                continue
-            pending_save = msg.prefix_cache.pending_save
-            checkpoint = pending_save.node.state_checkpoint
-            if checkpoint.frozen_block_id < 0:
-                continue
-            src_block_idx = pending_save.step // self.config.block_size
-            if src_block_idx >= len(msg.logical_blocks):
-                raise RuntimeError('SSM prefix-cache save source block is missing.')
-            logical_pairs.append((msg.logical_blocks[src_block_idx], checkpoint.frozen_block_id))
-
+        copy_plan = self.state_checkpoints.reserve_prefill_save_batch(messages, save_steps)
+        state_save_plan = _make_state_checkpoint_copy_plan(copy_plan.state_pairs)
         kv_save_plan = None
-        if logical_pairs:
-            kv_save_plan = self._make_kv_prefix_cache_copy_plan(logical_pairs)
+        if copy_plan.kv_block_pairs:
+            kv_save_plan = self._make_kv_prefix_cache_copy_plan(copy_plan.kv_block_pairs)
         return kv_save_plan, state_save_plan
 
     def _prepare_prefill_cache_inputs(self,
@@ -1000,10 +943,8 @@ def _make_decode_cache_inputs(self, valid_seqs: 'SeqList', delta: ModelInputsDel
         if (decode_state_interval <= 0 or self.spec_decoding or delta.max_q_seqlen != 1):
             return None
 
-        state_checkpoints = self.state_checkpoints
-        save_state_offsets = [state_checkpoints.reserve_decode_save(seq, decode_state_interval)
-                              for seq in valid_seqs]
-        state_save_plan = _make_state_prefix_cache_save_plan(valid_seqs, save_state_offsets)
+        copy_plan = self.state_checkpoints.reserve_decode_save_batch(valid_seqs, decode_state_interval)
+        state_save_plan = _make_state_checkpoint_copy_plan(copy_plan.state_pairs)
         if state_save_plan is None:
             return None
         return CacheCheckpointInputs(state_save_plan=state_save_plan)
@@ -1257,16 +1198,12 @@ def update_running_seqs(self, running: 'SeqList', inputs: 'ModelInputs|None'):
         else:
             self.running_seqs += running
 
-    def deactivate_evict_seqs(self):
-        """Deactivate and evict seqs."""
-        scheduler = self.scheduler
+    def preempt_invalid_decode_seqs(self):
+        """Return decode sequences rejected during prefetch to waiting."""
         to_evict_seqs = self.to_evict_seqs
         if len(to_evict_seqs) == 0:
             return
-        # deactivate seqs(running -> ready)
-        scheduler.deactivate_seqs(to_evict_seqs)
-        # ready to waiting
-        scheduler.evict_seqs(to_evict_seqs)
+        self.scheduler.preempt_seqs(to_evict_seqs)
         self.to_evict_seqs.clear()
 
     @torch.inference_mode()
@@ -1321,7 +1258,9 @@ def do_prefill_chunked(self):
         scheduler = self.scheduler
         return not scheduler.has_ready()
 
-    async def _send_next_inputs_impl(self, prefill: bool = None, enable_empty: bool = False):
+    async def _send_next_inputs_impl(self, enable_empty: bool = False):
+        self.scheduler.resume_completed_migrations()
+        prefill = self.do_prefill()
         forward_inputs = self._make_forward_inputs(prefill, enable_empty)
         if forward_inputs is None:
             return None, None
@@ -1338,14 +1277,12 @@ async def _send_next_inputs_impl(self, prefill: bool = None, enable_empty: bool
         return forward_inputs, next_running
 
     async def send_next_inputs(self):
-        prefill = self.do_prefill()
-        return await self._send_next_inputs_impl(prefill)
+        return await self._send_next_inputs_impl()
 
     async def prefetch_next_inputs(self):
-        prefill = self.do_prefill()
         # send next forward
         logger.debug('Prefetching next forward inputs.')
-        return await self._send_next_inputs_impl(prefill, True)
+        return await self._send_next_inputs_impl(enable_empty=True)
 
 
 def build_inputs_maker(engine: 'Engine'):
diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py
index d7b03f7927..1a3bc1f213 100644
--- a/lmdeploy/pytorch/messages.py
+++ b/lmdeploy/pytorch/messages.py
@@ -286,7 +286,7 @@ def __init__(self, session_id: int, seq_meta: SequenceMeta, lifecycle: 'Sequence
         self.session_id = session_id
         self.seq_meta = seq_meta
         self.sequences: SeqMap = dict()
-        self.lifecycle = lifecycle
+        self._lifecycle = lifecycle
 
     def add_sequence(self,
                      token_ids: Tensor,
@@ -301,7 +301,7 @@ def add_sequence(self,
         if sampling_param is None:
             sampling_param = SamplingParam()
 
-        seq_id = self.lifecycle.new_sequence_id()
+        seq_id = self._lifecycle.new_sequence_id()
         seq = self.seq_meta.strategy.make_sequence(seq_id=seq_id,
                                                    session=self,
                                                    sampling_param=sampling_param,
@@ -316,7 +316,7 @@ def add_sequence(self,
             mode=UpdateTokenMode.INPUTS,
         )
         status = MessageStatus.WAITING if migration_request is None else MessageStatus.MIGRATION_WAITING
-        self.lifecycle.add_sequence(seq, status)
+        self._lifecycle.add_sequence(seq, status)
 
         # metrics
         seq.record_event(EventType.QUEUED)
@@ -325,7 +325,11 @@ def add_sequence(self,
 
     def remove_sequence(self, seq: 'SchedulerSequence'):
         """Remove sequence."""
-        self.lifecycle.remove_sequence(seq)
+        self._lifecycle.remove_sequence(seq)
+
+    def end_sequence(self, seq: 'SchedulerSequence') -> None:
+        """Notify terminal completion and release the sequence."""
+        self._lifecycle.end_sequence(seq)
 
 
 def _div_up(x, n):
@@ -804,6 +808,14 @@ def set_state(self, state: 'StateBase'):
     def status(self):
         return self.state.status
 
+    def activate(self) -> None:
+        """Advance the sequence from its current resumable state."""
+        self.state.activate()
+
+    def finish(self) -> None:
+        """Finish the sequence's current running lifecycle."""
+        self.state.finish()
+
     @property
     def return_logits(self):
         return self.sampling_param.out_logits
diff --git a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
index ffac74c1d6..6b0e47bea6 100644
--- a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
+++ b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
@@ -14,6 +14,7 @@
 import heapq
 import time
 from collections.abc import Callable, Iterable
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -32,6 +33,14 @@
 logger = get_logger('lmdeploy')
 
 
+@dataclass(frozen=True)
+class CheckpointCopyPlan:
+    """Logical state and KV copies owned by one checkpoint batch."""
+
+    state_pairs: tuple[tuple[int, int], ...] = ()
+    kv_block_pairs: tuple[tuple[int, int], ...] = ()
+
+
 class StateCheckpointLifecycle:
     """Manage node-owned state and optional frozen partial-KV checkpoints.
 
@@ -61,6 +70,93 @@ def __init__(self,
         self._index = index
         self._snapshot_match_data = snapshot_match_data
 
+    def prepare_restore_batch(self, seqs: list[SchedulerSequence]) -> CheckpointCopyPlan:
+        """Pin selected restores and expose their logical copy pairs."""
+        self.pin_restores(seqs)
+        state_pairs = []
+        kv_block_pairs = []
+        for seq in seqs:
+            restore = seq.prefix_cache.restore
+            if not restore.is_selected:
+                continue
+            if not restore.pinned:
+                raise RuntimeError('Failed to acquire SSM prefix-cache restore checkpoint.')
+
+            state_pairs.append((int(restore.slot), int(seq.logical_state)))
+            checkpoint = restore.node.state_checkpoint
+            if checkpoint.frozen_block_id < 0:
+                continue
+            dst_block_idx = checkpoint.step // self._block_size
+            if dst_block_idx >= len(seq.logical_blocks):
+                raise RuntimeError('SSM prefix-cache restore destination block is missing.')
+            kv_block_pairs.append((
+                int(checkpoint.frozen_block_id),
+                int(seq.logical_blocks[dst_block_idx]),
+            ))
+        return CheckpointCopyPlan(
+            state_pairs=tuple(state_pairs),
+            kv_block_pairs=tuple(kv_block_pairs),
+        )
+
+    def reserve_prefill_save_batch(
+        self,
+        seqs: list[SchedulerSequence],
+        steps: tuple[int, ...] | None = None,
+    ) -> CheckpointCopyPlan:
+        """Reserve prefill checkpoints and expose their logical copy pairs."""
+        if steps is not None and len(steps) != len(seqs):
+            raise ValueError('steps must have one entry per prefill sequence.')
+        if steps is None:
+            state_offsets = [self.reserve_save(seq) for seq in seqs]
+        else:
+            state_offsets = [self.reserve_save(seq, step=step) for seq, step in zip(seqs, steps)]
+
+        state_pairs = []
+        kv_block_pairs = []
+        for seq, state_idx in zip(seqs, state_offsets):
+            if state_idx < 0:
+                continue
+            state_pairs.append((int(seq.logical_state), int(state_idx)))
+            pending_save = seq.prefix_cache.pending_save
+            checkpoint = pending_save.node.state_checkpoint
+            if checkpoint.frozen_block_id < 0:
+                continue
+            src_block_idx = pending_save.step // self._block_size
+            if src_block_idx >= len(seq.logical_blocks):
+                raise RuntimeError('SSM prefix-cache save source block is missing.')
+            kv_block_pairs.append((
+                int(seq.logical_blocks[src_block_idx]),
+                int(checkpoint.frozen_block_id),
+            ))
+        return CheckpointCopyPlan(
+            state_pairs=tuple(state_pairs),
+            kv_block_pairs=tuple(kv_block_pairs),
+        )
+
+    def reserve_decode_save_batch(
+        self,
+        seqs: list[SchedulerSequence],
+        interval: int,
+    ) -> CheckpointCopyPlan:
+        """Reserve replaceable decode checkpoints for one forward batch."""
+        state_pairs = []
+        for seq in seqs:
+            state_idx = self.reserve_decode_save(seq, interval)
+            if state_idx >= 0:
+                state_pairs.append((int(seq.logical_state), int(state_idx)))
+        return CheckpointCopyPlan(state_pairs=tuple(state_pairs))
+
+    def finish_forward_dispatch(
+        self,
+        seqs: list[SchedulerSequence],
+        *,
+        has_save_plan: bool,
+    ) -> None:
+        """Publish queued saves and release restore pins before prefetch."""
+        if has_save_plan:
+            self.publish_saves(seqs, pin_saves=True)
+        self.unpin_restores(seqs)
+
     def reserve_save(self, seq: SchedulerSequence, step: int = None, is_decode: bool = False):
         """Reserve a checkpoint at an exact safe prefill boundary."""
         self.discard_save(seq)
diff --git a/lmdeploy/pytorch/paging/kv_load_coordinator.py b/lmdeploy/pytorch/paging/kv_load_coordinator.py
index d9df94a525..b3a2fb9210 100644
--- a/lmdeploy/pytorch/paging/kv_load_coordinator.py
+++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py
@@ -549,7 +549,7 @@ def finish_deferred_loads_after_worker_drain(self) -> None:
 
     def _finish_deferred_end(self, seq: SchedulerSequence) -> None:
         session = seq.session
-        session.lifecycle.end_sequence(seq)
+        session.end_sequence(seq)
         if not session.sequences:
             self.sessions.pop(session.session_id, None)
 
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index b5bae8b72b..3bd3ed6ef8 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -3,6 +3,7 @@
 """Public paging scheduler and sequence-lifecycle facade."""
 
 from collections import OrderedDict
+from collections.abc import Iterable
 from contextlib import contextmanager
 from dataclasses import dataclass
 from itertools import chain
@@ -338,9 +339,7 @@ def schedule_running(self, running: SeqList, num_required_tokens: int = 1, preal
                 self.block_trie.allocate(seq)
                 continue
 
-            seq.state.deactivate()
-            self.kv_load_coordinator.release_tracking(seq)
-            seq.state.evict()
+            self.preempt_seqs((seq, ))
             valid_mask[idx] = False
         return valid_mask
 
@@ -383,7 +382,7 @@ def end_session(self, session_id: int):
                 continue
             # stop session so it won't get scheduled again
             seq.state.stop()
-            self.sequence_lifecycle.end_sequence(seq)
+            session.end_sequence(seq)
         if not session.sequences:
             self.sessions.pop(session_id)
 
@@ -498,39 +497,32 @@ def resolve_gpu_block_offsets(self, logical_block_ids):
         """Resolve paging-owned logical ids for a forward cache-copy plan."""
         return self.block_manager.resolve_gpu_block_offsets(logical_block_ids)
 
-    def evict_seqs(self, running: SeqList):
-        """Evict running sequences."""
+    def activate_seqs(self, running: SeqList):
+        """Mark a ready batch as running at the engine dispatch boundary."""
         for seq in running:
-            self.kv_load_coordinator.release_tracking(seq)
-            seq.state.evict()
-
-    def activate_seqs(self, running: SeqList, filter_status: MessageStatus = MessageStatus.READY):
-        """Lock running sequence."""
-        for seq in running:
-            if seq.status == filter_status:
+            if seq.status == MessageStatus.READY:
                 seq.state.activate()
 
-    def deactivate_seqs(self, running: SeqList, filter_status: MessageStatus = MessageStatus.RUNNING):
-        for seq in running:
-            if seq.status == filter_status:
+    def preempt_seqs(self, seqs: Iterable[SchedulerSequence]) -> None:
+        """Return invalid decode sequences to their evictable queue states."""
+        for seq in seqs:
+            if seq.status == MessageStatus.RUNNING:
                 seq.state.deactivate()
-
-    def activate_migration_seqs(self, running: SeqList):
-        """Lock running sequence."""
-        return self.activate_seqs(running, filter_status=MessageStatus.MIGRATION_READY)
-
-    def deactivate_migration_seqs(self, running: SeqList):
-        """Unlock running migration."""
-        return self.deactivate_seqs(running, filter_status=MessageStatus.MIGRATION_RUNNING)
+            self.kv_load_coordinator.release_tracking(seq)
+            seq.state.evict()
 
     @contextmanager
     def seqs_migration_activation(self, running: SeqList):
-        """Context manager to activate and deactivate sequences."""
-        self.activate_migration_seqs(running)
+        """Keep a migration batch running only while applying its output."""
+        for seq in running:
+            if seq.status == MessageStatus.MIGRATION_READY:
+                seq.state.activate()
         try:
             yield running
         finally:
-            self.deactivate_migration_seqs(running)
+            for seq in running:
+                if seq.status == MessageStatus.MIGRATION_RUNNING:
+                    seq.state.deactivate()
 
     def resume_completed_migrations(self):
         """Move completed migration sequences back to the waiting queue."""
diff --git a/lmdeploy/pytorch/strategies/ar/sequence.py b/lmdeploy/pytorch/strategies/ar/sequence.py
index 0d39b7e0ae..a0e800e806 100644
--- a/lmdeploy/pytorch/strategies/ar/sequence.py
+++ b/lmdeploy/pytorch/strategies/ar/sequence.py
@@ -148,4 +148,4 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode
             # fill token
             msg.update_token_ids(token, model_meta=model_meta, mode=update_mode, routed_experts=routed_experts)
             if stop:
-                msg.state.finish()
+                msg.finish()
diff --git a/lmdeploy/pytorch/strategies/ar_spec/sequence.py b/lmdeploy/pytorch/strategies/ar_spec/sequence.py
index 6494cb7dcf..78d2c45913 100644
--- a/lmdeploy/pytorch/strategies/ar_spec/sequence.py
+++ b/lmdeploy/pytorch/strategies/ar_spec/sequence.py
@@ -247,4 +247,4 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode
                                  routed_experts=routed_experts,
                                  stop_pos=stop_pos[idx])
             if stop:
-                msg.state.finish()
+                msg.finish()
diff --git a/lmdeploy/pytorch/strategies/dllm/sequence.py b/lmdeploy/pytorch/strategies/dllm/sequence.py
index 71d5639f94..4b3c243ed3 100644
--- a/lmdeploy/pytorch/strategies/dllm/sequence.py
+++ b/lmdeploy/pytorch/strategies/dllm/sequence.py
@@ -275,4 +275,4 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode
             msg.update_token_ids(token, dllm_mask=mask, model_meta=model_meta, mode=update_mode)
             if stop:
                 msg.set_stop_pos(stop_pos[idx])
-                msg.state.finish()
+                msg.finish()
diff --git a/tests/pytorch/engine/test_inputs_maker.py b/tests/pytorch/engine/test_inputs_maker.py
index 588c1c8e36..8393e00891 100644
--- a/tests/pytorch/engine/test_inputs_maker.py
+++ b/tests/pytorch/engine/test_inputs_maker.py
@@ -16,12 +16,12 @@
     InputsMakerAsync,
     InputsMakerConfig,
     LongContextChunker,
-    _make_state_prefix_cache_restore_plan,
-    _make_state_prefix_cache_save_plan,
+    _make_state_checkpoint_copy_plan,
 )
 from lmdeploy.pytorch.engine.model_agent.agent import BatchedOutputs
 from lmdeploy.pytorch.kv_connector import KVConnectorOutput
 from lmdeploy.pytorch.messages import MessageStatus, StateCheckpointRestore, StateCheckpointSaveReservation
+from lmdeploy.pytorch.paging.block_trie.checkpoint_lifecycle import CheckpointCopyPlan
 
 
 @dataclass
@@ -133,6 +133,9 @@ def __init__(self, running, waiting=None, num_ready=0, num_running=0):
     def has_kv_connector(self):
         return self.kv_connector is not None
 
+    def resume_completed_migrations(self):
+        pass
+
     def schedule(self,
                  is_prefill: bool,
                  prealloc_size: int,
@@ -217,14 +220,11 @@ def test_engine_loop_keeps_state_save_pinned_until_output_boundary():
     class _StateCheckpoints:
         pinned = False
 
-        def publish_saves(self, seqs, pin_saves=False):
-            events.append(('publish_saves', pin_saves))
-            assert pin_saves
+        def finish_forward_dispatch(self, seqs, *, has_save_plan):
+            events.append(('finish_dispatch', has_save_plan))
+            assert has_save_plan
             self.pinned = True
 
-        def unpin_restores(self, seqs):
-            events.append(('unpin_restores', self.pinned))
-
         def unpin_saves(self, seqs):
             events.append(('unpin_saves', self.pinned))
             self.pinned = False
@@ -267,8 +267,7 @@ async def get_output_async(self):
     assert next_running is None
     assert events == [
         'update_running',
-        ('publish_saves', True),
-        ('unpin_restores', True),
+        ('finish_dispatch', True),
         ('prefetch', True),
         ('get_output', True),
         ('unpin_saves', True),
@@ -311,13 +310,10 @@ def test_engine_loop_skips_prefetch_when_sleep_requested_but_unpins_state_save()
     class _StateCheckpoints:
         pinned = False
 
-        def publish_saves(self, seqs, pin_saves=False):
-            events.append(('publish_saves', pin_saves))
+        def finish_forward_dispatch(self, seqs, *, has_save_plan):
+            events.append(('finish_dispatch', has_save_plan))
             self.pinned = True
 
-        def unpin_restores(self, seqs):
-            events.append(('unpin_restores', self.pinned))
-
         def unpin_saves(self, seqs):
             events.append(('unpin_saves', self.pinned))
             self.pinned = False
@@ -353,8 +349,7 @@ async def get_output_async(self):
     assert next_running is None
     assert events == [
         'update_running',
-        ('publish_saves', True),
-        ('unpin_restores', True),
+        ('finish_dispatch', True),
         'get_output',
         ('unpin_saves', True),
     ]
@@ -369,9 +364,6 @@ class _Scheduler:
         def has_unfinished(self):
             return False
 
-        def resume_completed_migrations(self):
-            events.append('resume_completed_migrations')
-
     class _InputsMaker:
 
         def has_pending_long_context_chunk(self):
@@ -390,7 +382,56 @@ async def send_next_inputs(self):
     result = asyncio.run(asyncio.wait_for(loop._main_loop_try_send_next_inputs(), timeout=1.0))
 
     assert result == ('forward_inputs', ['long-seq'])
-    assert events == ['resume_completed_migrations', 'send_next_inputs']
+    assert events == ['send_next_inputs']
+
+
+@pytest.mark.parametrize(
+    ('method_name', 'enable_empty'),
+    [
+        ('send_next_inputs', False),
+        ('prefetch_next_inputs', True),
+    ],
+)
+def test_inputs_maker_resumes_completed_migrations_before_selecting_work(method_name, enable_empty):
+    events = []
+
+    class _Scheduler:
+
+        def resume_completed_migrations(self):
+            events.append('resume_migrations')
+
+        def tick(self):
+            events.append('tick')
+
+    class _Executor:
+
+        async def forward_async(self, forward_inputs):
+            events.append('forward_async')
+
+    def do_prefill():
+        events.append('select_work')
+        return True
+
+    def make_forward_inputs(prefill, enable_empty=False):
+        events.append(('make_forward_inputs', prefill, enable_empty))
+        return dict(running=['seq'], inputs=None, delta=None)
+
+    maker = InputsMakerAsync.__new__(InputsMakerAsync)
+    maker.scheduler = _Scheduler()
+    maker.executor = _Executor()
+    maker.do_prefill = do_prefill
+    maker._make_forward_inputs = make_forward_inputs
+
+    result = asyncio.run(getattr(maker, method_name)())
+
+    assert result == ({'inputs': None, 'delta': None}, ['seq'])
+    assert events == [
+        'resume_migrations',
+        'select_work',
+        ('make_forward_inputs', True, enable_empty),
+        'forward_async',
+        'tick',
+    ]
 
 
 def test_migration_loop_schedules_and_processes_ready_batch():
@@ -1373,22 +1414,11 @@ def test_do_prefill_default_forces_pending_last_chunk_prefill():
     assert maker.do_prefill_default()
 
 
-def test_state_prefix_cache_restore_plan_is_compact():
-    messages = [_state_seq(4, 11), _state_seq(5, -1), _state_seq(6, 13)]
-
-    plan = _make_state_prefix_cache_restore_plan(messages)
+def test_state_checkpoint_copy_plan_is_compact():
+    plan = _make_state_checkpoint_copy_plan(((11, 4), (13, 6)))
 
     assert plan == ((11, 13), (4, 6))
-    assert _make_state_prefix_cache_restore_plan([_state_seq(4)]) is None
-
-
-def test_state_prefix_cache_save_plan_is_compact():
-    messages = [_state_seq(4), _state_seq(5), _state_seq(6)]
-
-    plan = _make_state_prefix_cache_save_plan(messages, [-1, 21, 22])
-
-    assert plan == ((5, 6), (21, 22))
-    assert _make_state_prefix_cache_save_plan(messages, [-1, -1, -1]) is None
+    assert _make_state_checkpoint_copy_plan(()) is None
 
 
 def test_prepare_prefill_cache_inputs_groups_state_restore_and_save_plans():
@@ -1397,20 +1427,16 @@ def test_prepare_prefill_cache_inputs_groups_state_restore_and_save_plans():
 
     class _StateCheckpoints:
 
-        def pin_restores(self, seqs):
+        def prepare_restore_batch(self, seqs):
             events.append('pin_restores')
             for seq in seqs:
                 if seq.prefix_cache.restore.is_selected:
                     seq.prefix_cache.restore.pinned = True
+            return CheckpointCopyPlan(state_pairs=((11, 4), ))
 
-        def reserve_save(self, seq, step=None):
-            assert step is None
-            state_idx = {4: 21, 5: -1}[seq.logical_state]
-            if state_idx >= 0:
-                checkpoint = SimpleNamespace(step=0, frozen_block_id=-1)
-                node = SimpleNamespace(state_checkpoint=checkpoint)
-                seq.prefix_cache.pending_save.reserve(state_idx, 0, node, False)
-            return state_idx
+        def reserve_prefill_save_batch(self, seqs, steps=None):
+            assert steps is None
+            return CheckpointCopyPlan(state_pairs=((4, 21), ))
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
     maker.config = SimpleNamespace(is_ssm=True, enable_prefix_caching=True)
@@ -1430,16 +1456,12 @@ def test_prepare_prefill_cache_inputs_uses_explicit_chunk_end_step():
 
     class _StateCheckpoints:
 
-        def pin_restores(self, seqs):
-            for seq in seqs:
-                seq.prefix_cache.restore.pinned = True
+        def prepare_restore_batch(self, seqs):
+            return CheckpointCopyPlan(state_pairs=((11, 4), ))
 
-        def reserve_save(self, seq, step=None):
-            reserve_steps.append(step)
-            checkpoint = SimpleNamespace(step=step, frozen_block_id=-1)
-            node = SimpleNamespace(state_checkpoint=checkpoint)
-            seq.prefix_cache.pending_save.reserve(21, step, node, False)
-            return 21
+        def reserve_prefill_save_batch(self, seqs, steps=None):
+            reserve_steps.extend(steps)
+            return CheckpointCopyPlan(state_pairs=((4, 21), ))
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
     maker.config = SimpleNamespace(is_ssm=True, enable_prefix_caching=True)
@@ -1454,24 +1476,20 @@ def reserve_save(self, seq, step=None):
 
 def test_prepare_prefill_cache_inputs_groups_partial_kv_restore_and_save_plans():
     messages = [_state_seq(4, 11), _state_seq(5, 12)]
-    for msg, dst_block in zip(messages, (20, 21)):
-        checkpoint = SimpleNamespace(step=17, frozen_block_id=70)
-        msg.prefix_cache.restore.node = SimpleNamespace(state_checkpoint=checkpoint)
-        msg.logical_blocks = np.array([10, dst_block], dtype=np.int64)
 
     class _StateCheckpoints:
 
-        def pin_restores(self, seqs):
-            for seq in seqs:
-                seq.prefix_cache.restore.pinned = True
+        def prepare_restore_batch(self, seqs):
+            return CheckpointCopyPlan(
+                state_pairs=((11, 4), (12, 5)),
+                kv_block_pairs=((70, 20), (70, 21)),
+            )
 
-        def reserve_save(self, seq, step=None):
-            state_idx = {4: 21, 5: 22}[seq.logical_state]
-            frozen_block = {4: 80, 5: 81}[seq.logical_state]
-            checkpoint = SimpleNamespace(step=17, frozen_block_id=frozen_block)
-            node = SimpleNamespace(state_checkpoint=checkpoint)
-            seq.prefix_cache.pending_save.reserve(state_idx, 17, node, False)
-            return state_idx
+        def reserve_prefill_save_batch(self, seqs, steps=None):
+            return CheckpointCopyPlan(
+                state_pairs=((4, 21), (5, 22)),
+                kv_block_pairs=((20, 80), (21, 81)),
+            )
 
     scheduler = _CopyPlanScheduler()
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
@@ -1502,9 +1520,9 @@ def test_make_decode_cache_inputs_compacts_valid_state_saves():
 
     class _StateCheckpoints:
 
-        def reserve_decode_save(self, seq, interval):
+        def reserve_decode_save_batch(self, seqs, interval):
             assert interval == 16
-            return {4: 31, 5: -1}[seq.logical_state]
+            return CheckpointCopyPlan(state_pairs=((4, 31), ))
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
     maker.config = SimpleNamespace(is_ssm=True,
@@ -1535,7 +1553,7 @@ def test_make_decode_cache_inputs_respects_feature_gates(is_ssm, enabled, interv
 
     class _StateCheckpoints:
 
-        def reserve_decode_save(self, seq, interval):
+        def reserve_decode_save_batch(self, seqs, interval):
             raise AssertionError('disabled decode checkpoint path must not reserve state')
 
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
diff --git a/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py b/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
index 47920af6a2..96feb64fb7 100644
--- a/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
+++ b/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
@@ -494,6 +494,54 @@ def test_ssm_checkpoint_save_owns_and_releases_partial_tail(self, ssm_scheduler)
         assert block_mgr.get_num_free_gpu_blocks() == free_blocks
         assert ssm_scheduler.state_manager.get_num_free_checkpoint() == free_states
 
+    def test_prepare_restore_batch_owns_partial_checkpoint_copy_pairs(self, ssm_scheduler):
+        block_trie = ssm_scheduler.block_trie
+        block_size = ssm_scheduler.seq_meta.block_size
+        checkpoint_tokens = [1] * block_size * 2 + [2]
+        checkpoint_seq, checkpoint_node, state_idx = self._add_published_ssm_checkpoint(
+            ssm_scheduler,
+            checkpoint_tokens,
+        )
+        checkpoint_seq.session.remove_sequence(checkpoint_seq)
+        seq = ssm_scheduler.add_session(100).add_sequence(checkpoint_tokens + [3])
+
+        output = ssm_scheduler.schedule(is_prefill=True)
+        copy_plan = block_trie.state_checkpoints.prepare_restore_batch([seq])
+
+        checkpoint = checkpoint_node.state_checkpoint
+        assert output.running == [seq]
+        assert seq.prefix_cache.restore.pinned
+        assert copy_plan.state_pairs == ((state_idx, seq.logical_state), )
+        assert copy_plan.kv_block_pairs == ((checkpoint.frozen_block_id, seq.logical_blocks[2]), )
+
+        block_trie.state_checkpoints.finish_forward_dispatch([seq], has_save_plan=False)
+
+        assert not seq.prefix_cache.restore.is_selected
+        assert checkpoint.pin_count == 0
+
+    def test_reserve_prefill_save_batch_owns_partial_checkpoint_copy_pairs(self, ssm_scheduler):
+        block_mgr = ssm_scheduler.block_manager
+        block_trie = ssm_scheduler.block_trie
+        state_manager = ssm_scheduler.state_manager
+        block_size = ssm_scheduler.seq_meta.block_size
+        seq = ssm_scheduler.add_session(0).add_sequence([1] * block_size * 2 + [2])
+        block_mgr.allocate(seq)
+        block_trie.allocate(seq)
+        state_manager.allocate(seq)
+
+        copy_plan = block_trie.state_checkpoints.reserve_prefill_save_batch([seq])
+        checkpoint = seq.prefix_cache.pending_save.node.state_checkpoint
+
+        assert copy_plan.state_pairs == ((seq.logical_state, checkpoint.slot), )
+        assert copy_plan.kv_block_pairs == ((seq.logical_blocks[2], checkpoint.frozen_block_id), )
+
+        block_trie.state_checkpoints.finish_forward_dispatch([seq], has_save_plan=True)
+
+        assert checkpoint.published
+        assert seq.prefix_cache.producer_save_pin.is_acquired
+        block_trie.state_checkpoints.unpin_saves([seq])
+        assert checkpoint.pin_count == 0
+
     def test_ssm_checkpoint_partial_tail_allocation_failure_rolls_back_state(self, ssm_cache_config,
                                                                              scheduler_config, seq_meta):
         ssm_cache_config.num_gpu_blocks = 3
diff --git a/tests/pytorch/paging/test_scheduler_ssm.py b/tests/pytorch/paging/test_scheduler_ssm.py
index de01ded958..796565e682 100644
--- a/tests/pytorch/paging/test_scheduler_ssm.py
+++ b/tests/pytorch/paging/test_scheduler_ssm.py
@@ -3,7 +3,6 @@
 import torch
 
 from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
-from lmdeploy.pytorch.engine.inputs_maker import _make_state_prefix_cache_save_plan
 from lmdeploy.pytorch.messages import MessageStatus, SequenceMeta, UpdateTokenMode
 from lmdeploy.pytorch.paging.scheduler import Scheduler
 
@@ -123,17 +122,13 @@ def test_ssm_same_batch_duplicate_checkpoint_save_has_unique_dst_offsets():
     assert seq_a.logical_state != seq_b.logical_state
     assert seq_a.prefix_cache.trie_cursor is seq_b.prefix_cache.trie_cursor
 
-    save_state_offsets = [
-        scheduler.block_trie.state_checkpoints.reserve_save(seq) for seq in output.running
-    ]
-    save_plan = _make_state_prefix_cache_save_plan(output.running, save_state_offsets)
-    assert save_plan is not None
-    save_src_offsets, save_dst_offsets = save_plan
+    copy_plan = scheduler.block_trie.state_checkpoints.reserve_prefill_save_batch(output.running)
+    save_src_offsets, save_dst_offsets = zip(*copy_plan.state_pairs)
 
     assert save_src_offsets == (seq_a.logical_state, )
-    assert save_dst_offsets == (save_state_offsets[0], )
-    assert save_state_offsets[0] >= 0
-    assert save_state_offsets[1] == -1
+    assert save_dst_offsets == (seq_a.prefix_cache.pending_save.slot, )
+    assert save_dst_offsets[0] >= 0
+    assert not seq_b.prefix_cache.pending_save.is_pending
     assert len(save_dst_offsets) == len(set(save_dst_offsets))