[vlm] fix: contain multimodal feature transport failures (#37047)

This commit is contained in:
Mick
2026-09-01 13:46:38 +08:00
committed by GitHub
parent 33ed29a0ee
commit ae2bd5728b
16 changed files with 1408 additions and 96 deletions
+6
View File
@@ -105,6 +105,12 @@ class BaseBatchReq(msgspec.Struct, tag=True, kw_only=True, array_like=True):
return msgspec_struct_pydantic_core_schema(cls, handler)
class MMInputsProcessError(msgspec.Struct, frozen=True):
"""Request-local multimodal input failure produced after tokenizer fanout."""
message: str
class BeamSearchOutput(BaseBatchReq, kw_only=True):
sequences: List[BeamSearchSequence]
+98 -22
View File
@@ -36,6 +36,7 @@ from sglang.srt.managers.schedule_batch import (
CudaIpcTensorTransportProxy,
Modality,
MultimodalInputs,
MultimodalProcessorOutput,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.multimodal.transport import (
@@ -1285,6 +1286,9 @@ class ShmPointerMMData:
"""
def __init__(self, tensor: torch.Tensor, precomputed_hash: Optional[int] = None):
self._shm_handle = None
self.tensor = None
self._materialization_error = None
if not tensor.is_cpu:
tensor = tensor.cpu()
if not tensor.is_contiguous():
@@ -1311,7 +1315,6 @@ class ShmPointerMMData:
raise
self.shm_name = shm.name
shm.close()
self._shm_handle = None
def __getstate__(self):
return {
@@ -1326,27 +1329,78 @@ class ShmPointerMMData:
self.shape = state["shape"]
self.dtype = state["dtype"]
self.precomputed_hash = state.get("precomputed_hash")
self._shm_handle = shared_memory.SharedMemory(name=self.shm_name)
# Zero-copy view into shared memory (no clone, no unlink)
self.tensor = torch.frombuffer(self._shm_handle.buf, dtype=self.dtype).reshape(
self.shape
)
self._shm_handle = None
self.tensor = None
self._materialization_error = None
# keep deserialization infallible so all TP ranks finish the broadcast
handle = None
tensor = None
try:
handle = shared_memory.SharedMemory(name=self.shm_name)
tensor = torch.frombuffer(handle.buf, dtype=self.dtype)
self.tensor = tensor.reshape(self.shape)
self._shm_handle = handle
except Exception as error:
tensor = None
if handle is not None:
try:
handle.close()
except Exception:
logger.warning(
"Failed to close a malformed multimodal SHM handle",
exc_info=True,
)
self._materialization_error = f"{type(error).__name__}: {error}"
def materialize(self) -> torch.Tensor:
"""Clone tensor from shm to owned memory, then release shm handle."""
tensor = self.tensor.clone()
if self._shm_handle is not None:
self._shm_handle.close()
try:
if self._materialization_error is not None:
raise RuntimeError(self._materialization_error)
return self.tensor.clone()
finally:
self.close_and_unlink()
def close_and_unlink(self) -> None:
"""Release this rank's view and unlink the shared feature segment."""
handle = self._shm_handle
self._shm_handle = None
self.tensor = None
if handle is None:
try:
self._shm_handle.unlink()
handle = shared_memory.SharedMemory(name=self.shm_name)
except FileNotFoundError:
pass # Another rank already unlinked
self._shm_handle = None
return tensor
return
except OSError:
logger.warning(
"Failed to reopen a multimodal SHM segment for cleanup",
exc_info=True,
)
return
try:
try:
handle.unlink()
except FileNotFoundError:
pass
except OSError:
logger.warning(
"Failed to unlink a multimodal SHM segment",
exc_info=True,
)
finally:
try:
handle.close()
except Exception:
logger.warning(
"Failed to close a multimodal SHM handle",
exc_info=True,
)
def __del__(self):
# Only close; never unlink. Unlinking is materialize()'s job.
if getattr(self, "_shm_handle", None) is not None:
if self._shm_handle is not None:
self.tensor = None
self._shm_handle.close()
self._shm_handle = None
@@ -1427,10 +1481,9 @@ def has_shm_features(recv_reqs):
if isinstance(req, BaseBatchReq):
if has_shm_features(req.batch):
return True
elif (
isinstance(req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput))
and req.mm_inputs
):
elif isinstance(
req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)
) and isinstance(req.mm_inputs, (MultimodalProcessorOutput, MultimodalInputs)):
for item in req.mm_inputs.mm_items:
if _feature_has_shm(item.feature):
return True
@@ -1439,6 +1492,30 @@ def has_shm_features(recv_reqs):
return False
def _discard_tensor_or_list(value) -> None:
if isinstance(value, ShmPointerMMData):
value.close_and_unlink()
elif isinstance(value, (list, tuple)):
for tensor in value:
if isinstance(tensor, ShmPointerMMData):
tensor.close_and_unlink()
def discard_shm_features(obj) -> None:
"""Release SHM features that will not be consumed by this request."""
if isinstance(obj, BaseBatchReq):
for sub_obj in obj.batch:
discard_shm_features(sub_obj)
return
if not isinstance(obj, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)):
return
if not isinstance(obj.mm_inputs, (MultimodalProcessorOutput, MultimodalInputs)):
return
for item in obj.mm_inputs.mm_items:
_discard_tensor_or_list(item.feature)
_discard_tensor_or_list(item.precomputed_embeddings)
def _unwrap_tensor_or_list(value):
"""Restore ShmPointerMMData wrappers back into standard torch.Tensors."""
if isinstance(value, ShmPointerMMData):
@@ -1464,10 +1541,9 @@ def unwrap_shm_features(obj):
unwrap_shm_features(sub_obj)
return obj
# Handle single requests
if (
isinstance(obj, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput))
and obj.mm_inputs
):
if isinstance(
obj, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)
) and isinstance(obj.mm_inputs, (MultimodalProcessorOutput, MultimodalInputs)):
for item in obj.mm_inputs.mm_items:
if item.feature is not None:
item.feature = _unwrap_tensor_or_list(item.feature)
+52 -13
View File
@@ -509,6 +509,22 @@ class MultimodalDataItem(msgspec.Struct, kw_only=True, dict=True, array_like=Tru
)
self.feature.acknowledge_consumption(consumer_count)
def release_transport_proxies(self, consumer_count: int = 1) -> None:
"""Best-effort release of proxies left by an abandoned request."""
values = [self.feature, self.precomputed_embeddings]
values.extend(self.model_specific_data.values())
for value in values:
if not isinstance(value, CudaIpcTensorTransportProxy):
continue
count = self._resolve_transport_consumer_count(value, consumer_count)
try:
value.release_without_reconstruction(count)
except Exception:
logger.warning(
"Failed to release an abandoned multimodal transport proxy",
exc_info=True,
)
@staticmethod
def _resolve_transport_consumer_count(proxy, requested_count: int) -> int:
"""Clamp a group acknowledgement to the proxy's actual consumer set."""
@@ -643,7 +659,18 @@ class MultimodalInputs:
def release_features(self):
"""Release feature tensors to free GPU memory."""
for item in self.mm_items:
item.feature = None
try:
# A request can be rejected before a deferred GPU feature is
# reconstructed. Acknowledge that transport lease before the
# proxy is dropped so the tokenizer pool can reuse its slice.
item.acknowledge_deferred_cuda_ipc_feature()
except Exception:
logger.warning(
"Failed to release an unused multimodal feature transport",
exc_info=True,
)
finally:
item.feature = None
@staticmethod
def from_processor_output(obj: MultimodalProcessorOutput):
@@ -653,14 +680,19 @@ class MultimodalInputs:
# try reconstructing from cuda-ipc
reconstruct_device = None
for mm_item in mm_items:
if (
mm_item.has_cuda_ipc_proxy()
and not mm_item.can_defer_cuda_ipc_feature_reconstruction()
):
if reconstruct_device is None:
reconstruct_device = torch.cuda.current_device()
mm_item.reconstruct(reconstruct_device)
try:
for mm_item in mm_items:
if (
mm_item.has_cuda_ipc_proxy()
and not mm_item.can_defer_cuda_ipc_feature_reconstruction()
):
if reconstruct_device is None:
reconstruct_device = torch.cuda.current_device()
mm_item.reconstruct(reconstruct_device)
except BaseException:
for mm_item in mm_items:
mm_item.release_transport_proxies()
raise
if envs.SGLANG_MM_BUFFER_SIZE_MB.get() > 0:
# Multi-modal feature hashing optimization:
@@ -1892,9 +1924,18 @@ class Req(ReqDllmMixin):
logger.info(f"{prefix}: {self.time_stats.convert_to_duration()}")
self.has_log_time_stats = True
def set_finish_with_abort(self, error_msg: str):
def set_finish_with_abort(
self,
error_msg: str,
status_code: int = HTTPStatus.BAD_REQUEST,
err_type: str = "BadRequestError",
):
if get_parallel().tp_rank == 0:
logger.error(f"{error_msg}, {self.rid=}")
# Session requests share historical multimodal inputs with their prior
# request. The session owns and releases those features when it closes.
if self.multimodal_inputs is not None and self.session is None:
self.multimodal_inputs.release_features()
self.multimodal_inputs = None
self.grammar = None
self.origin_input_ids = array(
@@ -1902,9 +1943,7 @@ class Req(ReqDllmMixin):
) # set it to one token to skip the long prefill
self.return_logprob = False
self.logprob_start_len = -1
self.to_finish = FINISH_ABORT(
error_msg, HTTPStatus.BAD_REQUEST, "BadRequestError"
)
self.to_finish = FINISH_ABORT(error_msg, status_code, err_type)
def update_reasoning_tokens(self, token_id, think_end_ids):
if self._is_reasoning_over:
+166 -25
View File
@@ -151,6 +151,7 @@ from sglang.srt.managers.io_struct import (
LoadLoRAAdapterFromTensorsReqOutput,
LoadLoRAAdapterReqInput,
LoadLoRAAdapterReqOutput,
MMInputsProcessError,
OpenSessionReqInput,
PauseGenerationReqInput,
ProfileReq,
@@ -373,6 +374,16 @@ STEP_MAX_US = 2_000_000
LOAD_STALL_REFRESH_S = 0.05
@dataclasses.dataclass(frozen=True)
class _MultimodalInputBroadcast:
inputs: Optional[MultimodalInputs] = None
error: Optional[str] = None
class _MultimodalInputProcessingError(RuntimeError):
pass
def _accumulate_decode_moment(
totals: list[float],
batch_size: int,
@@ -1955,11 +1966,12 @@ class Scheduler(
def process_input_requests(self, recv_reqs: List):
now = time.monotonic()
self.session_controller.maybe_reap(now)
if get_mm().mm_feature_transport == "cuda_vmm":
for recv_req in recv_reqs:
self._materialize_cuda_vmm_inputs(recv_req)
for recv_req in recv_reqs:
vmm_errors = None
if get_mm().mm_feature_transport == "cuda_vmm":
vmm_errors = self._materialize_cuda_vmm_inputs(recv_req)
# Skip health check when server is busy — ongoing requests already carry health info.
if is_health_check_generate_req(recv_req) and not self.is_fully_idle(
for_health_check=True
@@ -1969,6 +1981,10 @@ class Scheduler(
)
continue
if vmm_errors is not None and any(vmm_errors):
self._dispatch_tokenized_mm_requests(recv_req, vmm_errors)
continue
output = self._request_dispatcher(recv_req)
if output is not None:
if self.rust_server is not None:
@@ -1986,26 +2002,87 @@ class Scheduler(
if self.external_corpus_manager is not None:
self.external_corpus_manager.check_pending_load()
def _materialize_cuda_vmm_inputs(self, recv_req):
"""Release VMM slices before request handling can reject the request."""
@staticmethod
def _tokenized_requests(recv_req):
if isinstance(
recv_req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)
):
tokenized_reqs = (recv_req,)
elif isinstance(
return (recv_req,)
if isinstance(
recv_req,
(BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput),
):
tokenized_reqs = recv_req
else:
return
return tuple(recv_req)
return ()
def _gather_vmm_materialization_errors(
self, local_error: Optional[str]
) -> List[Optional[str]]:
if not (
torch.distributed.is_available()
and torch.distributed.is_initialized()
and self.dp_tp_cpu_group is not None
):
return [local_error]
world_size = torch.distributed.get_world_size(group=self.dp_tp_cpu_group)
errors = [None] * world_size
torch.distributed.all_gather_object(
errors,
local_error,
group=self.dp_tp_cpu_group,
)
return errors
def _materialize_cuda_vmm_inputs(self, recv_req) -> Optional[List[Optional[str]]]:
"""Materialize each request and agree on failures across TP ranks."""
tokenized_reqs = self._tokenized_requests(recv_req)
if not tokenized_reqs:
return None
request_errors = []
for tokenized_req in tokenized_reqs:
if tokenized_req.mm_inputs is not None and not isinstance(
tokenized_req.mm_inputs, MultimodalInputs
):
tokenized_req.mm_inputs = MultimodalInputs.from_processor_output(
tokenized_req.mm_inputs
local_error = None
try:
if tokenized_req.mm_inputs is not None and not isinstance(
tokenized_req.mm_inputs, MultimodalInputs
):
tokenized_req.mm_inputs = MultimodalInputs.from_processor_output(
tokenized_req.mm_inputs
)
except Exception as error:
local_error = f"{type(error).__name__}: {error}"
rank_errors = self._gather_vmm_materialization_errors(local_error)
failed_ranks = [
rank for rank, error in enumerate(rank_errors) if error is not None
]
if failed_ranks:
details = "; ".join(
f"rank {rank}: {rank_errors[rank]}" for rank in failed_ranks
)
error_msg = f"Multimodal feature reconstruction failed ({details})"
logger.error(error_msg)
tokenized_req.mm_inputs = None
request_errors.append(error_msg)
else:
request_errors.append(None)
return request_errors
def _dispatch_tokenized_mm_requests(
self, recv_req, errors: List[Optional[str]]
) -> None:
tokenized_reqs = self._tokenized_requests(recv_req)
if len(tokenized_reqs) != len(errors):
raise RuntimeError("VMM materialization results do not match requests")
for tokenized_req, error in zip(tokenized_reqs, errors, strict=True):
if isinstance(tokenized_req, TokenizedGenerateReqInput):
self.handle_generate_request(tokenized_req, mm_input_error=error)
elif isinstance(tokenized_req, TokenizedEmbeddingReqInput):
self.handle_embedding_request(tokenized_req, mm_input_error=error)
else:
raise TypeError(
f"Unsupported tokenized request type: {type(tokenized_req).__name__}"
)
def init_profiler(self) -> None:
@@ -2346,6 +2423,11 @@ class Scheduler(
Returns:
MultimodalInputs | None
Raises:
_MultimodalInputProcessingError: The entry rank could not build the
request's multimodal inputs. The same error is broadcast to all
ranks before it is raised.
"""
if raw_mm_inputs is None:
return None
@@ -2371,18 +2453,29 @@ class Scheduler(
# Since the Scheduler is single-threaded, any large CPU cost will impact
# handling of other messages. For example, CPU hits 99.9% can significantly
# increase the CUDA kernel launch time.
result = None
if self.dp_tp_group.rank_in_group == 0:
# Only the entry rank materializes once from dict.
image_inputs = MultimodalInputs.from_processor_output(raw_mm_inputs)
# Broadcast to other TP ranks (use src=0 within the group).
try:
result = _MultimodalInputBroadcast(
inputs=MultimodalInputs.from_processor_output(raw_mm_inputs)
)
except Exception as error:
result = _MultimodalInputBroadcast(
error=(
"Multimodal input processing failed on the TP entry rank: "
f"{type(error).__name__}: {error}"
)
)
# Broadcast either the prepared inputs or the request-local error.
if group_world_size > 1:
obj_list = [image_inputs]
obj_list = [result]
torch.distributed.broadcast_object_list(
obj_list,
src=self.dp_tp_group.first_rank,
group=self.dp_tp_cpu_group,
)
image_inputs = obj_list[0]
result = obj_list[0]
else:
# Non-entry ranks: receive if group size > 1; otherwise materialize locally.
if group_world_size > 1:
@@ -2392,13 +2485,19 @@ class Scheduler(
src=self.dp_tp_group.first_rank,
group=self.dp_tp_cpu_group,
)
image_inputs = obj_list[0]
result = obj_list[0]
else:
image_inputs = MultimodalInputs.from_processor_output(raw_mm_inputs)
result = _MultimodalInputBroadcast(
inputs=MultimodalInputs.from_processor_output(raw_mm_inputs)
)
return image_inputs
if result.error is not None:
raise _MultimodalInputProcessingError(result.error)
return result.inputs
def _get_multimodal_inputs(self, mm_inputs):
if isinstance(mm_inputs, MMInputsProcessError):
raise _MultimodalInputProcessingError(mm_inputs.message)
if isinstance(mm_inputs, MultimodalInputs):
return mm_inputs
@@ -2487,6 +2586,8 @@ class Scheduler(
def handle_generate_request(
self,
recv_req: TokenizedGenerateReqInput,
*,
mm_input_error: Optional[str] = None,
):
# Route: normal request / session request / session-not-found
session_id = (
@@ -2635,6 +2736,16 @@ class Scheduler(
self._maybe_namespace_elastic_radix_cache(req)
if mm_input_error is not None:
req.set_finish_with_abort(
mm_input_error,
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
err_type="InternalServerError",
)
self.init_req_max_new_tokens(req)
self._add_request_to_queue(req)
return
if self.spec_algorithm.is_dflash_family():
error_msg = validate_dflash_request(req, self.enable_overlap)
if error_msg is not None:
@@ -2694,7 +2805,17 @@ class Scheduler(
# Handle multimodal inputs
if recv_req.mm_inputs is not None:
image_inputs = self._get_multimodal_inputs(recv_req.mm_inputs)
try:
image_inputs = self._get_multimodal_inputs(recv_req.mm_inputs)
except _MultimodalInputProcessingError as error:
req.set_finish_with_abort(
str(error),
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
err_type="InternalServerError",
)
self.init_req_max_new_tokens(req)
self._add_request_to_queue(req)
return
SessionController.adjust_mm_offsets(recv_req, req, image_inputs)
@@ -3025,6 +3146,8 @@ class Scheduler(
def handle_embedding_request(
self,
recv_req: TokenizedEmbeddingReqInput,
*,
mm_input_error: Optional[str] = None,
):
req = Req(
recv_req.rid,
@@ -3045,9 +3168,27 @@ class Scheduler(
req.tokenizer = self.tokenizer
self._maybe_namespace_elastic_radix_cache(req)
if mm_input_error is not None:
req.set_finish_with_abort(
mm_input_error,
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
err_type="InternalServerError",
)
self._add_request_to_queue(req)
return
# Handle multimodal inputs
if recv_req.mm_inputs is not None:
image_inputs = self._get_multimodal_inputs(recv_req.mm_inputs)
try:
image_inputs = self._get_multimodal_inputs(recv_req.mm_inputs)
except _MultimodalInputProcessingError as error:
req.set_finish_with_abort(
str(error),
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
err_type="InternalServerError",
)
self._add_request_to_queue(req)
return
# Expand a single image token into multiple dummy tokens for receiving image embeddings
# The `pad_input_ids_func` is model-specific and may be None for
# embedding models or models not requiring special padding.
@@ -1,5 +1,6 @@
from __future__ import annotations
import logging
from dataclasses import dataclass
from http import HTTPStatus
from typing import (
@@ -11,19 +12,22 @@ from typing import (
Union,
)
import torch
import zmq
from torch.distributed import barrier
from torch.distributed import ReduceOp, all_reduce, barrier
from sglang.srt.disaggregation.utils import prepare_abort
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import (
BatchTokenizedEmbeddingReqInput,
BatchTokenizedGenerateReqInput,
MMInputsProcessError,
TokenizedEmbeddingReqInput,
TokenizedGenerateReqInput,
sock_recv,
)
from sglang.srt.managers.mm_utils import (
discard_shm_features,
has_shm_features,
unwrap_shm_features,
)
@@ -44,6 +48,8 @@ if TYPE_CHECKING:
ScriptedTokenizerRecvProxy,
)
logger = logging.getLogger(__name__)
@dataclass(kw_only=True, slots=True, frozen=True)
class SchedulerRequestReceiver:
@@ -249,24 +255,63 @@ class SchedulerRequestReceiver:
return recv_reqs
def _finalize_shm_features(self, recv_reqs: Optional[List]) -> None:
# Unwrap shared memory features AFTER all broadcasts complete,
# so that ShmPointerMMData metadata (not full tensor data) is what
# gets serialized during broadcast_pyobj.
if recv_reqs:
if self.model_config.is_multimodal and has_shm_features(recv_reqs):
# The broadcast source returns with its original objects while
# peer ranks may still be unpickling ShmPointerMMData
# (-> shm_open). Synchronize the same CPU groups that carried
# SHM-backed work requests before materialize() unlinks them.
if get_parallel().enable_dp_attention:
if self.ps.attn_tp_size > 1:
barrier(group=self.attn_tp_cpu_group)
if self.ps.attn_cp_size > 1:
barrier(group=self.attn_cp_cpu_group)
elif self.ps.tp_size > 1:
barrier(group=self.tp_cpu_group)
for req in recv_reqs:
"""Materialize SHM features or mark the request failed on every rank."""
if not recv_reqs or not self.model_config.is_multimodal:
return
tokenized_reqs = []
for req in recv_reqs:
if isinstance(req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)):
tokenized_reqs.append(req)
elif isinstance(
req,
(BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput),
):
tokenized_reqs.extend(req.batch)
if not tokenized_reqs or not has_shm_features(tokenized_reqs):
return
# 1. wait until every rank has opened the shared feature segments
parallel = get_parallel()
if parallel.enable_dp_attention:
if self.ps.attn_tp_size > 1:
barrier(group=self.attn_tp_cpu_group)
if self.ps.attn_cp_size > 1:
barrier(group=self.attn_cp_cpu_group)
elif self.ps.tp_size > 1:
barrier(group=self.tp_cpu_group)
# 2. materialize independently so one bad VLM request does not stop the loop
failed = torch.zeros(len(tokenized_reqs), dtype=torch.int32)
for index, req in enumerate(tokenized_reqs):
if not has_shm_features([req]):
continue
try:
unwrap_shm_features(req)
except Exception:
logger.exception(
"Failed to materialize shared-memory multimodal features for rid=%s",
req.rid,
)
discard_shm_features(req)
failed[index] = 1
# 3. all ranks reject the same requests before entering model collectives
if parallel.enable_dp_attention:
if self.ps.attn_tp_size > 1:
all_reduce(failed, op=ReduceOp.MAX, group=self.attn_tp_cpu_group)
if self.ps.attn_cp_size > 1:
all_reduce(failed, op=ReduceOp.MAX, group=self.attn_cp_cpu_group)
elif self.ps.tp_size > 1:
all_reduce(failed, op=ReduceOp.MAX, group=self.tp_cpu_group)
error = MMInputsProcessError(
"Failed to materialize shared-memory multimodal features on a scheduler rank."
)
for index, req in enumerate(tokenized_reqs):
if failed[index].item():
discard_shm_features(req)
req.mm_inputs = error
def _split_work_and_control_reqs(self, recv_reqs: List):
work_reqs = [
@@ -38,6 +38,7 @@ from sglang.srt.multimodal.processors.executor import MultimodalProcessorExecuto
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,
)
@@ -1875,19 +1876,41 @@ class BaseMultimodalProcessor(ABC):
def _prepare_mm_items_for_transport(
self, mm_items: List[MultimodalDataItem]
) -> List[MultimodalDataItem]:
"""Wrap final GPU features for dispatch to the scheduler."""
"""Wrap final GPU features, rolling back every lease if one wrap fails."""
if not self.use_cuda_ipc:
return mm_items
# Pool misses fall back to plain CPU tensors. The scheduler copies out
# and releases each successful pool slice.
for item in mm_items:
if isinstance(item.feature, torch.Tensor):
item.feature = self._wrap_tensor_for_cuda_ipc(item.feature)
if isinstance(item.precomputed_embeddings, torch.Tensor):
item.precomputed_embeddings = self._wrap_tensor_for_cuda_ipc(
item.precomputed_embeddings
updates = []
try:
for item in mm_items:
fields = (
("feature", item.feature),
("precomputed_embeddings", item.precomputed_embeddings),
)
for field, tensor in fields:
if not isinstance(tensor, torch.Tensor):
continue
wrapped = self._wrap_tensor_for_cuda_ipc(tensor)
setattr(item, field, wrapped)
updates.append((item, field, tensor, wrapped))
except BaseException as error:
rollback_errors = []
for item, field, tensor, wrapped in reversed(updates):
try:
if isinstance(wrapped, CudaIpcTensorTransportProxy):
self.cudaipc_mmfeature_pool.cancel_proxy(wrapped)
except BaseException as rollback_error:
rollback_errors.append(rollback_error)
finally:
setattr(item, field, tensor)
if rollback_errors:
error.add_note(
f"{len(rollback_errors)} CUDA IPC rollback operation(s) also failed"
)
raise error from rollback_errors[0]
raise
return mm_items
async def process_and_combine_mm_data_async(
@@ -645,8 +645,6 @@ class KimiK3ImageProcessor(
model_specific_data=model_specific_data,
)
item.set_hash(artifact.feature_hash)
if self.use_cuda_ipc and isinstance(item.feature, torch.Tensor):
item.feature = self._wrap_tensor_for_cuda_ipc(item.feature)
if self.keep_mm_features_on_device and item.feature is not None:
item.model_specific_data[DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY] = (
True
@@ -655,7 +653,7 @@ class KimiK3ImageProcessor(
return MultimodalProcessorOutput(
input_ids=input_ids.tolist(),
mm_items=items,
mm_items=self._prepare_mm_items_for_transport(items),
im_token_id=self.mm_tokens.image_token_id,
)
@@ -582,14 +582,7 @@ class MossVLImageProcessor(SGLangBaseProcessor):
if mm_items and vision_token_info:
mm_items[0].set("vision_token_info", vision_token_info[0])
if self.use_cuda_ipc:
for item in mm_items:
if isinstance(item.feature, torch.Tensor):
item.feature = self._wrap_tensor_for_cuda_ipc(item.feature)
if isinstance(item.precomputed_embeddings, torch.Tensor):
item.precomputed_embeddings = self._wrap_tensor_for_cuda_ipc(
item.precomputed_embeddings
)
mm_items = self._prepare_mm_items_for_transport(mm_items)
return MultimodalProcessorOutput(
input_ids=input_ids.tolist(),
@@ -150,6 +150,17 @@ class MmItemMemoryPool:
use_pool_handle_cache=use_pool_handle_cache,
)
def cancel_proxy(self, proxy: "CudaIpcTensorTransportProxy") -> None:
"""Return a published slice when its request was never dispatched."""
ipc_extra = proxy.proxy_state["ipc_extra"]
if tuple(ipc_extra["pool_handle"]) != tuple(self._pool_ipc_handle):
raise RuntimeError("CUDA IPC proxy does not belong to this pool")
self._pool.cancel_lease(
ready_byte_offset=proxy.ready_byte_offset,
ack_byte_offset=proxy.ack_byte_offset,
generation=proxy.generation,
)
def _warn_pool_full_once(self, nbytes: int):
if self._pool_full_warned:
return
@@ -310,6 +321,10 @@ class CudaIpcTensorTransportProxy(StreamOrderedPoolConsumerMixin):
)
self._retain_storage_until_stream_completes(storage, device_id)
def release_without_reconstruction(self, consumer_count: int = 1) -> None:
"""Release a pool slice when its request abandons this proxy."""
self.acknowledge_consumption(consumer_count)
def reconstruct_on_target_device(
self,
rebuild_device_idx,
@@ -362,6 +362,50 @@ class StreamOrderedMmFeaturePool:
raise
return lease, destination
def cancel_lease(
self,
*,
ready_byte_offset: int,
ack_byte_offset: int,
generation: int,
) -> None:
"""Acknowledge every consumer for a lease that was not dispatched."""
slot_stride = self.control_words_per_slot * CONTROL_WORD_BYTES
if (
ready_byte_offset % slot_stride != 0
or ack_byte_offset != ready_byte_offset + CONTROL_WORD_BYTES
):
raise RuntimeError(f"Invalid {self.transport_name} pool lease offsets")
slot = ready_byte_offset // slot_stride
with self._lock:
lease = self._occupied.get(slot)
if (
lease is None
or lease.generation != generation
or lease.ready_byte_offset != ready_byte_offset
or lease.ack_byte_offset != ack_byte_offset
):
raise RuntimeError(
f"Cannot cancel inactive {self.transport_name} pool lease "
f"(slot={slot}, generation={generation})"
)
with torch.cuda.device(self.device_id):
stream_wait_value32(
self.device_id,
self.base_address + ready_byte_offset,
generation,
self.transport_name,
)
for rank in range(self.consumer_count):
stream_write_value32(
self.device_id,
self.base_address + ack_byte_offset + rank * CONTROL_WORD_BYTES,
generation,
self.transport_name,
)
def shutdown(self) -> None:
self._recycler_stop_event.set()
if self._recycle_thread.is_alive():
@@ -904,6 +904,13 @@ class CudaVmmPackedTensorTransportProxy(CudaVmmTensorTransportProxy):
"Packed CUDA VMM features must be reconstructed before release"
)
def release_without_reconstruction(self, consumer_count: int | None = None) -> None:
"""Release the shared packed allocation when its request is abandoned."""
if self._consumer_acknowledged:
return
self._packed_owner.acknowledge_consumption(consumer_count)
self._consumer_acknowledged = True
def reconstruct_on_target_device(
self, rebuild_device_idx, consumer_count: int | None = None
):