# SPDX-License-Identifier: Apache-2.0 """Blend V3: paged-aware CacheBlend as an EngineModule. Plugs into the unified MPCacheServer; standard REGISTER_KV_CACHE + CB_REGISTER_ROPE_V3 for setup; STORE wrapper registers fingerprints; retrieve scatters into the request's paged blocks. """ # Standard from dataclasses import dataclass from queue import Empty as QueueEmpty from queue import Queue from typing import TYPE_CHECKING, Any import threading import time if TYPE_CHECKING: # First Party from lmcache.v1.mp_coordinator.blend_client import ( BlendCoordinatorClient, RemoteMatch, ) # Third Party import numpy as np import torch # 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, TrimPolicy, ipc_key_to_object_keys, ) from lmcache.v1.distributed.storage_manager import PrefetchHandle from lmcache.v1.gpu_connector.gpu_ops import lmcache_memcpy_async_h2d from lmcache.v1.mp_coordinator.blend_client import PENDING from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.multiprocess.custom_types import ( CBMatchResult, CBUnifiedLookupResult, DeviceIPCWrapper, IPCCacheServerKey, ) from lmcache.v1.multiprocess.engine_context import MPCacheServerContext from lmcache.v1.multiprocess.engine_module import ( HandlerSpec, InstanceLivenessTarget, ThreadPoolType, ) from lmcache.v1.multiprocess.modules.lmcache_driven_transfer import ( LMCacheDrivenTransferModule, ) from lmcache.v1.multiprocess.modules.lookup import compute_extra_count from lmcache.v1.multiprocess.protocol import RequestType from lmcache.v1.multiprocess.token_hasher import ( TokenHasher, chunk_hash_windows_numba, rolling_hash_windows_numba, update_table_id_numba, ) from lmcache.v1.platform.base_cache_context import BaseCacheContext import lmcache.c_ops as lmc_ops logger = init_logger(__name__) @dataclass class _CBRopeState: """Per-instance RoPE state IPC-shared from vLLM; dangles on reallocate. Models with per-layer-type RoPE (distinct local/global theta) register one cache per distinct rope and a per-layer index into ``cos_sin_caches``. """ head_size: int is_neox_style: bool # NeoX = contiguous halves; else GPT-J. cos_sin_caches: list[torch.Tensor] group_to_cache: list[int] # engine group idx -> cache idx; empty = cache 0 def cache_for_group(self, engine_group_idx: int) -> torch.Tensor: """The cos/sin cache for one engine group. Engine groups partition layers by attention type, and rope follows attention type (sliding=local theta, full=global theta), so each engine group has exactly one cache. Args: engine_group_idx: The kernel group's engine group index. Returns: The group's cos/sin cache tensor. Raises: RuntimeError: If ``engine_group_idx`` is outside the map. """ if not self.group_to_cache: return self.cos_sin_caches[0] if engine_group_idx >= len(self.group_to_cache): raise RuntimeError( f"CB re-RoPE: engine group {engine_group_idx} has no rope " f"cache mapping (map covers {len(self.group_to_cache)} groups)." ) return self.cos_sin_caches[self.group_to_cache[engine_group_idx]] @dataclass class _CBUnifiedJob: """Per-request poll state for non-blocking cb_unified_lookup. Stashed across polls because the underlying status/found polls are consume-once. """ matches: list[CBMatchResult] num_tokens: int = 0 # Prefix leg (blend_v3-owned submit/poll). ``prefix_handle`` is None when # there is no GPU context / no full chunk (poll reports 0 coverage). prefix_handle: PrefetchHandle | None = None prefix_world_size: int = 1 prefix_chunks: int | None = None # stashed when the prefix poll completes retained_chunks: list[int] | None = None # SEGMENTED_PREFIX: full gapped set sparse_started: bool = False # prefix done -> sparse leg submitted/skipped handle: PrefetchHandle | None = None # sparse handle, None if no sparse leg non_prefix: list[CBMatchResult] | None = None per_hash_obj_keys: dict | None = None expanded_uidx: list[int] | None = None found_uidx: set[int] | None = None # stashed when the sparse poll completes l2_keys: int = 0 # sparse keys needing an L2 load (0 => no L2 read, span skipped) coord_submitted: bool = False # coordinator match query was issued coord_deadline: float = 0.0 # time.monotonic() wall-clock cutoff for the leg class BlendTokenRangeMatcherV3: """V3 matcher: token-level probe (any offset) + full-hash collision rejection. Self-contained (does not inherit a base matcher).""" _TABLE_BITS: int = 20 # 2^20 ~ 1 M entries _TABLE_SIZE: int = 1 << _TABLE_BITS _BASE: np.uint64 = np.uint64(0x9E3779B97F4A7C15) # Fibonacci-hashing const def __init__(self, chunk_size: int = 256): """Initialize the V3 matcher. Args: chunk_size (int): Tokens per non-overlapping fingerprint chunk. """ 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 token_hash (full bytes); None once evicted 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 (reverse lookup for 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() # V3 addition: compact_chunk_id -> full poly hash, for collision reject. self._chunk_poly_hash: list[int] = [] def on_new_token_hashes( self, token_ids: list[int], token_hashes: list[bytes], start_chunk_idx: int = 0, position_offset: int = 0, ) -> None: """Index a stored sequence's non-overlapping chunks into the matcher. Records each new chunk's poly hash + start position so a later match_sub_sequence can find it. Thread-safe (holds the matcher lock). Args: token_ids (list[int]): The stored sequence's token IDs. token_hashes (list[bytes]): Per-chunk content hashes (one per chunk), used as the dedup/eviction key. start_chunk_idx (int): First chunk to index; 1 skips chunk 0 (the prefix lookup leg owns it). position_offset (int): Added to each recorded start position (for indexing a tail-slice of a larger sequence). Returns: None. """ arr = np.array(token_ids, dtype=np.uint64) chunk_hashes = chunk_hash_windows_numba(arr, self.chunk_size, self._BASE) n = int(chunk_hashes.shape[0]) if n == 0 or start_chunk_idx >= n: return with self._lock: new_idxs = [ i for i in range(start_chunk_idx, 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] base_id = len(self._chunk_token_hash) if base_id + n_new > self._TABLE_SIZE: logger.error( "BlendTokenRangeMatcherV3 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( "BlendTokenRangeMatcherV3 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) update_table_id_numba(new_chunk_hashes, self._table_id, compact_ids) for k, orig_i in enumerate(new_idxs): th = token_hashes[orig_i] cid = int(compact_ids[k]) poly_hash = int(new_chunk_hashes[k]) slot = poly_hash & int(self._mask) self._chunk_token_hash.append(th) self._chunk_poly_hash.append(poly_hash) self._token_hash_to_start[th] = ( position_offset + 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 every registered chunk reused anywhere in a query sequence. Vectorized direct-address probe over all token positions, then a small verify loop over the surviving hits (a full poly-hash check rejects bucket collisions; evicted/unknown chunks are skipped). Thread-safe. Args: token_ids (list[int]): The query sequence's token IDs. Returns: list[CBMatchResult]: One result per unique reused chunk (cur_st = its first query position, old_st = its stored position). Empty if the query is shorter than one chunk or nothing matched. """ if len(token_ids) < self.chunk_size: return [] arr = np.array(token_ids, dtype=np.uint64) rolling = rolling_hash_windows_numba(arr, self.chunk_size, self._BASE) with self._lock: if not self._chunk_token_hash: return [] # Vectorized direct-address probe over all positions. The table is # sparse (TABLE_SIZE >> registered chunks), so only true matches and # a few bucket collisions reach the Python verify loop below. cids_at_pos = self._table_id[rolling & self._mask] hit_positions = np.nonzero(cids_at_pos >= 0)[0] seen_cids: set[int] = set() results: list[CBMatchResult] = [] for pos in hit_positions: pos = int(pos) cid = int(cids_at_pos[pos]) if cid in seen_cids: continue if int(rolling[pos]) != self._chunk_poly_hash[cid]: continue # bucket-only collision th = self._chunk_token_hash[cid] if th is None: continue # evicted old_st = self._token_hash_to_start.get(th) if old_st is None: continue seen_cids.add(cid) results.append( CBMatchResult( old_st=old_st, old_ed=old_st + self.chunk_size, cur_st=pos, cur_ed=pos + self.chunk_size, hash=th, ) ) logger.info( "[match_probe] n_tok=%d table_hits=%d matches=%d", len(token_ids), len(hit_positions), len(results), ) return results def remove_chunks(self, token_hashes: list[bytes]) -> None: """Evict the given chunks from the matcher. Clears each chunk's table slot + poly hash so later probes cannot match it. Thread-safe. Args: token_hashes (list[bytes]): Content hashes of the chunks to evict. """ with self._lock: for th in token_hashes: cid = self._token_hash_to_compact_id.get(th) if cid is None: continue 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 self._chunk_token_hash[cid] = None self._chunk_poly_hash[cid] = 0 self._token_hash_to_start.pop(th, None) del self._token_hash_to_compact_id[th] def _unique_token_coverage(results: list[CBMatchResult]) -> int: """Total token coverage, merging overlapping ranges (sliding-window probe can return overlaps; naive sum would double-count).""" 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 BlendV3Module(InstanceLivenessTarget): """Paged-aware V3 CacheBlend. Wraps LMCacheDrivenTransfer STORE to register fingerprints; serves CB rope/lookup/retrieve RPCs; reads cross-module GPU state via :class:`LMCacheDrivenTransferModule.cache_contexts`.""" def __init__( self, ctx: MPCacheServerContext, lmcache_driven_transfer: LMCacheDrivenTransferModule, coordinator: "BlendCoordinatorClient | None" = None, enable_segmented_prefix: bool = False, ): self._ctx = ctx self._transfer_module = lmcache_driven_transfer # Server config (--enable-segmented-prefix): retain the gapped prefix on # a mid-prefix L2 retrieve failure instead of truncating at the gap. self._segmented_prefix = enable_segmented_prefix # Optional bridge to the fleet-wide fingerprint directory. ``None`` => # purely local matching (publish/query paths skipped). self._coordinator = coordinator self._token_range_matcher = BlendTokenRangeMatcherV3(ctx.chunk_size) self._event_bus = ctx.event_bus self._cb_rope_state: dict[int, _CBRopeState] = {} # L2 opt: cache TP-expanded obj_keys at lookup, pop at retrieve. self._lookup_obj_keys_cache: dict[str, dict[bytes, list]] = {} self._lookup_obj_keys_lock = threading.Lock() # Non-blocking cb_unified_lookup poll state (submit-once, poll-on-recall) # so the handler never holds a worker thread across the L2->L1 loads. self._cb_jobs: dict[str, _CBUnifiedJob] = {} self._cb_jobs_lock = threading.Lock() # Async fingerprint registration: store enqueues, worker drains. _FpJob = tuple[list[int], list[bytes], int, int] self._fingerprint_queue: "Queue[_FpJob]" = Queue() self._fingerprint_stop = threading.Event() self._fingerprint_worker = threading.Thread( target=self._drain_fingerprint_queue, name="cb-fingerprint-worker", daemon=True, ) self._fingerprint_worker.start() # In-flight fingerprint hashes; storage_gate keeps these from eviction. self._pending_fp_hashes: set[bytes] = set() self._pending_fp_lock = threading.Lock() # Lazy eviction strikes; evict only at threshold so async re-store # can refresh the bucket first. self._stale_strike: dict[bytes, int] = {} self._STALE_STRIKE_THRESHOLD = 2 # ------------------------------------------------------------------ # EngineModule protocol # ------------------------------------------------------------------ @property def context(self) -> MPCacheServerContext: return self._ctx def get_handlers(self) -> list[HandlerSpec]: # STORE shadows LMCacheDrivenTransfer's; compositor registers V3 last. return [ HandlerSpec(RequestType.STORE, self.store, ThreadPoolType.AFFINITY), HandlerSpec( RequestType.CB_REGISTER_ROPE_V3, self.cb_register_rope, ThreadPoolType.SYNC, ), HandlerSpec( RequestType.CB_UNREGISTER_ROPE_V3, self.cb_unregister_rope, ThreadPoolType.SYNC, ), HandlerSpec( RequestType.CB_UNIFIED_LOOKUP, self.cb_unified_lookup, ThreadPoolType.NORMAL, ), HandlerSpec( RequestType.CB_RETRIEVE_PRE_COMPUTED_V3, self.cb_retrieve_pre_computed, ThreadPoolType.AFFINITY, ), ] def report_status(self) -> dict: # Meta is derived live from MP server gpu_transfe cache_contexts = self._transfer_module.context_entries_snapshot() def _meta(iid: int) -> "tuple[str, int] | None": entry = cache_contexts.get(iid) return (entry.model_name, entry.world_size) if entry is not None else None return { "registered_cb_rope_instances": list(self._cb_rope_state.keys()), "cb_rope_meta": {str(iid): _meta(iid) for iid in self._cb_rope_state}, "active_cb_lookups": len(self._cb_jobs), } def close(self) -> None: self._fingerprint_stop.set() if self._coordinator is not None: # Joins the client's daemon thread and closes its httpx.Client; # otherwise the coordinator leg leaks both on server shutdown. self._coordinator.close() self._cb_rope_state.clear() # ------------------------------------------------------------------ # V3 RPCs # ------------------------------------------------------------------ def cb_register_rope( self, instance_id: int, cos_sin_caches_ipc: list[DeviceIPCWrapper], head_size: int, is_neox_style: bool, group_to_cache: list[int], ) -> None: """Bolt CB re-RoPE state onto an already-registered KV-cache instance. Idempotent; ``REGISTER_KV_CACHE`` must precede this. Strips any YaRN/longrope mscale baked into each rope cache so re-RoPE stays a pure rotation. Args: instance_id (int): KV-cache instance to attach rope state to. cos_sin_caches_ipc (list[DeviceIPCWrapper]): IPC handles to vLLM's cos/sin rope cache(s) — one per distinct rope (dual-RoPE models send local/global); single-rope models send one. head_size (int): Rotary head dimension. is_neox_style (bool): True for NeoX (contiguous halves), else GPT-J. group_to_cache (list[int]): Per-engine-group index into the caches list; empty means every group uses cache 0. Raises: ValueError: If ``instance_id`` has no registered KV cache, the cache list is empty, or ``group_to_cache`` references a missing cache or does not cover every engine group of the registered model. """ entry = self._transfer_module.get_and_touch_context_entry(instance_id) if entry is None: raise ValueError( f"Instance {instance_id} has no paged KV cache registered; " "send REGISTER_KV_CACHE before CB_REGISTER_ROPE_V3." ) if not cos_sin_caches_ipc: raise ValueError("CB_REGISTER_ROPE_V3 requires >=1 cos/sin cache.") if group_to_cache: if min(group_to_cache) < 0 or max(group_to_cache) >= len( cos_sin_caches_ipc ): raise ValueError( f"group_to_cache {group_to_cache} contains indices outside " f"[0, {len(cos_sin_caches_ipc)}) for the sent cache(s)." ) # Fail at registration, not mid-retrieve: every engine group of # the registered model must have a cache mapping. max_eg_idx = max( ( g.engine_group_idx for g in entry.cache_context.kv_layer_groups_manager.kernel_groups ), default=-1, ) if len(group_to_cache) <= max_eg_idx: raise ValueError( f"group_to_cache covers {len(group_to_cache)} engine " f"group(s) but the registered model has engine groups up " f"to index {max_eg_idx}." ) cos_sin_caches: list[torch.Tensor] = [] for cache_idx, cache_ipc in enumerate(cos_sin_caches_ipc): cos_sin_cache = cache_ipc.to_tensor() # YaRN/longrope bake an mscale m into the rope cache # (cos²+sin²=m²≠1). vLLM already folds m into stored K, but CB # re-RoPE assumes a pure rotation, so an un-normalized m injects # an m² error per K element. _c32 = cos_sin_cache.to(torch.float32) _half = _c32.shape[1] // 2 _m = float((_c32[:, :_half] ** 2 + _c32[:, _half:] ** 2).mean().sqrt()) if abs(_m - 1.0) >= 1e-3: logger.info( "CB re-RoPE: cache %d: stripping rope-cache mscale=%.4f " "(m²=%.4f → K inflation if uncorrected) → unit magnitude", cache_idx, _m, _m * _m, ) cos_sin_cache = (_c32 / _m).to(cos_sin_cache.dtype) cos_sin_caches.append(cos_sin_cache) self._cb_rope_state[instance_id] = _CBRopeState( head_size=head_size, is_neox_style=is_neox_style, cos_sin_caches=cos_sin_caches, group_to_cache=list(group_to_cache), ) logger.info( "Registered CB rope state for instance %d " "(%d cache(s), shapes=%s dtype=%s, head_size=%d, is_neox=%s, " "group_map=%s)", instance_id, len(cos_sin_caches), [tuple(c.shape) for c in cos_sin_caches], cos_sin_caches[0].dtype, head_size, is_neox_style, "uniform" if not group_to_cache else str(group_to_cache), ) def cb_unregister_rope(self, instance_id: int) -> None: """Drop the instance's CB rope state; the paged KV cache is left intact. Args: instance_id (int): Instance whose rope state to remove (use ``UNREGISTER_KV_CACHE`` to free the KV cache itself). """ self._cb_rope_state.pop(instance_id, None) if self._transfer_module.get_and_touch_context_entry(instance_id) is None: logger.warning( "cb_unregister_rope: instance %d not registered", instance_id ) return logger.info("Unregistered CB rope state for instance %d", instance_id) def drop_instance_state(self, instance_id: int) -> None: """Drop blend state for a reaped instance (InstanceLivenessTarget hook). Only the CB rope state is held per instance; the GPU cache context is owned by ``LMCacheDrivenTransferModule`` (no mirror here), so reaping the GPU entry frees it directly. A no-op if no rope state is held. Args: instance_id: The reaped worker's instance ID. """ if self._cb_rope_state.pop(instance_id, None) is not None: logger.info("Dropped CB rope state for reaped instance %d", instance_id) # ------------------------------------------------------------------ # Unified lookup (CB_UNIFIED_LOOKUP) + shared helpers # ------------------------------------------------------------------ def _drain_fingerprints_sync(self) -> None: """Sync-drain pending fingerprint registrations (the async drainer races at low max_tokens).""" while True: try: job = self._fingerprint_queue.get_nowait() except QueueEmpty: break tokens_in_range, chunk_hashes, start_chunk_idx, position_offset = job try: self._token_range_matcher.on_new_token_hashes( tokens_in_range, chunk_hashes, start_chunk_idx=start_chunk_idx, position_offset=position_offset, ) except Exception: logger.exception("CB fingerprint registration failed (sync drain)") def _match_fingerprints(self, key: IPCCacheServerKey) -> list[CBMatchResult]: """Drain pending registrations and fingerprint-match sub-sequences. Returns the raw matches (any order, possibly overlapping); the caller applies the prefix filter + overlap dedup once via :meth:`_non_overlapping_after_prefix`. """ self._drain_fingerprints_sync() return self._token_range_matcher.match_sub_sequence(list(key.token_ids)) @staticmethod def _non_overlapping_after_prefix( matches: list[CBMatchResult], prefix_tokens: int ) -> list[CBMatchResult]: """Matches outside the prefix coverage, leftmost-greedy overlap-deduped. Drops matches the prefix leg already covers (``cur_st < prefix_tokens``), then keeps a left-to-right non-overlapping subset -- two matches over the same request range can't both scatter. Filtering precedes the dedup so a prefix-covered match cannot suppress a usable one in the greedy pass. Args: matches: Candidate matches in any order; ``cur_st``/``cur_ed`` are request token positions. prefix_tokens: Contiguous prefix coverage in tokens; matches starting before it are dropped. Pass ``0`` to keep all (dedup only). Returns: Non-overlapping matches in ascending ``cur_st`` order. """ kept: list[CBMatchResult] = [] covered_end = -1 for r in sorted( (r for r in matches if r.cur_st >= prefix_tokens), key=lambda r: r.cur_st, ): if r.cur_st >= covered_end: kept.append(r) covered_end = r.cur_ed return kept def _resolve_cb_layout_desc( self, model_name: str, world_size: int ) -> "MemoryLayoutDesc | None": """Find the CB KV buffer layout for ``(model_name, world_size)``. Reads the thread-safe ``layout_desc_registry`` (populated by ``lmcache_driven_transfer`` on KV-cache registration) rather than iterating ``cache_contexts``: iteration races concurrent register/unregister, and the registry holds the complete multi-group descriptor instead of a single-group manual reconstruction. Args: model_name (str): Model name to match. world_size (int): Tensor-parallel world size to match. Returns: MemoryLayoutDesc | None: The matching layout, or None if no registered CB context matches. """ return self._ctx.layout_desc_registry.find(model_name, world_size) def _sparse_prefetch_submit( self, key: IPCCacheServerKey, layout_desc: "MemoryLayoutDesc", matches: list[CBMatchResult], ) -> "tuple[PrefetchHandle, dict[bytes, list], list[int]]": """Coalesce all matches into one sparse L2->L1 prefetch and submit it. Non-blocking. Dedups object keys before submit (sparse keeps one read lock per loaded key, so a duplicate would leak). The caller polls ``query_prefetch_status(handle)`` then calls :meth:`_sparse_classify` with the found set. Args: key (IPCCacheServerKey): The request key. layout_desc (MemoryLayoutDesc): CB KV buffer layout for L1 alloc. matches (list[CBMatchResult]): Non-prefix matches to prefetch. Returns: tuple[PrefetchHandle, dict[bytes, list], list[int]]: the prefetch handle, per-hash TP-expanded object keys, and each expanded position's deduped-key index (maps the per-key found set back to every chunk). """ world_size = key.world_size per_hash_obj_keys: dict[bytes, list] = {} all_hashes = [r.hash for r in matches] all_obj_keys = ipc_key_to_object_keys(key, all_hashes, [0])[0] for i, h in enumerate(all_hashes): per_hash_obj_keys[h] = all_obj_keys[i * world_size : (i + 1) * world_size] # Dedup keys before submit (sparse keeps one read lock per loaded key; # a duplicate would leak). Map each expanded position to its deduped # index so the per-key found set resolves back to every chunk. uniq_keys: list = [] key_to_uidx: dict = {} expanded_uidx: list[int] = [] for k in all_obj_keys: uidx = key_to_uidx.get(k) if uidx is None: uidx = len(uniq_keys) key_to_uidx[k] = uidx uniq_keys.append(k) expanded_uidx.append(uidx) handle: PrefetchHandle = self._ctx.storage_manager.submit_prefetch_task( uniq_keys, layout_desc, external_request_id=key.request_id, policy=TrimPolicy.SPARSE, ) return handle, per_hash_obj_keys, expanded_uidx def _sparse_classify( self, key: IPCCacheServerKey, matches: list[CBMatchResult], found_uidx: set[int], per_hash_obj_keys: dict[bytes, list], expanded_uidx: list[int], ) -> list[CBMatchResult]: """Classify each prefetched chunk as found or stale, and finalize state. A chunk is found only if every TP rank's key loaded; stale chunks take an eviction strike (evicted at threshold, kept while still in-flight). Stashes the found chunks' obj_keys for the retrieve path. Args: key (IPCCacheServerKey): The request key. matches (list[CBMatchResult]): The submitted non-prefix matches. found_uidx (set[int]): Deduped-key indices that loaded. per_hash_obj_keys (dict[bytes, list]): Per-hash TP-expanded keys. expanded_uidx (list[int]): Each expanded position's deduped index. Returns: list[CBMatchResult]: The found subset, in cur_st order. """ world_size = key.world_size found_cb_match_result: list[CBMatchResult] = [] stale_hashes: list[bytes] = [] for j, r in enumerate(matches): base = j * world_size if all(expanded_uidx[base + t] in found_uidx for t in range(world_size)): found_cb_match_result.append(r) else: stale_hashes.append(r.hash) # Reset strikes for confirmed hashes. if found_cb_match_result: with self._pending_fp_lock: for r in found_cb_match_result: self._stale_strike.pop(r.hash, None) # Stale: in-flight keep; >= threshold strikes -> evict. if stale_hashes: with self._pending_fp_lock: truly_evict: list[bytes] = [] for h in stale_hashes: if h in self._pending_fp_hashes: continue n = self._stale_strike.get(h, 0) + 1 if n >= self._STALE_STRIKE_THRESHOLD: truly_evict.append(h) self._stale_strike.pop(h, None) else: self._stale_strike[h] = n if truly_evict: self._token_range_matcher.remove_chunks(truly_evict) self._event_bus.publish( Event( event_type=EventType.CB_CHUNKS_EVICTED, metadata={"num_chunks": len(stale_hashes)}, ) ) # Stash per-hash obj_keys for retrieve (L2 opt). if found_cb_match_result: cache_entry = { r.hash: per_hash_obj_keys[r.hash] for r in found_cb_match_result if r.hash in per_hash_obj_keys } with self._lookup_obj_keys_lock: self._lookup_obj_keys_cache[key.request_id] = cache_entry return found_cb_match_result def _submit_prefix_leg( self, key: IPCCacheServerKey, tp_size: int, policy: TrimPolicy, ) -> "tuple[PrefetchHandle | None, int]": """Submit the CB prefix prefetch (non-blocking). Opens the ``cb.prefix_lookup`` span (CB namespace — CB requests no longer feed the MP request / mp.lookup_prefetch spans or the MP hit-rate aggregate; the CB hit-rate metric carries prefix tokens via CB_LOOKUP_END) and writes the shared session (``set_tokens`` + ``lookup_ipc_key``) so ``end_session``'s L1 keep-alive touch still resolves the request's keys. Args: key (IPCCacheServerKey): Request key (token IDs, request_id, model, world_size). tp_size (int): Tensor-parallel size for MLA multi-reader locking. policy (TrimPolicy): ``PREFIX`` or ``SEGMENTED_PREFIX``. Returns: tuple: ``(handle, world_size)``. ``handle`` is None when there is no GPU context or no full chunk (the poll then reports 0 coverage). """ rid = key.request_id model_name, world_size = key.model_name, key.world_size self._event_bus.publish( Event(event_type=EventType.CB_PREFIX_LOOKUP_START, session_id=rid) ) layout_desc = self._resolve_cb_layout_desc(model_name, world_size) if layout_desc is None: logger.error( "No CB GPU context for model %s ws %d during prefix lookup!", model_name, world_size, ) return None, world_size chunk_hashes = self._ctx.token_hasher.compute_chunk_hashes(list(key.token_ids)) if not chunk_hashes: return None, world_size # Lookup-hash logger (chunk hashes, for debug); guarded so the metadata # dict is built only when a subscriber is listening. if self._event_bus.has_subscribers(EventType.MP_LOOKUP): self._event_bus.publish( Event( event_type=EventType.MP_LOOKUP, session_id=rid, metadata={ "request_id": rid, "chunk_hashes": chunk_hashes, "model_name": model_name, "chunk_size": self._ctx.chunk_size, "seq_len": len(key.token_ids), "dtypes": [str(d) for d in layout_desc.dtypes], "shapes": [list(s) for s in layout_desc.shapes], }, ) ) # Shared session: end_session reads lookup_ipc_key + the rolling hashes # to keep the request's KV alive in L1. session = self._ctx.session_manager.get_or_create(rid) session.set_tokens(list(key.token_ids)) session.lookup_ipc_key = key extra_count = compute_extra_count(tp_size, world_size) obj_keys = ipc_key_to_object_keys(key, chunk_hashes, [0])[0] handle = self._ctx.storage_manager.submit_prefetch_task( obj_keys, layout_desc, extra_count=extra_count, external_request_id=rid, policy=policy, ) return handle, world_size def _poll_prefix_leg( self, job: "_CBUnifiedJob", rid: str, segmented: bool ) -> "tuple[int, list[int] | None] | None": """Poll the CB prefix handle; on completion close the cb.prefix_lookup span. Consume-once: publishes CB_PREFIX_LOOKUP_END exactly once when the prefetch lands. The prefix hit tokens are accounted on the CB hit-rate metric at CB_LOOKUP_END, not here. For SEGMENTED_PREFIX also surfaces the gapped retained chunk set. Args: job (_CBUnifiedJob): Poll state holding the prefix handle + world size. rid (str): Request ID (event session_id). segmented (bool): SEGMENTED_PREFIX active -> also surface the gapped retained chunk set. Returns: tuple | None: ``(leading_chunks, retained_or_None)`` when resident; ``None`` while still loading. ``retained`` is the full gapped chunk set for SEGMENTED_PREFIX, else None. """ if job.prefix_handle is not None: bm = self._ctx.storage_manager.query_prefetch_status(job.prefix_handle) if bm is None: return None # still loading ws = job.prefix_world_size # NOTE(Kuntai): assumes uniform world size and prefix-ordered keys # that break at the first miss. leading = bm.count_leading_ones() // ws retained = ( sorted({ki // ws for ki in bm.get_indices_list()}) if segmented else None ) else: # No GPU context / no full chunk: nothing loaded. leading, retained = 0, ([] if segmented else None) self._event_bus.publish( Event( event_type=EventType.CB_PREFIX_LOOKUP_END, session_id=rid, metadata={"prefix_chunks": leading}, ) ) return leading, retained def cb_unified_lookup( self, key: IPCCacheServerKey, tp_size: int ) -> CBUnifiedLookupResult | None: """Non-blocking single-RPC CB lookup (submit-once, poll-on-recall). First call submits the prefix lookup + fingerprint match; later calls poll both legs, returning ``None`` until the prefix and the sparse non-prefix complement are both resident in L1 (so a worker thread never blocks on the L2->L1 loads). The prefix job's L1 read locks persist for the retrieve. Args: key (IPCCacheServerKey): Request key (token IDs, request_id, model, world_size). tp_size (int): Tensor-parallel size for the prefix lookup. Returns: CBUnifiedLookupResult | None: ``None`` while either leg is still loading (the caller re-issues to poll); on completion, the prefix coverage in tokens plus the found non-prefix segments. """ rid = key.request_id chunk_size = self._ctx.chunk_size with self._cb_jobs_lock: job = self._cb_jobs.get(rid) if job is None: # First invocation: start events + submit prefix + fingerprint match. self._event_bus.publish( Event(event_type=EventType.CB_REQUEST_START, session_id=rid) ) self._event_bus.publish( Event( event_type=EventType.CB_LOOKUP_START, session_id=rid, metadata={"num_tokens": len(key.token_ids)}, ) ) # SEGMENTED_PREFIX: request the contiguous prefix with gap-tolerant # retention so a mid-prefix L2 retrieve failure leaves the post-gap # chunks L1-resident (picked up by the sparse leg as L1 hits, hole # recomputed) instead of truncating the prefix at the gap. prefix_policy = ( TrimPolicy.SEGMENTED_PREFIX if self._segmented_prefix else TrimPolicy.PREFIX ) # Prefix leg: blend_v3 owns the submit + the cb.prefix_lookup span # (under cb.lookup); prefix hit tokens land on the CB hit-rate # metric via CB_LOOKUP_END below. prefix_handle, prefix_ws = self._submit_prefix_leg( key, tp_size, prefix_policy ) # Local and coordinator matching are mutually exclusive: with a # coordinator the fleet directory is the only source, so skip the # local matcher (and its span). The coordinator leg is async # (submitted below, resolved at poll) and is timed by cb.lookup. matches: list[CBMatchResult] if self._coordinator is not None: matches = [] else: # Local fingerprint match: CPU-bound, tight span. self._event_bus.publish( Event( event_type=EventType.CB_FINGERPRINT_MATCH_START, session_id=rid, ) ) matches = self._match_fingerprints(key) self._event_bus.publish( Event( event_type=EventType.CB_FINGERPRINT_MATCH_END, session_id=rid, metadata={"matches": len(matches)}, ) ) job = _CBUnifiedJob( matches=matches, num_tokens=len(key.token_ids), prefix_handle=prefix_handle, prefix_world_size=prefix_ws, ) job.coord_submitted = self._submit_coordinator_match(key) if job.coord_submitted and self._coordinator is not None: job.coord_deadline = time.monotonic() + self._coordinator.match_budget_s # Coordinator match leg: async span, ended on the resolving poll # (or deadline) in _poll_coordinator_match. self._event_bus.publish( Event( event_type=EventType.CB_COORDINATOR_MATCH_START, session_id=rid, ) ) with self._cb_jobs_lock: self._cb_jobs[rid] = job segmented = self._segmented_prefix # --- Prefix leg: poll (consume-once) until the L1+L2 prefix lands. --- if job.prefix_chunks is None: res = self._poll_prefix_leg(job, rid, segmented) if res is None: return None # prefix still loading -> defer job.prefix_chunks, prefix_retained = res if segmented: job.retained_chunks = prefix_retained # Poll above set it (or we returned); narrow for the arithmetic below. assert job.prefix_chunks is not None prefix_chunks: int = job.prefix_chunks # Prefix done: reconcile the complement outside the prefix coverage and # submit one sparse prefetch for it (once). Prefix-covered chunks never # enter the sparse prefetch, so they cannot leak a read lock. if not job.sparse_started: prefix_tokens = prefix_chunks * chunk_size if self._coordinator is not None: candidates = self._poll_coordinator_match(job, rid) if candidates is None: return None # coordinator still in flight (bounded by deadline) else: candidates = job.matches # Under SEGMENTED_PREFIX, a same-position match the prefix leg already # retained rides the segmented tail (prefix-class: pure load, no CHECK) # -- drop it here so it is not scattered twice. A same-position match # the tail does NOT cover is a genuine cross-context hit: keep it as # non-prefix (re-RoPE no-ops at delta 0, but it still needs CHECK). # Shifted (cur != old) matches are always kept. Then the single # prefix-filter + overlap dedup over the rest. if segmented: retained = set(job.retained_chunks or []) candidates = [ c for c in candidates if c.old_st != c.cur_st or (c.cur_st // chunk_size) not in retained ] job.non_prefix = self._non_overlapping_after_prefix( candidates, prefix_tokens ) if job.non_prefix: layout_desc = self._resolve_cb_layout_desc( key.model_name, key.world_size ) if layout_desc is not None: ( job.handle, job.per_hash_obj_keys, job.expanded_uidx, ) = self._sparse_prefetch_submit(key, layout_desc, job.non_prefix) # Only trace the span when the prefetch actually reads L2; # all-L1-resident matches do no L2 work worth a span. job.l2_keys = len(job.handle.l2_orig_indices) if job.l2_keys > 0: self._event_bus.publish( Event( event_type=EventType.CB_SPARSE_PREFETCH_START, session_id=rid, metadata={ "n_chunks": len(job.non_prefix), "world_size": key.world_size, "n_keys": len(job.non_prefix) * key.world_size, "l2_keys": job.l2_keys, }, ) ) else: logger.error( "No CB GPU context for model %s ws %d during cb_unified_lookup", key.model_name, key.world_size, ) job.non_prefix = [] job.sparse_started = True # --- Sparse leg: poll (consume-once) until the scattered chunks land. --- if job.handle is not None and job.found_uidx is None: bm = self._ctx.storage_manager.query_prefetch_status(job.handle) if bm is None: return None # sparse still loading -> defer job.found_uidx = set(bm.get_indices_list()) if job.l2_keys > 0: self._event_bus.publish( Event( event_type=EventType.CB_SPARSE_PREFETCH_END, session_id=rid, metadata={ "found_keys": len(job.found_uidx), "l2_keys": job.l2_keys, }, ) ) # --- BOTH legs ready: classify the complement + finalize. --- if job.handle is not None: found = self._sparse_classify( key, job.non_prefix or [], job.found_uidx or set(), job.per_hash_obj_keys or {}, job.expanded_uidx or [], ) else: found = [] prefix_tokens = prefix_chunks * chunk_size num_tokens = job.num_tokens # Segmented tail: post-gap chunks the SEGMENTED_PREFIX prefix leg kept # resident (retained index > the leading run). Delivered at their # original positions (old_st == cur_st) so the connector tags them # ``prefix`` (pure load, no recompute); only the gap is recomputed. The # storage key is the same prefix-chained chunk hash the prefix leg used, # so no fingerprint match is needed to retrieve them. segmented_tail: list[CBMatchResult] = [] if segmented and job.retained_chunks: chunk_hashes = self._ctx.token_hasher.compute_chunk_hashes( list(key.token_ids) ) for i in job.retained_chunks: if i < prefix_chunks or i >= len(chunk_hashes): continue # leading run (already prefix) / sub-chunk tail st = i * chunk_size segmented_tail.append( CBMatchResult( old_st=st, old_ed=st + chunk_size, cur_st=st, cur_ed=st + chunk_size, hash=chunk_hashes[i], ) ) seg_tail_tokens = _unique_token_coverage(segmented_tail) non_prefix_hit_tokens = _unique_token_coverage(found) self._event_bus.publish( Event( event_type=EventType.CB_LOOKUP_END, session_id=rid, metadata={ "num_tokens": num_tokens, "fingerprint_hits": len(found), "prefix_hits": job.prefix_chunks, "prefix_chunks": job.prefix_chunks, "storage_hits": len(found), "stale_chunks": len(job.non_prefix or []) - len(found), "no_gpu_context": False, "prefix_hit_tokens": prefix_tokens, "segmented_prefix_hit_tokens": seg_tail_tokens, "non_prefix_hit_tokens": non_prefix_hit_tokens, "hit_tokens": prefix_tokens + _unique_token_coverage(found + segmented_tail), "requested_tokens": (num_tokens // chunk_size) * chunk_size, }, ) ) with self._cb_jobs_lock: self._cb_jobs.pop(rid, None) return CBUnifiedLookupResult( prefix_coverage_tokens=prefix_tokens, non_prefix_segments=found, segmented_prefix_segments=segmented_tail, ) def store( self, key: IPCCacheServerKey, instance_id: int, gpu_block_ids: list[list[int]], event_ipc_handle: bytes, ) -> tuple[bytes, bool]: """Paged store, then register the stored chunks as match fingerprints. Delegates the KV write to ``LMCacheDrivenTransfer.store``, then (worker 0 only) enqueues the chunk hashes for async fingerprint registration ordered after the L1 commit. Chunk 0 of a position-0 store is skipped (owned by the standard prefix path). Fingerprint failures are logged, never raised — they do not affect store correctness. Args: key (IPCCacheServerKey): Store key (token IDs + ``[start, end)``). instance_id (int): Target KV-cache instance. gpu_block_ids (list[list[int]]): Per-layer-group paged block IDs. event_ipc_handle (bytes): IPC handle to the producer's CUDA event. Returns: tuple[bytes, bool]: The underlying ``LMCacheDrivenTransfer.store`` result (event handle, success). """ result = self._transfer_module.store( key, instance_id, gpu_block_ids, event_ipc_handle ) # The matcher is engine-shared; only worker 0 registers. if key.worker_id not in (0, None): return result # Enqueue on cupy_stream so CUDA FIFO ordering puts registration # after the L1-commit callback; otherwise lookups see the chunk as # not-yet-committed and drop the whole group as stale. try: session = self._ctx.session_manager.get_or_create(key.request_id) chunk_hashes = [ TokenHasher.hash_to_bytes(h) for h in session.get_hashes(key.start, key.end) ] if not chunk_hashes: return result tokens_in_range = list(key.token_ids)[key.start : key.end] start_chunk_idx = 1 if key.start == 0 else 0 job = (tokens_in_range, chunk_hashes, start_chunk_idx, key.start) with self._pending_fp_lock: self._pending_fp_hashes.update(chunk_hashes[start_chunk_idx:]) entry = self._transfer_module.get_and_touch_context_entry(instance_id) gpu_ctx = entry.cache_context if entry is not None else None if gpu_ctx is not None and gpu_ctx.cupy_stream is not None: gpu_ctx.cupy_stream.launch_host_func( self._fingerprint_queue.put_nowait, job ) else: self._fingerprint_queue.put_nowait(job) except Exception: logger.exception( "CB fingerprint enqueue failed for request %s " "(does not affect store correctness)", key.request_id, ) if self._coordinator is not None: self._publish_fingerprints(key, chunk_hashes, tokens_in_range) return result def _publish_fingerprints( self, key: IPCCacheServerKey, chunk_hashes: list[bytes], tokens_in_range: list[int], ) -> None: """Publish this stored range's chunk fingerprints to the coordinator. Best-effort and fire-and-forget (enqueue only): one wire ``ChunkFingerprint`` per stored chunk -- its content poly-hash (the same ``chunk_hash_windows_numba`` the match probes, with the fleet base), its shared-L2 ``object_key`` (the chunk storage key ``th``), and its token position. Never raises into the store path. Args: key: The store request key (model/scope/positions). chunk_hashes: Per-chunk storage keys (``th``) for the range. tokens_in_range: The stored tokens ``token_ids[start:end]``. """ coordinator = self._coordinator if coordinator is None or not chunk_hashes: return try: model_scope = key.model_name store_range = { "model_scope": model_scope, "tokens": list(tokens_in_range), "object_keys": [h.hex() for h in chunk_hashes], "old_st_base": key.start, } coordinator.enqueue_register([store_range]) except Exception: logger.warning( "CB coordinator publish build failed for request %s " "(does not affect store correctness)", key.request_id, ) def _submit_coordinator_match(self, key: IPCCacheServerKey) -> bool: """Issue a fleet directory match query for this request (best-effort). Args: key: The lookup request key. Returns: ``True`` if a query was submitted (so the finalize step should poll for it), ``False`` when there is no coordinator or submission failed. """ coordinator = self._coordinator if coordinator is None: return False try: tokens = list(key.token_ids) if len(tokens) < self._ctx.chunk_size: return False coordinator.submit_match(key.request_id, key.model_name, tokens) return True except Exception: logger.warning( "CB coordinator match submit failed for request %s", key.request_id ) return False def _poll_coordinator_match( self, job: "_CBUnifiedJob", rid: str ) -> "list[CBMatchResult] | None": """Poll the coordinator match result, deferring until it resolves. Mirrors the prefix/sparse legs: ``return None`` to defer while pending. A per-lookup wall-clock deadline (``job.coord_deadline``) bounds the total wait, including queue/pool time. Past the deadline the leg is abandoned and the lookup proceeds local-only (the client's later fill, if any, is dropped via ``take_match``). Args: job: The per-request poll state. rid: Request id. Returns: The global segments (possibly empty) once resolved or timed out, or ``None`` to defer (still in flight and within the deadline). """ coordinator = self._coordinator if coordinator is None or not job.coord_submitted: return [] poll = coordinator.poll_match(rid) if poll is PENDING: if time.monotonic() < job.coord_deadline: return None # defer; bounded by job.coord_deadline coordinator.take_match(rid) logger.warning( "CB coordinator match deadline exceeded for %s; local-only", rid ) self._event_bus.publish( Event( event_type=EventType.CB_COORDINATOR_MATCH_END, session_id=rid, metadata={"matches": 0, "timed_out": True}, ) ) return [] coordinator.take_match(rid) segments = self._build_global_segments(poll) if isinstance(poll, list) else [] self._event_bus.publish( Event( event_type=EventType.CB_COORDINATOR_MATCH_END, session_id=rid, metadata={"matches": len(segments), "timed_out": False}, ) ) return segments def _build_global_segments( self, matches: "list[RemoteMatch]" ) -> list[CBMatchResult]: """Convert coordinator matches into chunk-granular retrievable segments. Each coordinator ``object_key`` is the hex of the chunk's content hash (the same ``th`` a local ``CBMatchResult.hash`` holds), so the matches are returned as ``CBMatchResult`` directly: the retrieve path then expands ``hash`` to per-rank shared-L2 object keys via ``ipc_key_to_object_keys``, identical to local matches. Args: matches: Matched chunks returned by the coordinator client. Returns: One :class:`CBMatchResult` per matched chunk (request order). """ chunk_size = self._ctx.chunk_size return [ CBMatchResult( old_st=m.old_st, old_ed=m.old_st + chunk_size, cur_st=m.cur_st, cur_ed=m.cur_st + chunk_size, hash=bytes.fromhex(m.object_key), ) for m in matches ] def _drain_fingerprint_queue(self) -> None: """Best-effort background drainer for _fingerprint_queue.""" while not self._fingerprint_stop.is_set(): try: job = self._fingerprint_queue.get(timeout=0.1) except QueueEmpty: continue tokens_in_range, chunk_hashes, start_chunk_idx, position_offset = job try: self._token_range_matcher.on_new_token_hashes( tokens_in_range, chunk_hashes, start_chunk_idx=start_chunk_idx, position_offset=position_offset, ) except Exception: logger.exception("CB fingerprint registration failed (async)") finally: with self._pending_fp_lock: self._pending_fp_hashes.difference_update( chunk_hashes[start_chunk_idx:] ) def _apply_cb_rope_batched( self, gpu_context: BaseCacheContext, rope_state: _CBRopeState, batch_len: int, slots_to_rope: list[tuple[int, int, int]], ) -> None: """Re-RoPE the given tmp-pool slots in place (K-only, per kernel group). Args: gpu_context (GPUCacheContext): The instance's GPU cache context. rope_state (_CBRopeState): Cached cos/sin + head layout. batch_len (int): Number of tmp slots staged for this batch. slots_to_rope (list[tuple[int, int, int]]): ``(slot_idx, old_st, cur_st)`` per shifted slot — re-RoPE K from stored position ``old_st`` to new position ``cur_st``. Raises: RuntimeError: On a compressed (compress_ratio != 1) or MLA (kv_size != 2) layout, or a head_size/hidden_dim mismatch. """ if not slots_to_rope: return num_groups = gpu_context.kv_layer_groups_manager.num_kernel_groups for group_idx in range(num_groups): group = gpu_context.kv_layer_groups_manager.kernel_groups[group_idx] if group.tokens_per_block != group.slots_per_block: raise RuntimeError( f"CB v3: group {group_idx} is compressed " f"(tokens_per_block={group.tokens_per_block}, " f"slots_per_block={group.slots_per_block}); " f"compressed layouts unsupported." ) all_slots = [ gpu_context.get_temp_kernel_group_buffer(slot_idx, group_idx) for slot_idx in range(batch_len) ] if all_slots[0].shape[0] != 2: raise RuntimeError( f"CB v3: group {group_idx} has kv_size={all_slots[0].shape[0]}; " "MLA layouts unsupported." ) num_layers, slots, hidden_dim = all_slots[0].shape[1:] n_heads = hidden_dim // rope_state.head_size if n_heads * rope_state.head_size != hidden_dim: raise RuntimeError( f"CB rope: group {group_idx} hidden_dim ({hidden_dim}) " f"not a multiple of head_size ({rope_state.head_size})." ) # Per-group rope cache: dual-RoPE models rotate each # kernel group with its own theta's cos/sin. group_cos_sin = rope_state.cache_for_group(group.engine_group_idx) # slot ramp tiled across layers is invariant per (num_layers, # slots) — cache it; each shifted slot then just adds its offset. device = all_slots[0].device sp_key = (str(device), num_layers, slots) sp_cache = getattr(self, "_cb_sp_rep_cache", None) if sp_cache is None: sp_cache = {} self._cb_sp_rep_cache = sp_cache slot_positions_rep = sp_cache.get(sp_key) if slot_positions_rep is None: slot_positions_rep = torch.arange( slots, device=device, dtype=torch.long ).repeat(num_layers) sp_cache[sp_key] = slot_positions_rep for slot_idx, old_st, cur_st in slots_to_rope: # reshape returns an in-place view (tmp slots are contiguous). k_view = all_slots[slot_idx][0].reshape( num_layers * slots, n_heads, rope_state.head_size ) lmc_ops.rotary_embedding_k_fused( old_st + slot_positions_rep, cur_st + slot_positions_rep, k_view, rope_state.head_size, group_cos_sin, rope_state.is_neox_style, ) def cb_retrieve_pre_computed( self, key: IPCCacheServerKey, cb_match_result: list[CBMatchResult], gpu_block_ids: list[list[int]], instance_id: int, event_ipc_handle: bytes, ) -> tuple[bytes, bool]: """Scatter every matched token range into the request's paged KV. Reuses the lookup's prefetched chunks: fills tmp slots, K-only re-RoPEs the shifted (non-prefix) subset, then writes per-token via the slot kernel — so non-block-aligned matches and partial vLLM blocks shared with recomputed tokens are written correctly (no block-alignment trim). Only matches past the currently allocated slots are dropped (vLLM may call this twice: partial- then full-block alloc). Args: key (IPCCacheServerKey): The request key. cb_match_result (list[CBMatchResult]): Matched ranges to scatter (prefix-hit and shifted), any order. gpu_block_ids (list[list[int]]): This request's paged block table per engine (kernel) group; single-group models pass [[...]]. Mirrors the engine RETRIEVE/STORE per-group block-id contract. instance_id (int): Target KV-cache instance. event_ipc_handle (bytes): IPC handle to the forward's CUDA event. Returns: tuple[bytes, bool]: The scatter-complete event handle and whether the scatter ran (False if the prefetched objects were unavailable). Raises: ValueError: If the instance has no registered KV cache or rope state. MLA layouts are unsupported (raised during re-RoPE). """ entry = self._transfer_module.get_and_touch_context_entry(instance_id) if entry is None: raise ValueError( f"Instance {instance_id} not registered for paged KV cache" ) if instance_id not in self._cb_rope_state: raise ValueError( f"Instance {instance_id} has no CB rope state; " "send CB_REGISTER_ROPE_V3 before CB_RETRIEVE_PRE_COMPUTED_V3." ) gpu_context = entry.cache_context rope_state = self._cb_rope_state[instance_id] chunk_size = self._ctx.chunk_size _retrieve_t0 = time.perf_counter() cb_match_result = sorted(cb_match_result, key=lambda r: r.cur_st) # L2 opt: reuse lookup's obj_keys cache; fall back to re-resolve. with self._lookup_obj_keys_lock: cached = self._lookup_obj_keys_cache.pop(key.request_id, None) if cached is not None and all(r.hash in cached for r in cb_match_result): # The lookup cached all-ranks obj keys (world_size per hash). This # retrieve is per-worker, so select THIS rank's key -> M objects, not # M*world_size (else the zip below silently truncates and mispairs # ranks at TP>1). Mirrors the non-cached path's per-worker resolve. if key.worker_id is not None and key.world_size > 1: all_obj_keys = [cached[r.hash][key.worker_id] for r in cb_match_result] else: all_obj_keys = [k for r in cb_match_result for k in cached[r.hash]] else: all_obj_keys = ipc_key_to_object_keys( key, [r.hash for r in cb_match_result], [0] )[0] # Lookup read-locked the full found set, but the connector may have # dropped some matches (parent-covered / misaligned) before retrieve, # leaking their per-key read locks. Release those orphans now (disjoint # from all_obj_keys, which retrieve still consumes; needs the key cache). if cached is not None: retrieved_hashes = {r.hash for r in cb_match_result} orphan_keys = [ k for h, ks in cached.items() if h not in retrieved_hashes for k in ks ] if orphan_keys: self._ctx.storage_manager.finish_read_prefetched(orphan_keys) logger.debug( "CB V3 released %d prefetched-but-unretrieved keys (req=%s)", len(orphan_keys), key.request_id, ) # Non-prefix sparse hits split by re-rope need (not prefix coverage). n_non_shifted = sum(1 for r in cb_match_result if r.old_st == r.cur_st) n_shifted = len(cb_match_result) - n_non_shifted if not all_obj_keys: self._event_bus.publish( Event( event_type=EventType.CB_REQUEST_END, session_id=key.request_id, ) ) return event_ipc_handle, True logger.debug("CB V3 retrieving object keys: %s", all_obj_keys) # CB v3 only supports uncompressed single-block-id-space layouts # (enforced per group in ``_apply_cb_rope_batched``), so the first # kernel group's chunk geometry is representative. tokens_per_block = gpu_context.kv_layer_groups_manager.kernel_groups[ 0 ].tokens_per_block if chunk_size % tokens_per_block != 0: raise ValueError( f"chunk_size {chunk_size} must be a multiple of " f"tokens_per_block {tokens_per_block}" ) num_groups = gpu_context.kv_layer_groups_manager.num_kernel_groups with ( torch_dev.device(gpu_context.device), torch_dev.stream(gpu_context.stream), ): check_interprocess_event_support() event = torch_dev.Event(interprocess=True) # One staged block-id tensor per engine group, indexed by group. block_ids_per_group_gpu = gpu_context.stage_block_ids(gpu_block_ids) # Resolve each kernel group's block table + block size once. Select # by engine_group_idx (kernel groups may share one, e.g. MiniMax-M3). kgm = gpu_context.kv_layer_groups_manager resolved_groups: list[tuple[torch.Tensor, int]] = [] for group_idx in range(num_groups): eg_idx = kgm.kernel_groups[group_idx].engine_group_idx if eg_idx >= len(block_ids_per_group_gpu): # Engine groups have independent block tables under HMA; # substituting another group's table would scatter KV into # the wrong physical blocks (silent corruption). raise ValueError( f"CB retrieve: kernel group {group_idx} maps to engine " f"group {eg_idx}, but only " f"{len(block_ids_per_group_gpu)} block table(s) were " "provided." ) resolved_groups.append( ( block_ids_per_group_gpu[eg_idx], kgm.kernel_groups[group_idx].tokens_per_block, ) ) self._event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_RETRIEVE_START, session_id=key.request_id, metadata={ "num_chunks": len(cb_match_result), "model_name": key.model_name, }, ), ) 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) try: with self._ctx.storage_manager.read_prefetched_results( all_obj_keys ) as memory_objs: if memory_objs is None: return event_ipc_handle, False # Per-token scatter handles any cur_st; just bound the # matched range to the allocated slots. pairs: list[tuple[CBMatchResult, Any]] = [] # Bound by the smallest group: under HMA the sliding group # has fewer blocks than the full group, so [0] isn't safe. num_slots = min( int(block_ids.numel()) * group_bs for block_ids, group_bs in resolved_groups ) for r, memory_obj in zip(cb_match_result, memory_objs, strict=True): if r.cur_ed > num_slots: logger.warning( "Dropping CB match cur_st=%d cur_ed=%d: exceeds " "%d slots. Request %s.", r.cur_st, r.cur_ed, num_slots, key.request_id, ) continue pairs.append((r, memory_obj)) # cb.scatter span (GPU): the L1->paged write of every # applied match. Re-RoPE is folded in (n_shifted) — it is # interleaved per-batch, so not a separate span. self._event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_SCATTER_START, session_id=key.request_id, metadata={ "scattered_tokens": sum( r.cur_ed - r.cur_st for r, _ in pairs ), "n_prefix": sum( 1 for r, _ in pairs if r.old_st == r.cur_st ), "n_shifted": sum( 1 for r, _ in pairs if r.old_st != r.cur_st ), "dropped": len(cb_match_result) - len(pairs), }, ), ) # Consecutive matches → one batched scatter per group. runs: list[list[tuple[CBMatchResult, Any]]] = [] for r_obj in pairs: r = r_obj[0] if runs and runs[-1][-1][0].cur_ed == r.cur_st: runs[-1].append(r_obj) else: runs.append([r_obj]) max_batch = gpu_context.max_batch_size for run in runs: for batch_start in range(0, len(run), max_batch): batch = run[batch_start : batch_start + max_batch] batch_len = len(batch) # (a) H2D fill into per-chunk tmp slots. for slot_idx, (_, memory_obj) in enumerate(batch): # Single object group => object_group_idx=0. flat_slot = gpu_context.get_temp_object_group_buffer( slot_idx, 0 ) lmcache_memcpy_async_h2d(memory_obj, flat_slot) # (b) Re-RoPE shifted (non-prefix) slots in place. slots_to_rope = [ (slot_idx, r.old_st, r.cur_st) for slot_idx, (r, _) in enumerate(batch) if r.old_st != r.cur_st ] self._apply_cb_rope_batched( gpu_context, rope_state, batch_len, slots_to_rope ) # (c) Per-token slot scatter: partial vLLM blocks # shared with recomputed tokens stay disjoint. pos = torch.cat( [ torch.arange( r.cur_st, r.cur_ed, device=gpu_context.device, dtype=torch.long, ) for (r, _) in batch ] ) for group_idx in range(num_groups): # This group's block table + size (resolved above). group_block_ids, group_bs = resolved_groups[group_idx] # Per-group block count: under HMA the sliding # group has fewer blocks than the full group, so # gpu_context.num_blocks (group 0's) would # truncate the other groups' bounds check. page_buffer_size = ( kgm.kernel_groups[group_idx].shape_desc.nb * group_bs ) slot_mapping = group_block_ids[ pos // group_bs ] * group_bs + (pos % group_bs) tmp_buffers = [ gpu_context.get_temp_kernel_group_buffer( slot_idx, group_idx ) for slot_idx in range(batch_len) ] key_value = torch.cat(tmp_buffers, dim=2) lmc_ops.multi_layer_kv_transfer( key_value, gpu_context.get_kernel_group_kv_pointers(group_idx), slot_mapping, gpu_context.device, page_buffer_size, lmc_ops.TransferDirection.H2D, gpu_context.get_engine_kv_format(group_idx), block_size=group_bs, head_size=rope_state.head_size, ) self._event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_SCATTER_END, session_id=key.request_id, ), ) except Exception: logger.exception("Error during retrieving prefetched results") self._event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_RETRIEVE_END, session_id=key.request_id, metadata={"success": False}, ), ) self._event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_REQUEST_END, session_id=key.request_id, ), ) return event_ipc_handle, False event.record() self._event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_RETRIEVE_END, session_id=key.request_id, metadata={"success": True}, ), ) _scatter_ms = (time.perf_counter() - _retrieve_t0) * 1000 logger.info( "Retrieved pre-computed for %d match results into request %s " "paged blocks (scatter_ms=%.2f, non_shifted=%d shifted=%d)", len(cb_match_result), key.request_id, _scatter_ms, n_non_shifted, n_shifted, ) self._event_bus.publish_on_stream( gpu_context.cupy_stream, Event( event_type=EventType.CB_REQUEST_END, session_id=key.request_id, ), ) return event.ipc_handle(), True