Skip to content

Fix multimodal feature offload races #40621

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xiezhq-hermann merged 12 commits into from
Sep 25, 2026
Merged

Fix multimodal feature offload races #40621

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
692bf8b
Wait for multimodal feature copies before host reads
Sep 21, 2026
988d1ce
Allow encoders to keep raw multimodal inputs on CPU
Sep 21, 2026
d393c76
Protect multimodal host offload across streams
Sep 22, 2026
201c1fc
Keep multimodal offload changes focused on race safety
cctry Sep 22, 2026
c02ce4e
Inline multimodal offload at its only call site
cctry Sep 22, 2026
e14a461
Share multimodal offload readiness across encoder inputs
cctry Sep 22, 2026
769e614
Simplify multimodal offload waits
cctry Sep 22, 2026
db13d73
Reuse one worker event for multimodal host offload
cctry Sep 22, 2026
060de9a
Remove redundant multimodal offload event comment
cctry Sep 22, 2026
e0856b3
Acquire multimodal offload stream once per batch
cctry Sep 22, 2026
f1170d5
Prepare multimodal feature dtypes before cache lookup
cctry Sep 22, 2026
de6a7f7
Merge branch 'main' into fix/mm-feature-copy-readiness
xiezhq-hermann Sep 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 25 additions & 8 deletions python/sglang/srt/managers/mm_schedule.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
_is_xpu = is_xpu()

embedding_cache: Optional[MultiModalStaticCache] = None
host_offload_event: Optional[torch.cuda.Event] = None


