# SPDX-License-Identifier: Apache-2.0 """Blend (context-blend / cross-request KV reuse) module for MPCacheServer.""" # Standard from typing import Any import threading import time # Third Party import numpy as np # First Party from lmcache import torch_dev, torch_device_type from lmcache.logging import init_logger from lmcache.utils import check_interprocess_event_support from lmcache.v1.distributed.api import ( MemoryLayoutDesc, ObjectKey, PrefetchHandle, ipc_key_to_object_keys, ) from lmcache.v1.gpu_connector.gpu_ops import ( lmcache_memcpy_async_d2h, lmcache_memcpy_async_h2d, ) from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.multiprocess.custom_types import ( CBMatchResult, IPCCacheServerKey, KVCache, ) from lmcache.v1.multiprocess.engine_context import MPCacheServerContext from lmcache.v1.multiprocess.engine_module import ( HandlerSpec, ThreadPoolType, ) from lmcache.v1.multiprocess.protocols.base import RequestType from lmcache.v1.multiprocess.token_hasher import ( chunk_hash_windows_numba, rolling_hash_windows_numba, unique_hits_direct_id_numba, update_table_id_numba, ) from lmcache.v1.platform.cuda.cache_context import PlainGPUCacheContext logger = init_logger(__name__) class BlendTokenRangeMatcher: """Fast token-range matcher using polynomial rolling/chunk hashes and a direct-address lookup table. Table layout: poly_chunk_hash (u64) -> compact_chunk_id (i64, sequential 0...N-1). Because compact IDs are bounded by _TABLE_SIZE, unique_hits_direct_id_numba can use a fixed ``seen`` array of _TABLE_SIZE bytes (~1 MB) rather than one sized by an arbitrary max hash -- no memory explosion. Auxiliary storage: _chunk_token_hash[i] : token_hash for chunk i (None if evicted) _token_hash_to_start : token_hash -> start position in seq _compact_id_to_slot[i] : table slot for compact_id i _token_hash_to_compact_id : token_hash -> compact_chunk_id Methods: on_new_token_hashes -- register a sequence; builds fingerprints and writes compact IDs. match_sub_sequence -- sliding-window probe -> compact IDs -> token_hash -> start. Skips evicted entries. remove_chunks -- lazily evict stale entries. Clears the table slot and auxiliary maps. Args: chunk_size: Number of tokens per chunk for fingerprint computation. """ _TABLE_BITS: int = 20 # 2^20 ~ 1 M entries _TABLE_SIZE: int = 1 << _TABLE_BITS _BASE: np.uint64 = np.uint64(0x9E3779B97F4A7C15) # Fibonacci-hashing constant def __init__(self, chunk_size: int = 256): self.chunk_size = chunk_size # poly_chunk_hash -> compact_chunk_id; -1 = empty self._table_id = np.full(self._TABLE_SIZE, -1, dtype=np.int64) self._mask = np.uint64(self._TABLE_SIZE - 1) # compact_chunk_id -> caller-supplied token_hash (full bytes) self._chunk_token_hash: list[bytes | None] = [] # token_hash -> start position in its registered sequence self._token_hash_to_start: dict[bytes, int] = {} # compact_chunk_id -> table slot index (for reverse lookup during eviction) self._compact_id_to_slot = np.full(self._TABLE_SIZE, -1, dtype=np.int64) # token_hash -> compact_chunk_id (for eviction lookup) self._token_hash_to_compact_id: dict[bytes, int] = {} self._lock = threading.Lock() def on_new_token_hashes( self, token_ids: list[int], token_hashes: list[bytes], ) -> None: """Register a new token sequence and index its non-overlapping chunks. Args: token_ids: Raw token IDs for the full sequence (num_tokens elements). Used to compute polynomial chunk fingerprints that match the rolling hashes computed in match_sub_sequence. token_hashes: Per-chunk bytes hashes supplied by the caller (one per complete chunk of chunk_size tokens). Stored as the storage key returned in CBMatchResult.hash. """ arr = np.array(token_ids, dtype=np.uint64) # Polynomial fingerprints for non-overlapping chunks, built from raw # token IDs so they match the rolling hashes in match_sub_sequence chunk_hashes = chunk_hash_windows_numba(arr, self.chunk_size, self._BASE) n = int(chunk_hashes.shape[0]) if n == 0: return with self._lock: # Filter chunks already registered to avoid duplicate compact-ID # allocation. When both cb_store_pre_computed and cb_store_final # fire for the same token sequence they produce identical hashes; # registering twice orphans the first compact ID permanently since # _token_hash_to_compact_id is overwritten but the old list slot is # not freed. new_idxs = [ i for i in range(n) if token_hashes[i] not in self._token_hash_to_compact_id ] if not new_idxs: return n_new = len(new_idxs) new_chunk_hashes = chunk_hashes[new_idxs] # Compact sequential IDs: bounded by _TABLE_SIZE, safe for seen-array sizing # NOTE: base_id grows monotonically (evicted slots are not reused); the hard # limit is on total chunks ever registered, not active chunks. base_id = len(self._chunk_token_hash) if base_id + n_new > self._TABLE_SIZE: logger.error( "BlendTokenRangeMatcher compact-ID overflow: %d chunks " "registered, cannot add %d more (limit %d). Skipping.", base_id, n_new, self._TABLE_SIZE, ) return if base_id + n_new > int(self._TABLE_SIZE * 0.8): logger.warning( "BlendTokenRangeMatcher nearing capacity: %d/%d compact IDs used. " "Hash collision rate is rising; hit rate will degrade.", base_id + n_new, self._TABLE_SIZE, ) compact_ids = np.arange(base_id, base_id + n_new, dtype=np.int64) # Write table: poly_chunk_hash -> compact_chunk_id update_table_id_numba(new_chunk_hashes, self._table_id, compact_ids) # Persist compact_id -> token_hash, token_hash -> start, and reverse maps for k, orig_i in enumerate(new_idxs): th = token_hashes[orig_i] cid = int(compact_ids[k]) slot = int(new_chunk_hashes[k]) & int(self._mask) self._chunk_token_hash.append(th) self._token_hash_to_start[th] = orig_i * self.chunk_size self._compact_id_to_slot[cid] = slot self._token_hash_to_compact_id[th] = cid def match_sub_sequence( self, token_ids: list[int], ) -> list[CBMatchResult]: """Find stored chunks whose fingerprints appear anywhere in token_ids. Uses a sliding-window rolling hash so matches need not be aligned to chunk_size boundaries in the query. Entries previously evicted via remove_chunks (token_hash set to None) are silently skipped. Args: token_ids: Query token sequence to probe (raw token IDs as uint64). Returns: One CBMatchResult per unique stored chunk that was hit. old_st/old_ed : positions in the originally registered sequence cur_st/cur_ed : positions in the query (token_ids) where the match was found hash : token_hash bytes (from registration) for cache key lookup """ if len(token_ids) < self.chunk_size: return [] arr = np.array(token_ids, dtype=np.uint64) # Sliding-window polynomial hashes over the query rolling = rolling_hash_windows_numba(arr, self.chunk_size, self._BASE) with self._lock: if not self._chunk_token_hash: return [] # Probe table; seen array is _TABLE_SIZE bytes (~1 MB), fixed and safe hit_ids = unique_hits_direct_id_numba( rolling, self._table_id, self._mask, self._TABLE_SIZE ) if hit_ids.shape[0] == 0: return [] # For each hit compact_id, find the first query position where it matched hit_id_set = set(int(cid) for cid in hit_ids) cid_to_query_pos: dict[int, int] = {} for q_pos in range(rolling.shape[0]): idx = int(rolling[q_pos]) & int(self._mask) cid = int(self._table_id[idx]) if cid in hit_id_set and cid not in cid_to_query_pos: cid_to_query_pos[cid] = q_pos if len(cid_to_query_pos) == len(hit_id_set): break results: list[CBMatchResult] = [] for cid in hit_ids: cid_int = int(cid) th = self._chunk_token_hash[cid_int] if th is None: continue old_st = self._token_hash_to_start.get(th) cur_st = cid_to_query_pos.get(cid_int) if old_st is None or cur_st is None: continue results.append( CBMatchResult( old_st=old_st, old_ed=old_st + self.chunk_size, cur_st=cur_st, cur_ed=cur_st + self.chunk_size, hash=th, ) ) return results def remove_chunks(self, token_hashes: list[bytes]) -> None: """Evict stale entries whose backing data is no longer in storage. Args: token_hashes: Token hashes of chunks to remove from the table. """ with self._lock: for th in token_hashes: cid = self._token_hash_to_compact_id.get(th) if cid is None: continue # Clear the table slot slot = int(self._compact_id_to_slot[cid]) if slot < 0: logger.warning( "compact_id %d has no valid table slot; " "entry may have been evicted twice", cid, ) continue self._table_id[slot] = -1 self._compact_id_to_slot[cid] = -1 # Clean up auxiliary maps self._chunk_token_hash[cid] = None self._token_hash_to_start.pop(th, None) del self._token_hash_to_compact_id[th] def has_chunk(self, token_hash: bytes) -> bool: """Return True if token_hash is currently registered in the matcher. Used before lazy registration to avoid creating duplicate compact-ID entries for a hash that is already in the fingerprint table. Args: token_hash: The storage hash bytes for a single chunk (as returned by TokenHasher.compute_chunk_hashes). Returns: True if the chunk is registered and not evicted, False otherwise. """ with self._lock: return token_hash in self._token_hash_to_compact_id def _unique_token_coverage(results: list[CBMatchResult]) -> int: """Return the number of unique query tokens covered by a set of CBMatchResults. match_sub_sequence is a sliding-window probe, so two results from different registered chunks can have overlapping [cur_st, cur_ed) ranges. Summing chunk_size per result would double-count the overlapping tokens and produce hit_rate > 1. This function merges the intervals first. Args: results: Found CBMatchResult objects (each covers [cur_st, cur_ed) tokens). Returns: Total number of unique query-token positions covered. """ if not results: return 0 intervals = sorted((r.cur_st, r.cur_ed) for r in results) coverage = 0 cur_end = -1 for st, ed in intervals: if st >= cur_end: coverage += ed - st elif ed > cur_end: coverage += ed - cur_end cur_end = max(cur_end, ed) return coverage class BlendModule: """Handles blend (context-blend / cross-request KV reuse) operations. Owns CB-specific GPU context registrations and the token range matcher. Provides handlers for CB register, unregister, store, retrieve, and lookup. Args: ctx: The shared engine context. """ def __init__(self, ctx: MPCacheServerContext) -> None: self._ctx = ctx self._cb_gpu_contexts: dict[int, PlainGPUCacheContext] = {} self._cb_gpu_context_meta: dict[int, tuple[str, int]] = {} self._token_range_matcher = BlendTokenRangeMatcher(ctx.chunk_size) self._gpu_copy_lock = threading.Lock() @property def context(self) -> MPCacheServerContext: """Return the shared engine context. Exposed for testing only.""" return self._ctx def get_handlers(self) -> list[HandlerSpec]: """Return handler specs for all request types this module serves. Returns: A list of HandlerSpec entries mapping request types to their handler callables and thread pool assignments. """ return [ HandlerSpec( RequestType.CB_REGISTER_KV_CACHE, self.cb_register_kv_cache, ThreadPoolType.SYNC, ), HandlerSpec( RequestType.CB_UNREGISTER_KV_CACHE, self.cb_unregister_kv_cache, ThreadPoolType.SYNC, ), HandlerSpec( RequestType.CB_STORE_PRE_COMPUTED, self.cb_store_pre_computed, ThreadPoolType.AFFINITY, ), HandlerSpec( RequestType.CB_RETRIEVE_PRE_COMPUTED_V2, self.cb_retrieve_pre_computed, ThreadPoolType.AFFINITY, ), HandlerSpec( RequestType.CB_STORE_FINAL, self.cb_store_final, ThreadPoolType.AFFINITY, ), HandlerSpec( RequestType.CB_LOOKUP_PRE_COMPUTED_V2, self.cb_lookup_pre_computed, ThreadPoolType.NORMAL, ), ] def report_status(self) -> dict: """Return blend module status information. Returns: A dict containing registered CB GPU instance IDs and per-instance KV cache layout metadata. """ cb_gpu_context_meta: dict[str, dict] = {} for gpu_id, meta in self._cb_gpu_context_meta.items(): model_name, world_size = meta entry: dict = { "model_name": model_name, "world_size": world_size, } ctx = self._cb_gpu_contexts.get(gpu_id) if ctx is not None: # bytes per token = 2 (K+V) * num_layers * hidden_dim_size * # itemsize; num_tokens is the cache capacity, not a per-token # cost. cache_size_per_token = ( 2 * ctx.num_layers * ctx.hidden_dim_size * ctx.dtype.itemsize ) entry["kv_cache_layout"] = { "num_layers": ctx.num_layers, "num_tokens": ctx.num_tokens, "hidden_dim_size": ctx.hidden_dim_size, "dtype": str(ctx.dtype), "cache_size_per_token": cache_size_per_token, } cb_gpu_context_meta[str(gpu_id)] = entry return { "registered_cb_gpu_ids": list(self._cb_gpu_contexts.keys()), "cb_gpu_context_meta": cb_gpu_context_meta, } def close(self) -> None: """Release resources owned by this module.""" self._cb_gpu_contexts.clear() self._cb_gpu_context_meta.clear() def cb_register_kv_cache( self, instance_id: int, kv_caches: KVCache, model_name: str, world_size: int, ) -> None: """Register the KV cache buffer from the blend engine. Args: instance_id: Unique identifier for the blend engine instance. kv_caches: KVCache object containing the GPU buffer pointers. model_name: The name of the model associated with this KV cache. world_size: The world size associated with this KV cache. """ gpu_context = PlainGPUCacheContext(kv_caches, self._ctx.chunk_size) self._cb_gpu_contexts[instance_id] = gpu_context self._cb_gpu_context_meta[instance_id] = (model_name, world_size) layout_desc = MemoryLayoutDesc( shapes=[gpu_context.get_kv_buffer_shape(self._ctx.chunk_size)], dtypes=[gpu_context.dtype], ) self._ctx.layout_desc_registry.register(model_name, world_size, layout_desc) logger.info( "Registered CB KV cache for instance_id %d with %d layers", instance_id, gpu_context.num_layers, ) def cb_unregister_kv_cache(self, instance_id: int) -> None: """Unregister the KV cache buffer for the given instance_id. Args: instance_id: Unique identifier for the blend engine instance to unregister. """ if instance_id in self._cb_gpu_contexts: model_name, world_size = self._cb_gpu_context_meta[instance_id] del self._cb_gpu_contexts[instance_id] del self._cb_gpu_context_meta[instance_id] self._ctx.layout_desc_registry.unregister(model_name, world_size) logger.info("Unregistered CB KV cache for instance_id %d", instance_id) else: logger.warning( "Attempted to unregister non-existent CB KV cache for instance_id %d", instance_id, ) def cb_lookup_pre_computed(self, key: IPCCacheServerKey) -> list[CBMatchResult]: """Lookup the pre-computed chunks in the underlying storage. Uses BlendTokenRangeMatcher for a fast local pre-filter, then submits prefetch tasks for matched chunks using their stored hashes directly. Chunks that the fingerprint table matched but are no longer present in storage are lazily evicted from the matcher via remove_chunks. Args: key: IPCCacheServerKey containing the token ids to lookup. Returns: List of CBMatchResult for chunks that were actually found in storage, ready to be passed to cb_retrieve_pre_computed. """ num_tokens = len(key.token_ids) self._ctx.event_bus.publish( Event( event_type=EventType.CB_REQUEST_START, session_id=key.request_id, ) ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_LOOKUP_START, session_id=key.request_id, metadata={"num_tokens": num_tokens}, ) ) cb_match_result = self._token_range_matcher.match_sub_sequence( list(key.token_ids) ) if not cb_match_result: self._ctx.event_bus.publish( Event( event_type=EventType.CB_LOOKUP_END, session_id=key.request_id, metadata={ "num_tokens": num_tokens, "fingerprint_hits": 0, "prefix_hits": 0, "storage_hits": 0, "stale_chunks": 0, "no_gpu_context": False, "hit_tokens": 0, "requested_tokens": (num_tokens // self._ctx.chunk_size) * self._ctx.chunk_size, }, ) ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_REQUEST_END, session_id=key.request_id, ) ) return [] # Sort by query position cb_match_result.sort(key=lambda r: r.cur_st) # The sliding-window probe returns O(table_size) overlapping matches. # Greedy leftmost-first picks one chunk per slot; lossless when matches # are chunk-aligned (the CB case). deduped: list[CBMatchResult] = [] covered_end = -1 for r in cb_match_result: if r.cur_st >= covered_end: deduped.append(r) covered_end = r.cur_ed cb_match_result = deduped # Group consecutive matched chunks groups: list[list[CBMatchResult]] = [] for result in cb_match_result: if groups and groups[-1][-1].cur_ed == result.cur_st: groups[-1].append(result) else: groups.append([result]) prefetch_handles: list[PrefetchHandle] = [] found_cb_match_result: list[CBMatchResult] = [] model_name, world_size = key.model_name, key.world_size # Find the cb gpu context and calculate the layout desc layout_desc: MemoryLayoutDesc | None = None for gpu_id, (m_name, w_size) in self._cb_gpu_context_meta.items(): if m_name == model_name and w_size == world_size: cb_ctx = self._cb_gpu_contexts[gpu_id] layout_desc = MemoryLayoutDesc( shapes=[cb_ctx.get_kv_buffer_shape(self._ctx.chunk_size)], dtypes=[cb_ctx.dtype], ) break if layout_desc is None: logger.error( "No CB GPU context found for model %s with world size %d " "during cb_lookup_pre_computed!", model_name, world_size, ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_LOOKUP_END, session_id=key.request_id, metadata={ "num_tokens": num_tokens, "fingerprint_hits": 0, "prefix_hits": 0, "storage_hits": 0, "stale_chunks": 0, "no_gpu_context": True, "hit_tokens": 0, "requested_tokens": (num_tokens // self._ctx.chunk_size) * self._ctx.chunk_size, }, ) ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_REQUEST_END, session_id=key.request_id, ) ) return [] # Submit prefetch for each group. All candidates use the standard chunk # hash computed by token_hasher, which matches the hash used at store # time, so ipc_key_to_object_keys resolves correctly. for group in groups: chunk_hashes = [r.hash for r in group] obj_keys = ipc_key_to_object_keys(key, chunk_hashes, [0])[0] handle = self._ctx.storage_manager.submit_prefetch_task( obj_keys, layout_desc, external_request_id=key.request_id, ) prefetch_handles.append(handle) logger.debug( "Submitted prefetch for %d chunks starting at %d", len(group), group[0].cur_st, ) # Collect only the CBMatchResults for chunks actually found in storage stale_hashes: list[bytes] = [] for handle, group in zip(prefetch_handles, groups, strict=False): found = None while True: found = self._ctx.storage_manager.query_prefetch_status(handle) if found is not None: break time.sleep(0.001) # Real found count after dedup the TP found_count = found.count_leading_ones() // world_size start = group[0].cur_st end = group[-1].cur_ed if found_count > 0: found_cb_match_result.extend(group[:found_count]) # Chunks after found_count in the group are stale stale_hashes.extend(r.hash for r in group[found_count:]) logger.debug( "Found %d pre-computed chunks for range (%d, %d)", found_count, start, end, ) else: stale_hashes.extend(r.hash for r in group) logger.debug( "No pre-computed chunks found for range (%d, %d)", start, end, ) # Evict stale fingerprint entries; remove_chunks safely skips hashes that # were never registered (e.g. prefix-probe candidates not in storage). if stale_hashes: self._token_range_matcher.remove_chunks(stale_hashes) logger.debug( "Evicted %d stale chunks from fingerprint table", len(stale_hashes), ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_CHUNKS_EVICTED, metadata={"num_chunks": len(stale_hashes)}, ) ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_LOOKUP_END, session_id=key.request_id, metadata={ "num_tokens": num_tokens, "fingerprint_hits": len(found_cb_match_result), "prefix_hits": 0, "storage_hits": len(found_cb_match_result), "stale_chunks": len(stale_hashes), "no_gpu_context": False, "hit_tokens": _unique_token_coverage(found_cb_match_result), "requested_tokens": (num_tokens // self._ctx.chunk_size) * self._ctx.chunk_size, }, ) ) return found_cb_match_result def _cb_store_gpu_copy( self, obj_keys: list[ObjectKey], gpu_context: PlainGPUCacheContext, offset: int, event_ipc_handle: bytes, start_event: Event | None = None, ) -> tuple[Any, dict]: """Helper function to perform GPU-to-CPU copy operations for storing chunks. Args: obj_keys: List of object keys to store. gpu_context: GPU context for the blend engine instance. offset: The starting offset in the CB KV cache buffer. event_ipc_handle: The IPC handle for the GPU event that signals the completion of LLM inference. start_event: Optional event to publish on the stream after waiting for the vLLM GPU event, marking the true start of the store operation. Returns: A tuple of (event, reserved_dict) where event is the GPU event and reserved_dict is the dictionary of reserved memory objects. """ with ( torch_dev.device(gpu_context.device), torch_dev.stream(gpu_context.stream), ): # Not all backends support interprocess Events (CUDA IPC specific) check_interprocess_event_support() event = torch_dev.Event(interprocess=True) # Wait for vLLM event to finish # Not all backends support IPC event handles (CUDA IPC specific) if not hasattr(torch_dev.Event, "from_ipc_handle"): raise RuntimeError( f"Backend '{torch_device_type}' does not support IPC event " "handles (Event.from_ipc_handle not available). " "Multiprocess IPC requires CUDA." ) vllm_event = torch_dev.Event.from_ipc_handle( gpu_context.device, event_ipc_handle ) vllm_event.wait(stream=gpu_context.stream) if start_event is not None: self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, start_event ) # Prepare for the copy num_tokens = self._ctx.chunk_size cpu_shape = gpu_context.get_kv_buffer_shape(num_tokens) layout_desc = MemoryLayoutDesc( shapes=[cpu_shape], dtypes=[gpu_context.dtype] ) reserved_dict = self._ctx.storage_manager.reserve_write( obj_keys, layout_desc, "new" ) for idx, obj_key in enumerate(obj_keys): if obj_key in reserved_dict: memory_obj = reserved_dict[obj_key] else: continue offset_start = idx * self._ctx.chunk_size + offset offset_end = offset_start + self._ctx.chunk_size # Copy from GPU to CPU tmp_buffer = gpu_context.get_tmp_gpu_buffer(offset_end - offset_start) gpu_kv_slice = gpu_context.slice_kv_cache_on_tokens( offset_start, offset_end ) with self._gpu_copy_lock: tmp_buffer.copy_(gpu_kv_slice, non_blocking=True) lmcache_memcpy_async_d2h(tmp_buffer, memory_obj) event.record() # Call finish_write after the copy is done gpu_context.cupy_stream.launch_host_func( self._ctx.storage_manager.finish_write, list(reserved_dict.keys()), ) return event, reserved_dict def cb_store_pre_computed( self, key: IPCCacheServerKey, offset: int, instance_id: int, event_ipc_handle: bytes, ) -> tuple[bytes, bool]: """Store the pre-computed chunks in the underlying storage for later retrieval. Args: key: IPCCacheServerKey containing the token ids for which the pre-computed chunks are stored. offset: The starting offset in the CB KV cache buffer where the pre-computed chunks begin. instance_id: The instance_id of the blend engine instance to store the pre-computed chunks for. event_ipc_handle: The IPC handle for the CUDA event that signals the completion of LLM inference. Returns: IPC handle bytes for the event that signals the completion of storing the pre-computed chunks, and a boolean flag indicating if the store is successful. Raises: ValueError: If instance_id is not registered for CB KV cache. Note: The input tokens should not have any separator in it. It should just be one "paragraph". This function will discard the last partial chunk and only store the full chunks. """ num_tokens = key.end - key.start if instance_id not in self._cb_gpu_contexts: raise ValueError( f"Instance ID {instance_id} not registered for CB KV cache" ) gpu_context = self._cb_gpu_contexts[instance_id] # CPU-synchronous sentinel: GPU store is about to be enqueued. self._ctx.event_bus.publish( Event( event_type=EventType.CB_STORE_PRE_COMPUTED_SUBMITTED, session_id=key.request_id, metadata={"instance_id": instance_id}, ) ) self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_STORE_PRE_COMPUTED_START, session_id=key.request_id, metadata={"instance_id": instance_id, "num_tokens": num_tokens}, ), ) # Compute normal prefix hashes so these chunks are accessible both via # the CB lookup path and via the standard lookup/retrieve path. chunk_hashes = self._ctx.token_hasher.compute_chunk_hashes(list(key.token_ids)) # convert to object key obj_keys = ipc_key_to_object_keys(key, chunk_hashes, [0])[0] reserved_dict: dict = {} try: event, reserved_dict = self._cb_store_gpu_copy( obj_keys, gpu_context, offset, event_ipc_handle ) # Register chunk hashes with the local matcher for fast sub-sequence lookup token_hashes = list(chunk_hashes) # NOTE(Jiayi): We only register the token hashes for worker_id 0 or None # to avoid duplicate registration across workers. if key.worker_id in [0, None]: self._token_range_matcher.on_new_token_hashes( list(key.token_ids), token_hashes ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_FINGERPRINTS_REGISTERED, session_id=key.request_id, metadata={ "num_chunks": len(token_hashes), "num_tokens": len(list(key.token_ids)), }, ) ) logger.info( "Stored pre-computed doc with %d tokens, num stored chunks: %d", key.end - key.start, len(reserved_dict), ) except Exception: logger.exception("Cannot store pre-computed chunks due to exception") self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_STORE_PRE_COMPUTED_END, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_tokens": num_tokens, "stored_chunks": 0, "success": False, }, ), ) raise self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_STORE_PRE_COMPUTED_END, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_tokens": num_tokens, "stored_chunks": len(reserved_dict), "success": True, }, ), ) return event.ipc_handle(), True def cb_retrieve_pre_computed( self, key: IPCCacheServerKey, cb_match_result: list[CBMatchResult], offset: int, instance_id: int, event_ipc_handle: bytes, ) -> tuple[bytes, bool]: """Retrieve pre-computed chunks from storage and copy them to the CB KV buffer. Args: key: IPCCacheServerKey containing the token ids for which the pre-computed chunks are retrieved. cb_match_result: List of CBMatchResult returned by cb_lookup_pre_computed, containing the per-chunk hashes and query positions. offset: The starting offset in the CB KV cache buffer to copy the retrieved chunks to. instance_id: The instance_id of the blend engine instance to retrieve the pre-computed chunks for. event_ipc_handle: The IPC handle for the CUDA event that signals the completion of LLM inference. Returns: IPC handle bytes for the event that signals the completion of retrieving the pre-computed chunks, and a boolean flag indicating if the retrieval is successful. Raises: ValueError: If instance_id is not registered for CB KV cache. Note: cb_lookup_pre_computed must be called first before calling this function. """ if instance_id not in self._cb_gpu_contexts: raise ValueError( f"Instance ID {instance_id} not registered for CB KV cache" ) gpu_context = self._cb_gpu_contexts[instance_id] # One obj_key per match_result, in cur_st order cb_match_result = sorted(cb_match_result, key=lambda r: r.cur_st) num_chunks = len(cb_match_result) chunk_hashes = [r.hash for r in cb_match_result] all_obj_keys = ipc_key_to_object_keys(key, chunk_hashes, [0])[0] # CPU-synchronous sentinel: GPU retrieve is about to be enqueued. self._ctx.event_bus.publish( Event( event_type=EventType.CB_RETRIEVE_SUBMITTED, session_id=key.request_id, metadata={"instance_id": instance_id}, ) ) logger.debug("DEBUG object keys to retrieve: %s", all_obj_keys) with ( torch_dev.device(gpu_context.device), torch_dev.stream(gpu_context.stream), ): # Not all backends support interprocess Events (CUDA IPC specific) check_interprocess_event_support() event = torch_dev.Event(interprocess=True) self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_RETRIEVE_START, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_chunks": num_chunks, }, ), ) try: with self._ctx.storage_manager.read_prefetched_results( all_obj_keys ) as memory_objs: if memory_objs is None: logger.error("Some keys not found during CB retrieve!") self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_RETRIEVE_END, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_chunks": num_chunks, "success": False, }, ), ) return event.ipc_handle(), False for r, memory_obj in zip( cb_match_result, memory_objs, strict=False ): gpu_st = r.cur_st + offset gpu_ed = gpu_st + self._ctx.chunk_size tmp_buffer = gpu_context.get_tmp_gpu_buffer( self._ctx.chunk_size ) target_buffer = gpu_context.slice_kv_cache_on_tokens( gpu_st, gpu_ed ) with self._gpu_copy_lock: lmcache_memcpy_async_h2d(memory_obj, tmp_buffer) target_buffer.copy_(tmp_buffer, non_blocking=True) except Exception: logger.exception("Error during retrieving prefetched results") self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_RETRIEVE_END, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_chunks": num_chunks, "success": False, }, ), ) return event.ipc_handle(), False finally: event.record() # TODO: here we simply "unlock" all the keys, which may cause # double-unlock if error happens during read_prefetched_results. # We should consider not unlocking objects in read_prefetched_results # if error happens. gpu_context.cupy_stream.launch_host_func( self._ctx.storage_manager.finish_read_prefetched, all_obj_keys, ) logger.info( "Retrieved pre-computed for %d match results to GPU offset starting at %d", len(cb_match_result), offset, ) self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_RETRIEVE_END, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_chunks": num_chunks, "success": True, }, ), ) return event.ipc_handle(), True def cb_store_final( self, key: IPCCacheServerKey, offset: int, instance_id: int, event_ipc_handle: bytes, ) -> tuple[bytes, bool]: """Store the final chunks in the underlying storage after processing. The stored chunks should be accessible for normal mode LLMs. Args: key: IPCCacheServerKey containing the token ids for which the final chunks are stored. offset: The starting offset in the CB KV cache buffer where the final chunks are stored. instance_id: The instance_id of the blend engine instance to store the final chunks for. event_ipc_handle: The IPC handle for the CUDA event that signals the completion of LLM inference. Returns: IPC handle bytes for the event that signals the completion of storing the final chunks, and a boolean flag indicating if the store is successful. Raises: ValueError: If instance_id is not registered for CB KV cache. """ num_tokens = key.end - key.start # Get GPU context if instance_id not in self._cb_gpu_contexts: raise ValueError( f"Instance ID {instance_id} not registered for CB KV cache" ) gpu_context = self._cb_gpu_contexts[instance_id] # CPU-synchronous sentinels: SUBMITTED before SESSION_END so the # tracing subscriber's in-flight counter is non-zero when SESSION_END # arrives and correctly defers root span closure. self._ctx.event_bus.publish( Event( event_type=EventType.CB_STORE_FINAL_SUBMITTED, session_id=key.request_id, metadata={"instance_id": instance_id}, ) ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_REQUEST_END, session_id=key.request_id, ) ) # Compute normal hash for the keys chunk_hashes = self._ctx.token_hasher.compute_chunk_hashes(list(key.token_ids)) # convert to object key obj_keys = ipc_key_to_object_keys(key, chunk_hashes, [0])[0] reserved_dict: dict = {} try: event, reserved_dict = self._cb_store_gpu_copy( obj_keys, gpu_context, offset, event_ipc_handle, start_event=Event( event_type=EventType.CB_STORE_FINAL_START, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_tokens": num_tokens, }, ), ) # Register fingerprints so future CB lookups can find these chunks. # Mirrors cb_store_pre_computed; without this, chunks stored here are # invisible to cb_lookup_pre_computed, causing 0% hit rate on re-requests. if key.worker_id in [0, None]: self._token_range_matcher.on_new_token_hashes( list(key.token_ids), list(chunk_hashes) ) self._ctx.event_bus.publish( Event( event_type=EventType.CB_FINGERPRINTS_REGISTERED, session_id=key.request_id, metadata={ "num_chunks": len(chunk_hashes), "num_tokens": len(list(key.token_ids)), }, ) ) logger.info( "Stored final doc with %d tokens, num stored chunks: %d", key.end - key.start, len(reserved_dict), ) except Exception: logger.exception("Cannot store final chunks due to exception") self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_STORE_FINAL_END, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_tokens": num_tokens, "stored_chunks": 0, "success": False, }, ), ) raise self._ctx.event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_STORE_FINAL_END, session_id=key.request_id, metadata={ "instance_id": instance_id, "num_tokens": num_tokens, "stored_chunks": len(reserved_dict), "success": True, }, ), ) return event.ipc_handle(), True