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..c74d2d04d6 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,10 +412,11 @@ 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() + msgs[0].finish() else: self.end_session(session_id) resp_type = ResponseType.SUCCESS @@ -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 @@ -510,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 @@ -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: @@ -605,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)) @@ -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/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py index bd240f776e..343c873fb1 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 @@ -403,14 +406,12 @@ async def _main_loop_try_send_next_inputs(self): if self._sleep_requested: return None, None - self.scheduler.collect_migration_done() return await self.inputs_maker.send_next_inputs() async def _prefetch_next_inputs(self): - """Collect migration completions before prefetching the next batch.""" + """Prefetch the next batch unless sleep has started.""" if self._sleep_requested: return None, None - self.scheduler.collect_migration_done() return await self.inputs_maker.prefetch_next_inputs() async def _wait_for_schedulable_prefill(self): @@ -422,22 +423,11 @@ 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 - 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) - def _finish_forward_output(self, out: 'BatchedOutputs | None', running: 'SeqList', @@ -446,16 +436,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) @@ -481,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 @@ -528,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, @@ -546,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.""" @@ -558,7 +549,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 +606,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 +683,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..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 @@ -34,6 +35,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 @@ -57,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) @@ -95,6 +77,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 +105,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 +347,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 +681,7 @@ def __init__( self, executor: 'ExecutorBase', scheduler: 'Scheduler', + state_checkpoints: 'StateCheckpointLifecycle', adapter_manager: 'AdapterManager', engine_strategy: 'EngineStrategy', sampling_strategy: 'SamplingStrategy', @@ -697,11 +690,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 @@ -878,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)) @@ -893,38 +886,19 @@ 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]: """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.scheduler.block_trie.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.cache_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( @@ -933,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.scheduler.block_trie.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.cache_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, @@ -984,14 +939,12 @@ 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 - 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) @@ -1161,7 +1114,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) @@ -1245,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() @@ -1309,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 @@ -1326,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'): @@ -1342,6 +1291,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/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/messages.py b/lmdeploy/pytorch/messages.py index ab0056307f..1a3bc1f213 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 @@ -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] @@ -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,11 @@ 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 end_sequence(self, seq: 'SchedulerSequence') -> None: + """Notify terminal completion and release the sequence.""" + self._lifecycle.end_sequence(seq) def _div_up(x, n): @@ -818,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 @@ -966,7 +964,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/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/block_trie/checkpoint_lifecycle.py b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py index c5e09f19d3..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) @@ -281,6 +377,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/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/eviction_helper/__init__.py b/lmdeploy/pytorch/paging/eviction_helper/__init__.py index 6b5c44ff97..22a1834fa6 100644 --- a/lmdeploy/pytorch/paging/eviction_helper/__init__.py +++ b/lmdeploy/pytorch/paging/eviction_helper/__init__.py @@ -1,10 +1,29 @@ # 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 + from .recompute_eviction_helper import RecomputeEvictionHelper + 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, +) -> RecomputeEvictionHelper: """Build eviction helper.""" if eviction_type == 'copy': logger.warning('`copy` eviction has been deprecated, ' @@ -12,6 +31,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 deleted file mode 100644 index f075748f70..0000000000 --- a/lmdeploy/pytorch/paging/eviction_helper/base_eviction_helper.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. - -from ...messages import SchedulerSequence -from ..scheduler import Scheduler - -SeqList = list[SchedulerSequence] - - -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 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 bfa39c8919..fbc2a0290f 100644 --- a/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py +++ b/lmdeploy/pytorch/paging/eviction_helper/recompute_eviction_helper.py @@ -1,135 +1,130 @@ # Copyright (c) OpenMMLab. All rights reserved. +from __future__ import annotations -from ...messages import SchedulerSequence -from ..scheduler import Scheduler -from .base_eviction_helper import BaseEvictionHelper - - -class RecomputeEvictionHelper(BaseEvictionHelper): - """Recompute eviction.""" +from collections.abc import Iterable +from typing import TYPE_CHECKING - def __init__(self, scheduler: Scheduler): - super().__init__(scheduler) +from ...messages import SchedulerSequence - if len(self.cache_config.states_shapes) == 0: - self.evict_for_seq = self._evict_for_seq_default - else: - self.evict_for_seq = self._evict_for_ssm +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: + """Reclaim paging resources so a sequence can be recomputed.""" + + def __init__( + self, + *, + block_manager: BaseBlockManager, + block_trie: BlockTrie, + state_manager: StateManager, + load_coordinator: KVLoadCoordinator, + is_ssm: bool, + ) -> None: + 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 + + 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.scheduler.kv_load_coordinator.is_remote_ready(evict_seq): - continue - - # skip sequence with no blocks - if evict_seq.num_blocks == 0: + for evict_seq in evictable_seqs: + if not self._reclaim_candidate(evict_seq): continue - - 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.scheduler.kv_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: - 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. - """ + if self._try_make_block_capacity(num_required_blocks): + return True + + return self._try_make_block_capacity(num_required_blocks) + + 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.scheduler.kv_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.scheduler.kv_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: - 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 - return success + 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 + + 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 6bc485d610..b3a2fb9210 100644 --- a/lmdeploy/pytorch/paging/kv_load_coordinator.py +++ b/lmdeploy/pytorch/paging/kv_load_coordinator.py @@ -1,15 +1,16 @@ # 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 - 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. @@ -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.recompute_eviction_helper import RecomputeEvictionHelper + + +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): @@ -50,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. @@ -57,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: @@ -84,8 +118,20 @@ 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, + sessions: dict[int, SchedulerSession], + ) -> None: + self.lookup_enabled = lookup_enabled + self.connector = connector + self.block_manager = block_manager + self.block_trie = block_trie + 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 +150,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 +168,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 +199,168 @@ 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], + eviction_helper: RecomputeEvictionHelper, + ) -> 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, + eviction_helper=eviction_helper, + ) + + def _admit_load( + self, + seq: SchedulerSequence, + *, + num_external_tokens: int, + prealloc_size: int, + evictable_seqs: Iterable[SchedulerSequence], + eviction_helper: RecomputeEvictionHelper, + ) -> KVLoadAdmission: + """Admit the complete prefill, then allocate the remote interval.""" + 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, + eviction_helper, + ) + 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, + # 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 None + + 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], + 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.try_make_capacity_for( + seq, + list(evictable_seqs), + prealloc_size, + ) + if not full_prefill_fits: + seq.kv_token_limit = plan.original_kv_token_limit + return KVLoadAdmission.FULL_PREFILL_UNAVAILABLE + 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 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 = 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] + ) + connector.update_state_after_alloc( + seq, + load_block_ids, + 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=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 = plan.original_kv_token_limit + raise + seq.kv_token_limit = None + def start_load( self, seq: SchedulerSequence, @@ -166,10 +371,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: @@ -188,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 @@ -205,7 +409,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: @@ -217,20 +422,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 +446,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 +459,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``.""" @@ -265,12 +468,12 @@ 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: - self._remove_sequence(seq) - elif record.stop_requested: + if record.deferred_cleanup is _DeferredLoadCleanup.END: + 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 @@ -279,12 +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 - 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: + 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, @@ -292,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.end_requested = True + 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 @@ -311,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 @@ -325,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. @@ -336,22 +540,24 @@ 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) self._finish_cancelled_or_failed(record) - def _remove_sequence(self, seq: SchedulerSequence) -> None: - scheduler = self.scheduler - connector = scheduler.kv_connector - if connector is not None: - connector.request_finished(seq) - + def _finish_deferred_end(self, seq: SchedulerSequence) -> None: session = seq.session - session.remove_sequence(seq) + session.end_sequence(seq) if not session.sequences: - scheduler.sessions.pop(session.session_id, None) + self.sessions.pop(session.session_id, None) + + def shutdown(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/kv_save_coordinator.py b/lmdeploy/pytorch/paging/kv_save_coordinator.py index a1e7130060..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. """ @@ -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') @@ -59,10 +59,10 @@ 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.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 new file mode 100644 index 0000000000..97b08f111a --- /dev/null +++ b/lmdeploy/pytorch/paging/prefill_scheduler.py @@ -0,0 +1,892 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""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 +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.recompute_eviction_helper import RecomputeEvictionHelper + +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 _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.""" + + def __init__(self, prefill_scheduler: '_PrefillScheduler'): + self.prefill_scheduler = prefill_scheduler + self._info_cache: dict[int, _PrefillReorderInfo] = {} + + def reorder(self, + waiting: SeqList, + 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 + # 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 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. + 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 turn_policy.allows_nonfinal_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.""" + + 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_started(cls): + return cls(action=_PrefillAdmissionAction.LOAD_STARTED) + + +@dataclass(frozen=True, slots=True) +class _PrefixMatchStateSnapshot: + """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. + 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 _PrefixMatchPhase(enum.Enum): + """Lifecycle of one request-local tentative prefix transaction.""" + + IDLE = enum.auto() + TRACKING = enum.auto() + MATCHED = enum.auto() + + +class _TentativePrefixMatch: + """Manage one request's tentative ``BlockTrie.match`` side effects.""" + + __slots__ = ( + 'seq', + 'block_trie', + 'block_manager', + 'is_ssm', + '_preserve_existing_state', + '_stats_snapshot', + '_state_snapshot', + '_phase', + ) + + 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._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. + + 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._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._phase = _PrefixMatchPhase.TRACKING + + def match(self) -> None: + """Apply one tentative match after capturing its rollback boundary.""" + assert not self.is_matched + self.begin() + assert self._phase is _PrefixMatchPhase.TRACKING + self.block_trie.match(self.seq) + self._phase = _PrefixMatchPhase.MATCHED + + 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) -> None: + """Undo the tentative prefix-cache transaction.""" + if self._phase is _PrefixMatchPhase.IDLE: + return + + 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() + + 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.release_paging_resources() + 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._phase = _PrefixMatchPhase.IDLE + + +class _PrefillAdmissionAttempt: + """Own one waiting sequence's tentative admission side effects.""" + + def __init__(self, + prefill_scheduler: '_PrefillScheduler', + seq: SchedulerSequence, + stopped: SeqList, + evictable_waiting: SeqList, + prealloc_size: int, + batch_prefill_tokens: int, + batch_has_prefill: bool, + turn_policy: _PrefillTurnPolicy): + self.prefill_scheduler = prefill_scheduler + self.seq = seq + self.stopped = stopped + self.evictable_waiting = evictable_waiting + self.prealloc_size = prealloc_size + self.batch_prefill_tokens = batch_prefill_tokens + 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.turn_policy = turn_policy + 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._load_ready), + ) + self._resource_rollback_rejection: _PrefillAdmissionResult | None = None + + def run(self): + """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() + + gate_result = self._apply_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._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: + load_result = self._try_external_load() + if load_result is not None: + return load_result + if not self._prepare_and_make_capacity(): + return _PrefillAdmissionResult.stop() + return None + + def _admit_prefix_cache_resources(self): + """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.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(): + self._prefix_match.match() + + if self.load_coordinator.lookup_enabled: + 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(): + gate_rejection = self._rollback_match_after_resource_failure( + 'failed to pin SSM restore checkpoint') + if gate_rejection is not None: + return gate_rejection + + kv_result = self._admit_kv_resources(had_ssm_restore) + if kv_result is not None: + return kv_result + return self._admit_runtime_state() + + def _admit_kv_resources(self, had_ssm_restore: bool): + """Evict for KV, retrying once without a pinned SSM restore.""" + if self._prepare_and_make_capacity(): + return None + + reason = 'eviction failed' + if had_ssm_restore: + reason = 'eviction failed with pinned SSM restore' + 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. + if had_ssm_restore and self._prepare_and_make_capacity(): + return None + return _PrefillAdmissionResult.stop() + + def _admit_runtime_state(self): + """Admit an SSM runtime slot, retrying after match rollback.""" + prefill = self.prefill_scheduler + seq = self.seq + 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_make_capacity(): + return _PrefillAdmissionResult.stop() + if not state_checkpoints.make_runtime_state_available(): + seq.kv_token_limit = None + return _PrefillAdmissionResult.stop() + return None + + def _try_external_load(self): + """Map connector/paging admission to prefill queue policy.""" + prefill = self.prefill_scheduler + if self._load_ready: + return None + admission = self.load_coordinator.try_load( + 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 + 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_started() + 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.stopped) + yield from reversed(self.evictable_waiting) + + 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 False + if not prefill.block_trie.enabled: + return False + if self._has_private_local_tail(): + 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. + + 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 + seq = self.seq + return seq.num_blocks > int(seq.num_history_ids) // seq.block_size + + def _token_budget_rejection(self): + if self.turn_policy.allows_nonfinal_long_prefill: + return _PrefillAdmissionResult.stop() + return _PrefillAdmissionResult.skip() + + 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: + 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 + if (self.turn_policy.allows_nonfinal_long_prefill + or prefill._prefill_kv_token_limit(seq) is None): + return 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._accept_gate_enabling_match( + _PrefillAdmissionResult.skip()) + return None + + 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 + + rejection = self._token_budget_rejection() + if self._prefix_match.is_matched: + self._prefix_match.rollback('still exceeds prefill token budget') + return rejection + + if not self._try_match_prefix_for_prefill_gate(): + 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 rejection + + self._accept_gate_enabling_match(rejection) + return None + + 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._try_make_capacity(alloc_size): + return True + seq.kv_token_limit = None + return False + + 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.try_make_capacity_for( + self.seq, + list(self._evictable_sequences()), + alloc_size, + ) + + def _allocate_and_commit(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._effective_prealloc_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._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. + self.load_coordinator.mark_prefill_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: 'RecomputeEvictionHelper', + 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 _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, + *, + stopped: 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 = stopped + waiting + if not self.eviction_helper.try_make_capacity_for( + 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, + stopped: SeqList, + num_ready: int, + num_running: int, + turn_policy: _PrefillTurnPolicy, + prealloc_size: int = 0, + ): + """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 = [] + batch_prefill_tokens = 0 + + if max_batches <= 0 or not waiting: + return running + + waiting = _PrefillReorderer(self).reorder( + waiting, + turn_policy=turn_policy, + ) + skipped_waiting: SeqList = [] + while waiting and len(running) < max_batches: + seq = waiting.pop(0) + evictable_waiting = skipped_waiting + waiting + admission = _PrefillAdmissionAttempt( + self, + seq, + stopped=stopped, + evictable_waiting=evictable_waiting, + prealloc_size=prealloc_size, + batch_prefill_tokens=batch_prefill_tokens, + batch_has_prefill=bool(running), + turn_policy=turn_policy, + ).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 + 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: + break + + return running diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py index d5f2450ddc..3bd3ed6ef8 100644 --- a/lmdeploy/pytorch/paging/scheduler.py +++ b/lmdeploy/pytorch/paging/scheduler.py @@ -1,87 +1,34 @@ # 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 time -from collections import Counter, OrderedDict +"""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 typing import TYPE_CHECKING, Any +from itertools import chain +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 ..kv_connector import KVConnectorStepInput 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 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: from lmdeploy.pytorch.kv_connector.base import KVConnectorBase -logger = get_logger('lmdeploy') + from .block_trie.checkpoint_lifecycle import StateCheckpointLifecycle MapType = dict[int, int] SeqList = list[SchedulerSequence] @@ -89,7 +36,7 @@ @dataclass class SchedulerOutput: - """Output of schedule.""" + """Paging selection and connector snapshots for one model step.""" running: SeqList swap_in_map: MapType @@ -107,728 +54,12 @@ 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, scheduler: 'Scheduler'): - self.scheduler = 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.scheduler.kv_load_coordinator.is_remote_ready(seq) - ] - waiting = [ - seq for seq in waiting - if not self.scheduler.kv_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 - - scheduler = self.scheduler - chunk_limit = scheduler._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) - 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.""" - scheduler = self.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_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): - scheduler = self.scheduler - if scheduler._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 - - -@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 - waiting queue for asynchronous KV load. - """ - - admitted: bool - 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) - - @classmethod - def skip(cls): - return cls(admitted=False, should_skip=True) - - @classmethod - def stop(cls): - return cls(admitted=False) - - @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 - - -@dataclass(frozen=True) -class _PrefixMatchBaseline: - """Exact pre-attempt state used to undo 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 - # Matching mutates global hit statistics, so those are transactional too. - stats_snapshot: Any - - -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 scheduler loop still owns queue traversal and decides - whether a rejected candidate is skipped or ends the current prefill turn. - """ - - def __init__(self, - scheduler: 'Scheduler', - seq: SchedulerSequence, - evictable_waiting: SeqList, - prealloc_size: int, - token_count: int, - has_admitted: bool, - allow_long_prefill: bool): - self.scheduler = scheduler - self.seq = seq - self.evictable_waiting = evictable_waiting - 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.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 - - 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.scheduler._external_lookup_enabled and not self._remote_ready: - if self._lookup_is_pending(): - return self._check_result(_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) - - resource_result = self._admit_resources() - if resource_result is not None: - return self._check_result(resource_result) - - return self._check_result(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 _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. - - 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. - """ - scheduler = self.scheduler - connector = scheduler.kv_connector - assert connector is not None - if not connector.is_lookup_pending(self.seq.seq_id): - return False - 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() - if self.scheduler._external_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. - """ - 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() - # 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) - - if scheduler._external_lookup_enabled: - lookup_result = self._query_external_prefix() - if lookup_result is not None: - 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 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') - 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._rollback_gate(stats_snapshot, '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') - if result is not None: - return result - if not self._prepare_and_evict(): - return _PrefillAdmissionResult.stop() - if not scheduler._ensure_runtime_state_available(): - seq.kv_token_limit = None - return _PrefillAdmissionResult.stop() - - 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. - """ - 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( - self.seq, - self.seq.num_history_ids, - ) - if num_external_tokens is not None: - if num_external_tokens > 0: - 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, - ) - 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. - """ - 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 - 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) - 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 - return _PrefillAdmissionResult.load() - - 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 - or self._has_private_local_tail()): - return None - stats_snapshot = scheduler.block_trie.stats.snapshot() - scheduler.block_trie.match(self.seq) - return stats_snapshot - - 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. - """ - scheduler = self.scheduler - if not scheduler._external_lookup_enabled: - return False - 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() - return _PrefillAdmissionResult.skip() - - def _check_prefill_admission_gates(self): - """Apply prefill gates, tentatively matching only when it may help.""" - scheduler = self.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 - - if is_nonfinal_long_prefill and not self.allow_long_prefill: - stats_snapshot = self._match_prefix_for_prefill_gate() - if stats_snapshot 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') - return _PrefillAdmissionResult.skip() - self._keep_gate_prefix_match(stats_snapshot, _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: - 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()) - return None - self._rollback_prefix_match(stats_snapshot, 'still exceeds prefill token budget') - else: - self._rollback_prefix_match(self._gate_match_stats_snapshot, 'still exceeds prefill token budget') - return self._token_budget_rejection() - - def _prepare_and_evict(self): - """Apply chunk allocation limits and evict for this prefill.""" - scheduler = self.scheduler - seq = self.seq - alloc_size = scheduler._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.""" - 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) - - 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 - # 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._finish_prefix_cache_schedule(seq) - scheduler.kv_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. - scheduler.kv_load_coordinator.mark_scheduled(seq) - return _PrefillAdmissionResult.admit(prefill_token_count) - - 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__( @@ -842,9 +73,10 @@ def __init__( self.cache_config = cache_config self.sessions: dict[int, SchedulerSession] = OrderedDict() self.kv_connector = kv_connector - - # For Disaggregation - self.locked_sessions: dict[int, SchedulerSession] = OrderedDict() + 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) @@ -853,33 +85,58 @@ 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 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, 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 + # candidates from its caller; it does not reach back through Scheduler. + self.kv_load_coordinator = KVLoadCoordinator( + lookup_enabled=external_lookup_enabled, + connector=kv_connector, + block_manager=self.block_manager, + block_trie=self.block_trie, + 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, + 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.block_manager) - 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 - 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).""" @@ -893,148 +150,20 @@ def shutdown(self) -> None: """ connector = self.kv_connector self.kv_connector = None - self._external_lookup_enabled = False - self.kv_load_coordinator.clear() + self.sequence_lifecycle.disable_connector() + self.kv_load_coordinator.shutdown() self.kv_save_coordinator.clear() 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 _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.""" - 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) - - 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) + @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 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, @@ -1042,76 +171,64 @@ 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 + return self._prefill_scheduler.reserve_long_context_chunk( + seq, + stopped=self.hanging, + waiting=self.waiting, + chunk_size=chunk_size, + prealloc_size=prealloc_size, + is_last_chunk=is_last_chunk, + ) + + # 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. - 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 + # Sequence views. + @property + def waiting(self) -> SeqList: + return list(self.seq_manager.get_sequences(MessageStatus.WAITING).values()) - self.block_manager.allocate(seq, prealloc_size) - self.block_trie.allocate(seq) - return True + @property + def hanging(self) -> SeqList: + return list(self.seq_manager.get_sequences(MessageStatus.STOPPED).values()) - @staticmethod - def create_status_list_property(status: MessageStatus): - """Create status list property.""" + @property + def migration_waiting(self) -> SeqList: + return list(self.seq_manager.get_sequences(MessageStatus.MIGRATION_WAITING).values()) - def _get_status_list(self): - seq_map = self.seq_manager.get_sequences(status) - return list(seq_map.values()) + @property + def migration_done(self) -> SeqList: + return list(self.seq_manager.get_sequences(MessageStatus.MIGRATION_DONE).values()) - return property(_get_status_list) + # Sequence counts. + def num_waiting(self) -> int: + return self.seq_manager.num_sequences(MessageStatus.WAITING) - @staticmethod - def create_num_status_method(status: MessageStatus): - """Create num status method.""" + def num_remote_loading(self) -> int: + return self.seq_manager.num_sequences(MessageStatus.WAITING_FOR_REMOTE_KVS) - def _num_status(self): - return self.seq_manager.num_sequences(status) + def num_ready(self) -> int: + return self.seq_manager.num_sequences(MessageStatus.READY) - return _num_status + def num_running(self) -> int: + return self.seq_manager.num_sequences(MessageStatus.RUNNING) - @staticmethod - def create_has_status_method(status: MessageStatus): - """Create has status method.""" + # 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_status(self): - return self.seq_manager.num_sequences(status) > 0 + def has_remote_loading(self) -> bool: + return self.seq_manager.num_sequences(MessageStatus.WAITING_FOR_REMOTE_KVS) > 0 - return _has_status + def has_ready(self) -> bool: + return self.seq_manager.num_sequences(MessageStatus.READY) > 0 - # 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 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. @@ -1120,221 +237,110 @@ 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 _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 get_session(self, session_id: int) -> SchedulerSession | None: + """Return one session owner, if it exists.""" + return self.sessions.get(session_id) - def _reorder_migrating(): - """Reorder waiting.""" - return sorted(self.migration_waiting, key=lambda seq: seq.arrive_time) + def get_sessions(self) -> list[SchedulerSession]: + """Return a snapshot of current session owners.""" + return list(self.sessions.values()) - migration_waiting = _reorder_migrating() + def schedule_migration(self): + """Admit waiting migration sequences to paging resources.""" + migration_ready: SeqList = [] + 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.try_make_capacity_for( + seq, + evictable, + 0, + ): break # allocate session memory self.block_manager.allocate(seq) - self._finish_prefix_cache_schedule(seq) - _to_running(seq) - - 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.""" + self.block_trie.finalize_match(seq) 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.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: - skipped_waiting.append(seq) - continue - if admission.should_stop: - break - - _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.""" - - 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) + migration_ready.append(seq) - return self.ready[:self.scheduler_config.max_batches], swap_in_map, swap_out_map, copy_map + return migration_ready 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: - output = self._schedule_prefill(prealloc_size, allow_long_prefill, prefer_long_prefill) - else: - output = self._schedule_decoding(prealloc_size) - running, swap_in_map, swap_out_map, copy_map = output + """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') + + 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, + ) + 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 - 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) 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 - # running to ready - seq.state.deactivate() - # ready to waiting - self.kv_load_coordinator.release(seq) - seq.state.evict() + self.preempt_seqs((seq, )) valid_mask[idx] = False - valid_mask = list(reversed(valid_mask)) return valid_mask def stop_session(self, session_id: int): @@ -1352,7 +358,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() @@ -1372,13 +378,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() - if connector is not None: - connector.request_finished(seq) - session.remove_sequence(seq) + session.end_sequence(seq) if not session.sequences: self.sessions.pop(session_id) @@ -1427,7 +431,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 {}, @@ -1436,13 +440,17 @@ 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. 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. @@ -1452,14 +460,15 @@ 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_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 @@ -1475,54 +484,48 @@ 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) - 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(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() - - @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) - - 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 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/__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..f93e90689b 100644 --- a/lmdeploy/pytorch/paging/seq_states/states.py +++ b/lmdeploy/pytorch/paging/seq_states/states.py @@ -1,28 +1,88 @@ # 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.register_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.release_paging_resources(seq) + seq.session.sequences.pop(seq.seq_id) + self._seq_manager.unregister_sequence(seq) + + 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) + + 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 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) + 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 +95,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.""" @@ -70,9 +129,9 @@ def stop(self): """Stop the state.""" self.to_state(StoppedState) - def free(self): - """Free the state.""" - _free_seq(self.seq, self.scheduler) + 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}') @@ -87,10 +146,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 +257,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/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_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..70e6e19a01 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,7 +34,13 @@ def end_session(self, session_id): self.ended_sessions.append(session_id) self.sessions.pop(session_id) - def finish_deferred_kv_transfers_after_worker_drain(self): + def get_session(self, session_id): + return self.sessions.get(session_id) + + def get_sessions(self): + return list(self.sessions.values()) + + def finish_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_inputs_maker.py b/tests/pytorch/engine/test_inputs_maker.py index 985231e4b8..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 @@ -130,6 +130,12 @@ 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 resume_completed_migrations(self): + pass + def schedule(self, is_prefill: bool, prealloc_size: int, @@ -186,6 +192,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, @@ -203,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 @@ -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(resume_completed_migrations=lambda: None) + loop.state_checkpoints = state_checkpoints loop.inputs_maker = _InputsMaker(state_checkpoints) loop.executor = _Executor(state_checkpoints) loop._sleep_requested = False @@ -253,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), @@ -297,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 @@ -323,9 +333,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(resume_completed_migrations=lambda: None) + loop.state_checkpoints = state_checkpoints loop.inputs_maker = _InputsMaker() loop.executor = _Executor() loop._sleep_requested = True @@ -339,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), ] @@ -355,9 +364,6 @@ class _Scheduler: def has_unfinished(self): return False - def collect_migration_done(self): - events.append('collect_migration_done') - class _InputsMaker: def has_pending_long_context_chunk(self): @@ -376,42 +382,121 @@ 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 == ['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(): + 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 = [] - 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(): @@ -451,12 +536,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(), @@ -469,12 +555,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(), @@ -486,11 +573,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(), @@ -1326,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(): @@ -1350,25 +1427,20 @@ 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) - 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) @@ -1384,21 +1456,16 @@ 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) - 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, )) @@ -1409,31 +1476,26 @@ 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() - 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) @@ -1447,8 +1509,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=()) @@ -1459,15 +1520,15 @@ 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) - 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) @@ -1492,14 +1553,14 @@ 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) - 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/engine/test_kv_connector_wiring.py b/tests/pytorch/engine/test_kv_connector_wiring.py index aa7e1eb7e7..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,23 +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() - scheduler = Scheduler.__new__(Scheduler) - scheduler.kv_connector = connector - scheduler.kv_load_coordinator = Mock() - scheduler.kv_save_coordinator = Mock() - - scheduler.shutdown() - scheduler.shutdown() - - connector.shutdown.assert_called_once_with() - assert scheduler.kv_load_coordinator.clear.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/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_block_trie/test_checkpoint_lifecycle.py b/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py index 13eae3d708..96feb64fb7 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 @@ -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_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 new file mode 100644 index 0000000000..c7bd6d1060 --- /dev/null +++ b/tests/pytorch/paging/test_prefill_scheduler.py @@ -0,0 +1,661 @@ +# 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, 'try_make_capacity_for', 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_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) + + 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]) + 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) + + 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 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 + + +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.release_paging_resources() + + 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.try_make_capacity_for(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 4275c80366..58d6b9b8ae 100644 --- a/tests/pytorch/paging/test_scheduler.py +++ b/tests/pytorch/paging/test_scheduler.py @@ -1,116 +1,14 @@ -import time -from unittest.mock import Mock +# Copyright (c) OpenMMLab. All rights reserved. + +from types import SimpleNamespace import pytest import torch -import lmdeploy.pytorch.paging.scheduler as 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: @@ -163,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) @@ -210,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 @@ -255,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 @@ -272,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 @@ -284,307 +183,68 @@ 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 + 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'), + [ + ('waiting', MessageStatus.WAITING), + ('hanging', MessageStatus.STOPPED), + ('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), + ('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 = [] - seq_a.state.stop() - seq_b = scheduler.add_session(101).add_sequence([2] * block_size) + def num_sequences(actual_status): + queried.append(actual_status) + return 3 - output = scheduler.schedule(is_prefill=True) + monkeypatch.setattr(scheduler.seq_manager, 'num_sequences', num_sequences) - assert output.running == [seq_b] - assert seq_b.logical_state >= 0 - assert seq_a.logical_state == -1 - assert seq_a.status == MessageStatus.STOPPED + assert getattr(scheduler, name)() == expected + assert queried == [status] def test_schedule_migration_matches_current_sequence(): @@ -608,1400 +268,68 @@ 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 -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(): - 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) - - 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) - - scheduler._rollback_unscheduled_prefix_match(seq, stats_snapshot) - - 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_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 == [] + 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_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, '_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(scheduler_module._envs, 'opt_ttft_policy', 'fifo') - monkeypatch.setattr(scheduler_module._envs, 'opt_ttft_aging_sec', 0.25) - - 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 - - -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' - 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._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._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..fd7423f050 --- /dev/null +++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py @@ -0,0 +1,873 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from unittest.mock import Mock + +import pytest +import torch + +from lmdeploy.messages import KVTransferConfig +from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig +from lmdeploy.pytorch.kv_connector import ( + KVConnectorMetadata, + KVConnectorOutput, + KVConnectorResult, + KVConnectorStepInput, + 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.new_requests = [] + self.lookup_calls = [] + self.cancelled = [] + self.finished = [] + self.allocations = [] + self.shutdown_calls = 0 + + def on_new_request(self, request): + self.new_requests.append(request.seq_id) + + 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, step_input): + 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): + self.shutdown_calls += 1 + + +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_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_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) + 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) + 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) + + 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 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 + + +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( + 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_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, + enable_prefix_caching=False, + ) + 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 + 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_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 + self.step_input = None + + def build_connector_meta(self, step_input): + self.step_input = step_input + 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 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() + + # 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 step_input: metadata + + scheduler.build_connector_meta([seq], connector_token_lens=(8, )) + scheduler.block_manager.free(seq) + 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() diff --git a/tests/pytorch/paging/test_scheduler_ssm.py b/tests/pytorch/paging/test_scheduler_ssm.py new file mode 100644 index 0000000000..796565e682 --- /dev/null +++ b/tests/pytorch/paging/test_scheduler_ssm.py @@ -0,0 +1,405 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +import torch + +from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig +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 + + 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 == (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)) + + +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 _make_runtime_state_available_once_then_succeed(): + return next(ensure_results) + + 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]) + + 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( + 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)) + + ensure_results = iter([False, True]) + + def _make_runtime_state_available_once_then_succeed(): + return next(ensure_results) + + 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) + + 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()