diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index a59ef78..390b3e8 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -199,6 +199,11 @@ def _free_kv_cache_tensor(self, tensor: WorkerTensor) -> None: if worker is not None and worker.initialized: worker.free_tensor(tensor) + def materialize_kv_page_view(self, model_id: str): + """Return a page-contiguous view over runner-owned NPU KV cache.""" + worker = self._worker_for_runtime(self._kv_cache_runtime_name()) + return self.materialize_worker_page_view(model_id, worker) + @staticmethod def _validate_kv_cache_bounds( model: RuntimeModel, diff --git a/python/cli/main.py b/python/cli/main.py index 32342b9..8a77706 100644 --- a/python/cli/main.py +++ b/python/cli/main.py @@ -79,6 +79,12 @@ def build_parser() -> argparse.ArgumentParser: default=True, help="Enable chunked prefill (default: True). Use --no-enable-chunked-prefill to disable.", ) + parser.add_argument( + "--max-cpu-offload-blocks", + type=_non_negative_int, + default=0, + help="Maximum number of KV blocks that can be offloaded to CPU. Use 0 to disable CPU offload.", + ) # Misc parser.add_argument( @@ -114,6 +120,7 @@ def build_serving_engine_config(args: argparse.Namespace) -> EngineConfig: long_prefill_token_threshold=args.long_prefill_token_threshold, enable_prefix_cache=args.enable_prefix_caching, enable_chunk_prefill=args.enable_chunked_prefill, + max_cpu_offload_blocks=args.max_cpu_offload_blocks, ) @@ -144,6 +151,13 @@ def _build_executor_kwargs() -> dict[str, object]: return executor_kwargs +def _non_negative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("value must be non-negative") + return parsed + + def run_serve( config: EngineConfig, *, @@ -189,6 +203,7 @@ async def shutdown(): print(f" Chunked prefill threshold: {config.long_prefill_token_threshold}") print(f" Prefix cache: {'enabled' if config.enable_prefix_cache else 'disabled'}") print(f" Chunk prefill: {'enabled' if config.enable_chunk_prefill else 'disabled'}") + print(f" CPU KV offload blocks: {config.max_cpu_offload_blocks}") print(" Endpoints: /v1/completions, /v1/chat/completions, /v1/models, /health") uvicorn.run(app, host=host, port=port, log_level="info") diff --git a/python/core/async_engine.py b/python/core/async_engine.py index 5374389..e8db372 100644 --- a/python/core/async_engine.py +++ b/python/core/async_engine.py @@ -49,6 +49,7 @@ class EngineConfig: # Feature flags enable_prefix_cache: bool = True enable_chunk_prefill: bool = True + max_cpu_offload_blocks: int = 0 @dataclass @@ -93,6 +94,7 @@ def __init__( num_blocks=num_blocks, block_size=block_size, enable_prefix_cache=self.config.enable_prefix_cache, + max_cpu_offload_blocks=self.config.max_cpu_offload_blocks, ) scheduler_config = SchedulerConfig( @@ -102,6 +104,7 @@ def __init__( max_seq_len=runtime.max_seq_len, enable_prefix_cache=self.config.enable_prefix_cache, enable_chunk_prefill=self.config.enable_chunk_prefill, + max_cpu_offload_blocks=self.config.max_cpu_offload_blocks, ) self.scheduler = Scheduler(config=scheduler_config, kv_cache_manager=self.kv_cache_manager) @@ -258,6 +261,7 @@ async def _engine_loop(self) -> None: self._handle_step_error(scheduler_output) continue + self._process_transfer_outputs(step_output) with profile_span( "scheduler.process_step_output", cat="scheduler", @@ -305,6 +309,24 @@ def _process_step_output( ) ctx.queue.put_nowait(token_output) + def _process_transfer_outputs(self, step_output: StepOutput) -> None: + """Apply completed worker-side KV transfers to scheduler metadata.""" + if not step_output.completed_transfer_jobs: + return + with profile_span( + "scheduler.process_kv_transfer_output", + cat="scheduler", + args={"jobs": len(step_output.completed_transfer_jobs)}, + ): + for result in step_output.completed_transfer_jobs: + if not result.success: + logger.error( + "KV transfer job %s failed: %s", + result.job_id, + result.error or "unknown error", + ) + self.kv_cache_manager.complete_transfer_result(result) + def _handle_step_error(self, scheduler_output: SchedulerOutput) -> None: """On worker error, abort all requests in the failed batch.""" for sr in scheduler_output.scheduled_requests: diff --git a/python/core/kv_cache.py b/python/core/kv_cache.py index 9d0096e..391ef78 100644 --- a/python/core/kv_cache.py +++ b/python/core/kv_cache.py @@ -14,6 +14,16 @@ import torch +from .kv_offload import ( + CPULoadStoreSpec, + KVBlockLocation, + NPULoadStoreSpec, + OffloadKey, + SSDLoadStoreSpec, + TorchKVPageView, + TransferJob, + TransferResult, +) from .types import KvAllocation, ModelConfig, RuntimeConfig @@ -32,9 +42,24 @@ class KVCacheBlock: block_id: int ref_cnt: int = 0 block_hash: int | None = None + logical_block_id: int | None = None + physical_page_id: int | None = None + location: KVBlockLocation = KVBlockLocation.NPU + cpu_slot_id: int | None = None + ssd_slot_id: int | None = None + pending_job_id: int | None = None + offload_key: OffloadKey | None = None + dirty: bool = False + last_access_ts: float = 0.0 prev_free: "KVCacheBlock | None" = field(default=None, repr=False) next_free: "KVCacheBlock | None" = field(default=None, repr=False) + def __post_init__(self) -> None: + if self.logical_block_id is None: + self.logical_block_id = self.block_id + if self.physical_page_id is None: + self.physical_page_id = self.block_id + @dataclass(frozen=True) class KVCacheBlocks: @@ -123,22 +148,42 @@ def __init__( num_blocks: int | None = None, block_size: int = 64, enable_prefix_cache: bool = True, + max_cpu_offload_blocks: int = 0, ) -> None: """Create an empty registry of model-specific KV pools.""" + if max_cpu_offload_blocks < 0: + raise ValueError("max_cpu_offload_blocks must be non-negative") self._pools: dict[str, _CachePool] = {} self.block_size = block_size self.enable_prefix_cache = enable_prefix_cache + self.max_cpu_offload_blocks = max_cpu_offload_blocks self.blocks: list[KVCacheBlock] = [] self.free_queue = FreeKVCacheBlockQueue() self.hash_to_block: dict[int, KVCacheBlock] = {} self.request_blocks: dict[str, list[KVCacheBlock]] = {} + self.page_id_to_pending_jobs: dict[int, set[int]] = {} + self.cpu_slot_to_block_id: dict[int, int] = {} + self.ssd_slot_to_block_id: dict[int, int] = {} + self._num_physical_pages: int = 0 + self._free_physical_page_ids: list[int] = [] + self._free_cpu_slots: list[int] = [] + self._free_ssd_slots: list[int] = [] + self._next_cpu_slot_id: int = 0 + self._next_ssd_slot_id: int = 0 + self._next_transfer_job_id: int = 0 if num_blocks is not None: self._init_blocks(num_blocks, block_size) @property def num_free_blocks(self) -> int: """Return the number of immediately allocatable KV blocks.""" - return self.free_queue.count + npu_free_blocks = 0 + block = self.free_queue.head + while block is not None: + if block.physical_page_id is not None: + npu_free_blocks += 1 + block = block.next_free + return npu_free_blocks + len(self._free_physical_page_ids) @property def num_blocks(self) -> int: @@ -151,6 +196,7 @@ def _init_blocks(self, num_blocks: int, block_size: int) -> None: raise ValueError("KV block pool is already initialized with different dimensions") return self.block_size = block_size + self._num_physical_pages = num_blocks self.blocks = [KVCacheBlock(block_id=i) for i in range(num_blocks)] for block in self.blocks: self.free_queue.append(block) @@ -193,7 +239,7 @@ def allocate_for_prompt(self, model_id: str, request_id: str, prompt_len: int) - if blocks is None: raise RuntimeError("Insufficient KV cache blocks.") self.request_blocks[request_id] = blocks - page_ids = [block.block_id for block in blocks] + page_ids = [self._resident_physical_page_id(block) for block in blocks] return KvAllocation( request_id=request_id, model_id=model_id, @@ -211,13 +257,37 @@ def allocate_blocks(self, num_blocks: int) -> list[KVCacheBlock] | None: blocks: list[KVCacheBlock] = [] for _ in range(num_blocks): block = self.free_queue.popleft() + if block is None and self._free_physical_page_ids: + block = self._new_logical_block() if block is None: for allocated in blocks: self.release(allocated) return None + physical_page_id = block.physical_page_id + if physical_page_id is None: + physical_page_id = self._allocate_physical_page_id() + if physical_page_id is None: + for allocated in blocks: + self.release(allocated) + self.free_queue.append(block) + return None if block.block_hash is not None: self.hash_to_block.pop(block.block_hash, None) block.block_hash = None + if block.pending_job_id is not None: + for allocated in blocks: + self.release(allocated) + self.free_queue.append(block) + return None + if block.cpu_slot_id is not None: + self._release_cpu_slot(block) + if block.ssd_slot_id is not None: + self._release_ssd_slot(block) + block.location = KVBlockLocation.NPU + block.physical_page_id = physical_page_id + block.cpu_slot_id = None + block.ssd_slot_id = None + block.offload_key = None block.ref_cnt = 1 blocks.append(block) return blocks @@ -281,7 +351,7 @@ def release(self, block: KVCacheBlock) -> None: if block.ref_cnt <= 0: return block.ref_cnt -= 1 - if block.ref_cnt == 0: + if block.ref_cnt == 0 and block.pending_job_id is None: self.free_queue.append(block) def _iter_block_hashes(self, token_ids: list[int]): @@ -310,6 +380,378 @@ def compute_block_hashes(self, token_ids: list[int]) -> list[int]: """Compute chained hashes for all full blocks in the token sequence.""" return [block_hash for _, block_hash in self._iter_block_hashes(token_ids)] + def resident_block_ids(self, block_ids: list[int]) -> list[int]: + """Resolve logical block IDs to resident NPU physical page IDs.""" + physical_ids: list[int] = [] + for block_id in block_ids: + block = self.blocks[block_id] + if block.location != KVBlockLocation.NPU or block.physical_page_id is None: + raise RuntimeError(f"KV block {block_id} is not resident on NPU: {block.location.value}") + physical_ids.append(block.physical_page_id) + return physical_ids + + def mark_blocks_moving_to_ssd( + self, + block_ids: list[int], + ssd_slot_ids: list[int], + job_id: int, + *, + offload_keys: list[OffloadKey] | None = None, + ) -> None: + """Mark resident NPU blocks as being stored to SSD by one transfer job.""" + if len(block_ids) != len(ssd_slot_ids): + raise ValueError("block_ids and ssd_slot_ids must have the same length") + if offload_keys is not None and len(offload_keys) != len(block_ids): + raise ValueError("offload_keys and block_ids must have the same length") + for idx, block_id in enumerate(block_ids): + block = self.blocks[block_id] + if block.location != KVBlockLocation.NPU or block.physical_page_id is None: + raise RuntimeError(f"KV block {block_id} is not resident on NPU") + if block.pending_job_id is not None: + raise RuntimeError(f"KV block {block_id} already has pending job {block.pending_job_id}") + if block.ref_cnt == 0: + self.free_queue.remove(block) + block.location = KVBlockLocation.MOVING_TO_SSD + block.ssd_slot_id = ssd_slot_ids[idx] + block.pending_job_id = job_id + if offload_keys is not None: + block.offload_key = offload_keys[idx] + self.page_id_to_pending_jobs.setdefault(block.physical_page_id, set()).add(job_id) + + def mark_blocks_moving_to_cpu( + self, + block_ids: list[int], + cpu_slot_ids: list[int], + job_id: int, + *, + offload_keys: list[OffloadKey] | None = None, + ) -> None: + """Mark resident NPU blocks as being stored to CPU by one transfer job.""" + if len(block_ids) != len(cpu_slot_ids): + raise ValueError("block_ids and cpu_slot_ids must have the same length") + if offload_keys is not None and len(offload_keys) != len(block_ids): + raise ValueError("offload_keys and block_ids must have the same length") + for idx, block_id in enumerate(block_ids): + block = self.blocks[block_id] + if block.location != KVBlockLocation.NPU or block.physical_page_id is None: + raise RuntimeError(f"KV block {block_id} is not resident on NPU") + if block.pending_job_id is not None: + raise RuntimeError(f"KV block {block_id} already has pending job {block.pending_job_id}") + if block.ref_cnt == 0: + self.free_queue.remove(block) + block.location = KVBlockLocation.MOVING_TO_CPU + block.cpu_slot_id = cpu_slot_ids[idx] + block.pending_job_id = job_id + if offload_keys is not None: + block.offload_key = offload_keys[idx] + self.page_id_to_pending_jobs.setdefault(block.physical_page_id, set()).add(job_id) + + def complete_cpu_store_job(self, job_id: int, *, success: bool = True) -> None: + """Complete one NPU-to-CPU store job and publish or invalidate its blocks.""" + for block in self.blocks: + if block.pending_job_id != job_id or block.location != KVBlockLocation.MOVING_TO_CPU: + continue + self._clear_pending_page_job(block, job_id) + block.pending_job_id = None + if success: + block.location = KVBlockLocation.CPU + self._release_physical_page_id(block.physical_page_id) + block.physical_page_id = None + block.dirty = False + if block.ref_cnt == 0: + self.free_queue.append(block) + else: + block.location = KVBlockLocation.NPU + block.offload_key = None + self._release_cpu_slot(block) + if block.ref_cnt == 0: + self.free_queue.append(block) + + def complete_store_job(self, job_id: int, *, success: bool = True) -> None: + """Complete one NPU-to-SSD store job and publish or invalidate its blocks.""" + for block in self.blocks: + if block.pending_job_id != job_id or block.location != KVBlockLocation.MOVING_TO_SSD: + continue + self._clear_pending_page_job(block, job_id) + block.pending_job_id = None + if success: + block.location = KVBlockLocation.SSD + self._release_physical_page_id(block.physical_page_id) + block.physical_page_id = None + block.dirty = False + if block.ref_cnt == 0: + self.free_queue.append(block) + else: + block.location = KVBlockLocation.NPU + block.offload_key = None + self._release_ssd_slot(block) + if block.ref_cnt == 0: + self.free_queue.append(block) + + def mark_blocks_moving_to_npu( + self, + block_ids: list[int], + physical_page_ids: list[int], + job_id: int, + ) -> None: + """Mark SSD blocks as being loaded back to NPU pages.""" + if len(block_ids) != len(physical_page_ids): + raise ValueError("block_ids and physical_page_ids must have the same length") + for block_id in block_ids: + block = self.blocks[block_id] + if block.location != KVBlockLocation.SSD: + raise RuntimeError(f"KV block {block_id} is not stored on SSD") + if block.pending_job_id is not None: + raise RuntimeError(f"KV block {block_id} already has pending job {block.pending_job_id}") + for idx, block_id in enumerate(block_ids): + block = self.blocks[block_id] + if block.ref_cnt == 0: + self.free_queue.remove(block) + self._reserve_physical_page_id(physical_page_ids[idx]) + block.location = KVBlockLocation.MOVING_TO_NPU + block.physical_page_id = physical_page_ids[idx] + block.pending_job_id = job_id + self.page_id_to_pending_jobs.setdefault(physical_page_ids[idx], set()).add(job_id) + + def mark_cpu_blocks_moving_to_npu( + self, + block_ids: list[int], + physical_page_ids: list[int], + job_id: int, + ) -> None: + """Mark CPU-resident blocks as being loaded back to NPU pages.""" + if len(block_ids) != len(physical_page_ids): + raise ValueError("block_ids and physical_page_ids must have the same length") + for block_id in block_ids: + block = self.blocks[block_id] + if block.location != KVBlockLocation.CPU: + raise RuntimeError(f"KV block {block_id} is not stored on CPU") + if block.pending_job_id is not None: + raise RuntimeError(f"KV block {block_id} already has pending job {block.pending_job_id}") + for idx, block_id in enumerate(block_ids): + block = self.blocks[block_id] + if block.ref_cnt == 0: + self.free_queue.remove(block) + self._reserve_physical_page_id(physical_page_ids[idx]) + block.location = KVBlockLocation.MOVING_TO_NPU + block.physical_page_id = physical_page_ids[idx] + block.pending_job_id = job_id + self.page_id_to_pending_jobs.setdefault(physical_page_ids[idx], set()).add(job_id) + + def complete_load_job(self, job_id: int, *, success: bool = True) -> None: + """Complete one SSD-to-NPU load job and make loaded blocks resident.""" + for block in self.blocks: + if block.pending_job_id != job_id or block.location != KVBlockLocation.MOVING_TO_NPU: + continue + self._clear_pending_page_job(block, job_id) + block.pending_job_id = None + if success: + block.location = KVBlockLocation.NPU + block.dirty = False + if block.cpu_slot_id is not None: + self._release_cpu_slot(block) + if block.ssd_slot_id is not None: + self._release_ssd_slot(block) + if block.ref_cnt == 0: + self.free_queue.append(block) + else: + self._release_physical_page_id(block.physical_page_id) + block.location = KVBlockLocation.CPU if block.cpu_slot_id is not None else KVBlockLocation.SSD + block.physical_page_id = None + + def complete_cpu_load_job(self, job_id: int, *, success: bool = True) -> None: + """Complete one CPU-to-NPU load job and make loaded blocks resident.""" + self.complete_load_job(job_id, success=success) + + def pending_jobs_for_page(self, page_id: int) -> set[int]: + """Return transfer jobs that must be fenced before reusing one NPU page.""" + return set(self.page_id_to_pending_jobs.get(page_id, set())) + + def build_cpu_store_job( + self, + block_ids: list[int], + *, + request_id: str | None = None, + offload_keys: list[OffloadKey] | None = None, + ) -> TransferJob: + """Create and mark one NPU-to-CPU transfer job for logical blocks.""" + job_id = self._next_job_id() + page_ids = self.resident_block_ids(block_ids) + cpu_slot_ids = self._allocate_cpu_slots(block_ids) + try: + self.mark_blocks_moving_to_cpu( + block_ids, + cpu_slot_ids, + job_id, + offload_keys=offload_keys, + ) + except Exception: + for slot_id in cpu_slot_ids: + self.cpu_slot_to_block_id.pop(slot_id, None) + self._free_cpu_slots.append(slot_id) + raise + return TransferJob( + job_id=job_id, + request_id=request_id, + src=NPULoadStoreSpec(page_ids), + dst=CPULoadStoreSpec(cpu_slot_ids), + keys=tuple(offload_keys or ()), + ) + + def build_cpu_load_job( + self, + block_ids: list[int], + physical_page_ids: list[int] | None = None, + *, + request_id: str | None = None, + ) -> TransferJob: + """Create and mark one CPU-to-NPU transfer job for logical blocks.""" + cpu_slot_ids: list[int] = [] + for block_id in block_ids: + block = self.blocks[block_id] + if block.location != KVBlockLocation.CPU or block.cpu_slot_id is None: + raise RuntimeError(f"KV block {block_id} is not stored on CPU") + cpu_slot_ids.append(block.cpu_slot_id) + allocated_physical_pages = physical_page_ids is None + if physical_page_ids is None: + physical_page_ids = self._allocate_physical_page_ids(len(block_ids)) + if physical_page_ids is None: + raise RuntimeError("Insufficient free NPU pages for CPU KV load") + if len(block_ids) != len(physical_page_ids): + raise ValueError("block_ids and physical_page_ids must have the same length") + job_id = self._next_job_id() + try: + self.mark_cpu_blocks_moving_to_npu(block_ids, physical_page_ids, job_id) + except Exception: + if allocated_physical_pages: + for page_id in physical_page_ids: + self._release_physical_page_id(page_id) + raise + return TransferJob( + job_id=job_id, + request_id=request_id, + src=CPULoadStoreSpec(cpu_slot_ids), + dst=NPULoadStoreSpec(physical_page_ids), + ) + + def build_ssd_store_job( + self, + block_ids: list[int], + *, + request_id: str | None = None, + offload_keys: list[OffloadKey] | None = None, + ) -> TransferJob: + """Create and mark one NPU-to-SSD transfer job for logical blocks.""" + job_id = self._next_job_id() + page_ids = self.resident_block_ids(block_ids) + ssd_slot_ids = self._allocate_ssd_slots(block_ids) + try: + self.mark_blocks_moving_to_ssd( + block_ids, + ssd_slot_ids, + job_id, + offload_keys=offload_keys, + ) + except Exception: + for slot_id in ssd_slot_ids: + self.ssd_slot_to_block_id.pop(slot_id, None) + self._free_ssd_slots.append(slot_id) + raise + return TransferJob( + job_id=job_id, + request_id=request_id, + src=NPULoadStoreSpec(page_ids), + dst=SSDLoadStoreSpec(ssd_slot_ids), + keys=tuple(offload_keys or ()), + ) + + def build_ssd_load_job( + self, + block_ids: list[int], + physical_page_ids: list[int] | None = None, + *, + request_id: str | None = None, + ) -> TransferJob: + """Create and mark one SSD-to-NPU transfer job for logical blocks.""" + ssd_slot_ids: list[int] = [] + for block_id in block_ids: + block = self.blocks[block_id] + if block.location != KVBlockLocation.SSD or block.ssd_slot_id is None: + raise RuntimeError(f"KV block {block_id} is not stored on SSD") + ssd_slot_ids.append(block.ssd_slot_id) + allocated_physical_pages = physical_page_ids is None + if physical_page_ids is None: + physical_page_ids = self._allocate_physical_page_ids(len(block_ids)) + if physical_page_ids is None: + raise RuntimeError("Insufficient free NPU pages for SSD KV load") + if len(block_ids) != len(physical_page_ids): + raise ValueError("block_ids and physical_page_ids must have the same length") + job_id = self._next_job_id() + try: + self.mark_blocks_moving_to_npu(block_ids, physical_page_ids, job_id) + except Exception: + if allocated_physical_pages: + for page_id in physical_page_ids: + self._release_physical_page_id(page_id) + raise + return TransferJob( + job_id=job_id, + request_id=request_id, + src=SSDLoadStoreSpec(ssd_slot_ids), + dst=NPULoadStoreSpec(physical_page_ids), + ) + + def select_cpu_offload_candidates( + self, + num_blocks: int, + *, + excluded_block_ids: set[int] | None = None, + ) -> list[int]: + """Pick resident blocks that can be stored to CPU to free NPU pages.""" + if num_blocks <= 0: + return [] + excluded_block_ids = excluded_block_ids or set() + candidates = [ + block + for block in self.blocks + if block.block_id not in excluded_block_ids + and block.location == KVBlockLocation.NPU + and block.physical_page_id is not None + and block.pending_job_id is None + ] + candidates.sort(key=lambda block: (block.ref_cnt > 0, block.last_access_ts, block.block_id)) + return [block.block_id for block in candidates[:num_blocks]] + + def non_resident_block_ids(self, block_ids: list[int]) -> list[int]: + """Return blocks from the input list that are not currently resident on NPU.""" + return [ + block_id + for block_id in block_ids + if self.blocks[block_id].location != KVBlockLocation.NPU + or self.blocks[block_id].physical_page_id is None + ] + + def complete_transfer_result(self, result: TransferResult) -> None: + """Apply one worker-reported KV transfer completion to block metadata.""" + self.complete_transfer_job(result.job_id, success=result.success) + + def complete_transfer_job(self, job_id: int, *, success: bool = True) -> None: + """Complete one pending transfer job regardless of direction/medium.""" + matched = False + for block in self.blocks: + if block.pending_job_id != job_id: + continue + matched = True + if block.location == KVBlockLocation.MOVING_TO_CPU: + self.complete_cpu_store_job(job_id, success=success) + elif block.location == KVBlockLocation.MOVING_TO_SSD: + self.complete_store_job(job_id, success=success) + elif block.location == KVBlockLocation.MOVING_TO_NPU: + self.complete_load_job(job_id, success=success) + break + if not matched: + return + def ensure_one_more_slot(self, alloc: KvAllocation) -> int: """Ensure a request has capacity for one more token and return its slot.""" pool = self._pool(alloc.model_id) @@ -318,7 +760,7 @@ def ensure_one_more_slot(self, alloc: KvAllocation) -> int: if blocks is None: raise RuntimeError("Insufficient KV cache blocks.") self.request_blocks.setdefault(alloc.request_id, []).extend(blocks) - alloc.page_ids.extend(block.block_id for block in blocks) + alloc.page_ids.extend(self._resident_physical_page_id(block) for block in blocks) alloc.tokens_capacity = len(alloc.page_ids) * pool.page_size return self.slot_mapping_for_request(alloc, alloc.tokens_used) @@ -401,6 +843,15 @@ def materialize_single_layer_cache( pool.value_pages[layer_idx].reshape(-1, pool.head_dim), ) + def materialize_page_view(self, model_id: str) -> TorchKVPageView: + """Return a page-contiguous view over all K/V layers for CPU offload.""" + pool = self._pool(model_id) + components: list[torch.Tensor] = [] + for layer_idx in range(pool.num_layers): + components.append(pool.key_pages[layer_idx]) + components.append(pool.value_pages[layer_idx]) + return TorchKVPageView(components) + def free(self, alloc: KvAllocation) -> None: """Return an allocation's pages to the model pool.""" self.release_request(alloc.request_id) @@ -408,6 +859,113 @@ def free(self, alloc: KvAllocation) -> None: alloc.tokens_capacity = 0 alloc.tokens_used = 0 + def _clear_pending_page_job(self, block: KVCacheBlock, job_id: int) -> None: + page_id = block.physical_page_id + if page_id is None: + return + pending_jobs = self.page_id_to_pending_jobs.get(page_id) + if pending_jobs is None: + return + pending_jobs.discard(job_id) + if not pending_jobs: + del self.page_id_to_pending_jobs[page_id] + + def _next_job_id(self) -> int: + job_id = self._next_transfer_job_id + self._next_transfer_job_id += 1 + return job_id + + def _new_logical_block(self) -> KVCacheBlock: + block = KVCacheBlock(block_id=len(self.blocks), physical_page_id=None) + self.blocks.append(block) + return block + + def _resident_physical_page_id(self, block: KVCacheBlock) -> int: + if block.location != KVBlockLocation.NPU or block.physical_page_id is None: + raise RuntimeError(f"KV block {block.block_id} is not resident on NPU") + return block.physical_page_id + + def _allocate_physical_page_id(self) -> int | None: + if not self._free_physical_page_ids: + return None + return self._free_physical_page_ids.pop() + + def _allocate_physical_page_ids(self, count: int) -> list[int] | None: + if count <= 0: + return [] + if len(self._free_physical_page_ids) < count: + return None + return [self._free_physical_page_ids.pop() for _ in range(count)] + + def _release_physical_page_id(self, page_id: int | None) -> None: + if page_id is None: + return + if page_id < 0 or page_id >= self._num_physical_pages: + raise ValueError(f"Invalid NPU page id {page_id}") + if page_id in self._free_physical_page_ids: + return + self._free_physical_page_ids.append(page_id) + + def _reserve_physical_page_id(self, page_id: int) -> None: + if page_id < 0 or page_id >= self._num_physical_pages: + raise ValueError(f"Invalid NPU page id {page_id}") + try: + self._free_physical_page_ids.remove(page_id) + except ValueError: + pass + + def _allocate_cpu_slots(self, block_ids: list[int]) -> list[int]: + for block_id in block_ids: + block = self.blocks[block_id] + if block.cpu_slot_id is not None: + raise RuntimeError(f"KV block {block_id} already owns CPU slot {block.cpu_slot_id}") + new_slots_needed = max(0, len(block_ids) - len(self._free_cpu_slots)) + if ( + self._next_cpu_slot_id + new_slots_needed > self.max_cpu_offload_blocks + ): + raise RuntimeError("CPU KV offload capacity exceeded") + slots: list[int] = [] + for block_id in block_ids: + if self._free_cpu_slots: + slot_id = self._free_cpu_slots.pop() + else: + slot_id = self._next_cpu_slot_id + self._next_cpu_slot_id += 1 + self.cpu_slot_to_block_id[slot_id] = block_id + slots.append(slot_id) + return slots + + def _release_cpu_slot(self, block: KVCacheBlock) -> None: + slot_id = block.cpu_slot_id + if slot_id is None: + return + self.cpu_slot_to_block_id.pop(slot_id, None) + self._free_cpu_slots.append(slot_id) + block.cpu_slot_id = None + + def _allocate_ssd_slots(self, block_ids: list[int]) -> list[int]: + slots: list[int] = [] + for block_id in block_ids: + block = self.blocks[block_id] + if block.ssd_slot_id is not None: + raise RuntimeError(f"KV block {block_id} already owns SSD slot {block.ssd_slot_id}") + if self._free_ssd_slots: + slot_id = self._free_ssd_slots.pop() + else: + slot_id = self._next_ssd_slot_id + self._next_ssd_slot_id += 1 + self.ssd_slot_to_block_id[slot_id] = block_id + slots.append(slot_id) + return slots + + def _release_ssd_slot(self, block: KVCacheBlock) -> None: + slot_id = block.ssd_slot_id + if slot_id is None: + return + self.ssd_slot_to_block_id.pop(slot_id, None) + self._free_ssd_slots.append(slot_id) + block.ssd_slot_id = None + def _pool(self, model_id: str) -> _CachePool: """Return the registered cache pool for a model.""" if model_id not in self._pools: diff --git a/python/core/kv_offload.py b/python/core/kv_offload.py new file mode 100644 index 0000000..50b6678 --- /dev/null +++ b/python/core/kv_offload.py @@ -0,0 +1,458 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Protocol + +import torch + + +class KVBlockLocation(Enum): + """Residency state for one logical KV cache block.""" + + NPU = "npu" + CPU = "cpu" + SSD = "ssd" + MOVING_TO_CPU = "moving_to_cpu" + MOVING_TO_SSD = "moving_to_ssd" + MOVING_TO_NPU = "moving_to_npu" + INVALID = "invalid" + + +@dataclass(frozen=True) +class OffloadKey: + """Stable identifier for offloaded KV data, independent of NPU page ids.""" + + block_hash: int + group_id: int = 0 + layout_version: int = 0 + + +@dataclass(frozen=True) +class LoadStoreSpec: + """Base transfer endpoint descriptor.""" + + block_ids: tuple[int, ...] + medium: str + + +@dataclass(frozen=True) +class NPULoadStoreSpec(LoadStoreSpec): + """NPU KV page transfer endpoint.""" + + def __init__(self, page_ids: list[int] | tuple[int, ...]) -> None: + super().__init__(tuple(page_ids), "NPU") + + +@dataclass(frozen=True) +class SSDLoadStoreSpec(LoadStoreSpec): + """On-card SSD slot transfer endpoint.""" + + def __init__(self, slot_ids: list[int] | tuple[int, ...]) -> None: + super().__init__(tuple(slot_ids), "SSD") + + +@dataclass(frozen=True) +class CPULoadStoreSpec(LoadStoreSpec): + """CPU KV offload slot transfer endpoint.""" + + def __init__(self, slot_ids: list[int] | tuple[int, ...]) -> None: + super().__init__(tuple(slot_ids), "CPU") + + +@dataclass(frozen=True) +class TransferJob: + """One asynchronous KV cache transfer job.""" + + job_id: int + request_id: str | None + src: LoadStoreSpec + dst: LoadStoreSpec + keys: tuple[OffloadKey, ...] = () + + @property + def direction(self) -> tuple[str, str]: + """Return the source and destination media for handler dispatch.""" + return (self.src.medium, self.dst.medium) + + +@dataclass(frozen=True) +class TransferResult: + """Completion status for one transfer job.""" + + job_id: int + success: bool = True + error: str | None = None + + +class KvOffloadBackend(Protocol): + """Worker-side transfer backend interface.""" + + def submit(self, job: TransferJob) -> bool: + """Submit an asynchronous transfer job.""" + ... + + def poll(self) -> list[TransferResult]: + """Return transfer jobs completed since the previous poll.""" + ... + + def wait(self, job_ids: set[int]) -> list[TransferResult]: + """Block until the selected jobs are complete.""" + ... + + def shutdown(self) -> None: + """Release backend resources.""" + ... + + +class NoopKvOffloadBackend: + """Backend used when KV offload is disabled.""" + + def submit(self, job: TransferJob) -> bool: + return False + + def poll(self) -> list[TransferResult]: + return [] + + def wait(self, job_ids: set[int]) -> list[TransferResult]: + return [] + + def shutdown(self) -> None: + return None + + +class UnavailableKvOffloadBackend: + """Backend placeholder for transfer media whose runtime API is not wired yet.""" + + def __init__(self, medium: str, *, reason: str | None = None) -> None: + self.medium = medium + self.reason = reason or f"{medium} KV offload backend is not configured" + self.submitted_jobs: dict[int, TransferJob] = {} + self._finished: dict[int, TransferResult] = {} + + def submit(self, job: TransferJob) -> bool: + self.submitted_jobs[job.job_id] = job + self._finished[job.job_id] = TransferResult( + job_id=job.job_id, + success=False, + error=self.reason, + ) + return False + + def poll(self) -> list[TransferResult]: + results = list(self._finished.values()) + self._finished.clear() + return results + + def wait(self, job_ids: set[int]) -> list[TransferResult]: + results: list[TransferResult] = [] + for job_id in job_ids: + result = self._finished.pop(job_id, None) + if result is not None: + results.append(result) + return results + + def shutdown(self) -> None: + self.submitted_jobs.clear() + self._finished.clear() + + +class MockKvOffloadBackend: + """Deterministic in-memory backend for scheduler and state-machine tests.""" + + def __init__(self, *, complete_immediately: bool = True) -> None: + self.complete_immediately = complete_immediately + self.submitted_jobs: dict[int, TransferJob] = {} + self._finished: dict[int, TransferResult] = {} + self._pending: set[int] = set() + + def submit(self, job: TransferJob) -> bool: + self.submitted_jobs[job.job_id] = job + if self.complete_immediately: + self._finished[job.job_id] = TransferResult(job_id=job.job_id) + else: + self._pending.add(job.job_id) + return True + + def complete(self, job_id: int, *, success: bool = True, error: str | None = None) -> None: + """Mark one pending job complete.""" + if job_id not in self.submitted_jobs: + raise KeyError(f"Unknown transfer job {job_id}") + self._pending.discard(job_id) + self._finished[job_id] = TransferResult(job_id=job_id, success=success, error=error) + + def poll(self) -> list[TransferResult]: + results = list(self._finished.values()) + self._finished.clear() + return results + + def wait(self, job_ids: set[int]) -> list[TransferResult]: + for job_id in list(job_ids): + if job_id in self._pending: + self.complete(job_id) + results: list[TransferResult] = [] + for job_id in job_ids: + result = self._finished.pop(job_id, None) + if result is not None: + results.append(result) + return results + + def shutdown(self) -> None: + self.submitted_jobs.clear() + self._finished.clear() + self._pending.clear() + + +@dataclass +class OffloadStats: + """Counters used by mock and future real backends.""" + + stores_submitted: int = 0 + loads_submitted: int = 0 + stores_completed: int = 0 + loads_completed: int = 0 + failures: int = 0 + bytes_moved: int = 0 + extra: dict[str, int] = field(default_factory=dict) + + +class TorchKVPageView: + """Canonical page view over torch KV tensors. + + Each component tensor must expose pages on dimension 0. For the serving + manager this is typically one component per K/V tensor and layer, shaped + ``[num_pages, ...]``. Transfers copy the flattened bytes of every component + for each selected page. + """ + + def __init__(self, components: list[torch.Tensor]) -> None: + if not components: + raise ValueError("components must not be empty") + num_pages = int(components[0].shape[0]) + if num_pages <= 0: + raise ValueError("components must contain at least one page") + page_bytes = 0 + normalized: list[torch.Tensor] = [] + for component in components: + if int(component.shape[0]) != num_pages: + raise ValueError("all components must have the same num_pages") + if not component.is_contiguous(): + raise ValueError("canonical KV page components must be contiguous") + normalized.append(component) + page_bytes += int(component[0].nbytes) + self.components = normalized + self.num_pages = num_pages + self.page_size_bytes = page_bytes + + def copy_page_to(self, page_id: int, dst: torch.Tensor) -> None: + """Copy one page from this view into one flat uint8 destination row.""" + self._validate_page_id(page_id) + self._validate_cpu_row(dst) + offset = 0 + for component in self.components: + src = component[page_id].view(torch.uint8).reshape(-1) + end = offset + src.numel() + dst[offset:end].copy_(src.cpu()) + offset = end + + def copy_page_from(self, page_id: int, src: torch.Tensor) -> None: + """Copy one flat uint8 source row into one page in this view.""" + self._validate_page_id(page_id) + self._validate_cpu_row(src) + offset = 0 + for component in self.components: + dst = component[page_id].view(torch.uint8).reshape(-1) + end = offset + dst.numel() + dst.copy_(src[offset:end].to(dst.device)) + offset = end + + def _validate_page_id(self, page_id: int) -> None: + if page_id < 0 or page_id >= self.num_pages: + raise IndexError(f"page_id {page_id} is out of range for {self.num_pages} pages") + + def _validate_cpu_row(self, row: torch.Tensor) -> None: + if row.dtype != torch.uint8 or row.device.type != "cpu" or row.dim() != 1: + raise ValueError("CPU offload rows must be 1-D CPU torch.uint8 tensors") + if row.numel() != self.page_size_bytes: + raise ValueError(f"CPU offload row has {row.numel()} bytes, expected {self.page_size_bytes}") + + +class WorkerKVPageView: + """Canonical page view over runner-owned WorkerTensor KV caches.""" + + def __init__( + self, + *, + worker: Any, + key_pages: Any, + value_pages: Any, + num_layers: int, + num_pages: int, + num_kv_heads: int, + page_size: int, + head_dim: int, + ) -> None: + if num_layers <= 0 or num_pages <= 0 or num_kv_heads <= 0 or page_size <= 0 or head_dim <= 0: + raise ValueError("KV page view dimensions must be positive") + self.worker = worker + self.key_pages = key_pages + self.value_pages = value_pages + self.num_layers = num_layers + self.num_pages = num_pages + self.num_kv_heads = num_kv_heads + self.page_size = page_size + self.head_dim = head_dim + self.rows_per_page = num_kv_heads * page_size + element_size = torch.empty((), dtype=key_pages.torch_dtype).element_size() + if torch.empty((), dtype=value_pages.torch_dtype).element_size() != element_size: + raise ValueError("key_pages and value_pages must use dtypes with the same element size") + self._element_size = element_size + self.component_size_bytes = self.rows_per_page * head_dim * element_size + self.page_size_bytes = 2 * num_layers * self.component_size_bytes + + def copy_page_to(self, page_id: int, dst: torch.Tensor) -> None: + """Copy one worker-resident KV page into one flat CPU uint8 row.""" + self._validate_page_id(page_id) + self._validate_cpu_row(dst) + offset = 0 + for layer_idx in range(self.num_layers): + for tensor in (self.key_pages, self.value_pages): + src_ptr = tensor.data_ptr + self._component_offset(layer_idx, page_id) + end = offset + self.component_size_bytes + self.worker.copy_from( + dst[offset:end].data_ptr(), + src_ptr, + self.component_size_bytes, + worker_id=tensor.worker_id, + ) + offset = end + + def copy_page_from(self, page_id: int, src: torch.Tensor) -> None: + """Copy one flat CPU uint8 row into one worker-resident KV page.""" + self._validate_page_id(page_id) + self._validate_cpu_row(src) + offset = 0 + for layer_idx in range(self.num_layers): + for tensor in (self.key_pages, self.value_pages): + dst_ptr = tensor.data_ptr + self._component_offset(layer_idx, page_id) + end = offset + self.component_size_bytes + self.worker.copy_to( + dst_ptr, + src[offset:end].data_ptr(), + self.component_size_bytes, + worker_id=tensor.worker_id, + ) + offset = end + + def _component_offset(self, layer_idx: int, page_id: int) -> int: + row_offset = (layer_idx * self.num_pages + page_id) * self.rows_per_page + return row_offset * self.head_dim * self._element_size + + def _validate_page_id(self, page_id: int) -> None: + if page_id < 0 or page_id >= self.num_pages: + raise IndexError(f"page_id {page_id} is out of range for {self.num_pages} pages") + + def _validate_cpu_row(self, row: torch.Tensor) -> None: + if row.dtype != torch.uint8 or row.device.type != "cpu" or row.dim() != 1: + raise ValueError("CPU offload rows must be 1-D CPU torch.uint8 tensors") + if row.numel() != self.page_size_bytes: + raise ValueError(f"CPU offload row has {row.numel()} bytes, expected {self.page_size_bytes}") + + +class CPUKvOffloadBackend: + """Synchronous CPU offload backend for canonical torch KV page views.""" + + def __init__(self, page_view: Any, *, num_cpu_slots: int) -> None: + if num_cpu_slots <= 0: + raise ValueError("num_cpu_slots must be positive") + self.page_view = page_view + self.cpu_slots = torch.empty( + (num_cpu_slots, page_view.page_size_bytes), + dtype=torch.uint8, + device="cpu", + ) + self.submitted_jobs: dict[int, TransferJob] = {} + self._finished: dict[int, TransferResult] = {} + + @property + def num_cpu_slots(self) -> int: + """Return the number of CPU slots currently allocated.""" + return int(self.cpu_slots.shape[0]) + + def ensure_num_cpu_slots(self, num_cpu_slots: int) -> None: + """Grow the CPU slot tensor if a later transfer references more slots.""" + if num_cpu_slots <= self.num_cpu_slots: + return + new_slots = torch.empty( + (num_cpu_slots, self.page_view.page_size_bytes), + dtype=torch.uint8, + device="cpu", + ) + new_slots[: self.num_cpu_slots].copy_(self.cpu_slots) + self.cpu_slots = new_slots + + def submit(self, job: TransferJob) -> bool: + self.submitted_jobs[job.job_id] = job + try: + if job.direction == ("NPU", "CPU"): + self._copy_to_cpu(job) + elif job.direction == ("CPU", "NPU"): + self._copy_from_cpu(job) + else: + raise ValueError(f"unsupported CPU offload transfer direction: {job.direction}") + except Exception as exc: + self._finished[job.job_id] = TransferResult(job_id=job.job_id, success=False, error=str(exc)) + return False + self._finished[job.job_id] = TransferResult(job_id=job.job_id) + return True + + def poll(self) -> list[TransferResult]: + results = list(self._finished.values()) + self._finished.clear() + return results + + def wait(self, job_ids: set[int]) -> list[TransferResult]: + results: list[TransferResult] = [] + for job_id in job_ids: + result = self._finished.pop(job_id, None) + if result is not None: + results.append(result) + return results + + def shutdown(self) -> None: + self.submitted_jobs.clear() + self._finished.clear() + + def _copy_to_cpu(self, job: TransferJob) -> None: + if not isinstance(job.src, NPULoadStoreSpec) or not isinstance(job.dst, CPULoadStoreSpec): + raise TypeError("NPU -> CPU transfers require NPULoadStoreSpec and CPULoadStoreSpec") + self._validate_same_length(job) + for page_id, slot_id in zip(job.src.block_ids, job.dst.block_ids): + self._validate_slot_id(slot_id) + self.page_view.copy_page_to(page_id, self.cpu_slots[slot_id]) + + def _copy_from_cpu(self, job: TransferJob) -> None: + if not isinstance(job.src, CPULoadStoreSpec) or not isinstance(job.dst, NPULoadStoreSpec): + raise TypeError("CPU -> NPU transfers require CPULoadStoreSpec and NPULoadStoreSpec") + self._validate_same_length(job) + for slot_id, page_id in zip(job.src.block_ids, job.dst.block_ids): + self._validate_slot_id(slot_id) + self.page_view.copy_page_from(page_id, self.cpu_slots[slot_id]) + + @staticmethod + def _validate_same_length(job: TransferJob) -> None: + if len(job.src.block_ids) != len(job.dst.block_ids): + raise ValueError("source and destination specs must have the same number of blocks") + + def _validate_slot_id(self, slot_id: int) -> None: + if slot_id < 0 or slot_id >= self.cpu_slots.shape[0]: + raise IndexError(f"CPU slot {slot_id} is out of range for {self.cpu_slots.shape[0]} slots") diff --git a/python/core/model_runner.py b/python/core/model_runner.py index 1155303..571790f 100644 --- a/python/core/model_runner.py +++ b/python/core/model_runner.py @@ -17,6 +17,7 @@ from python.runtime.worker import WorkerTensor +from .kv_offload import WorkerKVPageView from .types import ( DecodeBatch, DecodeResult, @@ -34,6 +35,11 @@ class _KvCachePool: key_pages: WorkerTensor value_pages: WorkerTensor + num_layers: int + num_pages: int + num_kv_heads: int + page_size: int + head_dim: int class ModelRunner(ABC): @@ -67,6 +73,11 @@ def init_kv_cache( self._kv_caches[model_id] = _KvCachePool( key_pages=key_pages, value_pages=value_pages, + num_layers=config.num_hidden_layers, + num_pages=num_pages, + num_kv_heads=config.num_key_value_heads, + page_size=runtime.page_size, + head_dim=config.head_dim, ) def close_kv_cache(self) -> None: @@ -76,6 +87,22 @@ def close_kv_cache(self) -> None: self._free_kv_cache_tensor(pool.value_pages) self._kv_caches.clear() + def materialize_worker_page_view(self, model_id: str, worker) -> WorkerKVPageView: + """Return a page-contiguous worker KV view for CPU offload transfers.""" + if model_id not in self._kv_caches: + raise KeyError(f"KV cache for model {model_id!r} is not initialized") + pool = self._kv_caches[model_id] + return WorkerKVPageView( + worker=worker, + key_pages=pool.key_pages, + value_pages=pool.value_pages, + num_layers=pool.num_layers, + num_pages=pool.num_pages, + num_kv_heads=pool.num_kv_heads, + page_size=pool.page_size, + head_dim=pool.head_dim, + ) + @abstractmethod def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> WorkerTensor: """Allocate one worker-resident KV cache tensor.""" diff --git a/python/core/pypto_executor.py b/python/core/pypto_executor.py index 3b5df64..ab6bbb7 100644 --- a/python/core/pypto_executor.py +++ b/python/core/pypto_executor.py @@ -73,6 +73,14 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: ): return self._runners[model.config.model_id].run_decode(model, batch) + def materialize_kv_page_view(self, model_id: str): + """Return a worker-side canonical KV page view for offload backends.""" + runner = self._runners[model_id] + materialize = getattr(runner, "materialize_kv_page_view", None) + if not callable(materialize): + raise RuntimeError(f"Runner for model {model_id!r} does not expose a KV page view") + return materialize(model_id) + @contextlib.contextmanager def session(self): """Provide a generation lifecycle hook for PyPTO runtimes.""" diff --git a/python/core/scheduler.py b/python/core/scheduler.py index 6f9305c..ce406ef 100644 --- a/python/core/scheduler.py +++ b/python/core/scheduler.py @@ -15,6 +15,7 @@ from enum import Enum, auto from .kv_cache import KvCacheManager +from .kv_offload import KVBlockLocation, TransferJob class RequestStatus(Enum): @@ -45,6 +46,7 @@ class SchedulerConfig: # Feature flags enable_prefix_cache: bool = True enable_chunk_prefill: bool = True + max_cpu_offload_blocks: int = 0 @dataclass @@ -100,12 +102,18 @@ class ScheduledRequest: class SchedulerOutput: scheduled_requests: list[ScheduledRequest] = field(default_factory=list) preempted_requests: list[Request] = field(default_factory=list) + transfer_jobs: list[TransferJob] = field(default_factory=list) + jobs_to_flush: set[int] = field(default_factory=set) num_prefill_tokens: int = 0 num_decode_tokens: int = 0 @property def is_empty(self) -> bool: - return len(self.scheduled_requests) == 0 + return ( + len(self.scheduled_requests) == 0 + and len(self.transfer_jobs) == 0 + and len(self.jobs_to_flush) == 0 + ) @dataclass @@ -125,6 +133,10 @@ def __init__(self, config: SchedulerConfig, kv_cache_manager: KvCacheManager) -> self.waiting: deque[Request] = deque() self.running: list[Request] = [] self.requests: dict[str, Request] = {} + self.num_cpu_store_jobs: int = 0 + self.num_cpu_load_jobs: int = 0 + self.num_cpu_store_blocks: int = 0 + self.num_cpu_load_blocks: int = 0 def add_request(self, request: Request) -> None: if len(request.prompt_token_ids) > self.config.max_seq_len: @@ -164,11 +176,12 @@ def schedule(self) -> SchedulerOutput: # Phase 1: schedule RUNNING requests (decode or resumed prefill) scheduled_req_ids: set[str] = set() num_scheduled_tokens: dict[str, int] = {} - running_to_keep: list[Request] = [] - for request in self.running: + req_index = 0 + while req_index < len(self.running) and token_budget > 0: + request = self.running[req_index] num_new = request.num_new_tokens_needed if num_new <= 0: - running_to_keep.append(request) + req_index += 1 continue if self.config.enable_chunk_prefill and self.config.long_prefill_token_threshold > 0: @@ -176,32 +189,39 @@ def schedule(self) -> SchedulerOutput: num_new = min(num_new, token_budget) if num_new <= 0: - running_to_keep.append(request) + req_index += 1 continue num_blocks_needed = self._blocks_needed(request, num_new) if not self._try_allocate_blocks(request, num_blocks_needed): + preempted_request = self._pop_lowest_priority_running_request(request) + if preempted_request is None: + req_index += 1 + continue preempted = self._preempt_lowest_priority( - request, scheduled_req_ids, num_scheduled_tokens, output + preempted_request, + scheduled_req_ids, + num_scheduled_tokens, + output, ) if preempted is None: - running_to_keep.append(request) + self.running.insert(req_index, preempted_request) + req_index += 1 continue token_budget += preempted.get("returned_tokens", 0) output.preempted_requests.append(preempted["request"]) - if not self._try_allocate_blocks(request, num_blocks_needed): - running_to_keep.append(request) - continue + break is_prefill = request.is_prefill all_block_ids = request.cached_block_ids + request.allocated_block_ids + physical_block_ids = self.kv_cache_manager.resident_block_ids(all_block_ids) output.scheduled_requests.append( ScheduledRequest( request=request, num_new_tokens=num_new, is_prefill=is_prefill, num_computed_tokens=request.num_computed_tokens, - block_ids=list(all_block_ids), + block_ids=physical_block_ids, ) ) scheduled_req_ids.add(request.request_id) @@ -211,58 +231,71 @@ def schedule(self) -> SchedulerOutput: else: output.num_decode_tokens += num_new token_budget -= num_new - running_to_keep.append(request) + req_index += 1 - self.running = running_to_keep - - # Phase 2: schedule WAITING requests (new prefill) remaining_waiting: deque[Request] = deque() - while self.waiting and token_budget > 0: - if len(self.running) >= self.config.max_num_running_reqs: - break - - request = self.waiting.popleft() + # Phase 2: schedule WAITING requests (new prefill or resumed preemption). + if not output.preempted_requests: + while self.waiting and token_budget > 0: + if len(self.running) >= self.config.max_num_running_reqs: + break + + request = self.waiting.popleft() + + # Prefix cache lookup + existing_block_ids = request.cached_block_ids + request.allocated_block_ids + if self.config.enable_prefix_cache and not existing_block_ids: + cached_blocks = self.kv_cache_manager.get_computed_blocks(request.prompt_token_ids) + if cached_blocks: + request.cached_block_ids = [b.block_id for b in cached_blocks] + request.num_computed_tokens = len(cached_blocks) * self.kv_cache_manager.block_size + else: + cached_blocks = [] - # Prefix cache lookup - if self.config.enable_prefix_cache: - cached_blocks = self.kv_cache_manager.get_computed_blocks(request.prompt_token_ids) - if cached_blocks: - request.cached_block_ids = [b.block_id for b in cached_blocks] - request.num_computed_tokens = len(cached_blocks) * self.kv_cache_manager.block_size - else: - cached_blocks = [] + num_new = request.num_new_tokens_needed + if self.config.enable_chunk_prefill and self.config.long_prefill_token_threshold > 0: + num_new = min(num_new, self.config.long_prefill_token_threshold) + num_new = min(num_new, token_budget) - num_new = request.num_new_tokens_needed - if self.config.enable_chunk_prefill and self.config.long_prefill_token_threshold > 0: - num_new = min(num_new, self.config.long_prefill_token_threshold) - num_new = min(num_new, token_budget) - - if num_new <= 0: - remaining_waiting.append(request) - continue + if num_new <= 0: + remaining_waiting.append(request) + continue - num_blocks_needed = self._blocks_needed(request, num_new) - if not self._try_allocate_blocks(request, num_blocks_needed): - self.kv_cache_manager.release_cached_blocks(cached_blocks) - request.cached_block_ids = [] - request.num_computed_tokens = 0 - remaining_waiting.append(request) - break + all_block_ids = request.cached_block_ids + request.allocated_block_ids + missing_block_ids = self.kv_cache_manager.non_resident_block_ids(all_block_ids) + num_blocks_needed = self._blocks_needed(request, num_new) + if self.kv_cache_manager.num_free_blocks < len(missing_block_ids) + num_blocks_needed: + self.kv_cache_manager.release_cached_blocks(cached_blocks) + request.cached_block_ids = [] + request.num_computed_tokens = 0 + remaining_waiting.append(request) + break + if not self._ensure_blocks_resident(request, all_block_ids, output): + remaining_waiting.append(request) + continue - request.status = RequestStatus.RUNNING - self.running.append(request) - all_block_ids = request.cached_block_ids + request.allocated_block_ids - output.scheduled_requests.append( - ScheduledRequest( - request=request, - num_new_tokens=num_new, - is_prefill=True, - num_computed_tokens=request.num_computed_tokens, - block_ids=list(all_block_ids), + if not self._try_allocate_blocks(request, num_blocks_needed): + self.kv_cache_manager.release_cached_blocks(cached_blocks) + request.cached_block_ids = [] + request.num_computed_tokens = 0 + remaining_waiting.append(request) + break + + request.status = RequestStatus.RUNNING + self.running.append(request) + all_block_ids = request.cached_block_ids + request.allocated_block_ids + physical_block_ids = self.kv_cache_manager.resident_block_ids(all_block_ids) + output.scheduled_requests.append( + ScheduledRequest( + request=request, + num_new_tokens=num_new, + is_prefill=True, + num_computed_tokens=request.num_computed_tokens, + block_ids=physical_block_ids, + ) ) - ) - output.num_prefill_tokens += num_new - token_budget -= num_new + output.num_prefill_tokens += num_new + token_budget -= num_new remaining_waiting.extend(self.waiting) self.waiting = remaining_waiting @@ -350,25 +383,58 @@ def _try_allocate_blocks(self, request: Request, num_blocks: int) -> bool: request.allocated_block_ids.extend(block_ids) return True + def _ensure_blocks_resident( + self, + request: Request, + block_ids: list[int], + output: SchedulerOutput, + ) -> bool: + if any(self.kv_cache_manager.blocks[block_id].pending_job_id is not None for block_id in block_ids): + return False + missing_block_ids = self.kv_cache_manager.non_resident_block_ids(block_ids) + if not missing_block_ids: + return True + if self.config.max_cpu_offload_blocks <= 0: + return False + try: + job = self.kv_cache_manager.build_cpu_load_job( + missing_block_ids, + request_id=request.request_id, + ) + except RuntimeError: + return False + output.transfer_jobs.append(job) + output.jobs_to_flush.add(job.job_id) + self.num_cpu_load_jobs += 1 + self.num_cpu_load_blocks += len(missing_block_ids) + return False + + def _pop_lowest_priority_running_request(self, exclude: Request) -> Request | None: + """Remove and return the lowest-priority running request.""" + if not self.running: + return None + candidates = [r for r in self.running if r.request_id != exclude.request_id] + if candidates: + victim = max(candidates, key=lambda r: r.arrival_time) + self.running.remove(victim) + return victim + if exclude in self.running: + self.running.remove(exclude) + return exclude + return None + def _preempt_lowest_priority( self, - exclude: Request, + victim: Request, scheduled_req_ids: set[str], num_scheduled_tokens: dict[str, int], output: SchedulerOutput, ) -> dict | None: - """Preempt the lowest-priority running request to free blocks. + """Preempt a request that has already been removed from the running queue. If the victim was already scheduled in this iteration, it is removed from the scheduled output and its token budget is returned. """ - if not self.running: - return None - candidates = [r for r in self.running if r.request_id != exclude.request_id] - if not candidates: - return None - victim = max(candidates, key=lambda r: r.arrival_time) - returned_tokens = 0 if victim.request_id in scheduled_req_ids: scheduled_req_ids.discard(victim.request_id) @@ -381,15 +447,43 @@ def _preempt_lowest_priority( else: output.num_decode_tokens -= returned_tokens - self._free_request_blocks(victim) + if self.config.max_cpu_offload_blocks > 0: + resident_block_ids = [ + block_id for block_id in victim.cached_block_ids + victim.allocated_block_ids + if self.kv_cache_manager.blocks[block_id].location == KVBlockLocation.NPU + and self.kv_cache_manager.blocks[block_id].physical_page_id is not None + and self.kv_cache_manager.blocks[block_id].pending_job_id is None + and self.kv_cache_manager.blocks[block_id].ref_cnt == 1 + ] + if resident_block_ids: + try: + job = self.kv_cache_manager.build_cpu_store_job( + resident_block_ids, + request_id=victim.request_id, + ) + except RuntimeError: + job = None + if job is not None: + output.transfer_jobs.append(job) + output.jobs_to_flush.add(job.job_id) + self.num_cpu_store_jobs += 1 + self.num_cpu_store_blocks += len(resident_block_ids) + else: + self._reset_preempted_request_to_recompute(victim) + else: + self._reset_preempted_request_to_recompute(victim) + else: + self._reset_preempted_request_to_recompute(victim) victim.status = RequestStatus.PREEMPTED - victim.num_computed_tokens = 0 - victim.cached_block_ids = [] - victim.allocated_block_ids = [] - self.running = [r for r in self.running if r.request_id != victim.request_id] self.waiting.appendleft(victim) return {"request": victim, "returned_tokens": returned_tokens} + def _reset_preempted_request_to_recompute(self, request: Request) -> None: + self._free_request_blocks(request) + request.num_computed_tokens = 0 + request.cached_block_ids = [] + request.allocated_block_ids = [] + def _free_request_blocks(self, request: Request) -> None: self.kv_cache_manager.release_blocks_by_ids( request.cached_block_ids, diff --git a/python/core/serving_worker.py b/python/core/serving_worker.py index 0a7662f..e895fab 100644 --- a/python/core/serving_worker.py +++ b/python/core/serving_worker.py @@ -16,6 +16,15 @@ from typing import TYPE_CHECKING +from .kv_offload import ( + CPUKvOffloadBackend, + CPULoadStoreSpec, + KvOffloadBackend, + SSDLoadStoreSpec, + TransferJob, + TransferResult, + UnavailableKvOffloadBackend, +) from python.profile import get_profiler, profile_span from .types import ( DecodeBatch, @@ -52,6 +61,8 @@ def __init__( self.sampler = None self.model_record = None self._page_size: int = 64 + self._cpu_kv_offload_backend: CPUKvOffloadBackend | None = None + self._ssd_kv_offload_backend: KvOffloadBackend | None = None def init_device_and_model(self) -> None: from .model_loader import ModelLoader @@ -138,6 +149,7 @@ def _execute_step(self, scheduler_output) -> StepOutput: """Execute one batch step (may contain prefill + decode requests).""" runtime_model = self.model_record.runtime_model new_tokens: dict[str, int] = {} + completed_transfer_jobs = self._execute_kv_transfer_jobs(scheduler_output) prefill_requests = [ sr for sr in scheduler_output.scheduled_requests if sr.is_prefill @@ -157,7 +169,92 @@ def _execute_step(self, scheduler_output) -> StepOutput: if decode_requests: self._batch_decode(decode_requests, runtime_model, new_tokens) - return StepOutput(new_tokens=new_tokens) + return StepOutput(new_tokens=new_tokens, completed_transfer_jobs=completed_transfer_jobs) + + def _execute_kv_transfer_jobs(self, scheduler_output) -> list[TransferResult]: + """Execute worker-side KV transfer jobs carried by scheduler output.""" + transfer_jobs: list[TransferJob] = list(getattr(scheduler_output, "transfer_jobs", []) or []) + jobs_to_flush: set[int] = set(getattr(scheduler_output, "jobs_to_flush", set()) or set()) + if not transfer_jobs and not jobs_to_flush: + return [] + + completed: list[TransferResult] = [] + with profile_span( + "WorkerProcess.kv_transfer_jobs", + cat="worker", + args={"jobs": len(transfer_jobs), "flush": len(jobs_to_flush)}, + ): + for job in transfer_jobs: + try: + backend = self._ensure_kv_offload_backend(job) + backend.submit(job) + except Exception as exc: + completed.append(TransferResult(job_id=job.job_id, success=False, error=str(exc))) + + if self._cpu_kv_offload_backend is not None and jobs_to_flush: + completed.extend(self._cpu_kv_offload_backend.wait(jobs_to_flush)) + if self._ssd_kv_offload_backend is not None and jobs_to_flush: + completed.extend(self._ssd_kv_offload_backend.wait(jobs_to_flush)) + + if self._cpu_kv_offload_backend is not None: + completed.extend(self._cpu_kv_offload_backend.poll()) + if self._ssd_kv_offload_backend is not None: + completed.extend(self._ssd_kv_offload_backend.poll()) + + return completed + + def _ensure_kv_offload_backend(self, job: TransferJob) -> KvOffloadBackend: + """Return the worker backend responsible for one KV transfer job.""" + if self._max_cpu_slot_id(job) is not None: + return self._ensure_cpu_kv_offload_backend(job) + if self._uses_ssd(job): + return self._ensure_ssd_kv_offload_backend() + raise ValueError(f"unsupported KV transfer direction: {job.direction}") + + def _ensure_cpu_kv_offload_backend(self, job: TransferJob) -> CPUKvOffloadBackend: + """Create or grow the CPU KV offload backend required by one job.""" + max_slot_id = self._max_cpu_slot_id(job) + if max_slot_id is None: + raise ValueError(f"unsupported KV transfer direction: {job.direction}") + required_slots = max_slot_id + 1 + if self._cpu_kv_offload_backend is None: + materialize = getattr(self.executor, "materialize_kv_page_view", None) + if not callable(materialize): + raise RuntimeError("executor does not expose a worker KV page view") + page_view = materialize(self.config.model_id) + self._cpu_kv_offload_backend = CPUKvOffloadBackend(page_view, num_cpu_slots=required_slots) + else: + self._cpu_kv_offload_backend.ensure_num_cpu_slots(required_slots) + return self._cpu_kv_offload_backend + + def _ensure_ssd_kv_offload_backend(self) -> KvOffloadBackend: + """Return the SSD KV backend placeholder until the real NPU-SSD API is wired.""" + if self._ssd_kv_offload_backend is None: + self._ssd_kv_offload_backend = UnavailableKvOffloadBackend( + "SSD", + reason="NPU-attached SSD KV transfer API is not configured", + ) + return self._ssd_kv_offload_backend + + @staticmethod + def _max_cpu_slot_id(job: TransferJob) -> int | None: + cpu_specs = [ + spec for spec in (job.src, job.dst) + if isinstance(spec, CPULoadStoreSpec) + ] + if not cpu_specs: + return None + max_slot_id = -1 + for spec in cpu_specs: + if spec.block_ids: + max_slot_id = max(max_slot_id, max(spec.block_ids)) + if max_slot_id < 0: + raise ValueError("CPU transfer specs must include at least one slot") + return max_slot_id + + @staticmethod + def _uses_ssd(job: TransferJob) -> bool: + return isinstance(job.src, SSDLoadStoreSpec) or isinstance(job.dst, SSDLoadStoreSpec) def _batch_prefill( self, scheduled: list, runtime_model, new_tokens: dict[str, int] diff --git a/python/core/types.py b/python/core/types.py index 97d14a1..9ea27b5 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -14,6 +14,7 @@ import torch +from .kv_offload import TransferResult from .tokenizer import TokenizerAdapter @@ -235,4 +236,5 @@ class WorkerCommand: class StepOutput: """Result returned from worker process after executing a batch step.""" new_tokens: dict # {request_id: int} + completed_transfer_jobs: list[TransferResult] = field(default_factory=list) error: str | None = None diff --git a/tests/test_cli.py b/tests/test_cli.py index f88e5cb..ad562bc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -47,6 +47,7 @@ def test_build_serving_engine_config_uses_cli_args(tmp_path, monkeypatch): "--long-prefill-token-threshold", "64", "--no-enable-prefix-caching", "--no-enable-chunked-prefill", + "--max-cpu-offload-blocks", "12", ]) config = cli.build_serving_engine_config(args) @@ -71,6 +72,7 @@ def test_build_serving_engine_config_uses_cli_args(tmp_path, monkeypatch): assert config.long_prefill_token_threshold == 64 assert config.enable_prefix_cache is False assert config.enable_chunk_prefill is False + assert config.max_cpu_offload_blocks == 12 def test_parser_rejects_invalid_backend(tmp_path): @@ -89,6 +91,14 @@ def test_parser_rejects_removed_prompt_mode(tmp_path): _parse_args(["--model", str(model_dir), "--backend", "npu", "--prompt", "hello"]) +def test_parser_rejects_negative_cpu_offload_blocks(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + + with pytest.raises(SystemExit): + _parse_args(["--model", str(model_dir), "--max-cpu-offload-blocks", "-1"]) + + def test_main_starts_serving(tmp_path, monkeypatch): model_dir = tmp_path / "model" model_dir.mkdir()