hicache: type annotations for ruff ANN rules

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr

deepinfer/hicache/host_cache.py

@@ -33,6 +33,7 @@ import os
 import time
 from collections import OrderedDict, deque
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
@@ -41,6 +42,10 @@ from deepinfer import envs
 from deepinfer.logger import init_logger
 from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface
 
+if TYPE_CHECKING:
+    from deepinfer._C.engine import Sequence
+    from deepinfer.engine.llm_engine import LLMEngine
+
 logger = init_logger(__name__)
 
 
@@ -83,7 +88,7 @@ class Stats:
 
 
 class HostPrefixCache:
-    def __init__(self, engine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None:
+    def __init__(self, engine: LLMEngine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None:
         assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)"
         assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1"
         assert GPUCacheInterface.cache_info_ready, "cache engine not initialized"
@@ -241,7 +246,7 @@ class HostPrefixCache:
 
     # ------------------------------------------------------------------ load (hit)
 
-    def try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int:
+    def try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int:
         """Fill `seq` (not yet in the scheduler) from the host cache.
 
         Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not
@@ -256,7 +261,7 @@ class HostPrefixCache:
             self._fail()
             return 0
 
-    def _try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int:
+    def _try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int:
         bs = self.block_size
         n = len(prompt_ids)
         chain = self._chain(request_id, prompt_ids)
@@ -325,7 +330,7 @@ class HostPrefixCache:
 
     # ------------------------------------------------------------------ store (write-through at prefill end)
 
-    def on_prefill_step(self, scheduled, finished) -> None:
+    def on_prefill_step(self, scheduled: list[Sequence], finished: list[Sequence]) -> None:
         """Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed."""
         if not self.enabled:
             return
@@ -337,7 +342,7 @@ class HostPrefixCache:
             logger.exception("host prefix cache: store failed")
             self._fail()
 
-    def _store(self, seq) -> None:
+    def _store(self, seq: Sequence) -> None:
         chain = self._chains.get(seq.request_id)
         if chain is None:  # request did not go through try_resume: hash the prompt from the sequence itself
             chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32))

deepinfer/hicache/kernels.py

@@ -3,11 +3,12 @@
 from __future__ import annotations
 
 from pathlib import Path
+from types import ModuleType
 
 _module = None
 
 
-def copy_pages_module():
+def copy_pages_module() -> ModuleType:
     global _module
     if _module is None:
         from torch.utils.cpp_extension import load