Add KV cache CPU offload support

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,

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

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: