[EPD] Optimize multimodal global cache with paged embedding pool (#28441)

Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
Co-authored-by: liusy58 <liusy58@linux.alibaba.com>

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:
+    ) -> GlobalCacheEncodeContext:
-        # mm_inputs: dict
         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.
+        miss_d2h_handles = []
-        fallback_mask = torch.zeros(num_items, dtype=torch.int32)
+        if self.rank == 0 and new_slices:
-        cached_slices = []
+            miss_hashes = [ctx.str_mm_hashes[i] for i in missing_indices]
-
+            miss_d2h_handles = self.mm_global_cache.store_to_pool_async(
-        if self.rank == 0:
+                miss_hashes, new_slices, ctx.modality
-            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,
             )
 
-        # Step 5: All ranks run ViT for items that need fallback recomputation.
+        fallback_indices = await self._wait_global_cache_prefetch(
-        fallback_indices = [i for i in range(num_items) if fallback_mask[i].item() == 1]
+            ctx, hit_indices, hit_hashes
-        fallback_slices = None
+        )
+
+        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.
+            new_hashes = [ctx.str_mm_hashes[i] for i in missing_indices]
-        if self.rank == 0:
+            new_hashes += [ctx.str_mm_hashes[i] for i in fallback_indices]
-            final_slices = [None] * num_items
+            self._launch_global_cache_insert(
+                ctx,
+                new_hashes,
+                miss_d2h_handles + fallback_d2h_handles,
+            )
 
-            for i, idx in enumerate(missing_indices):
+            self.embedding_to_send[ctx.req_id] = EmbeddingData(
-                final_slices[idx] = new_slices[i]
+                ctx.req_id,
-
-            # 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,
                 num_parts,
                 part_idx,
-                grid_thw,
+                ctx.grid_thw,
-                modality,
+                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)
+            ctx = await self._prepare_global_cache_context(
-            grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
+                mm_items, modality, req_id, hashes
-            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
-                )
             )
 
-            # Rank 0: compute hashes
+            nbytes, total_tokens, embedding_dim, event = (
-            if self.rank == 0:
+                self._setup_mooncake_async_encode(
-                if hashes is None:
+                    ctx.req_id,
-                    mm_hashes = self._calculate_hashes_from_features(
+                    num_parts,
-                        mm_feature, grid_thw, modality
+                    part_idx,
-                    )
+                    ctx.grid_thw,
-                else:
+                    ctx.modality,
-                    mm_hashes = hashes
+                    ctx.aux_data,
-                str_mm_hashes = [str(h) for h in mm_hashes]
+                )
+            )
 
             # 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
+                    missing_indices, hit_indices = await self._lookup_global_cache(ctx)
-                    if self.rank == 0:
+                    hit_hashes = self._prefetch_global_cache_hits(ctx, hit_indices)
-                        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)
 
-                    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,
+                            ctx.mm_feature,
-                            mm_inputs,
+                            ctx.mm_inputs,
                             missing_indices,
-                            modality,
+                            ctx.modality,
-                            get_feature_fn,
+                            ctx.get_feature_fn,
-                            grid_thw,
+                            ctx.grid_thw,
                             keep_on_gpu=True,
                         )
 
-                    # Step 3: Rank 0 prefetches cache-hit embeddings and builds fallback_mask.
+                    fallback_indices = await self._wait_global_cache_prefetch(
-                    fallback_mask = torch.zeros(num_items, dtype=torch.int32)
+                        ctx, hit_indices, hit_hashes
-                    cached_slices = []
+                    )
 
-                    if self.rank == 0 and hit_indices:
+                    fallback_slices = []
-                        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
                     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,
+                            ctx.mm_feature,
-                            mm_inputs,
+                            ctx.mm_inputs,
                             fallback_indices,
-                            modality,
+                            ctx.modality,
-                            get_feature_fn,
+                            ctx.get_feature_fn,
-                            grid_thw,
+                            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):
+                        d2h_handles = []
-                            final_slices[idx] = new_slices[i]
+                        if new_slices:
-
+                            miss_hashes = [
-                        # Fill in successfully loaded cache-hit embeddings
+                                ctx.str_mm_hashes[i] for i in missing_indices
-                        if cached_slices:
+                            ]
-                            for i, idx in enumerate(hit_indices):
+                            miss_handles = self.mm_global_cache.store_to_pool_async(
-                                if cached_slices[i] is not None:
+                                miss_hashes, new_slices, ctx.modality
-                                    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
                             )
-                            for s in final_slices
+                            d2h_handles.extend(miss_handles)
-                        ]
+                        if fallback_slices:
-                        mm_embedding = torch.cat(final_slices, dim=0)
+                            fallback_hashes = [
-                        # Wait for any pending VIT / cat kernels to finish
+                                ctx.str_mm_hashes[i] for i in fallback_indices
-                        # 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
                             ]
-                            if loaded_hashes:
+                            fb_handles = self.mm_global_cache.store_to_pool_async(
-                                self.mm_global_cache.release_embeddings(loaded_hashes)
+                                fallback_hashes, fallback_slices, ctx.modality
+                            )
+                            d2h_handles.extend(fb_handles)
 
-                        # Background insert: store newly computed embeddings into global cache.
+                        mm_embedding = self._assemble_global_cache_gpu(
-                        # Includes both original misses and fallback-recomputed hits.
+                            ctx,
-                        all_new_hashes = [str_mm_hashes[i] for i in missing_indices]
+                            missing_indices,
-                        all_new_slices = list(new_slices)
+                            fallback_indices,
-                        if fallback_slices is not None:
+                            new_slices,
-                            all_new_hashes += [
+                            fallback_slices,
-                                str_mm_hashes[i] for i in fallback_indices
+                        )
-                            ]
-                            all_new_slices += list(fallback_slices)
-                        if all_new_hashes:
 
-                            async def _background_insert():
+                        new_hashes = [ctx.str_mm_hashes[i] for i in missing_indices]
-                                await asyncio.to_thread(
+                        new_hashes += [ctx.str_mm_hashes[i] for i in fallback_indices]
-                                    self.mm_global_cache.insert_batch,
+                        self._launch_global_cache_insert(
-                                    all_new_hashes,
+                            ctx,
-                                    all_new_slices,
+                            new_hashes,
-                                )
+                            d2h_handles,
+                        )
 
-                            insert_task = asyncio.create_task(_background_insert())
+                        self._forward_results[ctx.req_id]["embedding"] = mm_embedding
-                            self.background_tasks.add(insert_task)
-                            insert_task.add_done_callback(self.background_tasks.discard)
-
-                        self._forward_results[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"
                 )
 

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:
+    def get_key(self, mm_hash: str) -> str:
-        return f"emb_{image_hash}"
+        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