[Score] Send stacked GPU embedding overrides through a CUDA IPC pool

python/sglang/srt/entrypoints/engine_score_mixin.py

@@ -40,6 +40,7 @@ def score(
         score_extraction_token_id: Optional[int] = None,
         temperature: float = 1.0,
         return_token_logprobs: bool = False,
+        stacked_query_item_embed_overrides: Optional[List[torch.Tensor]] = None,
     ) -> ScoreResult:
         """
         Score items against a query using the loaded model.
@@ -84,6 +85,9 @@ def score(
                 softmax. Requires apply_softmax=True when different from 1.
             return_token_logprobs: Include uncalibrated full-vocabulary logprobs for
                 each item's candidates. Only supported for CausalLM, not raw logits.
+            stacked_query_item_embed_overrides: Alternative to query_/item_embed_overrides:
+                one already-stacked [N, hidden] tensor per output sequence, used without
+                a restack. See score_request.
 
         Returns:
             ScoreResult with scores (flat ``[num_items x num_labels]`` for pointwise
@@ -102,6 +106,7 @@ def score(
                 query_embed_overrides=query_embed_overrides,
                 item_embed_overrides=item_embed_overrides,
                 score_extraction_token_id=score_extraction_token_id,
+                stacked_query_item_embed_overrides=stacked_query_item_embed_overrides,
                 request=None,
                 return_pooled_hidden_states=return_pooled_hidden_states,
                 temperature=temperature,
@@ -123,6 +128,7 @@ async def async_score(
         score_extraction_token_id: Optional[int] = None,
         temperature: float = 1.0,
         return_token_logprobs: bool = False,
+        stacked_query_item_embed_overrides: Optional[List[torch.Tensor]] = None,
     ) -> ScoreResult:
         """Asynchronous version of score(). See score() for full documentation."""
         return await self.tokenizer_manager.score_request(
@@ -135,6 +141,7 @@ async def async_score(
             query_embed_overrides=query_embed_overrides,
             item_embed_overrides=item_embed_overrides,
             score_extraction_token_id=score_extraction_token_id,
+            stacked_query_item_embed_overrides=stacked_query_item_embed_overrides,
             request=None,
             return_pooled_hidden_states=return_pooled_hidden_states,
             temperature=temperature,

python/sglang/srt/managers/embed_types.py

@@ -23,6 +23,8 @@
 import msgspec
 import torch
 
+from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
+
 
 class PositionalEmbeds(msgspec.Struct, array_like=True):
     """Embeddings to place at specific token positions.
