from __future__ import annotations """ Copyright 2023-2025 SGLang Team Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ import logging import threading import time from queue import Empty, Full, Queue from typing import TYPE_CHECKING, List, NamedTuple, Optional import torch from sglang.srt.mem_cache.hicache_storage import ( HiCacheStorageConfig, HiCacheStorageExtraInfo, ) if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool_host import HostKVCache from sglang.srt.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from sglang.srt.layers.dp_attention import ( get_attention_dp_rank, get_attention_tp_rank, get_attention_tp_size, is_dp_attention_enabled, ) from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool from sglang.srt.utils import get_device_module logger = logging.getLogger(__name__) device_module = get_device_module() class LayerLoadingEvent: def __init__(self, num_layers: int): self._num_layers = num_layers self.load_events = [device_module.Event() for _ in range(num_layers)] self.start_event = device_module.Event() # start event on controller stream def complete(self, layer_index: int): assert 0 <= layer_index < self._num_layers self.load_events[layer_index].record() def wait(self, layer_index: int): device_module.current_stream().wait_event(self.load_events[layer_index]) @property def finish_event(self): return self.load_events[-1] class LayerDoneCounter: def __init__(self, num_layers: int): self.num_layers = num_layers # extra producer and consumer counters for overlap mode self.num_counters = 3 self.events = [LayerLoadingEvent(num_layers) for _ in range(self.num_counters)] self.producer_index = -1 self.consumer_index = -1 def update_producer(self): self.producer_index = (self.producer_index + 1) % self.num_counters assert self.events[ self.producer_index ].finish_event.query(), ( "Producer finish event should be ready before being reused." ) return self.producer_index def set_consumer(self, index: int): self.consumer_index = index def wait_until(self, threshold: int): if self.consumer_index < 0: return self.events[self.consumer_index].wait(threshold) def reset(self): self.producer_index = -1 self.consumer_index = -1 class CacheOperation: counter = 0 def __init__( self, host_indices: torch.Tensor, device_indices: torch.Tensor, node_id: int, priority: Optional[int] = None, ): self.host_indices = host_indices self.device_indices = device_indices self.node_ids = [node_id] self.data = None self.id = CacheOperation.counter CacheOperation.counter += 1 # default priority is the order of creation self.priority = priority if priority is not None else self.id @staticmethod def merge_ops(ops: List[CacheOperation]) -> CacheOperation: assert len(ops) > 0 if len(ops) == 1: return ops[0] host_indices = torch.cat([op.host_indices for op in ops]) device_indices = torch.cat([op.device_indices for op in ops]) node_ids = [] priority = min(op.priority for op in ops) for op in ops: node_ids.extend(op.node_ids) merged_op = CacheOperation(host_indices, device_indices, -1, priority) merged_op.node_ids = node_ids return merged_op def __lt__(self, other: CacheOperation): return self.priority < other.priority class HiCacheAck(NamedTuple): start_event: device_module.Event finish_event: device_module.Event node_ids: List[int] class TransferBuffer: """ Overlapping buffer preparation and transfer operations to improve throughput. """ def __init__( self, stop_event, buffer_count: int = 3, max_buffer_size: int = 1024 ) -> None: self.stop_event = stop_event self.buffers = Queue(maxsize=buffer_count) # todo: adjust the buffer size based on throughput profile of the system self.max_buffer_size = max_buffer_size def full(self) -> bool: return self.buffers.full() def empty(self) -> bool: return self.buffers.empty() def put(self, item, block=True, timeout=1) -> None: while not self.stop_event.is_set(): try: self.buffers.put(item, block=block, timeout=timeout) break except Full: if not block: break continue except Exception as e: logger.error(e) def get(self, block=True, timeout=1) -> Optional[CacheOperation]: try: return self.buffers.get(block=block, timeout=timeout) except Empty: return None except Exception as e: logger.error(e) def clear(self): self.buffers.queue.clear() class StorageOperation: counter = 0 def __init__( self, host_indices: torch.Tensor, token_ids: List[int], last_hash: Optional[str] = None, hash_value: Optional[List[str]] = None, prefix_keys: Optional[List[str]] = None, ): self.host_indices = host_indices self.token_ids = token_ids self.last_hash = last_hash self.completed_tokens = 0 self.hash_value = hash_value if hash_value is not None else [] self.prefix_keys = prefix_keys self.id = StorageOperation.counter StorageOperation.counter += 1 def __lt__(self, other: "StorageOperation"): return self.id < other.id class PrefetchOperation(StorageOperation): def __init__( self, request_id: str, host_indices: torch.Tensor, token_ids: List[int], last_hash: Optional[str] = None, prefix_keys: Optional[List[str]] = None, ): self.request_id = request_id self._lock = threading.Lock() self._terminated_flag = False self.start_time = time.monotonic() super().__init__(host_indices, token_ids, last_hash, prefix_keys=prefix_keys) def increment(self, num_tokens: int): with self._lock: if self._terminated_flag: return False self.completed_tokens += num_tokens return True def mark_terminate(self): with self._lock: self._terminated_flag = True def is_terminated(self) -> bool: return self._terminated_flag class HiCacheController: def __init__( self, token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, mem_pool_host: HostKVCache, page_size: int, tp_group: torch.distributed.ProcessGroup, load_cache_event: threading.Event, write_policy: str = "write_through_selective", io_backend: str = "", storage_backend: Optional[str] = None, prefetch_threshold: int = 256, model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, pp_rank: int = 0, pp_size: int = 1, ): self.tp_group = tp_group self.mem_pool_device_allocator = token_to_kv_pool_allocator self.mem_pool_device = token_to_kv_pool_allocator.get_kvcache() self.mem_pool_host = mem_pool_host self.write_policy = write_policy self.page_size = page_size self.io_backend = io_backend self.enable_storage = False self.storage_backend = None self.storage_backend_type = None self.pp_rank = pp_rank self.pp_size = pp_size # Default storage page IO functions (may be overridden by attach). self.page_get_func = self._generic_page_get self.page_set_func = self._generic_page_set # Dedicated stop event for storage background threads (prefetch/backup). # NOTE: Do NOT reuse `self.stop_event` here since it also guards core HiCache # transfer buffers (CPU<->GPU). We want to allow runtime attach/detach of # storage without stopping the whole controller. self.storage_stop_event = threading.Event() self.device = self.mem_pool_device.device self.layer_num = self.mem_pool_device.layer_num self.layer_done_counter = LayerDoneCounter(self.layer_num) self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter) if write_policy not in [ "write_through", "write_through_selective", "write_back", ]: raise ValueError(f"Invalid write policy: {write_policy}") # self.write_queue = PriorityQueue[CacheOperation]() self.load_queue: List[CacheOperation] = [] self.write_queue: List[CacheOperation] = [] self.ack_load_queue: List[HiCacheAck] = [] self.ack_write_queue: List[HiCacheAck] = [] self.stop_event = threading.Event() self.write_buffer = TransferBuffer(self.stop_event) self.load_buffer = TransferBuffer( self.stop_event, buffer_count=10, max_buffer_size=100 ) self.write_stream = device_module.Stream() self.load_stream = device_module.Stream() # If a storage backend is provided at startup, treat it as an implicit attach, # so init/runtime share the same lifecycle semantics and code paths. if storage_backend is not None: try: self.attach_storage_backend( storage_backend=storage_backend, prefetch_threshold=prefetch_threshold, model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, ) except ValueError as e: # Preserve the historical error shape on init for unknown backends. raise ValueError(f"Failed to create storage backend: {e}") from e def _start_storage_threads(self): """Start storage prefetch/backup threads and their queues. This is used by runtime attach, and also by reset when storage is enabled. """ assert self.enable_storage assert not self.storage_stop_event.is_set() self.prefetch_thread = threading.Thread( target=self.prefetch_thread_func, daemon=True ) self.backup_thread = threading.Thread( target=self.backup_thread_func, daemon=True ) self.prefetch_queue = Queue() self.backup_queue = Queue() self.prefetch_revoke_queue = Queue() self.ack_backup_queue = Queue() self.host_mem_release_queue = Queue() self.prefetch_thread.start() self.backup_thread.start() def _stop_storage_threads(self): """Stop storage prefetch/backup threads and drain internal queues. Caller should ensure no in-flight requests. """ # Always request stop. This is safe even when storage is already disabled, # and makes detach truly idempotent (previous partial detach may have left # threads alive). # NOTE: do NOT clear stop_event unless threads have fully stopped; otherwise # a still-alive thread may resume and touch released state. self.storage_stop_event.set() # Best-effort wakeups so threads exit promptly even if blocked on queues. try: if hasattr(self, "prefetch_queue"): self.prefetch_queue.put_nowait(None) if hasattr(self, "backup_queue"): self.backup_queue.put_nowait(None) if hasattr(self, "prefetch_buffer"): self.prefetch_buffer.put_nowait(None) except Exception: pass # Best-effort joins (threads are daemon, but join keeps state clean). threads = [] if hasattr(self, "prefetch_thread"): threads.append(self.prefetch_thread) if hasattr(self, "backup_thread"): threads.append(self.backup_thread) if hasattr(self, "prefetch_io_aux_thread"): threads.append(self.prefetch_io_aux_thread) for t in threads: try: t.join(timeout=10) except Exception: pass alive = [t for t in threads if getattr(t, "is_alive", lambda: False)()] if alive: logger.error( "Failed to stop HiCache storage threads cleanly: %s", [getattr(t, "name", repr(t)) for t in alive], ) raise RuntimeError("Failed to stop HiCache storage threads cleanly.") def attach_storage_backend( self, storage_backend: str, prefetch_threshold: int = 256, model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, ): """Attach (enable) storage backend at runtime. Requirement: no in-flight requests. This call is expected to run on the scheduler thread (control path), not concurrently with prefetch/backup. """ if self.enable_storage: raise RuntimeError("Storage backend already attached.") # Defensive: a previous partial detach may have flipped `enable_storage` but # left background threads alive. Attaching on top of them is unsafe. try: self._stop_storage_threads() except Exception as e: raise RuntimeError( "Cannot attach storage backend: previous detach did not stop storage threads cleanly." ) from e # Rollback-safe init: if creation fails, keep controller state consistent # for future attach attempts. self.storage_backend_type = storage_backend from sglang.srt.mem_cache.hicache_storage import get_hash_str self.get_hash_str = get_hash_str self.storage_config = self._generate_storage_config( model_name, storage_backend_extra_config ) # for MLA models, only one rank needs to backup the KV cache self.backup_skip = ( self.storage_config.is_mla_model # todo: load balancing and self.storage_config.tp_rank != 0 ) # Use storage backend factory for dynamic backend creation from sglang.srt.mem_cache.storage import StorageBackendFactory try: self.storage_backend = StorageBackendFactory.create_backend( storage_backend, self.storage_config, self.mem_pool_host ) self.storage_backend.register_mem_pool_host(self.mem_pool_host) self.enable_storage = True # todo: threshold policy for prefetching self.prefetch_threshold = max(prefetch_threshold, self.page_size) self.prefetch_capacity_limit = max( 0, int(0.8 * (self.mem_pool_host.size - self.mem_pool_device.size)) ) # granularity of batch storage IO operations, in number of pages self.storage_batch_size = 128 # tracking the number of tokens locked in prefetching, updated by the main scheduler thread self.prefetch_tokens_occupied = 0 # create a new communication group for synchronizing storage operations across TP workers self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group) if self.tp_world_size > 1: from sglang.srt.distributed.parallel_state import ( create_custom_parallel_group, ) group_ranks = torch.distributed.get_process_group_ranks(self.tp_group) self.prefetch_tp_group = create_custom_parallel_group( group_ranks=group_ranks, backend="gloo" ) # Select the get and set functions self.page_get_func = self._generic_page_get self.page_set_func = self._generic_page_set if (self.storage_backend_type in ["hf3fs", "mooncake", "eic", "nixl"]) or ( self.storage_backend_type == "dynamic" and bool(self.storage_config.extra_config.get("interface_v1", 0)) ): self.page_get_func = self._page_get_zero_copy self.page_set_func = self._page_set_zero_copy # Ensure stop_event is clear before starting threads. self.storage_stop_event.clear() self._start_storage_threads() except Exception: # Best-effort cleanup for partial init. try: self._stop_storage_threads() except Exception: pass try: if hasattr(self, "prefetch_tp_group"): try: torch.distributed.destroy_process_group(self.prefetch_tp_group) except Exception: pass self.prefetch_tp_group = None except Exception: pass try: if ( hasattr(self, "storage_backend") and self.storage_backend is not None ): if hasattr(self.storage_backend, "close"): self.storage_backend.close() except Exception: pass self.storage_backend = None self.storage_backend_type = None self.enable_storage = False self.page_get_func = self._generic_page_get self.page_set_func = self._generic_page_set raise def detach_storage_backend(self): """Detach (disable) storage backend at runtime. Requirement: no in-flight requests. This will stop storage threads and release the backend instance (best-effort close). """ # Idempotent cleanup: even if `enable_storage` is already False, # we may still have leftover resources (threads/backend/process group) from a # previous partial detach. We attempt cleanup whenever possible. try: self._stop_storage_threads() except Exception as e: # Do not proceed tearing down backend/process group if threads are not # fully stopped; otherwise still-alive threads may touch released state. # Caller can retry detach. logger.exception("Stop storage threads failed: %s", e) # IMPORTANT: Do not silently succeed. Upper layers rely on exceptions here # to avoid flipping `enable_storage` flags while threads are still alive. raise RuntimeError("Stop storage threads failed; detach aborted.") from e # Best-effort destroy process group created for storage ops. try: if ( hasattr(self, "prefetch_tp_group") and self.prefetch_tp_group is not None ): try: torch.distributed.destroy_process_group(self.prefetch_tp_group) except Exception: pass self.prefetch_tp_group = None except Exception: pass # Best-effort close (some backends rely on GC/destructor). try: if ( hasattr(self, "storage_backend") and self.storage_backend is not None and hasattr(self.storage_backend, "close") ): self.storage_backend.close() except Exception: logger.exception("Failed to close storage backend cleanly.") self.storage_backend = None self.storage_backend_type = None self.enable_storage = False self.page_get_func = self._generic_page_get self.page_set_func = self._generic_page_set # Now it's safe to clear the stop event for future re-attach. self.storage_stop_event.clear() def _generate_storage_config( self, model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, ): if is_dp_attention_enabled(): self.tp_rank = get_attention_tp_rank() self.tp_size = get_attention_tp_size() self.dp_rank = get_attention_dp_rank() else: self.tp_rank = get_tensor_model_parallel_rank() self.tp_size = get_tensor_model_parallel_world_size() self.dp_rank = 0 # Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool. is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool) return HiCacheStorageConfig( tp_rank=self.tp_rank, tp_size=self.tp_size, pp_rank=self.pp_rank, pp_size=self.pp_size, is_mla_model=is_mla_backend, is_page_first_layout=self.mem_pool_host.layout == "page_first", model_name=model_name, extra_config=storage_backend_extra_config, ) def reset(self): self.stop_event.set() self.storage_stop_event.set() self.write_queue.clear() self.load_queue.clear() self.write_buffer.clear() self.load_buffer.clear() self.ack_write_queue.clear() self.ack_load_queue.clear() if self.enable_storage: self.prefetch_thread.join() self.backup_thread.join() self.prefetch_queue.queue.clear() self.backup_queue.queue.clear() self.prefetch_revoke_queue.queue.clear() self.ack_backup_queue.queue.clear() self.stop_event.clear() self.storage_stop_event.clear() if self.enable_storage: self.prefetch_thread = threading.Thread( target=self.prefetch_thread_func, daemon=True ) self.backup_thread = threading.Thread( target=self.backup_thread_func, daemon=True ) self.prefetch_thread.start() self.backup_thread.start() def write( self, device_indices: torch.Tensor, priority: Optional[int] = None, node_id: int = -1, ) -> Optional[torch.Tensor]: """ Back up KV caches from device memory to host memory. """ host_indices = self.mem_pool_host.alloc(len(device_indices)) if host_indices is None: return None self.write_queue.append( CacheOperation(host_indices, device_indices, node_id, priority) ) self.start_writing() return host_indices def start_writing(self) -> None: if len(self.write_queue) == 0: return op = CacheOperation.merge_ops(self.write_queue) host_indices, device_indices = self.move_indices(op) self.write_queue.clear() start_event = device_module.Event() finish_event = device_module.Event() start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) self.mem_pool_host.backup_from_device_all_layer( self.mem_pool_device, host_indices, device_indices, self.io_backend ) finish_event.record() # NOTE: We must save the host indices and device indices here, # this is because we need to guarantee that these tensors are # still alive when the write stream is executing. if host_indices.is_cuda: host_indices.record_stream(self.write_stream) if device_indices.is_cuda: device_indices.record_stream(self.write_stream) self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) def load( self, host_indices: torch.Tensor, priority: Optional[int] = None, node_id: int = -1, ) -> Optional[torch.Tensor]: """ Load KV caches from host memory to device memory. """ device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) if device_indices is None: return None self.load_queue.append( CacheOperation(host_indices, device_indices, node_id, priority) ) return device_indices def move_indices(self, op: CacheOperation): host_indices, device_indices = op.host_indices, op.device_indices # move indices to GPU if using kernels, to host if using direct indexing if self.io_backend == "kernel": if not host_indices.is_cuda: host_indices = host_indices.to(self.device, non_blocking=True) return host_indices, device_indices elif self.io_backend == "direct": if self.mem_pool_host.layout == "layer_first": device_indices = device_indices.cpu() host_indices, idx = host_indices.sort() return host_indices, device_indices.index_select(0, idx) elif self.mem_pool_host.layout == "page_first_direct": return host_indices, device_indices.cpu() elif self.io_backend == "kernel_ascend": return host_indices, device_indices.cpu() else: raise ValueError(f"Unsupported io backend") def start_loading(self) -> int: if len(self.load_queue) == 0: return -1 producer_id = self.layer_done_counter.update_producer() op = CacheOperation.merge_ops(self.load_queue) host_indices, device_indices = self.move_indices(op) self.load_queue.clear() producer_event = self.layer_done_counter.events[producer_id] producer_event.start_event.record() with device_module.stream(self.load_stream): producer_event.start_event.wait(self.load_stream) for i in range(self.layer_num): self.mem_pool_host.load_to_device_per_layer( self.mem_pool_device, host_indices, device_indices, i, self.io_backend, ) producer_event.complete(i) # NOTE: We must save the host indices and device indices here, # this is because we need to guarantee that these tensors are # still alive when the load stream is executing. if host_indices.is_cuda: host_indices.record_stream(self.load_stream) if device_indices.is_cuda: device_indices.record_stream(self.load_stream) self.ack_load_queue.append( HiCacheAck( start_event=producer_event.start_event, finish_event=producer_event.finish_event, node_ids=op.node_ids, ) ) return producer_id def evict_device(self, device_indices: torch.Tensor) -> int: self.mem_pool_device_allocator.free(device_indices) return len(device_indices) def evict_host(self, host_indices: torch.Tensor, backup_only: bool = True) -> int: if not backup_only: raise ValueError("Other eviction policies are not supported yet.") self.mem_pool_host.free(host_indices) return len(host_indices) def prefetch( self, request_id: str, host_indices: torch.Tensor, new_input_tokens: List[int], last_hash: Optional[str] = None, prefix_keys: Optional[List[str]] = None, ) -> PrefetchOperation: """ Prefetch KV caches from storage backend to host memory. """ operation = PrefetchOperation( request_id, host_indices, new_input_tokens, last_hash, prefix_keys ) self.prefetch_queue.put(operation) return operation def terminate_prefetch(self, operation): operation.mark_terminate() return operation.completed_tokens, operation.hash_value def append_host_mem_release(self, host_indices: torch.Tensor): if host_indices.numel() == 0: return pages = host_indices.split(self.mem_pool_host.page_size) for page in pages: self.host_mem_release_queue.put(page) def _page_get_zero_copy( self, operation, hash_values, host_indices, extra_info=None ): results = self.storage_backend.batch_get_v1( hash_values, host_indices, extra_info ) inc = 0 for i in range(len(hash_values)): if not results[i]: logger.warning( f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}." ) break inc += self.page_size operation.increment(inc) # todo: deprecate def _generic_page_get(self, operation, hash_values, host_indices, extra_info=None): dummy_page_dst = [ self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values ] page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst) if page_data is None: return for i in range(len(hash_values)): if page_data[i] is None: logger.warning( f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}." ) break # Must set the data before increasing the completed tokens. # Otherwise this page may be read before being set. self.mem_pool_host.set_from_flat_data_page( host_indices[i * self.page_size], page_data[i], ) if not operation.increment(self.page_size): break # Operation terminated by controller def _page_transfer(self, operation): # Transfer batch by batch prefix_keys = operation.prefix_keys for i in range(0, len(operation.hash_value), self.storage_batch_size): batch_hashes = operation.hash_value[i : i + self.storage_batch_size] batch_host_indices = operation.host_indices[ i * self.page_size : (i + len(batch_hashes)) * self.page_size ] prev_completed_tokens = operation.completed_tokens # Get one batch token, and update the completed_tokens if succeed extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) self.page_get_func(operation, batch_hashes, batch_host_indices, extra_info) # Check termination if ( operation.completed_tokens != prev_completed_tokens + len(batch_hashes) * self.page_size ): operation.mark_terminate() break # Some operations fail or operation terminated by controller if prefix_keys and len(prefix_keys) > 0: prefix_keys += batch_hashes def prefetch_io_aux_func(self): """ Auxiliary function conducting IO operations for prefetching. """ while not self.storage_stop_event.is_set(): try: operation = self.prefetch_buffer.get(block=True, timeout=1) if operation is None: continue self._page_transfer(operation) # operation terminated by controller, release pre-allocated memory self.append_host_mem_release( operation.host_indices[operation.completed_tokens :] ) except Empty: continue def prefetch_rate_limited(self) -> bool: """ Rate limit the prefetching operations to avoid overwhelming the storage backend. """ # cancel prefetch if too much memory is occupied if self.prefetch_tokens_occupied >= self.prefetch_capacity_limit: return True # todo: more sophisticated rate limiting based on storage backend performance return False def _storage_hit_query(self, operation) -> tuple[list[str], int]: last_hash = operation.last_hash tokens_to_fetch = operation.token_ids prefix_keys = operation.prefix_keys.copy() if operation.prefix_keys else None storage_query_count = 0 hash_value = [] for start in range( 0, len(tokens_to_fetch), self.page_size * self.storage_batch_size ): end = min( start + self.page_size * self.storage_batch_size, len(tokens_to_fetch) ) batch_tokens = tokens_to_fetch[start:end] batch_hashes = [] for i in range(0, len(batch_tokens), self.page_size): last_hash = self.get_hash_str( batch_tokens[i : i + self.page_size], last_hash ) batch_hashes.append(last_hash) extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info) hash_value.extend(batch_hashes[:hit_page_num]) storage_query_count += hit_page_num * self.page_size if hit_page_num < len(batch_hashes): break if prefix_keys and len(prefix_keys) > 0: prefix_keys += batch_hashes return hash_value, storage_query_count def prefetch_thread_func(self): """ Manage prefetching operations from storage backend to host memory. """ self.prefetch_buffer = Queue() self.prefetch_io_aux_thread = threading.Thread( target=self.prefetch_io_aux_func, daemon=True ) self.prefetch_io_aux_thread.start() while (not self.storage_stop_event.is_set()) or not self.prefetch_queue.empty(): try: operation = self.prefetch_queue.get(block=True, timeout=1) if operation is None: continue hash_value, storage_hit_count = self._storage_hit_query(operation) if self.tp_world_size > 1: storage_hit_count_tensor = torch.tensor( storage_hit_count, dtype=torch.int ) torch.distributed.all_reduce( storage_hit_count_tensor, op=torch.distributed.ReduceOp.MIN, group=self.prefetch_tp_group, ) storage_hit_count = storage_hit_count_tensor.item() if storage_hit_count < self.prefetch_threshold: # not to prefetch if not enough benefits self.prefetch_revoke_queue.put(operation.request_id) self.append_host_mem_release(operation.host_indices) logger.debug( f"Revoking prefetch for request {operation.request_id} due to insufficient hits ({storage_hit_count})." ) else: operation.hash_value = hash_value[ : (storage_hit_count // self.page_size) ] # free the pre-allocated memory for pages that are not hit self.append_host_mem_release( operation.host_indices[storage_hit_count:] ) operation.host_indices = operation.host_indices[:storage_hit_count] logger.debug( f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}." ) self.prefetch_buffer.put(operation) except Empty: continue def write_storage( self, host_indices: torch.Tensor, token_ids: List[int], hash_value: Optional[List[str]] = None, prefix_keys: Optional[List[str]] = None, ) -> int: """ Write KV caches from host memory to storage backend. """ operation = StorageOperation( host_indices, token_ids, hash_value=hash_value, prefix_keys=prefix_keys ) self.backup_queue.put(operation) return operation.id # todo: deprecate def _generic_page_set(self, hash_values, host_indices, extra_info=None) -> bool: data = [ self.mem_pool_host.get_data_page(host_indices[i * self.page_size]) for i in range(len(hash_values)) ] return self.storage_backend.batch_set(hash_values, data) def _page_set_zero_copy(self, hash_values, host_indices, extra_info=None) -> bool: return all( self.storage_backend.batch_set_v1(hash_values, host_indices, extra_info) ) # Backup batch by batch def _page_backup(self, operation): # Backup batch by batch prefix_keys = operation.prefix_keys for i in range(0, len(operation.hash_value), self.storage_batch_size): batch_hashes = operation.hash_value[i : i + self.storage_batch_size] batch_host_indices = operation.host_indices[ i * self.page_size : (i + len(batch_hashes)) * self.page_size ] # Set one batch token, and record if success. # todo: allow partial success extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) success = self.page_set_func(batch_hashes, batch_host_indices, extra_info) if not success: logger.warning( f"Write page to storage: {len(batch_hashes)} pages failed." ) break if prefix_keys and len(prefix_keys) > 0: prefix_keys += batch_hashes operation.completed_tokens += self.page_size * len(batch_hashes) def backup_thread_func(self): """ Manage backup operations from host memory to storage backend. """ while not self.storage_stop_event.is_set(): try: operation = self.backup_queue.get(block=True, timeout=1) if operation is None: continue if not self.backup_skip: self._page_backup(operation) self.ack_backup_queue.put(operation) except Empty: continue