From 36b449af1911036fbc615e5edc5f2f384debce83 Mon Sep 17 00:00:00 2001 From: Yuang Chen <77919385+cccccya@users.noreply.github.com> Date: Tue, 7 Jul 2026 12:05:52 +0800 Subject: [PATCH] [EPD] Optimize multimodal global cache with paged embedding pool (#28441) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 晟海 Co-authored-by: liusy58 --- .../srt/disaggregation/encode_server.py | 694 +++++----- .../embedding_cache_controller.py | 1116 ++++++++++++----- .../mooncake_embedding_store.py | 46 +- .../test_embedding_cache_controller.py | 1037 ++++++--------- 4 files changed, 1584 insertions(+), 1309 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 4f2260707..a3bd73a98 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -13,8 +13,9 @@ import time import traceback import uuid from collections import defaultdict +from dataclasses import dataclass from http import HTTPStatus -from typing import Annotated, Dict, List, Optional, Set, Tuple, Union +from typing import Annotated, Any, Dict, List, Optional, Set, Tuple, Union import aiohttp import numpy as np @@ -132,6 +133,19 @@ class InternalError(MMError): super().__init__(message, code=HTTPStatus.INTERNAL_SERVER_ERROR) +@dataclass +class GlobalCacheEncodeContext: + req_id: str + modality: Modality + mm_inputs: dict + get_feature_fn: Any + grid_thw: List + mm_feature: Any + num_items: int + aux_data: dict + str_mm_hashes: Optional[List[str]] + + class TensorWrapper: """Wrapper to keep tensor alive while exposing buffer for zero-copy.""" @@ -328,6 +342,13 @@ class MMEncoder: ) self.background_tasks: Set[asyncio.Task] = set() + # Embedding dtype = model param dtype. Always available (both transfer + # backends and the global-cache pool rely on it). + self._embedding_dtype = next(self.model.parameters()).dtype + self._element_size = torch.tensor( + [], dtype=self._embedding_dtype + ).element_size() + if self.server_args.enable_mm_global_cache: from sglang.srt.mem_cache.storage.mooncake_store.embedding_cache_controller import ( EmbeddingCacheController, @@ -340,6 +361,7 @@ class MMEncoder: hidden_dims=hidden_dims, tp_group=get_tp_group().cpu_group, all_rank_get=False, + dtype=self._embedding_dtype, ) else: self.mm_global_cache = None @@ -347,10 +369,6 @@ class MMEncoder: # Pre-compute embedding metadata (needed by all ranks for mooncake) if self.server_args.encoder_transfer_backend == "mooncake": self._embedding_dims = self._infer_embedding_dims() - self._embedding_dtype = next(self.model.parameters()).dtype - self._element_size = torch.tensor( - [], dtype=self._embedding_dtype - ).element_size() if self.rank == 0: logger.info( @@ -870,16 +888,13 @@ class MMEncoder: sub_grids = [grid_thw[i] for i in indices] return self.slice_embedding(new_embeddings, sub_grids, modality) - async def encode_with_global_cache( + async def _prepare_global_cache_context( self, mm_items, modality: Modality, req_id: str, - num_parts: int, - part_idx: int, hashes: Optional[List[str]] = None, - ) -> torch.Tensor: - # mm_inputs: dict + ) -> GlobalCacheEncodeContext: mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality) grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) mm_feature = _convert(_get_mm_feature(mm_inputs, modality)) @@ -894,7 +909,7 @@ class MMEncoder: f"must be in grid space (1 per encoder grid entry)." ) - # Step 1: Rank 0 checks global cache and broadcasts hit/miss mask to all ranks. + str_mm_hashes = None if self.rank == 0: if hashes is None: mm_hashes = self._calculate_hashes_from_features( @@ -902,16 +917,22 @@ class MMEncoder: ) else: mm_hashes = hashes - # Convert hashes to strings (L2 cache expects string keys for Mooncake) + # L2 cache expects string keys for Mooncake. str_mm_hashes = [str(h) for h in mm_hashes] - exist_mask = await self.mm_global_cache.batch_is_exist(str_mm_hashes) - mask_tensor = torch.tensor( - [1 if e else 0 for e in exist_mask], dtype=torch.int32 - ) - else: - mm_hashes = None - mask_tensor = torch.zeros(num_items, dtype=torch.int32) + return GlobalCacheEncodeContext( + req_id=req_id, + modality=modality, + mm_inputs=mm_inputs, + get_feature_fn=get_feature_fn, + grid_thw=grid_thw, + mm_feature=mm_feature, + num_items=num_items, + aux_data=_build_mm_aux_data(mm_inputs, self.model_type), + str_mm_hashes=str_mm_hashes, + ) + + def _broadcast_global_cache_mask(self, mask_tensor: torch.Tensor): if self.server_args.tp_size > 1: torch.distributed.broadcast( mask_tensor, @@ -919,139 +940,312 @@ class MMEncoder: group=self.mm_global_cache.prefetch_tp_group, ) + async def _lookup_global_cache( + self, + ctx: GlobalCacheEncodeContext, + ) -> Tuple[List[int], List[int]]: + if self.rank == 0: + exist_mask = await self.mm_global_cache.batch_is_exist(ctx.str_mm_hashes) + mask_tensor = torch.tensor( + [1 if e else 0 for e in exist_mask], dtype=torch.int32 + ) + else: + mask_tensor = torch.zeros(ctx.num_items, dtype=torch.int32) + + self._broadcast_global_cache_mask(mask_tensor) + exist_mask = [m.item() == 1 for m in mask_tensor] missing_indices = [i for i, e in enumerate(exist_mask) if not e] hit_indices = [i for i, e in enumerate(exist_mask) if e] + return missing_indices, hit_indices + + def _prefetch_global_cache_hits( + self, + ctx: GlobalCacheEncodeContext, + hit_indices: List[int], + ) -> List[str]: + if self.rank != 0 or not hit_indices: + return [] + + hit_hashes = [ctx.str_mm_hashes[i] for i in hit_indices] + hit_tokens = [ + self.get_num_tokens(ctx.grid_thw[i], ctx.modality) for i in hit_indices + ] + self.mm_global_cache.prefetch(ctx.req_id, hit_hashes, hit_tokens, ctx.modality) + return hit_hashes + + async def _wait_global_cache_prefetch( + self, + ctx: GlobalCacheEncodeContext, + hit_indices: List[int], + hit_hashes: List[str], + ) -> List[int]: + fallback_mask = torch.zeros(ctx.num_items, dtype=torch.int32) + if self.rank == 0 and hit_indices: + try: + + async def _wait_prefetch(): + while not self.mm_global_cache.check_prefetch_progress(ctx.req_id): + await asyncio.sleep(0.005) + + await asyncio.wait_for(_wait_prefetch(), timeout=60.0) + + for i, idx in enumerate(hit_indices): + if not self.mm_global_cache.has_local_embedding(hit_hashes[i]): + fallback_mask[idx] = 1 + num_partial_fail = int(fallback_mask.sum().item()) + if num_partial_fail > 0: + logger.warning( + f"Req {ctx.req_id}: {num_partial_fail}/{len(hit_indices)} " + f"cache-hit items failed to load, falling back to ViT" + ) + except (asyncio.TimeoutError, Exception) as e: + logger.error( + f"Prefetch failed for req {ctx.req_id}: {e}. " + f"Falling back to ViT for {len(hit_indices)} hit items." + ) + for idx in hit_indices: + fallback_mask[idx] = 1 + + self._broadcast_global_cache_mask(fallback_mask) + fallback_indices = [ + i for i in range(ctx.num_items) if fallback_mask[i].item() == 1 + ] + return fallback_indices + + def _launch_global_cache_insert( + self, + ctx: GlobalCacheEncodeContext, + hashes: List[str], + d2h_handles: List[Any], + ): + if not hashes: + return + + async def _background_insert(): + await asyncio.to_thread( + self.mm_global_cache.wait_store_to_pool, + d2h_handles, + ) + await asyncio.to_thread( + self.mm_global_cache.insert_batch, + hashes, + ctx.modality, + ) + + task = asyncio.create_task(_background_insert()) + self.background_tasks.add(task) + task.add_done_callback(self.background_tasks.discard) + + @staticmethod + def _as_2d_tensor(tensor: torch.Tensor) -> torch.Tensor: + if tensor.ndim != 2: + tensor = tensor.reshape(-1, tensor.shape[-1]) + return tensor + + def _assemble_global_cache_cpu( + self, + ctx: GlobalCacheEncodeContext, + hit_indices: List[int], + missing_indices: List[int], + fallback_indices: List[int], + new_slices: List[torch.Tensor], + fallback_slices: List[torch.Tensor], + ) -> torch.Tensor: + miss_slice_pos = {idx: pos for pos, idx in enumerate(missing_indices)} + fallback_slice_pos = {idx: pos for pos, idx in enumerate(fallback_indices)} + fallback_index_set = set(fallback_indices) + token_counts = [ + self.get_num_tokens(grid, ctx.modality) for grid in ctx.grid_thw + ] + dim = self.mm_global_cache.get_embedding_dim(ctx.modality) + + mm_embedding = torch.empty( + (sum(token_counts), dim), + dtype=self._embedding_dtype, + pin_memory=True, + ) + + hit_view_hashes = [ + ctx.str_mm_hashes[idx] + for idx in hit_indices + if idx not in fallback_index_set + ] + hit_views = {} + try: + if hit_view_hashes: + cached_slice_lists = self.mm_global_cache.get_pool_views( + hit_view_hashes + ) + for h, slices in zip(hit_view_hashes, cached_slice_lists): + if slices is None: + raise InternalError( + f"Cached embedding {h} not available for req {ctx.req_id}" + ) + hit_views[h] = slices + + offset = 0 + for idx, num_tokens in enumerate(token_counts): + if idx in miss_slice_pos: + src = self._as_2d_tensor(new_slices[miss_slice_pos[idx]]) + mm_embedding[offset : offset + num_tokens].copy_( + src, non_blocking=True + ) + elif idx in fallback_slice_pos: + src = self._as_2d_tensor(fallback_slices[fallback_slice_pos[idx]]) + mm_embedding[offset : offset + num_tokens].copy_( + src, non_blocking=True + ) + else: + copied = 0 + for view in hit_views[ctx.str_mm_hashes[idx]]: + n = view.shape[0] + mm_embedding[offset + copied : offset + copied + n].copy_(view) + copied += n + offset += num_tokens + + torch.cuda.current_stream(self.device).synchronize() + return mm_embedding + finally: + if hit_view_hashes: + self.mm_global_cache.release_pool_views(hit_view_hashes) + + def _assemble_global_cache_gpu( + self, + ctx: GlobalCacheEncodeContext, + missing_indices: List[int], + fallback_indices: List[int], + new_slices: List[torch.Tensor], + fallback_slices: List[torch.Tensor], + ) -> torch.Tensor: + miss_slice_pos = {idx: pos for pos, idx in enumerate(missing_indices)} + fallback_slice_pos = {idx: pos for pos, idx in enumerate(fallback_indices)} + token_counts = [ + self.get_num_tokens(grid, ctx.modality) for grid in ctx.grid_thw + ] + embedding_dim = self.mm_global_cache.get_embedding_dim(ctx.modality) + mm_embedding = torch.empty( + (sum(token_counts), embedding_dim), + dtype=self._embedding_dtype, + device=self.device, + ) + + offset = 0 + copy_handles = [] + for idx, num_tokens in enumerate(token_counts): + if idx in miss_slice_pos: + mm_embedding[offset : offset + num_tokens].copy_( + new_slices[miss_slice_pos[idx]], + non_blocking=True, + ) + elif idx in fallback_slice_pos: + mm_embedding[offset : offset + num_tokens].copy_( + fallback_slices[fallback_slice_pos[idx]], + non_blocking=True, + ) + else: + handle = self.mm_global_cache.load_to_device_async( + ctx.str_mm_hashes[idx], mm_embedding, offset + ) + if handle is None: + raise InternalError( + f"Cached embedding {ctx.str_mm_hashes[idx]} disappeared " + f"during assembly for req {ctx.req_id}" + ) + copy_handles.append(handle) + offset += num_tokens + + self.mm_global_cache.wait_load_to_device(copy_handles) + torch.cuda.current_stream(mm_embedding.device).synchronize() + return mm_embedding + + async def encode_with_global_cache( + self, + mm_items, + modality: Modality, + req_id: str, + num_parts: int, + part_idx: int, + hashes: Optional[List[str]] = None, + ) -> torch.Tensor: + ctx = await self._prepare_global_cache_context( + mm_items, modality, req_id, hashes + ) + + missing_indices, hit_indices = await self._lookup_global_cache(ctx) + hit_hashes = self._prefetch_global_cache_hits(ctx, hit_indices) - # Step 2: All ranks run ViT together on cache-miss images. new_slices = [] if missing_indices: new_slices = self._encode_missing( - mm_feature, mm_inputs, missing_indices, modality, get_feature_fn + ctx.mm_feature, + ctx.mm_inputs, + missing_indices, + ctx.modality, + ctx.get_feature_fn, + ctx.grid_thw, + keep_on_gpu=True, ) - # Step 3: Rank 0 prefetches cache-hit embeddings and builds fallback_mask. - fallback_mask = torch.zeros(num_items, dtype=torch.int32) - cached_slices = [] - - if self.rank == 0: - if hit_indices: - hit_hashes = [str_mm_hashes[i] for i in hit_indices] - hit_tokens = [ - self.get_num_tokens(grid_thw[i], modality) for i in hit_indices - ] - self.mm_global_cache.prefetch(req_id, hit_hashes, hit_tokens, modality) - - try: - - async def _wait_prefetch(): - while not self.mm_global_cache.check_prefetch_progress(req_id): - await asyncio.sleep(0.005) - - await asyncio.wait_for(_wait_prefetch(), timeout=60.0) - - # Prefetch IO completed; check which items actually loaded. - cached_slices = self.mm_global_cache.get_embeddings(hit_hashes) - for i, idx in enumerate(hit_indices): - if cached_slices[i] is None: - fallback_mask[idx] = 1 - num_partial_fail = int(fallback_mask.sum().item()) - if num_partial_fail > 0: - logger.warning( - f"Req {req_id}: {num_partial_fail}/{len(hit_indices)} " - f"cache-hit items failed to load (pool full), " - f"falling back to ViT" - ) - except (asyncio.TimeoutError, Exception) as e: - logger.error( - f"Prefetch failed for req {req_id}: {e}. " - f"Falling back to ViT for {len(hit_indices)} hit items." - ) - for idx in hit_indices: - fallback_mask[idx] = 1 - - # Step 4: Broadcast fallback_mask to all ranks so they stay in sync. - if self.server_args.tp_size > 1: - torch.distributed.broadcast( - fallback_mask, - src=0, - group=self.mm_global_cache.prefetch_tp_group, + miss_d2h_handles = [] + if self.rank == 0 and new_slices: + miss_hashes = [ctx.str_mm_hashes[i] for i in missing_indices] + miss_d2h_handles = self.mm_global_cache.store_to_pool_async( + miss_hashes, new_slices, ctx.modality ) - # Step 5: All ranks run ViT for items that need fallback recomputation. - fallback_indices = [i for i in range(num_items) if fallback_mask[i].item() == 1] - fallback_slices = None + fallback_indices = await self._wait_global_cache_prefetch( + ctx, hit_indices, hit_hashes + ) + + fallback_slices = [] + fallback_d2h_handles = [] if fallback_indices: logger.info( - f"Req {req_id}: All ranks running ViT fallback " + f"Req {ctx.req_id}: All ranks running ViT fallback " f"for {len(fallback_indices)} items." ) fallback_slices = self._encode_missing( - mm_feature, mm_inputs, fallback_indices, modality, get_feature_fn + ctx.mm_feature, + ctx.mm_inputs, + fallback_indices, + ctx.modality, + ctx.get_feature_fn, + ctx.grid_thw, + keep_on_gpu=True, + ) + if self.rank == 0: + fallback_hashes = [ctx.str_mm_hashes[i] for i in fallback_indices] + fallback_d2h_handles = self.mm_global_cache.store_to_pool_async( + fallback_hashes, fallback_slices, ctx.modality + ) + + if self.rank == 0: + mm_embedding = self._assemble_global_cache_cpu( + ctx, + hit_indices, + missing_indices, + fallback_indices, + new_slices, + fallback_slices, ) - # Step 6: Rank 0 assembles final embedding and prepares for sending. - if self.rank == 0: - final_slices = [None] * num_items + new_hashes = [ctx.str_mm_hashes[i] for i in missing_indices] + new_hashes += [ctx.str_mm_hashes[i] for i in fallback_indices] + self._launch_global_cache_insert( + ctx, + new_hashes, + miss_d2h_handles + fallback_d2h_handles, + ) - for i, idx in enumerate(missing_indices): - final_slices[idx] = new_slices[i] - - # Fill in successfully loaded cache-hit embeddings - if cached_slices: - for i, idx in enumerate(hit_indices): - if cached_slices[i] is not None: - final_slices[idx] = cached_slices[i] - - # Fill in ViT fallback results for failed items - if fallback_slices is not None: - for i, idx in enumerate(fallback_indices): - final_slices[idx] = fallback_slices[i] - - mm_embedding = torch.cat(final_slices, dim=0) - - # Release embedding cache references now that torch.cat has - # copied the data into a new tensor. This allows the cache - # entries to be evicted under memory pressure. - if cached_slices: - loaded_hashes = [ - str_mm_hashes[idx] - for idx in hit_indices - if fallback_mask[idx].item() == 0 - ] - if loaded_hashes: - self.mm_global_cache.release_embeddings(loaded_hashes) - - # Background insert: store newly computed embeddings into global cache. - # Includes both original misses and fallback-recomputed hits. - all_new_hashes = [str_mm_hashes[i] for i in missing_indices] - all_new_slices = list(new_slices) - if fallback_slices is not None: - all_new_hashes += [str_mm_hashes[i] for i in fallback_indices] - all_new_slices += list(fallback_slices) - - if all_new_hashes: - - async def _background_insert(): - await asyncio.to_thread( - self.mm_global_cache.insert_batch, - all_new_hashes, - all_new_slices, - ) - - task = asyncio.create_task(_background_insert()) - self.background_tasks.add(task) - task.add_done_callback(self.background_tasks.discard) - - aux_data = _build_mm_aux_data(mm_inputs, self.model_type) - self.embedding_to_send[req_id] = EmbeddingData( - req_id, + self.embedding_to_send[ctx.req_id] = EmbeddingData( + ctx.req_id, num_parts, part_idx, - grid_thw, - modality, + ctx.grid_thw, + ctx.modality, mm_embedding, - **aux_data, + **ctx.aux_data, ) if self.profiler is not None: self.profiler.step() @@ -1079,28 +1273,20 @@ class MMEncoder: """Async encode with global cache for mooncake backend. All ranks participate in VIT forward; tp_size > 1 adds broadcasts for sync.""" try: - mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality) - grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) - mm_feature = _convert(_get_mm_feature(mm_inputs, modality)) - num_items = len(grid_thw) - aux_data = _build_mm_aux_data(mm_inputs) - - # Setup metadata and event management - nbytes, total_tokens, embedding_dim, event = ( - self._setup_mooncake_async_encode( - req_id, num_parts, part_idx, grid_thw, modality, aux_data - ) + ctx = await self._prepare_global_cache_context( + mm_items, modality, req_id, hashes ) - # Rank 0: compute hashes - if self.rank == 0: - if hashes is None: - mm_hashes = self._calculate_hashes_from_features( - mm_feature, grid_thw, modality - ) - else: - mm_hashes = hashes - str_mm_hashes = [str(h) for h in mm_hashes] + nbytes, total_tokens, embedding_dim, event = ( + self._setup_mooncake_async_encode( + ctx.req_id, + num_parts, + part_idx, + ctx.grid_thw, + ctx.modality, + ctx.aux_data, + ) + ) # All ranks: launch background task for cache check + VIT forward. # Do NOT use run_in_executor: get_feature_fn relies on a session @@ -1109,193 +1295,87 @@ class MMEncoder: # ThreadPoolExecutor worker thread. async def _run_forward_with_cache(): try: - # Step 1: Rank 0 checks cache, broadcast mask if TP > 1 - if self.rank == 0: - exist_mask = await self.mm_global_cache.batch_is_exist( - str_mm_hashes - ) - mask_tensor = torch.tensor( - [1 if e else 0 for e in exist_mask], - dtype=torch.int32, - ) - else: - mask_tensor = torch.zeros(num_items, dtype=torch.int32) + missing_indices, hit_indices = await self._lookup_global_cache(ctx) + hit_hashes = self._prefetch_global_cache_hits(ctx, hit_indices) - if self.server_args.tp_size > 1: - torch.distributed.broadcast( - mask_tensor, - src=0, - group=self.mm_global_cache.prefetch_tp_group, - ) - - exist_mask = [m.item() == 1 for m in mask_tensor] - missing_indices = [i for i, e in enumerate(exist_mask) if not e] - hit_indices = [i for i, e in enumerate(exist_mask) if e] - final_slices = [None] * num_items - - # Step 2: All ranks run VIT forward for cache misses - # (runs in event loop to preserve session context) new_slices = [] if missing_indices: new_slices = self._encode_missing( - mm_feature, - mm_inputs, + ctx.mm_feature, + ctx.mm_inputs, missing_indices, - modality, - get_feature_fn, - grid_thw, + ctx.modality, + ctx.get_feature_fn, + ctx.grid_thw, keep_on_gpu=True, ) - # Step 3: Rank 0 prefetches cache-hit embeddings and builds fallback_mask. - fallback_mask = torch.zeros(num_items, dtype=torch.int32) - cached_slices = [] + fallback_indices = await self._wait_global_cache_prefetch( + ctx, hit_indices, hit_hashes + ) - if self.rank == 0 and hit_indices: - hit_hashes = [str_mm_hashes[i] for i in hit_indices] - hit_tokens = [ - self.get_num_tokens(grid_thw[i], modality) - for i in hit_indices - ] - self.mm_global_cache.prefetch( - req_id, hit_hashes, hit_tokens, modality - ) - try: - - async def _wait_prefetch(): - while not self.mm_global_cache.check_prefetch_progress( - req_id - ): - await asyncio.sleep(0.005) - - await asyncio.wait_for(_wait_prefetch(), timeout=60.0) - - cached_slices = self.mm_global_cache.get_embeddings( - hit_hashes - ) - for i, idx in enumerate(hit_indices): - if cached_slices[i] is None: - fallback_mask[idx] = 1 - num_partial_fail = int(fallback_mask.sum().item()) - if num_partial_fail > 0: - logger.warning( - f"Req {req_id}: {num_partial_fail}/{len(hit_indices)} " - f"cache-hit items failed to load (pool full), " - f"falling back to ViT" - ) - except (asyncio.TimeoutError, Exception) as e: - logger.error( - f"Prefetch failed for {req_id}: {e}. " - f"Falling back to ViT for " - f"{len(hit_indices)} hit items." - ) - for idx in hit_indices: - fallback_mask[idx] = 1 - - # Step 4: Broadcast fallback_mask to all ranks so they stay in sync. - if self.server_args.tp_size > 1: - torch.distributed.broadcast( - fallback_mask, - src=0, - group=self.mm_global_cache.prefetch_tp_group, - ) - - # Step 5: All ranks run ViT for items that need fallback recomputation. - fallback_indices = [ - i for i in range(num_items) if fallback_mask[i].item() == 1 - ] - fallback_slices = None + fallback_slices = [] if fallback_indices: logger.info( - f"Req {req_id}: All ranks running ViT fallback " + f"Req {ctx.req_id}: All ranks running ViT fallback " f"for {len(fallback_indices)} items." ) fallback_slices = self._encode_missing( - mm_feature, - mm_inputs, + ctx.mm_feature, + ctx.mm_inputs, fallback_indices, - modality, - get_feature_fn, - grid_thw, + ctx.modality, + ctx.get_feature_fn, + ctx.grid_thw, keep_on_gpu=True, ) - # Step 6: Rank 0 assembles final embedding. if self.rank == 0: - for i, idx in enumerate(missing_indices): - final_slices[idx] = new_slices[i] - - # Fill in successfully loaded cache-hit embeddings - if cached_slices: - for i, idx in enumerate(hit_indices): - if cached_slices[i] is not None: - final_slices[idx] = cached_slices[i] - - # Fill in ViT fallback results for failed items - if fallback_slices is not None: - for i, idx in enumerate(fallback_indices): - final_slices[idx] = fallback_slices[i] - - # Move cached CPU slices to GPU and match model dtype - device = torch.device(f"cuda:{self.gpu_id}") - final_slices = [ - ( - s.to(device=device, dtype=self._embedding_dtype) - if s.device.type == "cpu" - else s + d2h_handles = [] + if new_slices: + miss_hashes = [ + ctx.str_mm_hashes[i] for i in missing_indices + ] + miss_handles = self.mm_global_cache.store_to_pool_async( + miss_hashes, new_slices, ctx.modality ) - for s in final_slices - ] - mm_embedding = torch.cat(final_slices, dim=0) - # Wait for any pending VIT / cat kernels to finish - # before publishing to /send: mooncake transfer_sync - # is a host-side RDMA read that bypasses CUDA streams - # and would otherwise race with in-flight kernels. - torch.cuda.current_stream(mm_embedding.device).synchronize() - - # Release cache refs after data is copied to GPU - if cached_slices: - loaded_hashes = [ - str_mm_hashes[idx] - for idx in hit_indices - if fallback_mask[idx].item() == 0 + d2h_handles.extend(miss_handles) + if fallback_slices: + fallback_hashes = [ + ctx.str_mm_hashes[i] for i in fallback_indices ] - if loaded_hashes: - self.mm_global_cache.release_embeddings(loaded_hashes) + fb_handles = self.mm_global_cache.store_to_pool_async( + fallback_hashes, fallback_slices, ctx.modality + ) + d2h_handles.extend(fb_handles) - # Background insert: store newly computed embeddings into global cache. - # Includes both original misses and fallback-recomputed hits. - all_new_hashes = [str_mm_hashes[i] for i in missing_indices] - all_new_slices = list(new_slices) - if fallback_slices is not None: - all_new_hashes += [ - str_mm_hashes[i] for i in fallback_indices - ] - all_new_slices += list(fallback_slices) - if all_new_hashes: + mm_embedding = self._assemble_global_cache_gpu( + ctx, + missing_indices, + fallback_indices, + new_slices, + fallback_slices, + ) - async def _background_insert(): - await asyncio.to_thread( - self.mm_global_cache.insert_batch, - all_new_hashes, - all_new_slices, - ) + new_hashes = [ctx.str_mm_hashes[i] for i in missing_indices] + new_hashes += [ctx.str_mm_hashes[i] for i in fallback_indices] + self._launch_global_cache_insert( + ctx, + new_hashes, + d2h_handles, + ) - insert_task = asyncio.create_task(_background_insert()) - self.background_tasks.add(insert_task) - insert_task.add_done_callback(self.background_tasks.discard) - - self._forward_results[req_id]["embedding"] = mm_embedding + self._forward_results[ctx.req_id]["embedding"] = mm_embedding logger.info( f"Global cache + VIT forward completed for " - f"{req_id}, shape={mm_embedding.shape}" + f"{ctx.req_id}, shape={mm_embedding.shape}" ) except Exception as e: logger.error( - f"Global cache + VIT forward failed for " f"{req_id}: {e}" + f"Global cache + VIT forward failed for {ctx.req_id}: {e}" ) if self.rank == 0: - self._forward_results[req_id]["error"] = str(e) + self._forward_results[ctx.req_id]["error"] = str(e) finally: if self.rank == 0: event.set() @@ -1306,7 +1386,7 @@ class MMEncoder: if self.rank == 0: logger.info( - f"Returning metadata immediately for {req_id}, " + f"Returning metadata immediately for {ctx.req_id}, " f"global cache + VIT forward running async" ) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/embedding_cache_controller.py b/python/sglang/srt/mem_cache/storage/mooncake_store/embedding_cache_controller.py index 98dcf37f1..938b092b2 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/embedding_cache_controller.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/embedding_cache_controller.py @@ -1,89 +1,285 @@ import asyncio import logging +import math import threading import time +from collections import OrderedDict +from dataclasses import dataclass +from enum import Enum, auto from queue import Empty, Queue -from typing import List, Optional +from typing import List, Optional, Tuple import torch +from sglang.srt.managers.schedule_batch import Modality from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import ( MooncakeEmbeddingStore, ) logger = logging.getLogger(__name__) +TARGET_PAGE_BYTES = 256 * 1024 +VISION_POOL_RATIO = 0.8 -class ContiguousMemoryAllocator: - """ - A simple allocator to manage variable-sized contiguous blocks - within a large pre-allocated flat buffer. - """ - def __init__(self, total_size_bytes: int): - self.total_size = total_size_bytes - # List of (offset, size) for free blocks - self.free_blocks = [(0, total_size_bytes)] - self.allocated_map = {} # {offset: size_bytes} - self.allocated_size = 0 # Running counter for O(1) get_allocated_size - self.lock = threading.Lock() +def _dtype_element_size(dtype: torch.dtype) -> int: + return torch.tensor([], dtype=dtype).element_size() - def allocate(self, size_bytes: int) -> Optional[int]: - with self.lock: - # Simple First-Fit allocation - for i, (offset, block_size) in enumerate(self.free_blocks): - if block_size >= size_bytes: - # Allocate from this block - remaining_size = block_size - size_bytes - if remaining_size > 0: - self.free_blocks[i] = (offset + size_bytes, remaining_size) - else: - self.free_blocks.pop(i) - self.allocated_map[offset] = size_bytes - self.allocated_size += size_bytes - return offset + +def compute_page_size(dim: int, element_size: int = 4) -> int: + return max(TARGET_PAGE_BYTES // (dim * element_size), 1) + + +class EntryState(Enum): + FILLING = auto() + READY = auto() + + +@dataclass(frozen=True) +class PageRun: + start: int + length: int + + @property + def end(self) -> int: + return self.start + self.length + + def page_ids(self) -> List[int]: + return list(range(self.start, self.end)) + + +class RangePageAllocator: + """Range-aware page allocator that prefers contiguous physical page runs.""" + + def __init__(self, num_pages: int): + self.free_ranges: List[Tuple[int, int]] = ( + [(0, num_pages)] if num_pages > 0 else [] + ) + + def allocate(self, num_tokens: int, page_size: int) -> Optional[List[PageRun]]: + required_pages = math.ceil(num_tokens / page_size) + if required_pages <= 0: + return [] + if self.free_pages < required_pages: return None - def free(self, offset: int, size_bytes: int): - with self.lock: - # Remove from allocated map and update counter - if offset in self.allocated_map: - self.allocated_size -= self.allocated_map[offset] - del self.allocated_map[offset] - - # Return block and merge adjacent free blocks - self.free_blocks.append((offset, size_bytes)) - self.free_blocks.sort() - - merged = [] - if not self.free_blocks: - return - - curr_offset, curr_size = self.free_blocks[0] - for next_offset, next_size in self.free_blocks[1:]: - if curr_offset + curr_size == next_offset: - curr_size += next_size + for idx, (start, length) in enumerate(self.free_ranges): + if length >= required_pages: + run = PageRun(start, required_pages) + if length == required_pages: + self.free_ranges.pop(idx) else: - merged.append((curr_offset, curr_size)) - curr_offset, curr_size = next_offset, next_size - merged.append((curr_offset, curr_size)) - self.free_blocks = merged + self.free_ranges[idx] = ( + start + required_pages, + length - required_pages, + ) + return [run] - def get_allocated_size(self) -> int: - """Return total allocated bytes. O(1) operation.""" - with self.lock: - return self.allocated_size + runs: List[PageRun] = [] + remaining = required_pages + while remaining > 0 and self.free_ranges: + start, length = self.free_ranges.pop(0) + take = min(length, remaining) + runs.append(PageRun(start, take)) + remaining -= take + if take < length: + self.free_ranges.insert(0, (start + take, length - take)) - def get_free_size(self) -> int: - """Return total free bytes.""" - with self.lock: - return sum(block_size for _, block_size in self.free_blocks) + if remaining != 0: + self.free(runs) + return None + return runs + + def free(self, runs: List[PageRun]): + if not runs: + return + for run in runs: + if run.length > 0: + self.free_ranges.append((run.start, run.length)) + self.free_ranges.sort() + + merged: List[Tuple[int, int]] = [] + for start, length in self.free_ranges: + if not merged: + merged.append((start, length)) + continue + prev_start, prev_length = merged[-1] + prev_end = prev_start + prev_length + if prev_end >= start: + merged[-1] = (prev_start, max(prev_end, start + length) - prev_start) + else: + merged.append((start, length)) + self.free_ranges = merged + + @property + def free_pages(self) -> int: + return sum(length for _, length in self.free_ranges) + + +class EvictableLRU: + """Per-pool eviction candidate queue. + + Only entries that can be freed immediately (READY, zero read pins) + belong in this queue. The controller is responsible for + adding and removing entries at the right state transitions; this + class does not inspect entry state. + """ + + def __init__(self): + self._lru: OrderedDict[str, float] = OrderedDict() + + def touch(self, mm_hash: str): + """Mark as recently used (move to tail of the queue).""" + self._lru.pop(mm_hash, None) + self._lru[mm_hash] = time.time() + + def remove(self, mm_hash: str): + """Remove from candidates.""" + self._lru.pop(mm_hash, None) + + def pop_oldest(self) -> Optional[str]: + """Pop the least-recently-used candidate. Returns None if empty.""" + if not self._lru: + return None + mm_hash, _ = self._lru.popitem(last=False) + return mm_hash + + def __len__(self) -> int: + return len(self._lru) + + def __contains__(self, mm_hash: str) -> bool: + return mm_hash in self._lru + + def keys(self): + return self._lru.keys() + + +@dataclass +class EmbeddingPool: + modality: str + dim: int + dtype: torch.dtype + page_size: int + tensor: torch.Tensor + num_pages: int + allocator: RangePageAllocator + page_bytes: int + pool_size_bytes: int + pin_memory: bool = True + + @classmethod + def create( + cls, + modality: str, + dim: int, + pool_size_bytes: int, + dtype: torch.dtype = torch.float32, + pin_memory: bool = True, + ) -> "EmbeddingPool": + element_size = _dtype_element_size(dtype) + page_size = compute_page_size(dim, element_size) + capacity_tokens = pool_size_bytes // (dim * element_size) + num_pages = capacity_tokens // page_size + total_tokens = num_pages * page_size + tensor = torch.empty( + (total_tokens, dim), + dtype=dtype, + pin_memory=pin_memory, + ) + page_bytes = page_size * dim * element_size + return cls( + modality=modality, + dim=dim, + dtype=dtype, + page_size=page_size, + tensor=tensor, + num_pages=num_pages, + allocator=RangePageAllocator(num_pages), + page_bytes=page_bytes, + pool_size_bytes=pool_size_bytes, + pin_memory=pin_memory, + ) + + +@dataclass +class EmbeddingCacheEntry: + hash: str + modality: object + num_tokens: int + dim: int + page_runs: List[PageRun] + state: EntryState + ref_count: int = 0 + + @property + def page_ids(self) -> List[int]: + return [page_id for run in self.page_runs for page_id in run.page_ids()] + + def pin(self): + self.ref_count += 1 + + def unpin(self): + if self.ref_count <= 0: + logger.warning("unpin called with ref_count=0 for %s", self.hash) + return + self.ref_count -= 1 + + def is_evictable(self) -> bool: + return self.state == EntryState.READY and self.ref_count == 0 + + +def build_transfer_buffers( + entry: EmbeddingCacheEntry, pool: EmbeddingPool +) -> Tuple[List[int], List[int]]: + """Build one pointer/size pair per physical page run.""" + if not entry.page_runs: + return [], [] + + ptrs: List[int] = [] + sizes: List[int] = [] + remaining_tokens = entry.num_tokens + element_size = _dtype_element_size(pool.dtype) + + for run in entry.page_runs: + if remaining_tokens <= 0: + break + valid_tokens = min(pool.page_size * run.length, remaining_tokens) + ptr = pool.tensor[run.start * pool.page_size].data_ptr() + size_bytes = valid_tokens * entry.dim * element_size + ptrs.append(ptr) + sizes.append(size_bytes) + remaining_tokens -= valid_tokens + return ptrs, sizes + + +@dataclass +class AsyncCopyHandle: + event: object + entry_hash: str + device: Optional[torch.device] = None + _src_ref: object = None + + def is_complete(self) -> bool: + if self.event is None: + return True + return bool(self.event.query()) + + def wait(self): + if self.event is not None: + self.event.synchronize() + self._src_ref = None class EmbeddingPrefetchOperation: """Groups all missing images of a request for a single batch GET.""" - def __init__(self, req_id: str, keys: List[str], ptrs: List[int], sizes: List[int]): + def __init__( + self, + req_id: str, + keys: List[str], + ptrs: List[List[int]], + sizes: List[List[int]], + ): self.req_id = req_id self.keys = keys self.ptrs = ptrs @@ -101,7 +297,7 @@ class EmbeddingPrefetchOperation: class EmbeddingInsertOperation: """Groups all newly computed images of a request for a single batch PUT.""" - def __init__(self, keys: List[str], ptrs: List[int], sizes: List[int]): + def __init__(self, keys: List[str], ptrs: List[List[int]], sizes: List[List[int]]): self.keys = keys self.ptrs = ptrs self.sizes = sizes @@ -118,43 +314,37 @@ class EmbeddingCacheController: all_rank_get=False, enable_eviction: bool = True, max_eviction_batch: int = 100, + dtype: torch.dtype = torch.float32, ): self.tp_world_size = tp_size self.tp_group = tp_group self.tp_rank = tp_rank self.all_rank_get = all_rank_get self.hidden_dims = hidden_dims or {} - self.element_size = torch.float32.itemsize + # Pool dtype must match the model's embedding dtype so that pool views, + # ViT output, and the final send buffer share one dtype — assembly then + # copies without any cast. Defaults to float32 for backward compat. + self.dtype = dtype + self.element_size = _dtype_element_size(self.dtype) self.enable_eviction = enable_eviction self.max_eviction_batch = max_eviction_batch - # 1. Mooncake Backend & Pinned Buffer self.mooncake_store = MooncakeEmbeddingStore() self.total_pool_size_bytes = int(max_pool_size_gb * 1024**3) - self.cpu_pool = torch.empty( - self.total_pool_size_bytes, dtype=torch.uint8, pin_memory=True - ) - self.mooncake_store.register_buffer(self.cpu_pool) + self.vision_pool, self.audio_pool = self._create_pools(pin_memory=True) + self.pools = { + "vision": self.vision_pool, + "audio": self.audio_pool, + } + self._register_pool_buffer(self.vision_pool) + self._register_pool_buffer(self.audio_pool) - # 2. Variable Size Memory Management - self.allocator = ContiguousMemoryAllocator(self.total_pool_size_bytes) - # {hash: (offset, num_tokens, embedding_dim, size_bytes, last_access_time)} - self.hash_to_metadata = {} + self.entries = {} + # self.lock protects entries, pool allocators, entry state/pins, + # and per-pool evictable LRUs. Do not mutate without holding self.lock. + self.vision_pool.evictable = EvictableLRU() + self.audio_pool.evictable = EvictableLRU() - # 3. LRU Tracking - # OrderedDict maintains insertion order, used as LRU cache - # hash -> access_time - self.access_order = {} - self.access_lock = threading.Lock() - - # 4. RDMA / read reference counting - # {hash: ref_count} — entries with ref_count > 0 cannot be evicted. - # Incremented when an RDMA transfer (GET/PUT) is in flight or when - # get_embeddings() returns a view into cpu_pool. Decremented after - # the RDMA completes or the caller releases the view. - self.ref_counts = {} - - # 5. Statistics self.stats = { "total_allocated": 0, "total_evicted": 0, @@ -162,8 +352,9 @@ class EmbeddingCacheController: "allocation_failures": 0, } - # 6. Task Tracking - self.ongoing_prefetch = {} # {req_id: EmbeddingPrefetchOperation} + self._copy_streams = {} + + self.ongoing_prefetch = {} self.prefetch_queue = Queue() self.insert_queue = Queue() @@ -186,275 +377,280 @@ class EmbeddingCacheController: else: self.prefetch_tp_group = None - def _update_access_time(self, image_hash: str): - """Update LRU access time for a hash.""" - with self.access_lock: - # Move to end (most recently used) - if image_hash in self.access_order: - del self.access_order[image_hash] - self.access_order[image_hash] = time.time() + def _create_pools(self, pin_memory: bool) -> Tuple[EmbeddingPool, EmbeddingPool]: + # vision pool uses IMAGE dim (IMAGE == VIDEO dim in all supported models) + vision_dim = self.hidden_dims.get(Modality.IMAGE) or self.hidden_dims.get( + Modality.VIDEO + ) + audio_dim = self.hidden_dims.get(Modality.AUDIO) or vision_dim + vision_bytes = int(self.total_pool_size_bytes * VISION_POOL_RATIO) + audio_bytes = self.total_pool_size_bytes - vision_bytes + return ( + EmbeddingPool.create( + "vision", vision_dim, vision_bytes, self.dtype, pin_memory + ), + EmbeddingPool.create( + "audio", audio_dim, audio_bytes, self.dtype, pin_memory + ), + ) - def _protect_hash(self, image_hash: str): - """Increment ref count to prevent eviction during RDMA or active read. + def _register_pool_buffer(self, pool: EmbeddingPool): + if pool.tensor.numel() == 0: + logger.warning( + f"[Rank {self.tp_rank}] {pool.modality} embedding pool has zero pages; " + f"dim={pool.dim}, budget={pool.pool_size_bytes} bytes" + ) + return + self.mooncake_store.register_buffer(pool.tensor) + logger.info( + f"[Rank {self.tp_rank}] Registered {pool.modality} embedding pool: " + f"dim={pool.dim}, pages={pool.num_pages}, " + f"page_tokens={pool.page_size}, " + f"capacity={pool.num_pages * pool.page_bytes / 1024**2:.2f} MB" + ) - NOTE: Caller must hold self.lock. + def _get_pool(self, modality: Modality) -> Optional[EmbeddingPool]: + if modality == Modality.AUDIO: + return self.audio_pool + if modality in (Modality.IMAGE, Modality.VIDEO): + return self.vision_pool + return None + + # --- LRU and state helpers (caller must hold self.lock) --- + + def _lru_touch(self, mm_hash: str): + """Mark an evictable entry as recently used in its pool's LRU.""" + entry = self.entries.get(mm_hash) + if entry is None or not entry.is_evictable(): + return + pool = self._get_pool(entry.modality) + if pool is not None: + pool.evictable.touch(mm_hash) + + def _mark_ready(self, entry: EmbeddingCacheEntry): + """Transition a FILLING entry to READY.""" + entry.state = EntryState.READY + pool = self._get_pool(entry.modality) + if pool is not None and entry.is_evictable(): + pool.evictable.touch(entry.hash) + + def _pin_read(self, entry: EmbeddingCacheEntry): + """Pin a READY entry for a read transfer.""" + if entry.ref_count == 0: + pool = self._get_pool(entry.modality) + if pool is not None: + pool.evictable.remove(entry.hash) + entry.pin() + + def _unpin_read(self, entry: EmbeddingCacheEntry): + """Release a read transfer pin.""" + entry.unpin() + if entry.is_evictable(): + pool = self._get_pool(entry.modality) + if pool is not None: + pool.evictable.touch(entry.hash) + + # --- Eviction --- + + def _evict_entry(self, mm_hash: str, remove_lru: bool = True): + """Free one entry and remove metadata. + + Caller must hold self.lock. """ - self.ref_counts[image_hash] = self.ref_counts.get(image_hash, 0) + 1 + entry = self.entries.get(mm_hash) + if entry is None: + return + pool = self._get_pool(entry.modality) + if remove_lru and pool is not None: + pool.evictable.remove(mm_hash) + pool.allocator.free(entry.page_runs) + self.stats["total_evicted"] += entry.num_tokens * entry.dim * self.element_size + del self.entries[mm_hash] - def _release_hash(self, image_hash: str): - """Decrement ref count after RDMA completes or caller releases a view. + def _evict_for_pool(self, pool: EmbeddingPool, required_pages: int): + """Evict oldest entries from this pool until enough pages are free. - NOTE: Caller must hold self.lock. + Caller must hold self.lock. """ - if image_hash in self.ref_counts: - self.ref_counts[image_hash] -= 1 - if self.ref_counts[image_hash] <= 0: - del self.ref_counts[image_hash] + if pool.allocator.free_pages >= required_pages: + return - def _select_eviction_candidates(self, required_bytes: int) -> List[str]: - """Select LRU candidates to free up at least required_bytes. - - NOTE: Caller must hold self.lock before calling this method. - """ - candidates = [] - freed_bytes = 0 - - with self.access_lock: - # Sort by access time (oldest first) - # Python dicts are insertion-ordered; the first keys are the oldest. - sorted_hashes = list(self.access_order.items()) - - for image_hash, _ in sorted_hashes: - if image_hash not in self.hash_to_metadata: - with self.access_lock: - self.access_order.pop(image_hash, None) - continue - if self.ref_counts.get(image_hash, 0) > 0: - continue - metadata = self.hash_to_metadata[image_hash] - size_bytes = metadata[3] if len(metadata) > 3 else 0 - candidates.append(image_hash) - freed_bytes += size_bytes - - if freed_bytes >= required_bytes: + evicted = 0 + while pool.allocator.free_pages < required_pages: + if evicted >= self.max_eviction_batch: break - - if len(candidates) >= self.max_eviction_batch: + mm_hash = pool.evictable.pop_oldest() + if mm_hash is None: break - - return candidates - - def _evict_hashes(self, hashes_to_evict: List[str]) -> int: - """Evict specified hashes and free their memory. Returns freed bytes. - - NOTE: Caller must hold self.lock before calling this method. - """ - total_freed = 0 - - # NOTE: self.lock should be held by the caller (e.g., insert_batch, - # prefetch). Do NOT acquire it here to avoid reentrant deadlock. - for image_hash in hashes_to_evict: - if image_hash not in self.hash_to_metadata: + entry = self.entries.get(mm_hash) + if entry is None: continue + self._evict_entry(mm_hash, remove_lru=False) + evicted += 1 - # Safety check: skip entries with in-flight RDMA or active reads - if self.ref_counts.get(image_hash, 0) > 0: - logger.warning( - f"[Rank {self.tp_rank}] Skipping eviction of {image_hash}: " - f"ref_count={self.ref_counts[image_hash]} (in-flight RDMA or active read)" - ) - continue - - offset, num_tokens, dim, size_bytes = self.hash_to_metadata[image_hash][:4] - - # Free memory in allocator - self.allocator.free(offset, size_bytes) - - # Remove from metadata and ref counts - del self.hash_to_metadata[image_hash] - self.ref_counts.pop(image_hash, None) - - # Remove from access order - with self.access_lock: - self.access_order.pop(image_hash, None) - - total_freed += size_bytes - self.stats["total_evicted"] += size_bytes - - if total_freed > 0: + if evicted > 0: self.stats["eviction_count"] += 1 - - if total_freed > 0: logger.info( - f"[Rank {self.tp_rank}] Evicted {len(hashes_to_evict)} embeddings, " - f"freed {total_freed / 1024**2:.2f} MB" + f"[Rank {self.tp_rank}] Evicted {evicted} embeddings from " + f"{pool.modality} pool" ) - return total_freed + def _allocate_with_eviction( + self, pool: EmbeddingPool, num_tokens: int + ) -> Optional[List[PageRun]]: + """Allocate pages, evicting LRU entries from the same pool if needed. - def _allocate_with_eviction(self, size_bytes: int) -> Optional[int]: - """Try to allocate memory, evicting old entries if necessary.""" - # First try direct allocation - offset = self.allocator.allocate(size_bytes) - if offset is not None: - self.stats["total_allocated"] += size_bytes - return offset - - # If failed and eviction is enabled, try eviction - if not self.enable_eviction: + Caller must hold self.lock. + """ + required_pages = math.ceil(num_tokens / pool.page_size) + if required_pages > pool.num_pages: self.stats["allocation_failures"] += 1 return None - # Select candidates to evict - candidates = self._select_eviction_candidates(size_bytes) - if not candidates: - n_protected = sum(1 for v in self.ref_counts.values() if v > 0) - logger.warning( - f"[Rank {self.tp_rank}] Cannot allocate {size_bytes / 1024**2:.2f} MB: " - f"pool full ({self.allocator.get_allocated_size() / 1024**2:.1f}/" - f"{self.total_pool_size_bytes / 1024**2:.1f} MB used), " - f"no evictable candidates " - f"({len(self.hash_to_metadata)} entries, {n_protected} protected)" - ) - self.stats["allocation_failures"] += 1 - return None + if self.enable_eviction: + self._evict_for_pool(pool, required_pages) - # Evict and try again - freed = self._evict_hashes(candidates) - if freed < size_bytes: - logger.warning( - f"[Rank {self.tp_rank}] Could not free enough memory: " - f"needed {size_bytes / 1024**2:.2f} MB, freed {freed / 1024**2:.2f} MB" - ) - - # Try allocation again after eviction - offset = self.allocator.allocate(size_bytes) - if offset is not None: - self.stats["total_allocated"] += size_bytes + page_runs = pool.allocator.allocate(num_tokens, pool.page_size) + if page_runs is not None: + self.stats["total_allocated"] += num_tokens * pool.dim * self.element_size else: self.stats["allocation_failures"] += 1 - - return offset + logger.warning( + f"[Rank {self.tp_rank}] Cannot allocate {required_pages} pages " + f"in {pool.modality} pool: free={pool.allocator.free_pages}" + ) + return page_runs def prefetch( self, req_id: str, - image_hashes: List[str], + mm_hashes: List[str], expected_tokens: List[int], modality=None, ): - """Issues ONE batch GET for all missing images in the request.""" - dim = self.hidden_dims.get(modality) if modality is not None else None - if not dim: - logger.warning( - f"Req {req_id}: Unknown dim for modality={modality}, skipping prefetch (will fallback to ViT)." - ) + """Issues ONE batch GET for cache-hit embeddings that are not local yet.""" + pool = self._get_pool(modality) + if pool is None: + logger.warning(f"prefetch: unknown modality {modality}; skipping.") return - keys, ptrs, sizes = [], [], [] + + keys, all_ptrs, all_sizes = [], [], [] with self.lock: - for h, num_tokens in zip(image_hashes, expected_tokens): - if h in self.hash_to_metadata: - # Update access time for LRU - self._update_access_time(h) - logger.debug( - f"Req {req_id}: Hash already in local metadata, skipping prefetch." - ) + for mm_hash, num_tokens in zip(mm_hashes, expected_tokens): + entry = self.entries.get(mm_hash) + if entry is not None: + if entry.state == EntryState.READY: + self._lru_touch(mm_hash) + else: + logger.debug( + f"Req {req_id}: {mm_hash} is FILLING; " f"treating as miss." + ) continue - size_bytes = num_tokens * dim * self.element_size - offset = self._allocate_with_eviction(size_bytes) - if offset is None: + page_runs = self._allocate_with_eviction(pool, int(num_tokens)) + if page_runs is None: logger.warning( - f"Req {req_id}: Failed to allocate {size_bytes / 1024**2:.2f} MB " - f"for prefetch, skipping this image." + f"Req {req_id}: Failed to allocate {num_tokens} tokens " + f"in {pool.modality} pool; falling back to encoder." ) continue - self.hash_to_metadata[h] = (offset, num_tokens, dim, size_bytes) - self._update_access_time(h) - self._protect_hash(h) - keys.append(h) - ptrs.append(self.cpu_pool.data_ptr() + offset) - sizes.append(size_bytes) + entry = EmbeddingCacheEntry( + hash=mm_hash, + modality=modality, + num_tokens=int(num_tokens), + dim=pool.dim, + page_runs=page_runs, + state=EntryState.FILLING, + ) + self.entries[mm_hash] = entry + keys.append(mm_hash) + entry_ptrs, entry_sizes = build_transfer_buffers(entry, pool) + all_ptrs.append(entry_ptrs) + all_sizes.append(entry_sizes) if not keys: return logger.info( - f"Req {req_id}: Starting global fetch for {len(keys)} images from Mooncake." + f"Req {req_id}: Starting global fetch for {len(keys)} " + f"embeddings from Mooncake." ) - op = EmbeddingPrefetchOperation(req_id, keys, ptrs, sizes) + op = EmbeddingPrefetchOperation(req_id, keys, all_ptrs, all_sizes) self.ongoing_prefetch[req_id] = op self.prefetch_queue.put(op) def insert_batch( - self, image_hashes: List[str], embedding_tensors: List[torch.Tensor] + self, + mm_hashes: List[str], + modality: Modality = None, ): - """Issues ONE batch PUT for all embeddings computed by this request. + """Issues ONE batch PUT for embeddings already in the host pool. - Note: Even if the embedding exists locally, we still push to Mooncake - to ensure multi-node cache consistency. Mooncake's batch_put has - built-in deduplication to avoid redundant transfers. + Only READY entries are pushed to Mooncake for multi-node sharing. + If an entry was never stored (e.g. store_to_pool_async allocation failed), + it is silently skipped. """ - keys, ptrs, sizes = [], [], [] - local_hit_count = 0 - new_count = 0 + pool = self._get_pool(modality) + if pool is None: + logger.warning(f"insert_batch: unknown modality {modality}; skipping.") + return + + keys, all_ptrs, all_sizes = [], [], [] skipped_count = 0 with self.lock: - for h, tensor in zip(image_hashes, embedding_tensors): - if h in self.hash_to_metadata: - # Update access time for existing entry - self._update_access_time(h) - self._protect_hash(h) - # Local cache hit: ensure Mooncake has it - offset, num_tokens, dim, size_bytes = self.hash_to_metadata[h][:4] - - # Still push to Mooncake for multi-node sharing - # (Mooncake batch_put will deduplicate if already exists) - keys.append(h) - ptrs.append(self.cpu_pool.data_ptr() + offset) - sizes.append(size_bytes) - local_hit_count += 1 - continue - - # Local cache miss: allocate and copy - num_tokens, dim = tensor.shape[0], tensor.shape[1] - size_bytes = num_tokens * dim * self.element_size - offset = self._allocate_with_eviction(size_bytes) - if offset is None: - logger.warning( - f"Failed to allocate {size_bytes / 1024**2:.2f} MB for insert, " - f"skipping this embedding." - ) + for mm_hash in mm_hashes: + entry = self.entries.get(mm_hash) + if entry is None or entry.state != EntryState.READY: skipped_count += 1 continue - # Copy to pinned pool for RDMA - target_view = ( - self.cpu_pool[offset : offset + size_bytes] - .view(torch.float32) - .view(num_tokens, dim) - ) - target_view.copy_(tensor.cpu()) - self.hash_to_metadata[h] = (offset, num_tokens, dim, size_bytes) - self._update_access_time(h) - self._protect_hash(h) - - keys.append(h) - ptrs.append(self.cpu_pool.data_ptr() + offset) - sizes.append(size_bytes) - new_count += 1 + self._pin_read(entry) + keys.append(mm_hash) + entry_ptrs, entry_sizes = build_transfer_buffers(entry, pool) + all_ptrs.append(entry_ptrs) + all_sizes.append(entry_sizes) if keys: logger.info( - f"Global Cache: Inserting {len(keys)} embeddings into Mooncake cluster " - f"({new_count} new, {local_hit_count} existing for replication, " - f"{skipped_count} skipped due to allocation failure)" + f"Global Cache: Inserting {len(keys)} embeddings into " + f"Mooncake cluster ({skipped_count} skipped)" ) - self.insert_queue.put(EmbeddingInsertOperation(keys, ptrs, sizes)) + self.insert_queue.put( + EmbeddingInsertOperation(keys, all_ptrs, all_sizes) + ) + + def _finish_get(self, op: EmbeddingPrefetchOperation, results: List[bool]): + with self.lock: + for mm_hash, success in zip(op.keys, results): + entry = self.entries.get(mm_hash) + if entry is None: + continue + if success: + if entry.state == EntryState.FILLING: + self._mark_ready(entry) + else: + pool = self._get_pool(entry.modality) + pool.evictable.remove(mm_hash) + pool.allocator.free(entry.page_runs) + del self.entries[mm_hash] + op.mark_done(all(results)) + + def _finish_put(self, op: EmbeddingInsertOperation, results: List[bool]): + with self.lock: + for mm_hash, success in zip(op.keys, results): + entry = self.entries.get(mm_hash) + if entry is None: + continue + if not success: + logger.warning( + f"[Rank {self.tp_rank}] Mooncake PUT failed for " + f"{mm_hash}; keeping local cache entry." + ) + self._unpin_read(entry) def _io_loop(self): """Asynchronous worker handling both Batch GET and Batch PUT.""" @@ -463,16 +659,19 @@ class EmbeddingCacheController: try: op = self.prefetch_queue.get_nowait() - results = self.mooncake_store.batch_get(op.keys, op.ptrs, op.sizes) + try: + results = self.mooncake_store.batch_get_into_multi_buffers( + op.keys, op.ptrs, op.sizes + ) + except Exception: + logger.exception("Mooncake multi-buffer GET failed") + results = [False] * len(op.keys) success_count = sum(results) logger.info( - f"Mooncake GET Finished: Req {op.req_id}, Successfully fetched {success_count}/{len(op.keys)} images." + f"Mooncake GET Finished: Req {op.req_id}, " + f"Successfully fetched {success_count}/{len(op.keys)} embeddings." ) - op.mark_done(all(results)) - # Release ref counts now that RDMA GET is complete - with self.lock: - for h in op.keys: - self._release_hash(h) + self._finish_get(op, results) self.prefetch_queue.task_done() processed_any = True except Empty: @@ -480,14 +679,18 @@ class EmbeddingCacheController: try: op = self.insert_queue.get_nowait() - self.mooncake_store.batch_put(op.keys, op.ptrs, op.sizes) + try: + results = self.mooncake_store.batch_put_from_multi_buffers( + op.keys, op.ptrs, op.sizes + ) + except Exception: + logger.exception("Mooncake multi-buffer PUT failed") + results = [False] * len(op.keys) + self._finish_put(op, results) logger.info( - f"Mooncake PUT Finished: Successfully stored {len(op.keys)} keys in cluster." + f"Mooncake PUT Finished: Stored {sum(results)}/{len(op.keys)} " + f"embeddings in cluster." ) - # Release ref counts now that RDMA PUT is complete - with self.lock: - for h in op.keys: - self._release_hash(h) self.insert_queue.task_done() processed_any = True except Empty: @@ -505,7 +708,7 @@ class EmbeddingCacheController: else: op = self.ongoing_prefetch[req_id] if op.is_finished: - local_ready = op.success + local_ready = True if self.all_rank_get and self.tp_world_size > 1: ready_tensor = torch.tensor( @@ -524,64 +727,285 @@ class EmbeddingCacheController: return True return False - def get_embeddings(self, image_hashes: List[str]) -> List[torch.Tensor]: - """Final reconstruction for model input. + def load_to_device_async( + self, mm_hash: str, dst_tensor: torch.Tensor, dst_token_offset: int + ) -> Optional[AsyncCopyHandle]: + """Async host-pool → device copy for a single READY entry. - Returns views into the pinned cpu_pool. Callers MUST call - release_embeddings() once they no longer need the returned - tensors (e.g. after .to(device) or torch.cat) so that the - entries can be evicted. + Returns an AsyncCopyHandle on success, or None if the entry is + missing/not ready. Pass returned handles to wait_load_to_device(). """ with self.lock: - tensors = [] - for h in image_hashes: - if h not in self.hash_to_metadata: - logger.warning(f"Hash {h} not found in local cache") - tensors.append(None) + entry = self.entries.get(mm_hash) + if entry is None: + logger.warning(f"Hash {mm_hash} not found in local cache") + return None + if entry.state != EntryState.READY: + logger.warning(f"Hash {mm_hash} is not ready; state={entry.state.name}") + return None + self._pin_read(entry) + + pool = self._get_pool(entry.modality) + try: + device = dst_tensor.device + copy_stream = self._get_copy_stream(device) + event = torch.cuda.Event() + copied = 0 + with torch.cuda.stream(copy_stream): + for run in entry.page_runs: + valid_tokens = min( + pool.page_size * run.length, entry.num_tokens - copied + ) + if valid_tokens <= 0: + break + src_start = run.start * pool.page_size + dst_start = dst_token_offset + copied + dst_tensor[dst_start : dst_start + valid_tokens].copy_( + pool.tensor[src_start : src_start + valid_tokens], + non_blocking=True, + ) + copied += valid_tokens + event.record(copy_stream) + return AsyncCopyHandle(event, mm_hash, device=torch.device(device)) + except Exception: + with self.lock: + self._unpin_read(entry) + raise + + def wait_load_to_device(self, handles: List[AsyncCopyHandle]): + """Wait for async GPU copies and release pool pins.""" + for handle in handles: + handle.wait() + with self.lock: + for handle in handles: + entry = self.entries.get(handle.entry_hash) + if entry is not None: + self._unpin_read(entry) + + def get_pool_views( + self, mm_hashes: List[str] + ) -> List[Optional[List[torch.Tensor]]]: + """Return zero-copy slice lists into the host pool for READY entries. + + Each element is a list of tensor views (one per page run) or None + if the entry is missing/not ready. The caller should flatten these + into a single list and do one torch.cat at the end. + Call release_pool_views() when done. + """ + results: List[Optional[List[torch.Tensor]]] = [] + with self.lock: + for mm_hash in mm_hashes: + entry = self.entries.get(mm_hash) + if entry is None or entry.state != EntryState.READY: + results.append(None) continue - # Update access time for LRU - self._update_access_time(h) - self._protect_hash(h) - offset, num_tokens, dim, size_bytes = self.hash_to_metadata[h][:4] - tensors.append( - self.cpu_pool[offset : offset + size_bytes] - .view(torch.float32) - .view(num_tokens, dim) - ) - return tensors + self._pin_read(entry) + pool = self._get_pool(entry.modality) + slices = [] + copied = 0 + for run in entry.page_runs: + valid = min(pool.page_size * run.length, entry.num_tokens - copied) + if valid <= 0: + break + start = run.start * pool.page_size + slices.append(pool.tensor[start : start + valid]) + copied += valid + results.append(slices) + return results - def release_embeddings(self, image_hashes: List[str]): - """Release reference counts on embeddings after the caller is done. - - Must be called once for every successful get_embeddings() call, - after the caller no longer needs the returned tensor views - (e.g. after .to(device) or torch.cat has copied the data). - """ + def release_pool_views(self, mm_hashes: List[str]): + """Release pins acquired by get_pool_views.""" with self.lock: - for h in image_hashes: - self._release_hash(h) + for mm_hash in mm_hashes: + entry = self.entries.get(mm_hash) + if entry is not None: + self._unpin_read(entry) + + def _get_copy_stream(self, device: torch.device) -> "torch.cuda.Stream": + key = str(device) + stream = self._copy_streams.get(key) + if stream is None: + stream = torch.cuda.Stream(device=device) + self._copy_streams[key] = stream + return stream + + def has_local_embedding(self, mm_hash: str) -> bool: + with self.lock: + entry = self.entries.get(mm_hash) + return entry is not None and entry.state == EntryState.READY + + def store_to_pool_async( + self, + mm_hashes: List[str], + tensors: List[torch.Tensor], + modality=None, + ) -> List[Tuple["EmbeddingCacheEntry", "AsyncCopyHandle"]]: + """Launch async D2H copies into host paged pool. + + Allocates pages and launches D2H copies on a side stream but does + NOT wait for completion. Entries remain in FILLING state. + + Items that cannot be stored (FILLING conflict, allocation failure, + dim mismatch) are silently skipped — the caller's assembly step + falls back to the GPU tensor for those items. + + Returns: pending D2H handles; pass to wait_store_to_pool(). + """ + for tensor in tensors: + if tensor.device.type != "cuda": + raise ValueError( + f"store_to_pool_async expects CUDA tensors, " + f"got device={tensor.device}" + ) + + pool = self._get_pool(modality) + if pool is None: + return [] + + pending: List[Tuple[torch.Tensor, EmbeddingCacheEntry]] = [] + with self.lock: + for mm_hash, tensor in zip(mm_hashes, tensors): + if tensor.ndim != 2: + tensor = tensor.reshape(-1, tensor.shape[-1]) + num_tokens, actual_dim = int(tensor.shape[0]), int(tensor.shape[1]) + + if actual_dim != pool.dim: + logger.warning( + f"[Rank {self.tp_rank}] Embedding dim mismatch for " + f"{mm_hash}: pool.dim={pool.dim}, tensor.dim={actual_dim}; " + f"skipping pool store" + ) + continue + + entry = self.entries.get(mm_hash) + if entry is not None: + if entry.state == EntryState.READY: + self._lru_touch(mm_hash) + continue + if entry.state == EntryState.FILLING: + continue + self._evict_entry(mm_hash) + + page_runs = self._allocate_with_eviction(pool, num_tokens) + if page_runs is None: + continue + + entry = EmbeddingCacheEntry( + hash=mm_hash, + modality=modality, + num_tokens=num_tokens, + dim=actual_dim, + page_runs=page_runs, + state=EntryState.FILLING, + ) + self.entries[mm_hash] = entry + pending.append((tensor, entry)) + + handles: List[Tuple[EmbeddingCacheEntry, AsyncCopyHandle]] = [] + for tensor, entry in pending: + handle = self._copy_tensor_to_pool(tensor, entry, pool) + handles.append((entry, handle)) + + return handles + + def wait_store_to_pool( + self, + handles: List[Tuple["EmbeddingCacheEntry", "AsyncCopyHandle"]], + ): + """Wait for async D2H copies and mark entries READY.""" + for entry, handle in handles: + handle.wait() + with self.lock: + for entry, handle in handles: + current = self.entries.get(entry.hash) + if current is entry and current.state == EntryState.FILLING: + self._mark_ready(current) + + def _copy_tensor_to_pool( + self, tensor: torch.Tensor, entry: EmbeddingCacheEntry, pool: EmbeddingPool + ) -> AsyncCopyHandle: + """Async D2H copy of a CUDA tensor into pool pages.""" + src = tensor.detach() + if src.ndim != 2: + src = src.reshape(-1, src.shape[-1]) + if not src.is_contiguous(): + src = src.contiguous() + + device = src.device + producer_stream = torch.cuda.current_stream(device) + copy_stream = self._get_copy_stream(device) + copy_stream.wait_stream(producer_stream) + src.record_stream(copy_stream) + event = torch.cuda.Event() + copied = 0 + with torch.cuda.stream(copy_stream): + for run in entry.page_runs: + valid_tokens = min( + pool.page_size * run.length, entry.num_tokens - copied + ) + if valid_tokens <= 0: + break + start = run.start * pool.page_size + pool.tensor[start : start + valid_tokens].copy_( + src[copied : copied + valid_tokens], + non_blocking=True, + ) + copied += valid_tokens + event.record(copy_stream) + return AsyncCopyHandle( + event=event, + entry_hash=entry.hash, + device=torch.device(device), + _src_ref=src, + ) + + def get_embedding_dim(self, modality=None) -> int: + return self._get_pool(modality).dim def get_stats(self) -> dict: """Return cache statistics.""" with self.lock: + allocated_bytes = sum( + sum(run.length for run in entry.page_runs) + * self._get_pool(entry.modality).page_bytes + for entry in self.entries.values() + ) + free_bytes = sum( + pool.allocator.free_pages * pool.page_bytes + for pool in self.pools.values() + ) return { **self.stats, - "num_cached": len(self.hash_to_metadata), - "num_protected": sum(1 for v in self.ref_counts.values() if v > 0), - "allocated_mb": self.allocator.get_allocated_size() / 1024**2, - "free_mb": self.allocator.get_free_size() / 1024**2, - "total_mb": self.total_pool_size_bytes / 1024**2, + "num_cached": len(self.entries), + "num_pinned": sum( + 1 for entry in self.entries.values() if entry.ref_count > 0 + ), + "allocated_mb": allocated_bytes / 1024**2, + "free_mb": free_bytes / 1024**2, + "total_mb": sum( + pool.num_pages * pool.page_bytes for pool in self.pools.values() + ) + / 1024**2, + "vision_free_pages": self.vision_pool.allocator.free_pages, + "audio_free_pages": self.audio_pool.allocator.free_pages, } - async def batch_is_exist(self, image_hashes: List[str]) -> List[bool]: + async def batch_is_exist(self, mm_hashes: List[str]) -> List[bool]: with self.lock: - local_results = [h in self.hash_to_metadata for h in image_hashes] + local_results = [] + for h in mm_hashes: + entry = self.entries.get(h) + if entry is not None and entry.state == EntryState.READY: + self._lru_touch(h) + local_results.append(True) + else: + local_results.append(False) local_hit_count = sum(local_results) global_hit_count = 0 if not all(local_results): missing_indices = [i for i, res in enumerate(local_results) if not res] - missing_hashes = [image_hashes[i] for i in missing_indices] + missing_hashes = [mm_hashes[i] for i in missing_indices] global_exists = await asyncio.to_thread( self.mooncake_store.batch_is_exist, missing_hashes @@ -591,7 +1015,7 @@ class EmbeddingCacheController: for i, exists in zip(missing_indices, global_exists): local_results[i] = exists - total = len(image_hashes) + total = len(mm_hashes) miss_count = total - local_hit_count - global_hit_count logger.info( f"=== Multi-Level Cache Check === " diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py index f358d97dc..578259f7c 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_embedding_store.py @@ -30,8 +30,8 @@ class MooncakeEmbeddingStore(MooncakeBaseStore): logger.info("Mooncake Embedding Store initialized successfully.") - def get_key(self, image_hash: str) -> str: - return f"emb_{image_hash}" + def get_key(self, mm_hash: str) -> str: + return f"emb_{mm_hash}" def batch_get( self, hashes: List[str], ptrs: List[int], sizes: List[int] @@ -66,3 +66,45 @@ class MooncakeEmbeddingStore(MooncakeBaseStore): keys = [self.get_key(h) for h in hashes] results = self.store.batch_is_exist(keys) return [res == 1 for res in results] + + def batch_get_into_multi_buffers( + self, + hashes: List[str], + ptrs: List[List[int]], + sizes: List[List[int]], + ) -> List[bool]: + keys = [self.get_key(h) for h in hashes] + results = self.store.batch_get_into_multi_buffers(keys, ptrs, sizes) + return [res > 0 for res in results] + + def batch_put_from_multi_buffers( + self, + hashes: List[str], + ptrs: List[List[int]], + sizes: List[List[int]], + ) -> List[bool]: + keys = [self.get_key(h) for h in hashes] + + # Skip keys that already exist in Mooncake + exists = self.store.batch_is_exist(keys) + put_keys = [] + put_ptrs = [] + put_sizes = [] + put_indices = [] + success_map = [True] * len(hashes) + + for i, status in enumerate(exists): + if status != 1: + put_keys.append(keys[i]) + put_ptrs.append(ptrs[i]) + put_sizes.append(sizes[i]) + put_indices.append(i) + + if not put_keys: + return success_map + + results = self.store.batch_put_from_multi_buffers(put_keys, put_ptrs, put_sizes) + for i, res in enumerate(results): + success_map[put_indices[i]] = res == 0 + + return success_map diff --git a/test/registered/unit/mem_cache/test_embedding_cache_controller.py b/test/registered/unit/mem_cache/test_embedding_cache_controller.py index 9b65cfe3a..59585a0ca 100644 --- a/test/registered/unit/mem_cache/test_embedding_cache_controller.py +++ b/test/registered/unit/mem_cache/test_embedding_cache_controller.py @@ -1,715 +1,444 @@ -"""Unit tests for EmbeddingCacheController — LRU eviction and RDMA ref counting.""" - -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") +"""Unit tests for EmbeddingCacheController paged host pool behavior.""" import threading -import time import unittest +from queue import Queue from unittest.mock import MagicMock import torch +from sglang.srt.managers.schedule_batch import Modality from sglang.srt.mem_cache.storage.mooncake_store.embedding_cache_controller import ( - ContiguousMemoryAllocator, EmbeddingCacheController, - EmbeddingInsertOperation, + EmbeddingCacheEntry, + EmbeddingPool, + EntryState, + EvictableLRU, + PageRun, + RangePageAllocator, + build_transfer_buffers, ) +from sglang.test.ci.ci_register import register_cpu_ci -# --------------------------------------------------------------------------- -# ContiguousMemoryAllocator tests -# --------------------------------------------------------------------------- +register_cpu_ci(est_time=5, suite="base-a-test-cpu") -class TestContiguousMemoryAllocator(unittest.TestCase): - def test_basic_alloc_free(self): - alloc = ContiguousMemoryAllocator(1024) - a = alloc.allocate(256) - self.assertIsNotNone(a) - self.assertEqual(a, 0) - b = alloc.allocate(256) - self.assertEqual(b, 256) - alloc.free(a, 256) - c = alloc.allocate(128) - self.assertEqual(c, 0) # reused freed block - - def test_alloc_fails_when_full(self): - alloc = ContiguousMemoryAllocator(256) - a = alloc.allocate(256) - self.assertIsNotNone(a) - b = alloc.allocate(1) - self.assertIsNone(b) - - def test_free_merges_adjacent(self): - alloc = ContiguousMemoryAllocator(512) - a = alloc.allocate(128) - b = alloc.allocate(128) - c = alloc.allocate(256) - alloc.free(a, 128) - alloc.free(b, 128) - # The two 128-byte blocks should merge into one 256-byte free block - d = alloc.allocate(256) - self.assertIsNotNone(d) - self.assertEqual(d, 0) - - def test_allocated_size_tracking(self): - alloc = ContiguousMemoryAllocator(1024) - self.assertEqual(alloc.get_allocated_size(), 0) - a = alloc.allocate(300) - self.assertEqual(alloc.get_allocated_size(), 300) - b = alloc.allocate(200) - self.assertEqual(alloc.get_allocated_size(), 500) - alloc.free(a, 300) - self.assertEqual(alloc.get_allocated_size(), 200) - alloc.free(b, 200) - self.assertEqual(alloc.get_allocated_size(), 0) - - def test_free_size_tracking(self): - alloc = ContiguousMemoryAllocator(1024) - self.assertEqual(alloc.get_free_size(), 1024) - alloc.allocate(400) - self.assertEqual(alloc.get_free_size(), 624) - - def test_double_free_is_safe(self): - alloc = ContiguousMemoryAllocator(256) - a = alloc.allocate(128) - alloc.free(a, 128) - # Second free of same offset — should not corrupt state - alloc.free(a, 128) - self.assertEqual(alloc.get_allocated_size(), 0) +def _make_pool(num_pages=16, dim=4, page_size=2, modality="vision"): + total_tokens = num_pages * page_size + tensor = torch.empty((total_tokens, dim), dtype=torch.float32) + return EmbeddingPool( + modality=modality, + dim=dim, + dtype=torch.float32, + page_size=page_size, + tensor=tensor, + num_pages=num_pages, + allocator=RangePageAllocator(num_pages), + page_bytes=page_size * dim * torch.float32.itemsize, + pool_size_bytes=total_tokens * dim * torch.float32.itemsize, + pin_memory=False, + ) -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_controller( - pool_mb=1.0, enable_eviction=True, hidden_dims=None, max_eviction_batch=10 -): - """Create an EmbeddingCacheController with a mocked MooncakeEmbeddingStore.""" +def _make_controller(num_pages=16, dim=4, page_size=2, enable_eviction=True): ctrl = EmbeddingCacheController.__new__(EmbeddingCacheController) ctrl.tp_world_size = 1 ctrl.tp_group = None ctrl.tp_rank = 0 ctrl.all_rank_get = False - ctrl.hidden_dims = hidden_dims or {"image": 1024} + ctrl.hidden_dims = { + Modality.IMAGE: dim, + Modality.VIDEO: dim, + Modality.AUDIO: dim, + } + ctrl.dtype = torch.float32 ctrl.element_size = torch.float32.itemsize ctrl.enable_eviction = enable_eviction - ctrl.max_eviction_batch = max_eviction_batch - - # Small pool for testing (1 MB by default) - ctrl.total_pool_size_bytes = int(pool_mb * 1024**2) - ctrl.cpu_pool = torch.empty( - ctrl.total_pool_size_bytes, dtype=torch.uint8, pin_memory=False - ) - - # Mock the mooncake store — no real RDMA + ctrl.max_eviction_batch = 10 ctrl.mooncake_store = MagicMock() - ctrl.mooncake_store.register_buffer = MagicMock() - - ctrl.allocator = ContiguousMemoryAllocator(ctrl.total_pool_size_bytes) - ctrl.hash_to_metadata = {} - ctrl.access_order = {} - ctrl.access_lock = threading.Lock() - ctrl.ref_counts = {} - + ctrl.total_pool_size_bytes = num_pages * page_size * dim * torch.float32.itemsize + ctrl.vision_pool = _make_pool(num_pages, dim, page_size) + ctrl.audio_pool = _make_pool(num_pages, dim, page_size, modality="audio") + ctrl.pools = {"vision": ctrl.vision_pool, "audio": ctrl.audio_pool} + ctrl.entries = {} + ctrl.vision_pool.evictable = EvictableLRU() + ctrl.audio_pool.evictable = EvictableLRU() ctrl.stats = { "total_allocated": 0, "total_evicted": 0, "eviction_count": 0, "allocation_failures": 0, } - ctrl.ongoing_prefetch = {} - ctrl.prefetch_queue = MagicMock() - ctrl.insert_queue = MagicMock() - + ctrl.prefetch_queue = Queue() + ctrl.insert_queue = Queue() ctrl.lock = threading.Lock() ctrl.stop_event = threading.Event() - - # Do NOT start the IO thread — tests drive _io_loop logic manually ctrl.io_thread = MagicMock() - ctrl.prefetch_tp_group = None + ctrl._copy_streams = {} return ctrl -def _embedding_bytes(num_tokens, dim): - return num_tokens * dim * torch.float32.itemsize - - -# --------------------------------------------------------------------------- -# LRU eviction tests -# --------------------------------------------------------------------------- - - -class TestLRUEviction(unittest.TestCase): - def test_evict_oldest_first(self): - dim = 64 - size = _embedding_bytes(1, dim) # 256 bytes - # Pool exactly fits 3 entries (768 bytes) - pool_bytes = size * 3 - ctrl = _make_controller(pool_mb=pool_bytes / (1024**2)) - - # Insert 3 entries — fills the pool - for i in range(3): - h = f"hash_{i}" - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - self.assertIsNotNone(offset) - ctrl.hash_to_metadata[h] = (offset, 1, dim, size) - ctrl._update_access_time(h) - - # Pool is full. Inserting a 4th should evict hash_0 (oldest). - h_new = "hash_new" - with ctrl.lock: - offset = ctrl._allocate_with_eviction(size) - - self.assertIsNotNone(offset) - with ctrl.lock: - self.assertNotIn("hash_0", ctrl.hash_to_metadata) - self.assertIn("hash_1", ctrl.hash_to_metadata) - self.assertIn("hash_2", ctrl.hash_to_metadata) - - def test_eviction_disabled(self): - dim = 64 - size = _embedding_bytes(1, dim) # 256 bytes - # Pool exactly fits 1 entry - ctrl = _make_controller(pool_mb=size / (1024**2), enable_eviction=False) - - # Fill the pool - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - self.assertIsNotNone(offset) - - # Try to allocate more — should fail without eviction - with ctrl.lock: - offset2 = ctrl._allocate_with_eviction(size) - self.assertIsNone(offset2) - - def test_access_time_updates_prevent_eviction(self): - ctrl = _make_controller(pool_mb=0.01) - dim = 64 - size = _embedding_bytes(1, dim) - - # Insert 2 entries - hashes = [] - for i in range(2): - h = f"hash_{i}" - hashes.append(h) - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata[h] = (offset, 1, dim, size) - ctrl._update_access_time(h) - - # Touch hash_0 to make it recently used - time.sleep(0.01) - with ctrl.lock: - ctrl._update_access_time("hash_0") - - # Trigger eviction — hash_1 should be evicted (older) - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(size) - self.assertIn("hash_1", candidates) - self.assertNotIn("hash_0", candidates) - - def test_eviction_stats(self): - ctrl = _make_controller(pool_mb=0.01) - dim = 64 - size = _embedding_bytes(1, dim) - - # Insert and then evict - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) - ctrl._update_access_time("h") - - with ctrl.lock: - freed = ctrl._evict_hashes(["h"]) - self.assertGreater(freed, 0) - self.assertEqual(ctrl.stats["eviction_count"], 1) - self.assertGreater(ctrl.stats["total_evicted"], 0) - - def test_evict_nonexistent_hash(self): - ctrl = _make_controller() - with ctrl.lock: - freed = ctrl._evict_hashes(["nonexistent"]) - self.assertEqual(freed, 0) - - def test_max_eviction_batch(self): - ctrl = _make_controller(pool_mb=1.0, max_eviction_batch=2) - dim = 64 - size = _embedding_bytes(1, dim) - - # Insert many small entries - for i in range(10): - h = f"hash_{i}" - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - if offset is None: - break - ctrl.hash_to_metadata[h] = (offset, 1, dim, size) - ctrl._update_access_time(h) - - # _select_eviction_candidates should return at most 2 - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(size * 100) - self.assertLessEqual(len(candidates), 2) - - -# --------------------------------------------------------------------------- -# Ref counting tests -# --------------------------------------------------------------------------- - - -class TestRefCounting(unittest.TestCase): - def test_protect_prevents_eviction(self): - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) - - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) - ctrl._update_access_time("h") - ctrl._protect_hash("h") - - # Should not be selected for eviction - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(size) - self.assertNotIn("h", candidates) - - # Release and retry - with ctrl.lock: - ctrl._release_hash("h") - candidates = ctrl._select_eviction_candidates(size) - self.assertIn("h", candidates) - - def test_evict_hashes_skips_protected(self): - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) - - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) - ctrl._update_access_time("h") - ctrl._protect_hash("h") - - # Attempt to evict — should be skipped - with ctrl.lock: - freed = ctrl._evict_hashes(["h"]) - self.assertEqual(freed, 0) - with ctrl.lock: - self.assertIn("h", ctrl.hash_to_metadata) - - def test_ref_count_multiple_protects(self): - ctrl = _make_controller() - with ctrl.lock: - ctrl._protect_hash("h") - ctrl._protect_hash("h") - self.assertEqual(ctrl.ref_counts["h"], 2) - - ctrl._release_hash("h") - self.assertEqual(ctrl.ref_counts["h"], 1) - - # Still protected - candidates = ctrl._select_eviction_candidates(1) - self.assertNotIn("h", candidates) - - ctrl._release_hash("h") - self.assertNotIn("h", ctrl.ref_counts) - - def test_release_nonexistent_is_safe(self): - ctrl = _make_controller() - with ctrl.lock: - ctrl._release_hash("nonexistent") # should not raise - - def test_prefetch_sets_ref_count(self): - """prefetch() should set ref_count=1 for each new entry.""" - ctrl = _make_controller() - dim = 64 - ctrl.hidden_dims = {"image": dim} - h = "img_hash_1" - - ctrl.prefetch("req1", [h], [1], modality="image") - - with ctrl.lock: - self.assertEqual(ctrl.ref_counts.get(h), 1) - - # Simulate RDMA completion - with ctrl.lock: - ctrl._release_hash(h) - self.assertNotIn(h, ctrl.ref_counts) - - def test_insert_batch_sets_ref_count(self): - """insert_batch() should set ref_count=1 for each new entry.""" - ctrl = _make_controller() - dim = 64 - h = "img_hash_1" - tensor = torch.randn(1, dim) - - ctrl.insert_batch([h], [tensor]) - - with ctrl.lock: - self.assertEqual(ctrl.ref_counts.get(h), 1) - - # Simulate RDMA completion - with ctrl.lock: - ctrl._release_hash(h) - self.assertNotIn(h, ctrl.ref_counts) - - def test_get_embeddings_sets_ref_count(self): - """get_embeddings() should set ref_count=1 per hash.""" - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) - - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) - ctrl._update_access_time("h1") - - tensors = ctrl.get_embeddings(["h1"]) - self.assertIsNotNone(tensors[0]) - - with ctrl.lock: - self.assertEqual(ctrl.ref_counts.get("h1"), 1) - - # Protected — eviction should skip - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(size) - self.assertNotIn("h1", candidates) - - # Release - ctrl.release_embeddings(["h1"]) - with ctrl.lock: - self.assertNotIn("h1", ctrl.ref_counts) - - def test_get_embeddings_missing_hash(self): - """Missing hashes should return None and not set ref_count.""" - ctrl = _make_controller() - tensors = ctrl.get_embeddings(["nonexistent"]) - self.assertIsNone(tensors[0]) - with ctrl.lock: - self.assertNotIn("nonexistent", ctrl.ref_counts) - - def test_release_embeddings_missing_hash_is_safe(self): - """Releasing a hash that was never protected should be a no-op.""" - ctrl = _make_controller() - ctrl.release_embeddings(["nonexistent"]) # should not raise - - def test_io_loop_releases_prefetch_ref(self): - """_io_loop should release ref_count after batch_get completes.""" - ctrl = _make_controller() - dim = 64 - ctrl.hidden_dims = {"image": dim} - h = "img_hash_1" - - ctrl.prefetch("req1", [h], [1], modality="image") - - with ctrl.lock: - self.assertEqual(ctrl.ref_counts.get(h), 1) - - # Simulate _io_loop completing the RDMA GET - op = ctrl.ongoing_prefetch.get("req1") - self.assertIsNotNone(op) - ctrl.mooncake_store.batch_get = MagicMock(return_value=[True]) - - # Manually execute what _io_loop does for prefetch - results = ctrl.mooncake_store.batch_get(op.keys, op.ptrs, op.sizes) - op.mark_done(all(results)) - with ctrl.lock: - for k in op.keys: - ctrl._release_hash(k) - - with ctrl.lock: - self.assertNotIn(h, ctrl.ref_counts) - - def test_io_loop_releases_insert_ref(self): - """_io_loop should release ref_count after batch_put completes.""" - ctrl = _make_controller() - dim = 64 - h = "img_hash_1" - tensor = torch.randn(1, dim) - - ctrl.insert_batch([h], [tensor]) - - with ctrl.lock: - self.assertEqual(ctrl.ref_counts.get(h), 1) - - # Get the enqueued insert operation - ctrl.insert_queue.put.assert_called_once() - insert_op = ctrl.insert_queue.put.call_args[0][0] - self.assertIsInstance(insert_op, EmbeddingInsertOperation) - - # Simulate _io_loop completing the RDMA PUT - ctrl.mooncake_store.batch_put = MagicMock() - ctrl.mooncake_store.batch_put(insert_op.keys, insert_op.ptrs, insert_op.sizes) - with ctrl.lock: - for k in insert_op.keys: - ctrl._release_hash(k) - - with ctrl.lock: - self.assertNotIn(h, ctrl.ref_counts) - - -# --------------------------------------------------------------------------- -# Race condition prevention tests -# --------------------------------------------------------------------------- - - -class TestRDMAEvictionRacePrevention(unittest.TestCase): - def test_eviction_during_prefetch_is_blocked(self): - """An entry with in-flight RDMA GET cannot be evicted.""" - ctrl = _make_controller(pool_mb=0.01) - dim = 64 - ctrl.hidden_dims = {"image": dim} - size = _embedding_bytes(1, dim) - - # Prefetch sets ref_count=1 - ctrl.prefetch("req1", ["h1"], [1], modality="image") - - # Now try to evict to make room — should skip h1 - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(size * 10) - self.assertNotIn("h1", candidates) - - # Direct eviction should also be skipped - with ctrl.lock: - freed = ctrl._evict_hashes(["h1"]) - self.assertEqual(freed, 0) - - def test_eviction_during_insert_is_blocked(self): - """An entry with in-flight RDMA PUT cannot be evicted.""" - ctrl = _make_controller(pool_mb=0.01) - dim = 64 - tensor = torch.randn(1, dim) - - ctrl.insert_batch(["h1"], [tensor]) - - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(999999) - self.assertNotIn("h1", candidates) - - with ctrl.lock: - freed = ctrl._evict_hashes(["h1"]) - self.assertEqual(freed, 0) - - def test_eviction_during_get_embeddings_is_blocked(self): - """An entry returned by get_embeddings() cannot be evicted.""" - ctrl = _make_controller(pool_mb=0.01) - dim = 64 - size = _embedding_bytes(1, dim) - - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) - ctrl._update_access_time("h1") - - tensors = ctrl.get_embeddings(["h1"]) - self.assertIsNotNone(tensors[0]) - - # Try to evict — should be blocked - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(999999) - self.assertNotIn("h1", candidates) - - # Release and verify eviction is now possible - ctrl.release_embeddings(["h1"]) - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(999999) - self.assertIn("h1", candidates) - - def test_concurrent_eviction_while_reading(self): - """Simulate a concurrent eviction attempt while a read holds a ref.""" - ctrl = _make_controller(pool_mb=0.05) - dim = 64 - size = _embedding_bytes(1, dim) - num_entries = 10 - - # Insert entries - for i in range(num_entries): - h = f"hash_{i}" - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - if offset is None: - break - ctrl.hash_to_metadata[h] = (offset, 1, dim, size) - ctrl._update_access_time(h) - - # Simulate get_embeddings holding refs on hash_0..hash_4 - held_hashes = [f"hash_{i}" for i in range(5)] - tensors = ctrl.get_embeddings(held_hashes) - - # Try to evict all — only unprotected entries should be candidates - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(999999) - for h in held_hashes: - self.assertNotIn(h, candidates) - - # Unprotected entries should be candidates - for i in range(5, num_entries): - h = f"hash_{i}" - if h in ctrl.hash_to_metadata: - self.assertIn(h, candidates) - - # Release refs - ctrl.release_embeddings(held_hashes) - with ctrl.lock: - candidates = ctrl._select_eviction_candidates(999999) - for h in held_hashes: - if h in ctrl.hash_to_metadata: - self.assertIn(h, candidates) - - def test_evict_hashes_cleans_up_ref_counts(self): - """After eviction, ref_counts for the evicted hash should be removed.""" - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) - - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) - ctrl._update_access_time("h") - # Stale ref_count (shouldn't happen normally, but test cleanup) - ctrl.ref_counts["h"] = 0 - - with ctrl.lock: - # ref_count is 0, so eviction should proceed - freed = ctrl._evict_hashes(["h"]) - self.assertGreater(freed, 0) - with ctrl.lock: - self.assertNotIn("h", ctrl.hash_to_metadata) - self.assertNotIn("h", ctrl.ref_counts) - - -# --------------------------------------------------------------------------- -# get_embeddings view safety tests -# --------------------------------------------------------------------------- - - -class TestGetEmbeddingsViewSafety(unittest.TestCase): - def test_get_embeddings_returns_view(self): - """get_embeddings returns a view into cpu_pool, not a copy.""" - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) - - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) - ctrl._update_access_time("h1") - - tensors = ctrl.get_embeddings(["h1"]) - self.assertIsNotNone(tensors[0]) - self.assertEqual(tensors[0].shape, (1, dim)) - - # Verify it's a view into cpu_pool (shares storage) - self.assertTrue( - tensors[0].storage().data_ptr() == ctrl.cpu_pool.storage().data_ptr() +class TestRangePageAllocator(unittest.TestCase): + def test_prefers_single_contiguous_run(self): + allocator = RangePageAllocator(num_pages=8) + + runs = allocator.allocate(num_tokens=6, page_size=2) + + self.assertEqual(runs, [PageRun(start=0, length=3)]) + self.assertEqual(allocator.free_ranges, [(3, 5)]) + + def test_free_merges_adjacent_ranges(self): + allocator = RangePageAllocator(num_pages=8) + first = allocator.allocate(num_tokens=4, page_size=2) + second = allocator.allocate(num_tokens=4, page_size=2) + + allocator.free(first) + allocator.free(second) + + self.assertEqual(allocator.free_ranges, [(0, 8)]) + + def test_scatter_fallback_returns_physical_order(self): + allocator = RangePageAllocator(num_pages=4) + a = allocator.allocate(num_tokens=2, page_size=2) + b = allocator.allocate(num_tokens=2, page_size=2) + c = allocator.allocate(num_tokens=2, page_size=2) + d = allocator.allocate(num_tokens=2, page_size=2) + allocator.free(a) + allocator.free(c) + + runs = allocator.allocate(num_tokens=4, page_size=2) + + self.assertEqual(runs, [PageRun(start=0, length=1), PageRun(start=2, length=1)]) + self.assertEqual([run.start for run in runs], sorted(run.start for run in runs)) + self.assertEqual(b, [PageRun(start=1, length=1)]) + self.assertEqual(d, [PageRun(start=3, length=1)]) + + def test_allocate_fails_when_total_free_pages_are_insufficient(self): + allocator = RangePageAllocator(num_pages=2) + allocator.allocate(num_tokens=4, page_size=2) + + self.assertIsNone(allocator.allocate(num_tokens=2, page_size=2)) + self.assertEqual(allocator.free_pages, 0) + + +class TestEntryStateAndPins(unittest.TestCase): + def test_ready_entry_with_no_pins_is_evictable(self): + entry = EmbeddingCacheEntry( + hash="h", + modality=Modality.IMAGE, + num_tokens=2, + dim=4, + page_runs=[PageRun(0, 1)], + state=EntryState.READY, ) - # Release - ctrl.release_embeddings(["h1"]) + self.assertTrue(entry.is_evictable()) - def test_data_preserved_while_ref_held(self): - """Data should remain intact as long as ref_count > 0.""" - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) + def test_filling_entry_is_not_evictable(self): + entry = EmbeddingCacheEntry( + hash="h", + modality=Modality.IMAGE, + num_tokens=2, + dim=4, + page_runs=[PageRun(0, 1)], + state=EntryState.FILLING, + ) - # Write known data - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) - ctrl._update_access_time("h1") - view = ( - ctrl.cpu_pool[offset : offset + size].view(torch.float32).view(1, dim) - ) - view.copy_(torch.ones(1, dim)) + self.assertFalse(entry.is_evictable()) - # Read via get_embeddings (holds ref) - tensors = ctrl.get_embeddings(["h1"]) - self.assertTrue(torch.all(tensors[0] == 1.0)) + def test_ready_entry_with_pin_is_not_evictable(self): + entry = EmbeddingCacheEntry( + hash="h", + modality=Modality.IMAGE, + num_tokens=2, + dim=4, + page_runs=[PageRun(0, 1)], + state=EntryState.READY, + ) - # Verify data is still valid - self.assertTrue(torch.all(tensors[0] == 1.0)) + entry.pin() + self.assertFalse(entry.is_evictable()) + entry.unpin() + self.assertTrue(entry.is_evictable()) - # Release - ctrl.release_embeddings(["h1"]) + def test_multiple_pins_require_all_unpins(self): + entry = EmbeddingCacheEntry( + hash="h", + modality=Modality.IMAGE, + num_tokens=2, + dim=4, + page_runs=[PageRun(0, 1)], + state=EntryState.READY, + ) + + entry.pin() + entry.pin() + + self.assertEqual(entry.ref_count, 2) + self.assertFalse(entry.is_evictable()) -# --------------------------------------------------------------------------- -# Stats tests -# --------------------------------------------------------------------------- +class TestEvictableLruInvariant(unittest.TestCase): + def _insert_entry( + self, + ctrl, + mm_hash, + modality=Modality.IMAGE, + state=EntryState.READY, + ): + pool = ctrl._get_pool(modality) + page_runs = pool.allocator.allocate(2, pool.page_size) + entry = EmbeddingCacheEntry( + hash=mm_hash, + modality=modality, + num_tokens=2, + dim=pool.dim, + page_runs=page_runs, + state=state, + ) + ctrl.entries[mm_hash] = entry + return entry + def test_filling_entry_is_not_in_evictable_lru_until_ready(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + entry = self._insert_entry( + ctrl, + "h", + state=EntryState.FILLING, + ) -class TestGetStats(unittest.TestCase): - def test_stats_include_num_protected(self): - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) + self.assertNotIn("h", ctrl.vision_pool.evictable) with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) - ctrl._update_access_time("h1") - ctrl._protect_hash("h1") - ctrl._protect_hash("h2") # h2 not in metadata, but has ref + ctrl._mark_ready(entry) - stats = ctrl.get_stats() - self.assertEqual(stats["num_protected"], 2) + self.assertEqual(list(ctrl.vision_pool.evictable.keys()), ["h"]) - def test_stats_eviction_tracking(self): - ctrl = _make_controller(pool_mb=0.01) - dim = 64 - size = _embedding_bytes(1, dim) + def test_first_read_pin_removes_candidate_and_last_release_reinserts(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + entry = self._insert_entry(ctrl, "h") + with ctrl.lock: + ctrl._lru_touch("h") with ctrl.lock: - offset = ctrl.allocator.allocate(size) - ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) - ctrl._update_access_time("h") + ctrl._pin_read(entry) + ctrl._pin_read(entry) + + self.assertNotIn("h", ctrl.vision_pool.evictable) with ctrl.lock: - ctrl._evict_hashes(["h"]) + ctrl._unpin_read(entry) + self.assertNotIn("h", ctrl.vision_pool.evictable) - stats = ctrl.get_stats() - self.assertEqual(stats["eviction_count"], 1) - self.assertGreater(stats["total_evicted"], 0) - self.assertEqual(stats["num_cached"], 0) - - -# --------------------------------------------------------------------------- -# _select_eviction_candidates iterator safety test -# --------------------------------------------------------------------------- - - -class TestEvictionCandidateIteratorSafety(unittest.TestCase): - def test_list_snapshot_prevents_concurrent_mutation(self): - """sorted_hashes should be a list snapshot, not a live dict view.""" - ctrl = _make_controller() - dim = 64 - size = _embedding_bytes(1, dim) - - # Insert entries - for i in range(5): - h = f"hash_{i}" - with ctrl.lock: - offset = ctrl.allocator.allocate(size) - if offset is None: - break - ctrl.hash_to_metadata[h] = (offset, 1, dim, size) - ctrl._update_access_time(h) - - # _select_eviction_candidates should work even if access_order - # is modified during iteration (the snapshot via list() prevents this) with ctrl.lock: - candidates = ctrl._select_eviction_candidates(size) - # Should return candidates without RuntimeError - self.assertIsInstance(candidates, list) + ctrl._unpin_read(entry) + self.assertEqual(list(ctrl.vision_pool.evictable.keys()), ["h"]) + + def test_evict_for_pool_pops_only_that_pool_candidates(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + self._insert_entry(ctrl, "vision_h", modality=Modality.IMAGE) + self._insert_entry(ctrl, "audio_h", modality=Modality.AUDIO) + with ctrl.lock: + ctrl._lru_touch("vision_h") + ctrl._lru_touch("audio_h") + + required_pages = ctrl.vision_pool.allocator.free_pages + 1 + with ctrl.lock: + ctrl._evict_for_pool(ctrl.vision_pool, required_pages) + + self.assertNotIn("vision_h", ctrl.entries) + self.assertIn("audio_h", ctrl.entries) + self.assertEqual(list(ctrl.audio_pool.evictable.keys()), ["audio_h"]) + + +class TestStoreToPool(unittest.TestCase): + def test_store_to_pool_async_raises_on_cpu_tensor(self): + ctrl = _make_controller(num_pages=8, dim=4, page_size=2) + tensor = torch.empty((2, 4), dtype=torch.float32) + + with self.assertRaises(ValueError): + ctrl.store_to_pool_async(["h"], [tensor], Modality.IMAGE) + + +def _insert_ready_entry(ctrl, mm_hash, tensor, modality=Modality.IMAGE): + """Manually write tensor into pool pages and create a READY entry.""" + pool = ctrl._get_pool(modality) + if tensor.ndim != 2: + tensor = tensor.reshape(-1, tensor.shape[-1]) + num_tokens = int(tensor.shape[0]) + page_runs = pool.allocator.allocate(num_tokens, pool.page_size) + copied = 0 + for run in page_runs: + valid = min(pool.page_size * run.length, num_tokens - copied) + start = run.start * pool.page_size + pool.tensor[start : start + valid].copy_(tensor[copied : copied + valid]) + copied += valid + entry = EmbeddingCacheEntry( + hash=mm_hash, + modality=modality, + num_tokens=num_tokens, + dim=int(tensor.shape[1]), + page_runs=page_runs, + state=EntryState.READY, + ) + ctrl.entries[mm_hash] = entry + pool.evictable.touch(mm_hash) + return entry + + +class TestMooncakeLifecycle(unittest.TestCase): + def test_prefetch_creates_filling_entry_and_get_success_marks_ready(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + + ctrl.prefetch("req", ["h"], [2], Modality.IMAGE) + op = ctrl.ongoing_prefetch["req"] + ctrl._finish_get(op, [True]) + + entry = ctrl.entries["h"] + self.assertEqual(entry.state, EntryState.READY) + + def test_prefetch_get_failure_frees_entry(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + + ctrl.prefetch("req", ["h"], [2], Modality.IMAGE) + op = ctrl.ongoing_prefetch["req"] + ctrl._finish_get(op, [False]) + + self.assertNotIn("h", ctrl.entries) + self.assertEqual(ctrl.vision_pool.allocator.free_pages, 4) + + def test_insert_batch_pins_and_releases_on_put(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4) + _insert_ready_entry(ctrl, "h", tensor) + + ctrl.insert_batch(["h"], Modality.IMAGE) + op = ctrl.insert_queue.get_nowait() + + entry = ctrl.entries["h"] + self.assertEqual(entry.ref_count, 1) + self.assertNotIn("h", ctrl.vision_pool.evictable) + + ctrl._finish_put(op, [True]) + + self.assertEqual(entry.ref_count, 0) + self.assertEqual(entry.state, EntryState.READY) + self.assertIn("h", ctrl.vision_pool.evictable) + + +class TestGetPoolViews(unittest.TestCase): + def test_get_pool_views_returns_none_for_filling_entry(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + ctrl.entries["h"] = EmbeddingCacheEntry( + hash="h", + modality=Modality.IMAGE, + num_tokens=2, + dim=4, + page_runs=[PageRun(0, 1)], + state=EntryState.FILLING, + ) + + views = ctrl.get_pool_views(["h"]) + self.assertIsNone(views[0]) + + def test_get_pool_views_returns_slices_and_release_unpins(self): + ctrl = _make_controller(num_pages=4, dim=4, page_size=2) + tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4) + _insert_ready_entry(ctrl, "h", tensor) + + views = ctrl.get_pool_views(["h"]) + self.assertIsNotNone(views[0]) + entry = ctrl.entries["h"] + self.assertEqual(entry.ref_count, 1) + + flat = torch.cat(views[0], dim=0) + self.assertTrue(torch.equal(flat, tensor)) + + ctrl.release_pool_views(["h"]) + self.assertEqual(entry.ref_count, 0) + self.assertIn("h", ctrl.vision_pool.evictable) + + +class TestTransferBuffers(unittest.TestCase): + def test_build_transfer_buffers_for_single_run(self): + pool = _make_pool(num_pages=8, dim=4, page_size=2) + entry = EmbeddingCacheEntry( + hash="h", + modality=Modality.IMAGE, + num_tokens=5, + dim=4, + page_runs=[PageRun(2, 3)], + state=EntryState.READY, + ) + + ptrs, sizes = build_transfer_buffers(entry, pool) + + self.assertEqual(ptrs, [pool.tensor[4].data_ptr()]) + self.assertEqual(sizes, [5 * 4 * torch.float32.itemsize]) + + def test_build_transfer_buffers_for_multiple_runs(self): + pool = _make_pool(num_pages=8, dim=4, page_size=2) + entry = EmbeddingCacheEntry( + hash="h", + modality=Modality.IMAGE, + num_tokens=5, + dim=4, + page_runs=[PageRun(0, 1), PageRun(3, 2)], + state=EntryState.READY, + ) + + ptrs, sizes = build_transfer_buffers(entry, pool) + + self.assertEqual(ptrs, [pool.tensor[0].data_ptr(), pool.tensor[6].data_ptr()]) + self.assertEqual( + sizes, + [ + 2 * 4 * torch.float32.itemsize, + 3 * 4 * torch.float32.itemsize, + ], + ) + + +class TestMooncakeEmbeddingStoreWrappers(unittest.TestCase): + def test_batch_put_multi_buffers_deduplicates_existing_keys(self): + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import ( + MooncakeEmbeddingStore, + ) + + store = MooncakeEmbeddingStore.__new__(MooncakeEmbeddingStore) + store.store = MagicMock() + store.store.batch_is_exist.return_value = [1, 0] + store.store.batch_put_from_multi_buffers.return_value = [0] + + results = store.batch_put_from_multi_buffers( + ["a", "b"], + [[11], [22]], + [[4], [4]], + ) + + self.assertEqual(results, [True, True]) + store.store.batch_put_from_multi_buffers.assert_called_once_with( + ["emb_b"], [[22]], [[4]] + ) + + def test_batch_get_multi_buffers_maps_positive_result_to_true(self): + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import ( + MooncakeEmbeddingStore, + ) + + store = MooncakeEmbeddingStore.__new__(MooncakeEmbeddingStore) + store.store = MagicMock() + store.store.batch_get_into_multi_buffers.return_value = [8, -1] + + results = store.batch_get_into_multi_buffers( + ["a", "b"], + [[11], [22]], + [[4], [4]], + ) + + self.assertEqual(results, [True, False]) if __name__ == "__main__":