@@ -36,7 +38,8 @@ class PositionalEmbeds(msgspec.Struct, array_like=True):
         positions: List of positions where embeddings should be injected.
     """
 
-    embeds: torch.Tensor
+    # torch.Tensor or CudaIpcTensorTransportProxy; msgspec cannot type that union.
+    embeds: object
     positions: List[int]
 
     def __post_init__(self):
@@ -51,8 +54,15 @@ def __post_init__(self):
             else:
                 # [1, hidden_dim] (already has the leading dim) → plain concat.
                 self.embeds = torch.cat(self.embeds, dim=0)
-        if self.embeds.shape[0] != len(self.positions):
+        if isinstance(self.embeds, torch.Tensor) and self.embeds.shape[0] != len(
+            self.positions
+        ):
             raise ValueError(
                 f"embeds length ({self.embeds.shape[0]}) != "
                 f"positions length ({len(self.positions)})"
             )
+
+    def reconstruct(self, target_device: int) -> None:
+        """Copy embeds out of the CUDA IPC pool onto target_device and release the slot."""
+        if isinstance(self.embeds, CudaIpcTensorTransportProxy):
+            self.embeds = self.embeds.reconstruct_on_target_device(target_device)

python/sglang/srt/managers/positional_embeds_transport.py

@@ -0,0 +1,94 @@
+"""Send CUDA override embeds to the scheduler through a CUDA IPC pool, not by value."""
+
+import logging
+from typing import Iterable, List, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.managers.embed_types import PositionalEmbeds
+from sglang.srt.multimodal.transport.cuda_ipc import (
+    MM_FEATURE_CACHE_SIZE,
+    MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
+    CudaIpcTensorTransportProxy,
+    MmItemMemoryPool,
+    get_mm_feature_pool_size_per_worker,
+)
+from sglang.srt.multimodal.transport.memory_pool import get_mm_feature_consumer_count
+from sglang.srt.runtime_context import get_parallel, get_serving
+from sglang.srt.utils import is_cuda
+
+logger = logging.getLogger(__name__)
+
+
+class PositionalEmbedsTransport:
+    def __init__(self, base_gpu_id: int, mm_processor) -> None:
+        parallel = get_parallel()
+        # By value otherwise: IPC needs one node, every pipeline stage would ack the
+        # slot, and the msgpack IPC encoder cannot encode the proxy.
+        self.enabled = (
+            is_cuda()
+            and parallel.nnodes == 1
+            and parallel.pp_size == 1
+            and envs.SGLANG_USE_PICKLE_IPC.get()
+        )
+        self._base_gpu_id = base_gpu_id
+        # A multimodal model with --mm-feature-transport=cuda_ipc already has a pool.
+        self._pool: Optional[MmItemMemoryPool] = getattr(
+            mm_processor, "cudaipc_mmfeature_pool", None
+        )
+        self._use_pool_handle_cache = envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.get()
+
+    def _get_pool(self) -> MmItemMemoryPool:
+        if self._pool is None:
+            pool_size = get_mm_feature_pool_size_per_worker(
+                MM_FEATURE_CACHE_SIZE, get_serving().tokenizer_worker_num
+            )
+            logger.info(
+                "Creating a %.0f MiB CUDA IPC pool on GPU %d for GPU-resident "
+                "embedding overrides (SGLANG_MM_FEATURE_CACHE_MB).",
+                pool_size / (1024 * 1024),
+                self._base_gpu_id,
+            )
+            self._pool = MmItemMemoryPool(
+                pool_size,
+                MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
+                self._base_gpu_id,
+                get_mm_feature_consumer_count(),
+            )
+        return self._pool
+
+    def prepare_for_dispatch(
+        self, tokenized_objs: Iterable
+    ) -> List[CudaIpcTensorTransportProxy]:
+        """Returns the proxies to cancel if dispatch fails; a full pool means by value."""
+        if not self.enabled:
+            return []
+        proxies = []
+        try:
+            for tokenized_obj in tokenized_objs:
+                overrides = tokenized_obj.positional_embed_overrides
+                if overrides is None or not (
+                    isinstance(overrides.embeds, torch.Tensor)
+                    and overrides.embeds.is_cuda
+                    and overrides.embeds.numel() > 0
+                ):
+                    continue
+                proxy = self._get_pool().wrap_tensor(
+                    overrides.embeds, use_pool_handle_cache=self._use_pool_handle_cache
+                )
+                if proxy is None:
+                    continue
+                proxies.append(proxy)
+                # A new object: requests may share the caller's PositionalEmbeds.
+                tokenized_obj.positional_embed_overrides = PositionalEmbeds(
+                    embeds=proxy, positions=overrides.positions
+                )
+        except BaseException:
+            self.cancel_for_dispatch(proxies)
+            raise
+        return proxies
+
+    def cancel_for_dispatch(self, proxies: List[CudaIpcTensorTransportProxy]) -> None:
+        for proxy in proxies:
+            self._pool.cancel_proxy(proxy)

python/sglang/srt/managers/schedule_batch.py

@@ -1367,7 +1367,11 @@ def __init__(
         # retracted request is rebootstrapped. Set in pause_generation(retract)
         # and consumed in the decode transfer commit; never plumbed to prefill.
         self.pd_rebootstrap_forced_output_id: Optional[int] = None
-        self.skip_radix_cache_insert = bootstrap_host == FAKE_BOOTSTRAP_HOST
+        # Override embeds give KV that the token ids alone do not determine.
+        self.skip_radix_cache_insert = (
+            bootstrap_host == FAKE_BOOTSTRAP_HOST
+            or positional_embed_overrides is not None
+        )
         # Set when the request's in-flight forward ran on KV that never landed:
         # that forward's output is dropped and the request finishes with it.
         self.discard_output_reason: Optional[FINISH_ABORT] = None

python/sglang/srt/managers/scheduler.py

@@ -2140,6 +2140,13 @@ def process_input_requests(self, recv_reqs: List):
         self.session_controller.maybe_reap(now)
 
         for recv_req in recv_reqs:
+            # Before any early exit, so every request releases its CUDA IPC pool slot.
+            for tokenized_req in self._tokenized_requests(recv_req):
+                if tokenized_req.positional_embed_overrides is not None:
+                    tokenized_req.positional_embed_overrides.reconstruct(
+                        get_device().gpu_id
+                    )
+
             vmm_errors = None
             if get_mm().mm_feature_transport == "cuda_vmm":
                 vmm_errors = self._materialize_cuda_vmm_inputs(recv_req)

python/sglang/srt/managers/tokenizer_manager.py

@@ -110,6 +110,7 @@
 from sglang.srt.managers.load_snapshot import create_load_snapshot_reader
 from sglang.srt.managers.mm_utils import wrap_shm_features
 from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
+from sglang.srt.managers.positional_embeds_transport import PositionalEmbedsTransport
 from sglang.srt.managers.schedule_batch import (
     MultimodalDataItem,
     get_request_return_hidden_states_mode,
@@ -539,6 +540,9 @@ def __init__(
         # Init request dispatcher
         self.init_request_dispatcher()
 
+        # Init the CUDA IPC transport for GPU-resident embedding overrides
+        self.init_positional_embeds_transport()
+
         # Construct this last so later initialization failures cannot orphan
         # the transport's recycler thread.
         self.cuda_vmm_feature_transport = CudaVmmFeatureTransport(
@@ -837,6 +841,11 @@ def set_startup_time(self, startup_time: Dict[str, Any]) -> None:
         if self.enable_metrics:
             self.metrics_collector.emit_startup_time(startup_time)
 
+    def init_positional_embeds_transport(self):
+        self.positional_embeds_transport = PositionalEmbedsTransport(
+            self.server_args.base_gpu_id, self.mm_processor
+        )
+
     def init_request_dispatcher(self):
         self._result_dispatcher = TypeBasedDispatcher(
             [
@@ -1730,13 +1739,17 @@ async def _send_one_request(
         tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput],
     ):
         prepared_mm_items = []
+        prepared_embeds = []
         dispatched = False
         try:
             prepared_mm_items = (
                 await self.cuda_vmm_feature_transport.prepare_for_dispatch_async(
                     (tokenized_obj.mm_inputs,)
                 )
             )
+            prepared_embeds = self.positional_embeds_transport.prepare_for_dispatch(
+                (tokenized_obj,)
+            )
             tokenized_obj.time_stats.set_api_server_dispatch_time()
             tokenized_obj = wrap_shm_features(tokenized_obj)
             time_stats = tokenized_obj.time_stats
@@ -1752,6 +1765,7 @@ async def _send_one_request(
         finally:
             if not dispatched:
                 self.cuda_vmm_feature_transport.cancel_for_dispatch(prepared_mm_items)
+                self.positional_embeds_transport.cancel_for_dispatch(prepared_embeds)
 
     def _mark_state_dispatched(self, rid: str):
         """Record that *rid* reached the scheduler.
@@ -1771,13 +1785,17 @@ async def _send_batch_request(
     ):
         """Send a batch of tokenized requests as a single batched request to the scheduler."""
         prepared_mm_items = []
+        prepared_embeds = []
         dispatched = False
         try:
             prepared_mm_items = (
                 await self.cuda_vmm_feature_transport.prepare_for_dispatch_async(
                     tokenized_obj.mm_inputs for tokenized_obj in tokenized_objs
                 )
             )
+            prepared_embeds = self.positional_embeds_transport.prepare_for_dispatch(
+                tokenized_objs
+            )
 
             set_time_batch(tokenized_objs, "set_api_server_dispatch_time")
             time_stats = [tokenized_obj.time_stats for tokenized_obj in tokenized_objs]
@@ -1799,6 +1817,7 @@ async def _send_batch_request(
         finally:
             if not dispatched:
                 self.cuda_vmm_feature_transport.cancel_for_dispatch(prepared_mm_items)
+                self.positional_embeds_transport.cancel_for_dispatch(prepared_embeds)
 
     def _coalesce_streaming_chunks(
         self,