def init_mm_embedding_cache(max_size: int = 0):
Expand Down Expand Up @@ -83,6 +84,12 @@ def _get_precomputed_embedding(
If some but not all have precomputed_embeddings, raise NotImplementedError.
If none have precomputed_embeddings, return None.
"""
if host_offload_event is not None and any(
isinstance(item.precomputed_embeddings, torch.Tensor)
and item.precomputed_embeddings.is_cpu
for item in items
):
host_offload_event.synchronize()
precomputed_embeddings = []
max_iterations = min(len(items_size) - 1, len(prefix_length))

Expand Down Expand Up @@ -209,9 +216,22 @@ def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) ->


def _move_items_to_device(
items: List[MultimodalDataItem], device: torch.device
items: List[MultimodalDataItem],
device: torch.device,
data_embedding_func: DataEmbeddingFunc,
) -> None:
"""Move item features to the target device (in-place, non-blocking)."""
"""Wait for feature readiness and upload unless the encoder defers the move."""
defer_move = _can_skip_pre_embed_feature_move(data_embedding_func)
if host_offload_event is not None and any(
isinstance(item.feature, torch.Tensor) and item.feature.is_cpu for item in items
):
if defer_move or device.type != "cuda":
# Deferred encoders can read CPU subsets of a mixed-device batch.
host_offload_event.synchronize()
else:
torch.cuda.current_stream(device).wait_event(host_offload_event)
if defer_move:
return
for item in items:
if isinstance(item.feature, torch.Tensor) and item.feature.device != device:
item.feature = item.feature.to(device, non_blocking=True)
Expand Down Expand Up @@ -270,8 +290,7 @@ def _get_chunked_embedding_full(
embedding_per_req = None

if embedding_per_req is None:
if not _can_skip_pre_embed_feature_move(data_embedding_func):
_move_items_to_device(embedding_items_per_req, device)
_move_items_to_device(embedding_items_per_req, device, data_embedding_func)
embedding = data_embedding_func(embedding_items_per_req)
if isinstance(embedding, list):
# This path caches the combined per-request embedding, so the
Expand Down Expand Up @@ -383,8 +402,7 @@ def _batch_encode_per_image_misses(
miss_items = [unique_misses[key][0] for key in ordered_cache_keys]
token_counts = [unique_misses[key][1] for key in ordered_cache_keys]

if not _can_skip_pre_embed_feature_move(data_embedding_func):
_move_items_to_device(miss_items, device)
_move_items_to_device(miss_items, device, data_embedding_func)
all_miss_embedding = data_embedding_func(miss_items)

if isinstance(all_miss_embedding, list):
Expand Down Expand Up @@ -461,8 +479,7 @@ def _get_chunked_embedding_by_item(

if miss_items:
miss_item_list = [item for _, item, _, _ in miss_items]
if not _can_skip_pre_embed_feature_move(data_embedding_func):
_move_items_to_device(miss_item_list, device)
_move_items_to_device(miss_item_list, device, data_embedding_func)
all_miss_embedding = data_embedding_func(miss_item_list)

if isinstance(all_miss_embedding, list):
Expand Down
68 changes: 49 additions & 19 deletions python/sglang/srt/managers/mm_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

from sglang.kernels.ops.memory.gpu_tensor_hash import gpu_tensor_hash
from sglang.srt.environ import envs
from sglang.srt.managers import mm_schedule
from sglang.srt.managers.io_struct import (
BaseBatchReq,
TokenizedEmbeddingReqInput,
Expand Down Expand Up @@ -661,6 +662,7 @@ def general_mm_embed_routine(
data_embedding_funcs: Dict[Modality, DataEmbeddingFunc] = None,
placeholder_tokens: Optional[dict[Modality, List[int]]] = None,
use_deepstack: Dict[Modality, bool] = {},
feature_dtypes: Optional[Dict[Modality, torch.dtype]] = None,
**kwargs,
) -> torch.Tensor:
"""
Expand All @@ -673,6 +675,7 @@ def general_mm_embed_routine(
data_embedding_funcs: A dictionary mapping from modality type to the corresponding embedding function.
placeholder_tokens: Token IDs for multimodal placeholders
use_deepstack: Whether to use deepstack embeddings for each modality, default False
feature_dtypes: Optional per-modality dtypes for preparing CUDA features before cache lookup.
**kwargs: Additional arguments passed to language model

Returns:
Expand All @@ -689,6 +692,23 @@ def general_mm_embed_routine(
mm_inputs_list = [
mm_input for mm_input in forward_batch.mm_inputs if mm_input is not None
]
if feature_dtypes:
for mm_input in mm_inputs_list:
for item in mm_input.mm_items:
dtype = feature_dtypes.get(item.modality)
feature = item.feature
if (
dtype is not None
and isinstance(feature, torch.Tensor)
and feature.is_cuda
and feature.dtype != dtype
):
# The cast may still read the source after replacing it.
feature.record_stream(
torch.cuda.current_stream(feature.device)
)
item.feature = feature.to(dtype=dtype)
del feature
extend_prefix_lens = [
prefix_len
for i, prefix_len in enumerate(forward_batch.extend_prefix_lens_cpu)
Expand Down Expand Up @@ -742,27 +762,37 @@ def general_mm_embed_routine(
# best-effort, offloading to CPU ensures we have a reliable fallback
# if a cache miss occurs in subsequent chunks, while still freeing up
# critical GPU memory.
if mm_inputs_list:
if mm_inputs_list and input_embeds.is_cuda:
language_only = get_disagg().language_only
stream = torch.cuda.current_stream(input_embeds.device)
offloaded = False
for mm_input_obj in mm_inputs_list:
if mm_input_obj and hasattr(mm_input_obj, "mm_items"):
for mm_item in mm_input_obj.mm_items:
feature = getattr(mm_item, "feature", None)
if isinstance(feature, torch.Tensor) and feature.is_cuda:
mm_item.feature = feature.to("cpu", non_blocking=True)
if get_disagg().language_only:
precomputed_embeddings = getattr(
mm_item, "precomputed_embeddings", None
for mm_item in mm_input_obj.mm_items:
feature = mm_item.feature
if isinstance(feature, torch.Tensor) and feature.is_cuda:
# The transport reconstructed this block on the scheduler stream
# and nothing else pins it to ours: without record_stream the
# allocator recycles it into the next scheduler allocation while
# the encoder read and this copy are still queued here.
feature.record_stream(stream)
mm_item.feature = feature.to("cpu", non_blocking=True)
offloaded = True
if language_only:
precomputed = mm_item.precomputed_embeddings
if (
isinstance(precomputed, torch.Tensor)
and precomputed.is_cuda
and not mm_item.keep_device_embedding
):
precomputed.record_stream(stream)
mm_item.precomputed_embeddings = precomputed.to(
"cpu", non_blocking=True
)
if (
isinstance(precomputed_embeddings, torch.Tensor)
and precomputed_embeddings.is_cuda
and not mm_item.keep_device_embedding
):
mm_item.precomputed_embeddings = (
precomputed_embeddings.to(
"cpu", non_blocking=True
)
)
offloaded = True
if offloaded:
if mm_schedule.host_offload_event is None:
mm_schedule.host_offload_event = torch.cuda.Event()
mm_schedule.host_offload_event.record(stream)
forward_batch.mm_inputs = None
forward_batch.mm_input_embeds = (
input_embeds.clone()
Expand Down
Loading