feat: add shared-memory radix-tree prefix cache reuse for host paged KV

.gitignore

@@ -4,6 +4,7 @@
 log_*.log
 log
 *.log
+tmp/
 
 *.pt
 runner.sh

batchgen/config/config.py

@@ -176,6 +176,45 @@ class HostPagedKVConfig:
     num_v_heads: int = 0 # Zero for MLA.
     v_head_dim: int = 0
     kv_dtype: str = "bfloat16" # "bfloat16 or float8_e4m3fn"
+    enable_prefix_reuse: bool = False
+    prefix_min_reuse_pages: int = 1
+    prefix_min_store_pages: int = 2
+    # Capacity knobs for prefix cache internals.
+    # 0 means "auto": values are derived from num_pages_per_layer at runtime.
+    sequence_page_node_capacity: int = 0
+    radix_node_capacity: int = 0
+    radix_edge_capacity: int = 0
+    prefix_entry_capacity: int = 0
+    prefix_page_ref_capacity: int = 0
+    prefix_page_budget: int = 0
+
+    def __post_init__(self) -> None:
+        self.apply_runtime_defaults()
+
+    def apply_runtime_defaults(self) -> None:
+        """Fill auto (0) capacity fields using num_pages_per_layer."""
+        num_pages = max(int(self.num_pages_per_layer), 0)
+        if num_pages == 0:
+            return
+
+        if self.sequence_page_node_capacity <= 0:
+            self.sequence_page_node_capacity = max(num_pages, num_pages * 4)
+        if self.radix_node_capacity <= 0:
+            self.radix_node_capacity = max(4096, num_pages // 2)
+        if self.radix_edge_capacity <= 0:
+            self.radix_edge_capacity = max(
+                self.radix_node_capacity * 2,
+                self.radix_node_capacity + 1,
+            )
+        if self.prefix_entry_capacity <= 0:
+            self.prefix_entry_capacity = max(1024, num_pages // 64)
+        if self.prefix_page_budget <= 0:
+            self.prefix_page_budget = max(128, num_pages // 2)
+        if self.prefix_page_ref_capacity <= 0:
+            self.prefix_page_ref_capacity = max(
+                num_pages,
+                self.prefix_entry_capacity * 8,
+            )
 
 @dataclass
 class DevicePagedKVConfig:

batchgen/config/engine_config_parser.py

@@ -209,6 +209,7 @@ def _parse_host_paged_kv_config(
         if key not in valid_fields:
             raise ValueError(f"Unknown key in Host_Paged_KV_Config: {key}")
         setattr(host_config, key, value)
+    host_config.apply_runtime_defaults()
 
 
 def _parse_device_paged_kv_config(

batchgen/kv_cache/dual_host_kv_coordinator.py

@@ -34,7 +34,9 @@ def _try_set_logger_name(config, name: str) -> bool:
 		return False
 
 
-def _build_host_config_from_profile(profile, shm_name: str, num_pages: int) -> Any:
+def _build_host_config_from_profile(
+	profile, shm_name: str, num_pages: int, enable_prefix_reuse: bool = True
+) -> Any:
 	"""Build a bg_lib.HostPagedKVConfig from a _HostKVModelProfile."""
 	from batchgen.kv_cache.host_kv_mananger_config import _dtype_size_bytes
 
@@ -55,6 +57,9 @@ def _build_host_config_from_profile(profile, shm_name: str, num_pages: int) -> A
 		profile.sequence_table_capacity or config.num_pages
 	)
 	config.alignment_bytes = profile.alignment_bytes
+	config.enable_prefix_reuse = bool(enable_prefix_reuse)
+	config.prefix_min_reuse_pages = 1
+	config.prefix_min_store_pages = 2
 	return config
 
 
@@ -107,6 +112,7 @@ def from_budget(
 		model_name: str,
 		host_kv_cache_size: int,
 		core_engine_module,
+		enable_prefix_reuse: bool = True,
 		enable_memfd: bool = False,
 		memfd_creator_pid: int = -1,
 		memfd_fd: int = -1,
@@ -127,10 +133,10 @@ def from_budget(
 		)
 
 		primary_config = _build_host_config_from_profile(
-			primary_profile, HOST_KV_SHM_NAME, num_pages,
+			primary_profile, HOST_KV_SHM_NAME, num_pages, enable_prefix_reuse,
 		)
 		aux_config = _build_host_config_from_profile(
-			aux_profile, HOST_KV_AUX_SHM_NAME, num_pages,
+			aux_profile, HOST_KV_AUX_SHM_NAME, num_pages, enable_prefix_reuse,
 		)
 
 		# Set distinct logger names to avoid C++ logger name collision
@@ -177,6 +183,7 @@ def create_managers(
 		cls,
 		model_name: str,
 		host_kv_cache_size: int,
+		enable_prefix_reuse: bool = True,
 		enable_memfd: bool = False,
 	) -> Optional[Tuple[Any, Any]]:
 		"""Server-side factory: create and initialize both host KV managers.
@@ -194,10 +201,10 @@ def create_managers(
 		)
 
 		primary_config = _build_host_config_from_profile(
-			primary_profile, HOST_KV_SHM_NAME, num_pages,
+			primary_profile, HOST_KV_SHM_NAME, num_pages, enable_prefix_reuse,
 		)
 		aux_config = _build_host_config_from_profile(
-			aux_profile, HOST_KV_AUX_SHM_NAME, num_pages,
+			aux_profile, HOST_KV_AUX_SHM_NAME, num_pages, enable_prefix_reuse,
 		)
 
 		# Set distinct logger names to avoid C++ logger name collision

batchgen/kv_cache/host_kv_mananger_config.py

@@ -227,7 +227,37 @@ def _resolve_profile(model_name: str) -> _HostKVModelProfile:
 	return _PROFILE_REGISTRY[_PROFILE_ALIASES[alias]]
 
 
-def build_host_kv_config(model_name: str, host_kv_cache_size: int) -> Any:
+def _apply_capacity_defaults(config: Any) -> None:
+	"""Populate HostPagedKVConfig capacity fields when they are zero."""
+	num_pages = int(config.num_pages)
+	if num_pages <= 0:
+		return
+
+	if int(config.sequence_page_node_capacity) <= 0:
+		config.sequence_page_node_capacity = max(num_pages, num_pages * 4)
+	if int(config.radix_node_capacity) <= 0:
+		config.radix_node_capacity = max(4096, num_pages // 2)
+	if int(config.radix_edge_capacity) <= 0:
+		config.radix_edge_capacity = max(
+			int(config.radix_node_capacity) * 2,
+			int(config.radix_node_capacity) + 1,
+		)
+	if int(config.prefix_entry_capacity) <= 0:
+		config.prefix_entry_capacity = max(1024, num_pages // 64)
+	if int(config.prefix_page_budget) <= 0:
+		config.prefix_page_budget = max(128, num_pages // 2)
+	if int(config.prefix_page_ref_capacity) <= 0:
+		config.prefix_page_ref_capacity = max(
+			num_pages,
+			int(config.prefix_entry_capacity) * 8,
+		)
+
+
+def build_host_kv_config(
+	model_name: str,
+	host_kv_cache_size: int,
+	enable_prefix_reuse: bool = True,
+) -> Any:
 	"""Builds a core HostPagedKVConfig for the given model and host budget."""
 
 	if host_kv_cache_size is None:
@@ -271,6 +301,10 @@ def build_host_kv_config(model_name: str, host_kv_cache_size: int) -> Any:
 		profile.sequence_table_capacity or config.num_pages
 	)
 	config.alignment_bytes = profile.alignment_bytes
+	config.enable_prefix_reuse = bool(enable_prefix_reuse)
+	config.prefix_min_reuse_pages = 1
+	config.prefix_min_store_pages = 2
+	_apply_capacity_defaults(config)
 	return config
 
 
@@ -356,7 +390,11 @@ def build_gpu_kv_config_aux(
 	)
 
 
-def build_host_kv_config_aux(model_name: str, host_kv_cache_size: int) -> Any | None:
+def build_host_kv_config_aux(
+	model_name: str,
+	host_kv_cache_size: int,
+	enable_prefix_reuse: bool = True,
+) -> Any | None:
 	"""Builds a HostPagedKVConfig for the DSA indexer host cache, or None."""
 
 	profile = _resolve_indexer_profile(model_name)
@@ -386,6 +424,10 @@ def build_host_kv_config_aux(model_name: str, host_kv_cache_size: int) -> Any |
 		profile.sequence_table_capacity or config.num_pages
 	)
 	config.alignment_bytes = profile.alignment_bytes
+	config.enable_prefix_reuse = bool(enable_prefix_reuse)
+	config.prefix_min_reuse_pages = 1
+	config.prefix_min_store_pages = 2
+	_apply_capacity_defaults(config)
 	return config
 
 

batchgen/server/server_args.py

@@ -92,6 +92,7 @@ class ServerArgs:
     disable_cuda_graphs: bool = True  # Disable CUDA graph capture for decode attention (128K+ crash: corrupted num_tokens_per_rank)
     cuda_graph_max_bucket_size: int = 128  # Max batch size per rank for CUDA graph capture
     cuda_graph_num_buckets: int = 16  # Number of CUDA graph bucket sizes
+    enable_prefix_cache: bool = True  # Enable host KV prefix cache reuse
     detokenization_include_special_tokens: bool = False  # When True, include special tokens in detokenized output
     # Dynamic host KV reservation settings
     host_kv_chunk_size: int = 8192  # Initial host KV chunk size in tokens (default: 8K)
@@ -349,6 +350,19 @@ def _build_parser() -> argparse.ArgumentParser:
         default=16,
         help="Maximum number of CUDA graph bucket sizes (default: 16). More buckets = longer capture time but less padding waste.",
     )
+    parser.set_defaults(enable_prefix_cache=True)
+    parser.add_argument(
+        "--enable-prefix-cache",
+        dest="enable_prefix_cache",
+        action="store_true",
+        help="Enable host KV prefix cache reuse (default: enabled)",
+    )
+    parser.add_argument(
+        "--disable-prefix-cache",
+        dest="enable_prefix_cache",
+        action="store_false",
+        help="Disable host KV prefix cache reuse",
+    )
     parser.add_argument(
         "--detokenization-include-special-tokens",
         action="store_true",
@@ -546,6 +560,7 @@ def prepare_server_args(argv: Optional[list[str]] = None) -> ServerArgs:
         cuda_graph_max_bucket_size=parsed.cuda_graph_max_bucket_size,
         cuda_graph_num_buckets=parsed.cuda_graph_num_buckets,
         detokenization_include_special_tokens=parsed.detokenization_include_special_tokens,
+        enable_prefix_cache=parsed.enable_prefix_cache,
         host_kv_chunk_size=parsed.host_kv_chunk_size,
         host_kv_eviction_watermark=parsed.host_kv_eviction_watermark,
         enable_host_kv_eviction=parsed.enable_host_kv_eviction,

batchgen/server/worker_manager.py

@@ -35,6 +35,14 @@
 PARAMETER_SERVER_ENDPOINT_ENV = "BATCHGEN_PARAMETER_SERVER_ENDPOINT"
 
 
+def _load_server_worker_main():
+    # Delay this import to avoid a package-init cycle when worker subprocesses
+    # import `batchgen.server.process_utils` via `batchgen.server_worker_main_loop`.
+    from batchgen.server_worker_main_loop import server_worker_main
+
+    return server_worker_main
+
+
 def _validate_shmem_enabled() -> None:
     """Check that THP shmem is enabled for --fast-init. Raises RuntimeError if not."""
     import re
@@ -204,6 +212,7 @@ def _diag(msg):
                 _diag(">>> allocate_host_kv_cache")
                 result = self.allocate_host_kv_cache(
                     self.args.host_kv_cache_size, self.args.model,
+                    enable_prefix_cache=self.args.enable_prefix_cache,
                     enable_memfd=self.args.fast_init,
                 )
                 _diag("<<< allocate_host_kv_cache")
@@ -629,6 +638,7 @@ def _spawn_workers(self) -> None:
             disable_cuda_graphs=self.args.disable_cuda_graphs,
             cuda_graph_max_bucket_size=self.args.cuda_graph_max_bucket_size,
             cuda_graph_num_buckets=self.args.cuda_graph_num_buckets,
+            enable_prefix_cache=self.args.enable_prefix_cache,
             detokenization_include_special_tokens=self.args.detokenization_include_special_tokens,
             host_kv_chunk_size=self.args.host_kv_chunk_size,
             enable_host_kv_eviction=self.args.enable_host_kv_eviction,
@@ -648,7 +658,7 @@ def _spawn_workers(self) -> None:
         )
         from batchgen.server_worker_main_loop import server_worker_main
         self.worker_process = mp.spawn(
-            server_worker_main,
+            _load_server_worker_main(),
             args=(
                 self.request_queue,
                 self.response_queue,
@@ -1005,6 +1015,7 @@ def _configure_host_kv_cache_budget(self) -> None:
     @staticmethod
     def allocate_host_kv_cache(
         host_kv_cache_size_gb: int, model_name: str,
+        enable_prefix_cache: bool = True,
         enable_memfd: bool = False,
     ) -> Any:
         from batchgen.kv_cache.dual_host_kv_coordinator import DualHostKVCoordinator
@@ -1013,6 +1024,7 @@ def allocate_host_kv_cache(
         dual = DualHostKVCoordinator.create_managers(
             model_name=model_name,
             host_kv_cache_size=int(host_kv_cache_size_gb * (1024**3)),
+            enable_prefix_reuse=enable_prefix_cache,
             enable_memfd=enable_memfd,
         )
         if dual is not None:
@@ -1025,6 +1037,7 @@ def allocate_host_kv_cache(
         config = build_host_kv_config(
             host_kv_cache_size=host_kv_cache_size_gb * (1024**3),
             model_name=model_name,
+            enable_prefix_reuse=enable_prefix_cache,
         )
         if enable_memfd:
             config.enable_memfd = True