From 36d254eeb0dc8a1631c49301f9959c461c389e91 Mon Sep 17 00:00:00 2001 From: superxf <1208713646@qq.com> Date: Mon, 15 Jun 2026 15:40:35 +0800 Subject: [PATCH] Add KV cache CPU offload support --- .../model/qwen3_14b/runner/npu_executor.py | 6 - examples/model/qwen3_14b/runner/npu_runner.py | 21 ++ python/cli/main.py | 29 +- python/core/async_engine.py | 78 +++++ python/core/kv_cache.py | 226 ++++++++++++++- python/core/kv_offload.py | 267 ++++++++++++++++++ python/core/model_runner.py | 15 + python/core/pypto_executor.py | 5 + python/core/scheduler.py | 74 ++++- python/core/serving_worker.py | 97 ++++++- python/core/types.py | 5 +- tests/test_cli.py | 26 +- 12 files changed, 827 insertions(+), 22 deletions(-) create mode 100644 python/core/kv_offload.py diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index c9b0ba5..68648a7 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -31,7 +31,6 @@ _VOCAB_PAD_MULTIPLE = 512 # must be a multiple of lm_head.VOCAB_CHUNK (64) -_QWEN14B_PAGE_SIZE = 128 _QWEN14B_BLOCK_DIM = 24 @@ -577,8 +576,3 @@ def _validate_supported_shape(model: RuntimeModel) -> None: raise ValueError( "Bundled kernels under model/ currently support Qwen3-14B layer shapes only: " + mismatch ) - if model.runtime.page_size != _QWEN14B_PAGE_SIZE: - raise ValueError( - "PyPTO Qwen3-14B kernels require runtime page_size " - f"{_QWEN14B_PAGE_SIZE}, got {model.runtime.page_size}." - ) diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 56dc597..7b4b16f 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -15,6 +15,7 @@ import torch from pypto.runtime import DeviceTensor +from python.core.kv_offload import WorkerKVPageView from python.core.model_runner import ModelRunner from python.core.types import ( DecodeBatch, @@ -201,6 +202,26 @@ def _materialize_kv_cache(self, model: RuntimeModel) -> Any: ModelRunner.init_kv_cache(self, model.config.model_id, spec[0], spec[1]) return self._kv_caches[model.config.model_id] + def materialize_kv_page_view(self, model_id: str) -> WorkerKVPageView: + """Return a byte-level transfer view over this model's runner-owned KV pages.""" + kv_cache = self._kv_caches.get(model_id) + if kv_cache is None: + spec = self._pending_kv_cache_specs.get(model_id) + if spec is None: + raise RuntimeError(f"KV cache for model {model_id!r} is not initialized") + ModelRunner.init_kv_cache(self, model_id, spec[0], spec[1]) + kv_cache = self._kv_caches[model_id] + return WorkerKVPageView( + worker=self._shared_l3_worker(), + key_pages=kv_cache.key_pages, + value_pages=kv_cache.value_pages, + num_layers=kv_cache.num_layers, + num_pages=kv_cache.num_pages, + num_kv_heads=kv_cache.num_kv_heads, + page_size=kv_cache.page_size, + head_dim=kv_cache.head_dim, + ) + @staticmethod def _validate_kv_cache_bounds( model: RuntimeModel, diff --git a/python/cli/main.py b/python/cli/main.py index 32342b9..6b9c753 100644 --- a/python/cli/main.py +++ b/python/cli/main.py @@ -28,6 +28,13 @@ _VALID_BACKENDS = {"npu"} +def _supported_block_size(value: str) -> int: + parsed = int(value) + if parsed != 128: + raise argparse.ArgumentTypeError("only block size 128 is currently supported") + return parsed + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="pypto-serving", @@ -48,7 +55,12 @@ def build_parser() -> argparse.ArgumentParser: # Runtime parser.add_argument("--max-model-len", type=int, default=512, help="Maximum sequence length (default: 512).") - parser.add_argument("--block-size", type=int, default=128, help="KV cache block size (default: 128).") + parser.add_argument( + "--block-size", + type=_supported_block_size, + default=128, + help="KV cache block size. Currently only 128 is supported (default: 128).", + ) # Generation parser.add_argument("--max-new-tokens", type=int, default=32, help="Maximum new tokens to generate (default: 32).") @@ -64,8 +76,8 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--long-prefill-token-threshold", type=int, - default=2048, - help="Chunked prefill threshold in serving mode (default: 2048).", + default=64, + help="Chunked prefill threshold in serving mode (default: 64).", ) parser.add_argument( "--enable-prefix-caching", @@ -79,6 +91,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=int, + default=0, + help="Maximum number of KV blocks to keep in CPU offload storage. 0 disables CPU offload (default: 0).", + ) # Misc parser.add_argument( @@ -114,6 +132,8 @@ 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, + enable_kv_cpu_offload=args.max_cpu_offload_blocks > 0, + max_cpu_offload_blocks=args.max_cpu_offload_blocks, ) @@ -189,6 +209,9 @@ 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" KV CPU offload: {'enabled' if config.enable_kv_cpu_offload else 'disabled'}") + if config.enable_kv_cpu_offload: + print(f" Max CPU 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..89ee0a1 100644 --- a/python/core/async_engine.py +++ b/python/core/async_engine.py @@ -49,6 +49,8 @@ class EngineConfig: # Feature flags enable_prefix_cache: bool = True enable_chunk_prefill: bool = True + enable_kv_cpu_offload: bool = False + max_cpu_offload_blocks: int = 0 @dataclass @@ -93,6 +95,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 +105,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, + enable_kv_cpu_offload=self.config.enable_kv_cpu_offload, ) self.scheduler = Scheduler(config=scheduler_config, kv_cache_manager=self.kv_cache_manager) @@ -228,6 +232,20 @@ async def _engine_loop(self) -> None: await asyncio.sleep(self.config.engine_loop_interval) continue + if scheduler_output.kv_transfer_jobs: + transfer_ok = await self._execute_kv_transfer_jobs( + scheduler_output.kv_transfer_jobs, + profile_name="scheduler.load_kv_transfer", + timeout_message="Worker response timed out during load KV transfer (300s)", + error_prefix="Worker returned load KV transfer error", + ) + if not transfer_ok: + self._handle_step_error(scheduler_output) + continue + + if not scheduler_output.scheduled_requests: + continue + finished_ids = self._pending_free_ids.copy() self._pending_free_ids.clear() with profile_span( @@ -264,13 +282,73 @@ async def _engine_loop(self) -> None: args={"new_tokens": len(step_output.new_tokens)}, ): self._process_step_output(scheduler_output, step_output) + await self._drain_store_kv_transfer_jobs() logger.info("Engine loop stopped") + async def _drain_store_kv_transfer_jobs(self) -> None: + """Execute NPU->CPU stores queued by requests that just finished.""" + jobs = self.scheduler.pop_store_kv_transfer_jobs() + if not jobs: + return + await self._execute_kv_transfer_jobs( + jobs, + profile_name="scheduler.store_kv_transfer", + timeout_message="Worker response timed out during store KV transfer (300s)", + error_prefix="Worker returned store KV transfer error", + ) + + async def _execute_kv_transfer_jobs( + self, + jobs, + *, + profile_name: str, + timeout_message: str, + error_prefix: str, + ) -> bool: + """Send KV transfer jobs to the worker and apply their completions.""" + with profile_span( + f"{profile_name}.queue", + cat="scheduler", + args={"jobs": len(jobs)}, + ): + self._input_queue.put( + WorkerCommand( + type="kv_transfer", + kv_transfer_jobs=jobs, + ) + ) + + try: + with profile_span(f"{profile_name}.wait", cat="scheduler"): + step_output: StepOutput = await asyncio.to_thread( + self._output_queue.get, timeout=300 + ) + except queue.Empty: + logger.error(timeout_message) + return False + + if step_output.error: + logger.error(f"{error_prefix}: {step_output.error}") + return False + + try: + self.scheduler.complete_transfer_results(step_output.kv_transfer_results) + except RuntimeError as exc: + logger.error(f"KV transfer completion failed: {exc}") + return False + return True + def _process_step_output( self, scheduler_output: SchedulerOutput, step_output: StepOutput ) -> None: """Process worker results: update scheduler state, push tokens to request queues.""" + try: + self.scheduler.complete_transfer_results(step_output.kv_transfer_results) + except RuntimeError as exc: + logger.error(f"KV transfer completion failed: {exc}") + self._handle_step_error(scheduler_output) + return request_outputs = self.scheduler.update_from_output( scheduler_output, step_output.new_tokens ) diff --git a/python/core/kv_cache.py b/python/core/kv_cache.py index 9d0096e..ceb083f 100644 --- a/python/core/kv_cache.py +++ b/python/core/kv_cache.py @@ -14,6 +14,13 @@ import torch +from .kv_offload import ( + CPULoadStoreSpec, + KVBlockLocation, + NPULoadStoreSpec, + TransferJob, + TransferResult, +) from .types import KvAllocation, ModelConfig, RuntimeConfig @@ -32,6 +39,10 @@ class KVCacheBlock: block_id: int ref_cnt: int = 0 block_hash: int | None = None + location: KVBlockLocation = KVBlockLocation.NPU + physical_page_id: int | None = None + cpu_slot_id: int | None = None + cpu_last_access: int = 0 prev_free: "KVCacheBlock | None" = field(default=None, repr=False) next_free: "KVCacheBlock | None" = field(default=None, repr=False) @@ -123,6 +134,7 @@ 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.""" self._pools: dict[str, _CachePool] = {} @@ -132,6 +144,11 @@ def __init__( self.free_queue = FreeKVCacheBlockQueue() self.hash_to_block: dict[int, KVCacheBlock] = {} self.request_blocks: dict[str, list[KVCacheBlock]] = {} + self.max_cpu_offload_blocks = int(max_cpu_offload_blocks) + self._next_transfer_job_id = 0 + self._cpu_access_clock = 0 + self._free_cpu_slots: list[int] = list(range(self.max_cpu_offload_blocks)) + self._pending_transfer_blocks: dict[int, list[int]] = {} if num_blocks is not None: self._init_blocks(num_blocks, block_size) @@ -145,13 +162,21 @@ def num_blocks(self) -> int: """Return the total number of physical KV blocks.""" return len(self.blocks) + @property + def num_free_cpu_slots(self) -> int: + """Return the number of CPU offload slots available for new stores.""" + return len(self._free_cpu_slots) + def _init_blocks(self, num_blocks: int, block_size: int) -> None: if self.blocks: if len(self.blocks) != num_blocks or self.block_size != block_size: raise ValueError("KV block pool is already initialized with different dimensions") return self.block_size = block_size - self.blocks = [KVCacheBlock(block_id=i) for i in range(num_blocks)] + self.blocks = [ + KVCacheBlock(block_id=i, physical_page_id=i) + for i in range(num_blocks) + ] for block in self.blocks: self.free_queue.append(block) @@ -204,21 +229,44 @@ def allocate_for_prompt(self, model_id: str, request_id: str, prompt_len: int) - def allocate_blocks(self, num_blocks: int) -> list[KVCacheBlock] | None: """Allocate physical KV blocks, evicting stale prefix hashes as needed.""" + if num_blocks <= 0: + return [] + return self.allocate_preferred_blocks(num_blocks, []) + + def allocate_preferred_blocks( + self, + num_blocks: int, + preferred_block_ids: list[int], + ) -> list[KVCacheBlock] | None: + """Allocate blocks, trying specific free block IDs before FIFO fallback.""" if num_blocks <= 0: return [] if self.num_free_blocks < num_blocks: return None blocks: list[KVCacheBlock] = [] + used_preferred: set[int] = set() + for block_id in preferred_block_ids: + if len(blocks) >= num_blocks: + break + if block_id in used_preferred or block_id < 0 or block_id >= len(self.blocks): + continue + block = self.blocks[block_id] + if not self._is_free(block): + continue + self.free_queue.remove(block) + self._prepare_allocated_block(block) + blocks.append(block) + used_preferred.add(block_id) + for _ in range(num_blocks): + if len(blocks) >= num_blocks: + break block = self.free_queue.popleft() if block is None: for allocated in blocks: self.release(allocated) return None - if block.block_hash is not None: - self.hash_to_block.pop(block.block_hash, None) - block.block_hash = None - block.ref_cnt = 1 + self._prepare_allocated_block(block) blocks.append(block) return blocks @@ -229,6 +277,17 @@ def allocate_block_ids(self, num_blocks: int) -> list[int] | None: return None return [block.block_id for block in blocks] + def allocate_preferred_block_ids( + self, + num_blocks: int, + preferred_block_ids: list[int], + ) -> list[int] | None: + """Allocate physical KV blocks and return their IDs, preferring given IDs.""" + blocks = self.allocate_preferred_blocks(num_blocks, preferred_block_ids) + if blocks is None: + return None + return [block.block_id for block in blocks] + def release_blocks_by_ids(self, *block_id_groups: list[int]) -> None: """Release request references for one or more groups of physical block IDs.""" for block_ids in block_id_groups: @@ -256,6 +315,8 @@ def get_cached_block(self, block_hash: int) -> KVCacheBlock | None: if block.ref_cnt == 0: self.free_queue.remove(block) block.ref_cnt += 1 + if block.location == KVBlockLocation.CPU: + self._touch_cpu_block(block) return block def cache_block(self, block: KVCacheBlock, block_hash: int) -> None: @@ -276,6 +337,108 @@ def cache_block_ids(self, block_ids: list[int], block_hashes: list[int], start: break self.cache_block(self.blocks[block_ids[idx]], block_hashes[idx]) + def build_cpu_store_job(self, block_ids: list[int], request_id: str | None = None) -> TransferJob: + """Build a transfer job that stores resident NPU blocks into CPU slots.""" + blocks = [self.blocks[block_id] for block_id in block_ids] + for block in blocks: + if block.ref_cnt != 0: + raise RuntimeError(f"Cannot offload KV block {block.block_id} with active references.") + if len(blocks) > self.max_cpu_offload_blocks: + raise RuntimeError("Insufficient CPU KV offload slots.") + while len(self._free_cpu_slots) < len(blocks): + if self._evict_lru_cpu_block() is None: + raise RuntimeError("Insufficient CPU KV offload slots.") + + page_ids: list[int] = [] + slot_ids: list[int] = [] + for block in blocks: + 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.") + if block.ref_cnt == 0: + self.free_queue.remove(block) + slot_id = self._free_cpu_slots.pop(0) + block.location = KVBlockLocation.TRANSFERRING + block.cpu_slot_id = slot_id + block.cpu_last_access = 0 + page_ids.append(block.physical_page_id) + slot_ids.append(slot_id) + + job = TransferJob( + self._next_job_id(), + request_id, + NPULoadStoreSpec(page_ids), + CPULoadStoreSpec(slot_ids), + ) + self._pending_transfer_blocks[job.job_id] = [block.block_id for block in blocks] + return job + + def build_cpu_load_job(self, block_ids: list[int], request_id: str | None = None) -> TransferJob: + """Build a transfer job that loads CPU-offloaded blocks back to NPU pages.""" + blocks = [self.blocks[block_id] for block_id in block_ids] + page_ids: list[int] = [] + slot_ids: list[int] = [] + for block in blocks: + if block.location != KVBlockLocation.CPU or block.cpu_slot_id is None: + raise RuntimeError(f"KV block {block.block_id} is not resident on CPU.") + block.location = KVBlockLocation.TRANSFERRING + block.physical_page_id = block.block_id + block.cpu_last_access = 0 + page_ids.append(block.physical_page_id) + slot_ids.append(block.cpu_slot_id) + + job = TransferJob( + self._next_job_id(), + request_id, + CPULoadStoreSpec(slot_ids), + NPULoadStoreSpec(page_ids), + ) + self._pending_transfer_blocks[job.job_id] = [block.block_id for block in blocks] + return job + + def complete_transfer_result(self, result: TransferResult) -> None: + """Apply a completed CPU offload transfer to block residency metadata.""" + block_ids = self._pending_transfer_blocks.pop(result.job_id, []) + if not block_ids: + return + blocks = [self.blocks[block_id] for block_id in block_ids] + if not result.success: + if isinstance(result.src, NPULoadStoreSpec) and isinstance(result.dst, CPULoadStoreSpec): + for block in blocks: + block.location = KVBlockLocation.NPU + if block.cpu_slot_id is not None: + self._free_cpu_slots.append(block.cpu_slot_id) + block.cpu_slot_id = None + block.cpu_last_access = 0 + if block.ref_cnt == 0: + self.free_queue.append(block) + elif isinstance(result.src, CPULoadStoreSpec) and isinstance(result.dst, NPULoadStoreSpec): + for block in blocks: + block.location = KVBlockLocation.CPU + block.physical_page_id = None + block.cpu_last_access = 0 + else: + raise TypeError("Unsupported KV offload transfer result") + raise RuntimeError(result.error or "KV offload transfer failed") + + if isinstance(result.src, NPULoadStoreSpec) and isinstance(result.dst, CPULoadStoreSpec): + for block, slot_id in zip(blocks, result.dst.slot_ids, strict=True): + block.location = KVBlockLocation.CPU + block.physical_page_id = None + block.cpu_slot_id = slot_id + self._touch_cpu_block(block) + if block.ref_cnt == 0: + self.free_queue.append(block) + elif isinstance(result.src, CPULoadStoreSpec) and isinstance(result.dst, NPULoadStoreSpec): + for block, page_id in zip(blocks, result.dst.page_ids, strict=True): + block.location = KVBlockLocation.NPU + block.physical_page_id = page_id + if block.cpu_slot_id is not None: + self._free_cpu_slots.append(block.cpu_slot_id) + block.cpu_slot_id = None + block.cpu_last_access = 0 + else: + raise TypeError("Unsupported KV offload transfer result") + def release(self, block: KVCacheBlock) -> None: """Release one request reference to a block.""" if block.ref_cnt <= 0: @@ -284,6 +447,59 @@ def release(self, block: KVCacheBlock) -> None: if block.ref_cnt == 0: self.free_queue.append(block) + def _next_job_id(self) -> int: + job_id = self._next_transfer_job_id + self._next_transfer_job_id += 1 + return job_id + + def _is_free(self, block: KVCacheBlock) -> bool: + return ( + block.ref_cnt == 0 + and ( + block == self.free_queue.head + or block == self.free_queue.tail + or block.prev_free is not None + or block.next_free is not None + ) + ) + + def _prepare_allocated_block(self, block: KVCacheBlock) -> None: + if block.block_hash is not None: + self.hash_to_block.pop(block.block_hash, None) + block.block_hash = None + if block.cpu_slot_id is not None: + self._free_cpu_slots.append(block.cpu_slot_id) + block.location = KVBlockLocation.NPU + block.physical_page_id = block.block_id + block.cpu_slot_id = None + block.cpu_last_access = 0 + block.ref_cnt = 1 + + def _touch_cpu_block(self, block: KVCacheBlock) -> None: + self._cpu_access_clock += 1 + block.cpu_last_access = self._cpu_access_clock + + def _evict_lru_cpu_block(self) -> KVCacheBlock | None: + candidates = [ + block + for block in self.blocks + if block.location == KVBlockLocation.CPU + and block.ref_cnt == 0 + and block.cpu_slot_id is not None + ] + if not candidates: + return None + victim = min(candidates, key=lambda block: block.cpu_last_access) + if victim.block_hash is not None: + self.hash_to_block.pop(victim.block_hash, None) + victim.block_hash = None + self._free_cpu_slots.append(victim.cpu_slot_id) + victim.cpu_slot_id = None + victim.cpu_last_access = 0 + victim.location = KVBlockLocation.NPU + victim.physical_page_id = victim.block_id + return victim + def _iter_block_hashes(self, token_ids: list[int]): """Yield (block_index, block_hash) for each full block in the token sequence.""" parent_hash = NONE_HASH diff --git a/python/core/kv_offload.py b/python/core/kv_offload.py new file mode 100644 index 0000000..89bf2d8 --- /dev/null +++ b/python/core/kv_offload.py @@ -0,0 +1,267 @@ +# 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 + +import itertools +from dataclasses import dataclass +from enum import Enum, auto +from typing import Any + +import torch +from pypto.runtime import DeviceTensor + + +class KVBlockLocation(Enum): + """Current residency of one logical KV block.""" + + NPU = auto() + CPU = auto() + TRANSFERRING = auto() + + +@dataclass(frozen=True) +class NPULoadStoreSpec: + """A list of physical NPU page IDs used as a transfer endpoint.""" + + page_ids: list[int] + + +@dataclass(frozen=True) +class CPULoadStoreSpec: + """A list of CPU offload slot IDs used as a transfer endpoint.""" + + slot_ids: list[int] + + +@dataclass(frozen=True) +class TransferJob: + """One page-level KV transfer between NPU cache pages and CPU slots.""" + + job_id: int + request_id: str | None + src: NPULoadStoreSpec | CPULoadStoreSpec + dst: NPULoadStoreSpec | CPULoadStoreSpec + + +@dataclass(frozen=True) +class TransferResult: + """Completion status for one KV transfer job.""" + + job_id: int + request_id: str | None + src: NPULoadStoreSpec | CPULoadStoreSpec + dst: NPULoadStoreSpec | CPULoadStoreSpec + success: bool + error: str | None = None + + +class WorkerKVPageView: + """Byte-level view over runner-owned paged K/V DeviceTensors. + + The runner stores KV as a flat view of + ``[num_layers, num_pages, num_kv_heads, page_size, head_dim]``. A logical + page is not one contiguous range across layers, so copies are split into + one contiguous K segment and one contiguous V segment per layer. + """ + + def __init__( + self, + *, + worker: Any, + key_pages: DeviceTensor, + value_pages: DeviceTensor, + num_layers: int, + num_pages: int, + num_kv_heads: int, + page_size: int, + head_dim: int, + ) -> None: + self.worker = worker + self.key_pages = key_pages + self.value_pages = value_pages + self.num_layers = int(num_layers) + self.num_pages = int(num_pages) + self.num_kv_heads = int(num_kv_heads) + self.page_size = int(page_size) + self.head_dim = int(head_dim) + self.dtype = key_pages.dtype + if value_pages.dtype != self.dtype: + raise ValueError("key_pages and value_pages must have the same dtype") + expected_shape = ( + self.num_layers * self.num_pages * self.num_kv_heads * self.page_size, + self.head_dim, + ) + if tuple(key_pages.shape) != expected_shape or tuple(value_pages.shape) != expected_shape: + raise ValueError( + "KV page tensors do not match the expected flat paged layout: " + f"expected={expected_shape}, key={key_pages.shape}, value={value_pages.shape}" + ) + + @property + def _element_size(self) -> int: + return torch.empty((), dtype=self.dtype).element_size() + + @property + def _layer_page_bytes(self) -> int: + return self.num_kv_heads * self.page_size * self.head_dim * self._element_size + + @property + def page_size_bytes(self) -> int: + """Total bytes for one logical page, including K and V for all layers.""" + + return 2 * self.num_layers * self._layer_page_bytes + + def copy_page_to(self, page_id: int, dst: torch.Tensor) -> None: + """Copy one NPU page into a contiguous uint8 CPU tensor.""" + + dst_u8 = self._validate_host_buffer(dst) + offset = 0 + for tensor in (self.key_pages, self.value_pages): + for layer_idx in range(self.num_layers): + self.worker.copy_from( + dst_u8.data_ptr() + offset, + self._device_segment_ptr(tensor, page_id, layer_idx), + self._layer_page_bytes, + ) + offset += self._layer_page_bytes + + def copy_page_from(self, page_id: int, src: torch.Tensor) -> None: + """Copy one contiguous uint8 CPU tensor into an NPU page.""" + + src_u8 = self._validate_host_buffer(src) + offset = 0 + for tensor in (self.key_pages, self.value_pages): + for layer_idx in range(self.num_layers): + self.worker.copy_to( + self._device_segment_ptr(tensor, page_id, layer_idx), + src_u8.data_ptr() + offset, + self._layer_page_bytes, + ) + offset += self._layer_page_bytes + + def _device_segment_ptr(self, tensor: DeviceTensor, page_id: int, layer_idx: int) -> int: + if page_id < 0 or page_id >= self.num_pages: + raise ValueError(f"page_id {page_id} is outside [0, {self.num_pages})") + if layer_idx < 0 or layer_idx >= self.num_layers: + raise ValueError(f"layer_idx {layer_idx} is outside [0, {self.num_layers})") + row = (layer_idx * self.num_pages + page_id) * self.num_kv_heads * self.page_size + return tensor.data_ptr + row * self.head_dim * self._element_size + + def _validate_host_buffer(self, tensor: torch.Tensor) -> torch.Tensor: + if tensor.device.type != "cpu": + raise ValueError("KV offload host buffer must be on CPU") + if tensor.dtype != torch.uint8: + raise ValueError("KV offload host buffer must be torch.uint8") + if not tensor.is_contiguous(): + raise ValueError("KV offload host buffer must be contiguous") + if tensor.numel() != self.page_size_bytes: + raise ValueError( + f"KV offload host buffer has {tensor.numel()} bytes, expected {self.page_size_bytes}" + ) + return tensor + + +class CPUKvOffloadBackend: + """Synchronous CPU offload backend for runner-owned NPU KV pages.""" + + def __init__( + self, + page_view: WorkerKVPageView, + *, + num_cpu_slots: int, + cpu_slots: torch.Tensor | None = None, + ) -> None: + self.page_view = page_view + if cpu_slots is None: + self.cpu_slots = torch.empty( + (int(num_cpu_slots), page_view.page_size_bytes), + dtype=torch.uint8, + device="cpu", + ).share_memory_() + else: + if tuple(cpu_slots.shape) != (int(num_cpu_slots), page_view.page_size_bytes): + raise ValueError( + "Preallocated CPU KV slots have wrong shape: " + f"got={tuple(cpu_slots.shape)}, expected={(int(num_cpu_slots), page_view.page_size_bytes)}" + ) + if cpu_slots.dtype != torch.uint8 or cpu_slots.device.type != "cpu" or not cpu_slots.is_contiguous(): + raise ValueError("Preallocated CPU KV slots must be contiguous CPU uint8") + if not cpu_slots.is_shared(): + raise ValueError("Preallocated CPU KV slots must be in shared memory") + self.cpu_slots = cpu_slots + self._completed: dict[int, TransferResult] = {} + self._job_ids = itertools.count() + + def ensure_num_cpu_slots(self, num_cpu_slots: int) -> None: + """Validate that the fixed shared CPU slot pool is large enough.""" + + num_cpu_slots = int(num_cpu_slots) + if num_cpu_slots <= self.cpu_slots.shape[0]: + return + raise RuntimeError( + "CPU KV offload slots cannot be resized after initialization: " + f"available={self.cpu_slots.shape[0]}, requested={num_cpu_slots}" + ) + + def next_job_id(self) -> int: + """Return a backend-local monotonically increasing job ID.""" + + return next(self._job_ids) + + def submit(self, job: TransferJob) -> bool: + """Run one transfer synchronously and record its result.""" + + try: + if isinstance(job.src, NPULoadStoreSpec) and isinstance(job.dst, CPULoadStoreSpec): + self._copy_npu_to_cpu(job.src.page_ids, job.dst.slot_ids) + elif isinstance(job.src, CPULoadStoreSpec) and isinstance(job.dst, NPULoadStoreSpec): + self._copy_cpu_to_npu(job.src.slot_ids, job.dst.page_ids) + else: + raise TypeError("KV offload only supports NPU<->CPU transfers") + except Exception as exc: + self._completed[job.job_id] = TransferResult( + job.job_id, + job.request_id, + job.src, + job.dst, + False, + str(exc), + ) + return False + self._completed[job.job_id] = TransferResult( + job.job_id, + job.request_id, + job.src, + job.dst, + True, + ) + return True + + def wait(self, job_ids: set[int]) -> list[TransferResult]: + """Return completed results for the requested job IDs.""" + + missing = job_ids - self._completed.keys() + if missing: + raise RuntimeError(f"Unknown or incomplete KV offload job IDs: {sorted(missing)}") + return [self._completed.pop(job_id) for job_id in job_ids] + + def _copy_npu_to_cpu(self, page_ids: list[int], slot_ids: list[int]) -> None: + if len(page_ids) != len(slot_ids): + raise ValueError("NPU page ids and CPU slot ids must have the same length") + self.ensure_num_cpu_slots(max(slot_ids, default=-1) + 1) + for page_id, slot_id in zip(page_ids, slot_ids, strict=True): + self.page_view.copy_page_to(page_id, self.cpu_slots[slot_id]) + + def _copy_cpu_to_npu(self, slot_ids: list[int], page_ids: list[int]) -> None: + if len(slot_ids) != len(page_ids): + raise ValueError("CPU slot ids and NPU page ids must have the same length") + self.ensure_num_cpu_slots(max(slot_ids, default=-1) + 1) + for slot_id, page_id in zip(slot_ids, page_ids, strict=True): + self.page_view.copy_page_from(page_id, self.cpu_slots[slot_id]) diff --git a/python/core/model_runner.py b/python/core/model_runner.py index 28830cc..70e5f7e 100644 --- a/python/core/model_runner.py +++ b/python/core/model_runner.py @@ -17,6 +17,7 @@ from pypto.runtime import DeviceTensor +from .kv_offload import WorkerKVPageView from .types import ( DecodeBatch, DecodeResult, @@ -34,6 +35,11 @@ class _KvCachePool: key_pages: DeviceTensor value_pages: DeviceTensor + 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,10 @@ def close_kv_cache(self) -> None: self._free_kv_cache_tensor(pool.value_pages) self._kv_caches.clear() + def materialize_kv_page_view(self, model_id: str) -> WorkerKVPageView: + """Return a byte-level view over runner-owned KV pages.""" + raise NotImplementedError("This runner does not expose KV page transfers.") + @abstractmethod def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> DeviceTensor: """Allocate one worker-resident KV cache tensor.""" diff --git a/python/core/pypto_executor.py b/python/core/pypto_executor.py index 6785d39..a289ab5 100644 --- a/python/core/pypto_executor.py +++ b/python/core/pypto_executor.py @@ -15,6 +15,7 @@ from .executor import ModelExecutor from .kv_cache import KvCacheManager +from .kv_offload import WorkerKVPageView from .model_runner import ModelRunner from python.profile import profile_span from .types import ( @@ -77,6 +78,10 @@ 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) -> WorkerKVPageView: + """Return a byte-level transfer view for one model's runner-owned KV cache.""" + return self._runners[model_id].materialize_kv_page_view(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..7ac0c72 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, TransferResult class RequestStatus(Enum): @@ -45,6 +46,7 @@ class SchedulerConfig: # Feature flags enable_prefix_cache: bool = True enable_chunk_prefill: bool = True + enable_kv_cpu_offload: bool = False @dataclass @@ -100,12 +102,13 @@ class ScheduledRequest: class SchedulerOutput: scheduled_requests: list[ScheduledRequest] = field(default_factory=list) preempted_requests: list[Request] = field(default_factory=list) + kv_transfer_jobs: list[TransferJob] = field(default_factory=list) 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.kv_transfer_jobs) == 0 @dataclass @@ -125,6 +128,7 @@ def __init__(self, config: SchedulerConfig, kv_cache_manager: KvCacheManager) -> self.waiting: deque[Request] = deque() self.running: list[Request] = [] self.requests: dict[str, Request] = {} + self._pending_store_kv_transfer_jobs: list[TransferJob] = [] def add_request(self, request: Request) -> None: if len(request.prompt_token_ids) > self.config.max_seq_len: @@ -151,14 +155,24 @@ def finish_request(self, request_id: str, status: RequestStatus) -> None: if request is None: return request.status = status + release_block_ids = request.cached_block_ids + request.allocated_block_ids self._free_request_blocks(request) + self._queue_cpu_store_jobs(release_block_ids, request.request_id) self.running = [r for r in self.running if r.request_id != request_id] def has_work(self) -> bool: - return len(self.running) > 0 or len(self.waiting) > 0 + return ( + len(self.running) > 0 + or len(self.waiting) > 0 + or len(self._pending_store_kv_transfer_jobs) > 0 + ) def schedule(self) -> SchedulerOutput: output = SchedulerOutput() + if self._pending_store_kv_transfer_jobs: + output.kv_transfer_jobs.extend(self._pending_store_kv_transfer_jobs) + self._pending_store_kv_transfer_jobs = [] + return output token_budget = self.config.max_num_scheduled_tokens # Phase 1: schedule RUNNING requests (decode or resumed prefill) @@ -229,6 +243,19 @@ def schedule(self) -> SchedulerOutput: 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 + if self.config.enable_kv_cpu_offload: + cpu_block_ids = [ + block.block_id + for block in cached_blocks + if block.location == KVBlockLocation.CPU + ] + if cpu_block_ids: + output.kv_transfer_jobs.append( + self.kv_cache_manager.build_cpu_load_job( + cpu_block_ids, + request_id=request.request_id, + ) + ) else: cached_blocks = [] @@ -269,6 +296,17 @@ def schedule(self) -> SchedulerOutput: return output + def complete_transfer_results(self, results: list[TransferResult]) -> None: + """Apply KV transfer completions produced by the worker process.""" + for result in results: + self.kv_cache_manager.complete_transfer_result(result) + + def pop_store_kv_transfer_jobs(self) -> list[TransferJob]: + """Return and clear NPU->CPU stores queued after request completion.""" + jobs = self._pending_store_kv_transfer_jobs + self._pending_store_kv_transfer_jobs = [] + return jobs + def update_from_output( self, scheduler_output: SchedulerOutput, new_token_ids: dict[str, int] ) -> list[RequestOutput]: @@ -317,7 +355,9 @@ def update_from_output( for req_id in finished_ids: request = self.requests.get(req_id) if request is not None: + release_block_ids = request.cached_block_ids + request.allocated_block_ids self._free_request_blocks(request) + self._queue_cpu_store_jobs(release_block_ids, request.request_id) self.running = [r for r in self.running if r.request_id != req_id] return outputs @@ -344,7 +384,15 @@ def _try_allocate_blocks(self, request: Request, num_blocks: int) -> bool: return True if self.kv_cache_manager.num_free_blocks < num_blocks: return False - block_ids = self.kv_cache_manager.allocate_block_ids(num_blocks) + current_block_ids = request.cached_block_ids + request.allocated_block_ids + preferred_block_ids: list[int] = [] + if current_block_ids: + start = current_block_ids[-1] + 1 + preferred_block_ids = list(range(start, start + num_blocks)) + block_ids = self.kv_cache_manager.allocate_preferred_block_ids( + num_blocks, + preferred_block_ids, + ) if block_ids is None: return False request.allocated_block_ids.extend(block_ids) @@ -411,3 +459,23 @@ def _cache_completed_blocks(self, request: Request) -> None: already_cached, total_blocks_computed, ) + + def _queue_cpu_store_jobs(self, block_ids: list[int], request_id: str) -> None: + """Queue CPU stores for completed cached blocks that are no longer referenced.""" + if not self.config.enable_kv_cpu_offload: + return + store_block_ids: list[int] = [] + for block_id in block_ids: + block = self.kv_cache_manager.blocks[block_id] + if block.block_hash is None: + continue + if block.location != KVBlockLocation.NPU or block.ref_cnt != 0: + continue + store_block_ids.append(block_id) + + for block_id in store_block_ids: + try: + job = self.kv_cache_manager.build_cpu_store_job([block_id], request_id=request_id) + except RuntimeError: + break + self._pending_store_kv_transfer_jobs.append(job) diff --git a/python/core/serving_worker.py b/python/core/serving_worker.py index 0a7662f..56f9e52 100644 --- a/python/core/serving_worker.py +++ b/python/core/serving_worker.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING +from .kv_offload import CPUKvOffloadBackend, TransferJob, TransferResult from python.profile import get_profiler, profile_span from .types import ( DecodeBatch, @@ -52,6 +53,8 @@ def __init__( self.sampler = None self.model_record = None self._page_size: int = 64 + self._kv_offload_backend: CPUKvOffloadBackend | None = None + self._kv_offload_cpu_slots: torch.Tensor | None = None def init_device_and_model(self) -> None: from .model_loader import ModelLoader @@ -94,6 +97,10 @@ def init_device_and_model(self) -> None: ) self._page_size = loaded.runtime_model.runtime.page_size + self._preallocate_kv_offload_cpu_slots( + loaded.config, + loaded.runtime_model.runtime, + ) register_model = getattr(self.executor, "register_model", None) if callable(register_model): @@ -124,17 +131,29 @@ def busy_loop(self) -> None: pass # No allocation cleanup needed try: - result = self._execute_step(cmd.scheduler_output) + result = self._execute_step( + cmd.scheduler_output, + ) self.output_queue.put(result) except Exception as e: logger.error(f"Worker step failed: {e}", exc_info=True) self.output_queue.put(StepOutput(new_tokens={}, error=str(e))) + elif cmd.type == "kv_transfer": + try: + result = self._execute_kv_transfer_step(cmd.kv_transfer_jobs or []) + self.output_queue.put(result) + except Exception as e: + logger.error(f"Worker KV transfer failed: {e}", exc_info=True) + self.output_queue.put(StepOutput(new_tokens={}, error=str(e))) else: logger.warning(f"Worker received unknown command: {cmd.type}") logger.info("Worker exiting") - def _execute_step(self, scheduler_output) -> StepOutput: + 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] = {} @@ -159,6 +178,80 @@ def _execute_step(self, scheduler_output) -> StepOutput: return StepOutput(new_tokens=new_tokens) + def _execute_kv_transfer_step(self, kv_transfer_jobs: list[TransferJob]) -> StepOutput: + """Execute a transfer-only worker command without running model kernels.""" + transfer_results: list[TransferResult] = [] + with profile_span( + "WorkerProcess.execute_kv_transfer_step", + cat="worker", + args={"jobs": len(kv_transfer_jobs)}, + ): + if kv_transfer_jobs: + transfer_results.extend(self._execute_kv_transfer_jobs(kv_transfer_jobs)) + failed_transfer = next((result for result in transfer_results if not result.success), None) + if failed_transfer is not None: + return StepOutput( + new_tokens={}, + kv_transfer_results=transfer_results, + error=failed_transfer.error or "KV offload transfer failed", + ) + return StepOutput(new_tokens={}, kv_transfer_results=transfer_results) + + def _preallocate_kv_offload_cpu_slots(self, model_config, runtime_config) -> None: + """Create CPU offload slots before the L3 worker forks chip children.""" + max_cpu_slots = int(getattr(self.config, "max_cpu_offload_blocks", 0)) + if max_cpu_slots <= 0: + return + kv_dtype = getattr(torch, runtime_config.kv_dtype) + element_size = torch.empty((), dtype=kv_dtype).element_size() + page_size_bytes = ( + 2 + * int(model_config.num_hidden_layers) + * int(model_config.num_key_value_heads) + * int(runtime_config.page_size) + * int(model_config.head_dim) + * element_size + ) + self._kv_offload_cpu_slots = torch.empty( + (max_cpu_slots, page_size_bytes), + dtype=torch.uint8, + device="cpu", + ).share_memory_() + + def _execute_kv_transfer_jobs(self, jobs: list[TransferJob]) -> list[TransferResult]: + """Run CPU KV offload transfers before model execution.""" + backend = self._get_kv_offload_backend(jobs) + results: list[TransferResult] = [] + for job in jobs: + backend.submit(job) + results.extend(backend.wait({job.job_id})) + return results + + def _get_kv_offload_backend(self, jobs: list[TransferJob]) -> CPUKvOffloadBackend: + if self._kv_offload_backend is not None: + return self._kv_offload_backend + max_slot_id = -1 + for job in jobs: + for endpoint in (job.src, job.dst): + slot_ids = getattr(endpoint, "slot_ids", None) + if slot_ids: + max_slot_id = max(max_slot_id, max(slot_ids)) + configured_slots = int(getattr(self.config, "max_cpu_offload_blocks", 0)) + num_cpu_slots = max(configured_slots, max_slot_id + 1, 1) + page_view = self.executor.materialize_kv_page_view(self.config.model_id) + cpu_slots = self._kv_offload_cpu_slots + if cpu_slots is not None and cpu_slots.shape[0] < num_cpu_slots: + raise RuntimeError( + "Preallocated CPU KV slots are smaller than requested transfer slots: " + f"preallocated={cpu_slots.shape[0]}, requested={num_cpu_slots}" + ) + self._kv_offload_backend = CPUKvOffloadBackend( + page_view, + num_cpu_slots=num_cpu_slots, + cpu_slots=cpu_slots, + ) + return self._kv_offload_backend + def _batch_prefill( self, scheduled: list, runtime_model, new_tokens: dict[str, int] ) -> None: diff --git a/python/core/types.py b/python/core/types.py index bb3343e..fd1ec9d 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -14,6 +14,7 @@ import torch +from .kv_offload import TransferJob, TransferResult from .tokenizer import TokenizerAdapter @@ -223,13 +224,15 @@ class GenerateResult: @dataclass class WorkerCommand: """Command sent from main process to worker process.""" - type: str # "step" | "shutdown" + type: str # "step" | "kv_transfer" | "shutdown" scheduler_output: object | None = None finished_request_ids: list | None = None + kv_transfer_jobs: list[TransferJob] | None = None @dataclass class StepOutput: """Result returned from worker process after executing a batch step.""" new_tokens: dict # {request_id: int} + kv_transfer_results: list[TransferResult] = field(default_factory=list) error: str | None = None diff --git a/tests/test_cli.py b/tests/test_cli.py index f88e5cb..4143f0c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -38,7 +38,7 @@ def test_build_serving_engine_config_uses_cli_args(tmp_path, monkeypatch): "--platform", "a5", "--device", "2", "--max-model-len", "1024", - "--block-size", "64", + "--block-size", "128", "--dtype", "bfloat16", "--kv-cache-dtype", "auto", "--max-new-tokens", "8", @@ -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", "3", ]) config = cli.build_serving_engine_config(args) @@ -60,7 +61,7 @@ def test_build_serving_engine_config_uses_cli_args(tmp_path, monkeypatch): "pypto_root": "/tmp/pypto", "save_kernels_dir": "/tmp/kernels", } - assert config.runtime_config.page_size == 64 + assert config.runtime_config.page_size == 128 assert config.runtime_config.max_batch_size == 4 assert config.runtime_config.max_seq_len == 1024 assert config.runtime_config.kv_dtype == "bfloat16" @@ -71,6 +72,8 @@ 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.enable_kv_cpu_offload is True + assert config.max_cpu_offload_blocks == 3 def test_parser_rejects_invalid_backend(tmp_path): @@ -81,6 +84,25 @@ def test_parser_rejects_invalid_backend(tmp_path): _parse_args(["--model", str(model_dir), "--backend", "cpu"]) +def test_parser_rejects_unsupported_block_size(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + + with pytest.raises(SystemExit): + _parse_args(["--model", str(model_dir), "--block-size", "64"]) + + +def test_cpu_offload_defaults_to_disabled(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + args = _parse_args(["--model", str(model_dir)]) + + config = cli.build_serving_engine_config(args) + + assert config.enable_kv_cpu_offload is False + assert config.max_cpu_offload_blocks == 0 + + def test_parser_rejects_removed_prompt_mode(tmp_path): model_dir = tmp_path / "model" model_dir.mkdir()