Add KV cache CPU offload support

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}."
-            )

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,

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")

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
         )