From 99709f734d8362282d4cec44d3b62465a87c9678 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 5 Aug 2026 20:24:12 +0800 Subject: [PATCH] [VLM] split multimodal scheduling from mm_utils (#32415) --- python/sglang/srt/managers/mm_schedule.py | 658 +++++++++++++++++ python/sglang/srt/managers/mm_utils.py | 665 +----------------- .../sglang/srt/mem_cache/kv_cache_builder.py | 2 +- .../test_mm_chunked_embedding_unit.py | 34 +- 4 files changed, 691 insertions(+), 668 deletions(-) create mode 100644 python/sglang/srt/managers/mm_schedule.py diff --git a/python/sglang/srt/managers/mm_schedule.py b/python/sglang/srt/managers/mm_schedule.py new file mode 100644 index 000000000..f08b4b1fe --- /dev/null +++ b/python/sglang/srt/managers/mm_schedule.py @@ -0,0 +1,658 @@ +"""Multimodal embedding scheduling and cache coordination.""" + +from dataclasses import dataclass, field +from typing import Callable, Dict, List, Optional, Tuple + +import torch + +from sglang.srt.managers.schedule_batch import MultimodalDataItem +from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache +from sglang.srt.multimodal.evs import EVSEmbeddingResult +from sglang.srt.runtime_context import get_parallel, get_schedule +from sglang.srt.utils import is_hip, is_npu +from sglang.utils import logger + +_is_hip = is_hip() +_is_npu = is_npu() + +embedding_cache: Optional[MultiModalStaticCache] = None + + +def init_mm_embedding_cache(max_size: int = 0): + global embedding_cache + embedding_cache = MultiModalStaticCache(max_size) + + +def get_embedding_chunk( + embedding: torch.Tensor, + extend_prefix_len: int, + extend_seq_len: int, + items_offset: List[Tuple[int, int]], +) -> Tuple[torch.Tensor, int, int]: + """ + Extract a chunk of embeddings based on the specified prefix length, sequence length, and offset ranges. + + Args: + embedding: The full embedding tensor to extract a chunk from + extend_prefix_len: The starting position (prefix length) for extraction + extend_seq_len: The number of tokens to extract + items_offset: List of [start, end] offset ranges for multimodal items in the input sequence + + Returns: + A tuple containing: + - The extracted embedding chunk as a tensor + - The start index used for extraction + - The end index used for extraction + + Note: + If there's no overlap between the requested range and the offset ranges, + an empty tensor is returned with zeros for start and end indices. + """ + start_index, end_index = 0, 0 + extend_start_index = extend_prefix_len + extend_end_index = extend_prefix_len + extend_seq_len - 1 + + for start, end in items_offset: + if extend_start_index >= start and extend_start_index <= end: + start_index += extend_start_index - start + elif extend_start_index > end: + start_index += end - start + 1 + + if extend_end_index >= start and extend_end_index <= end: + end_index += extend_end_index - start + 1 + elif extend_end_index > end: + end_index += end - start + 1 + # some models' embedding is 3-dim, reshape it to 2-dim + embedding = embedding.reshape(-1, embedding.shape[-1]) + embedding_chunk = embedding[start_index:end_index] + return embedding_chunk, start_index, end_index + + +def _get_precomputed_embedding( + items: List[MultimodalDataItem], + items_size: List[int], + prefix_length: List[int], + extend_length: List[int], + items_offset_list: List[List[Tuple[int, int]]], +) -> Optional[torch.Tensor]: + """ + If all items have precomputed_embeddings, return their concatenation. + If some but not all have precomputed_embeddings, raise NotImplementedError. + If none have precomputed_embeddings, return None. + """ + precomputed_embeddings = [] + max_iterations = min(len(items_size) - 1, len(prefix_length)) + + for i in range(max_iterations): + if items_size[i] == items_size[i + 1]: + continue + + items_per_req = items[items_size[i] : items_size[i + 1]] + extend_len = extend_length[i] if i < len(extend_length) else 0 + items_offset = items_offset_list[i] + + if any(item.precomputed_embeddings is None for item in items_per_req): + chunk = None + else: + req_embeddings = torch.concat( + [item.precomputed_embeddings for item in items_per_req] + ) + chunk, _, _ = get_embedding_chunk( + embedding=req_embeddings, + extend_prefix_len=prefix_length[i], + extend_seq_len=extend_len, + items_offset=items_offset, + ) + + if chunk is None and len(items_per_req) > 1: + return None + precomputed_embeddings.append(chunk) + + if any(feature is not None for feature in precomputed_embeddings): + if not all(feature is not None for feature in precomputed_embeddings): + raise NotImplementedError( + "MM inputs where only some items are precomputed." + ) + + # Normalize device across chunks before concat. + target_device = next( + (t.device for t in precomputed_embeddings if t.is_cuda), + precomputed_embeddings[0].device, + ) + precomputed_embeddings = [ + t if t.device == target_device else t.to(target_device, non_blocking=True) + for t in precomputed_embeddings + ] + result = torch.concat(precomputed_embeddings) + # some models embedding is 3-dim, reshape it to 2-dim (similar to get_embedding_chunk) + result = result.reshape(-1, result.shape[-1]) + return result + return None + + +# A modality's embedding function. May return the combined [tokens, hidden] +# tensor, an EVSEmbeddingResult, or one tensor per input item. The per-item +# form lets encoders that naturally produce per-item outputs (e.g. a wav +# AutoEncoder looping over clips) skip an encoder-side torch.cat that +# per-item consumers (_get_chunked_embedding_by_item) would immediately +# split back apart — and each cached entry then owns its storage instead of +# being a view pinning the concatenated buffer. +DataEmbeddingFunc = Callable[ + [List[MultimodalDataItem]], + torch.Tensor | List[torch.Tensor] | EVSEmbeddingResult, +] + + +def _flatten_embedding_result( + embedding: torch.Tensor | List[torch.Tensor], +) -> torch.Tensor: + """Normalize a DataEmbeddingFunc result to one [tokens, hidden] tensor.""" + if isinstance(embedding, list): + if not embedding: + raise ValueError( + "DataEmbeddingFunc returned an empty per-item list; expected " + "one entry per input item" + ) + flat = [e.reshape(-1, e.shape[-1]) for e in embedding] + return flat[0] if len(flat) == 1 else torch.cat(flat, dim=0) + return embedding + + +def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) -> bool: + """Models that materialize and batch visual features inside their encoder. + + instead of performing multiple H2D for each mm feature from all mm_items (followed by concatenation on device), + for some models which internally performs H2D on concated mm feature, these small H2D calls could be replaced with a single big H2D + """ + owner = getattr(data_embedding_func, "__self__", None) + if owner is None: + return False + if getattr(data_embedding_func, "__name__", None) not in ( + "get_image_feature", + "get_video_feature", + ): + return False + return owner.__class__.__name__ in { + "Qwen3VLForConditionalGeneration", + "Qwen3VLMoeForConditionalGeneration", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "KimiK25ForConditionalGeneration", + } + + +def _move_items_to_device( + items: List[MultimodalDataItem], device: torch.device +) -> None: + """Move item features to the target device (in-place, non-blocking).""" + for item in items: + if isinstance(item.feature, torch.Tensor) and item.feature.device != device: + item.feature = item.feature.to(device, non_blocking=True) + + +def _acknowledge_deferred_cuda_ipc_cache_hits( + items: List[MultimodalDataItem], +) -> None: + """Release lazy Kimi IPC slices when a cached embedding skips ViT. + + On an encoder-DP miss, exactly one rank copies an image and acknowledges + the full TP group. On a cache hit no rank copies it, so rank zero performs + the equivalent single acknowledgement. This preserves the fixed-pool + lifecycle without reintroducing an unnecessary GPU-to-GPU copy. + """ + parallel = get_parallel() + if parallel.attn_tp_rank != 0: + return + # The pool's recycler counts the whole TP group, so the acknowledgement must + # match that count even when an attention subgroup is smaller. + consumer_count = max(parallel.tp_size, 1) + for item in items: + item.acknowledge_deferred_cuda_ipc_feature(consumer_count) + + +def _get_chunked_embedding_full( + data_embedding_func: DataEmbeddingFunc, + embedding_items_per_req: List[MultimodalDataItem], + items_offset: List[Tuple[int, int]], + extend_prefix_len: int, + extend_seq_len: int, + input_ids: torch.Tensor, + device: torch.device, +) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + """ + Fallback: encode all items at once, cache combined result, extract chunk. + Used for non-bundled items or EVS results. + """ + item_hashes = [item.hash for item in embedding_items_per_req] + embedding_items_hash = MultiModalStaticCache.combine_hashes(item_hashes) + embedding_per_req = embedding_cache.get(item_hashes) + + 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) + embedding = data_embedding_func(embedding_items_per_req) + if isinstance(embedding, list): + # This path caches the combined per-request embedding, so the + # per-item form is flattened here. + embedding = _flatten_embedding_result(embedding) + embedding_per_req = ( + EmbeddingResult(embedding=embedding) + if isinstance(embedding, torch.Tensor) + else embedding + ) + embedding_cache.set(embedding_items_hash, embedding_per_req) + else: + _acknowledge_deferred_cuda_ipc_cache_hits(embedding_items_per_req) + + if isinstance(embedding_per_req, EVSEmbeddingResult): + item = embedding_items_per_req[0] + input_ids, items_offset = ( + embedding_per_req.redistribute_pruned_frames_placeholders( + input_ids, + items_offset, + item=item, + extend_prefix_len=extend_prefix_len, + extend_seq_len=extend_seq_len, + ) + ) + + embedding_per_req_chunk, _, _ = get_embedding_chunk( + embedding=embedding_per_req.embedding, + extend_prefix_len=extend_prefix_len, + extend_seq_len=extend_seq_len, + items_offset=items_offset, + ) + return embedding_per_req_chunk, input_ids + + +@dataclass +class PerImageRequestInfo: + """Metadata for a single request using the per-image encoding path.""" + + req_idx: int + items: List[MultimodalDataItem] + items_offset: List[Tuple[int, int]] + extend_prefix_len: int + extend_seq_len: int + overlapping: List[Tuple[int, MultimodalDataItem, int, int]] = field( + default_factory=list + ) + + +def _batch_encode_per_image_misses( + data_embedding_func: DataEmbeddingFunc, + per_image_requests: List[PerImageRequestInfo], + device: torch.device, +) -> Dict[int, torch.Tensor]: + """ + Collect cache misses across ALL per-image requests, deduplicate by hash, + encode in a single ViT call, and populate the cache. + + Returns: + hash_to_embedding: mapping from item.hash to its full embedding tensor. + """ + unique_misses: Dict[int, Tuple[MultimodalDataItem, int]] = {} + hash_to_embedding: Dict[int, torch.Tensor] = {} + + # Phase 1a: find overlapping items per request and collect cache misses + for req_info in per_image_requests: + chunk_start = req_info.extend_prefix_len + chunk_end = chunk_start + req_info.extend_seq_len # exclusive + overlapping = [] + if req_info.extend_seq_len > 0: + for idx, (item, (start, end)) in enumerate( + zip(req_info.items, req_info.items_offset) + ): + if end >= chunk_start and start < chunk_end: + overlapping.append((idx, item, start, end)) + req_info.overlapping = overlapping + + for _idx, item, start, end in overlapping: + if item.hash in hash_to_embedding: + continue + cached = embedding_cache.get_single(item.hash) + if cached is not None: + hash_to_embedding[item.hash] = cached.embedding + elif item.hash not in unique_misses: + token_count = end - start + 1 + unique_misses[item.hash] = (item, token_count) + + # Phase 1b: single ViT call for all unique cache misses + if unique_misses: + ordered_hashes = list(unique_misses.keys()) + miss_items = [unique_misses[h][0] for h in ordered_hashes] + token_counts = [unique_misses[h][1] for h in ordered_hashes] + + if not _can_skip_pre_embed_feature_move(data_embedding_func): + _move_items_to_device(miss_items, device) + all_miss_embedding = data_embedding_func(miss_items) + + if isinstance(all_miss_embedding, list): + # Per-item embeddings: no split needed, and each cache entry owns + # its storage (a torch.split view would pin the whole concatenated + # buffer for as long as any single item stays cached). Mirrors + # _get_chunked_embedding_by_item. + assert len(all_miss_embedding) == len(miss_items), ( + f"per-item embedding count {len(all_miss_embedding)} != " + f"cache-miss item count {len(miss_items)}" + ) + split_embeddings = [ + emb.reshape(-1, emb.shape[-1]) for emb in all_miss_embedding + ] + else: + all_miss_embedding = all_miss_embedding.reshape( + -1, all_miss_embedding.shape[-1] + ) + split_embeddings = torch.split(all_miss_embedding, token_counts, dim=0) + for h, emb in zip(ordered_hashes, split_embeddings): + embedding_cache.set(h, EmbeddingResult(embedding=emb)) + # Keep a local ref (no extra GPU memory) so assembly never fails due to LRU eviction. + hash_to_embedding[h] = emb + + return hash_to_embedding + + +def _get_chunked_embedding_by_item( + data_embedding_func: DataEmbeddingFunc, + embedding_items_per_req: List[MultimodalDataItem], + items_offset: List[Tuple[int, int]], + extend_prefix_len: int, + extend_seq_len: int, + device: torch.device, +) -> Optional[torch.Tensor]: + """ + Per-image chunk-aware encoding for one request. + Items must already be split per-image (each item has exactly one offset). + """ + chunk_start = extend_prefix_len + chunk_end = extend_prefix_len + extend_seq_len # exclusive + + if extend_seq_len <= 0: + return None + + overlapping = [] + for idx, (item, (start, end)) in enumerate( + zip(embedding_items_per_req, items_offset) + ): + if end >= chunk_start and start < chunk_end: + overlapping.append((idx, item, start, end)) + + if not overlapping: + return None + + cached_embeddings = {} + miss_items = [] + for idx, item, start, end in overlapping: + cached = embedding_cache.get_single(item.hash) + if cached is not None: + cached_embeddings[idx] = cached.embedding + _acknowledge_deferred_cuda_ipc_cache_hits([item]) + else: + miss_items.append((idx, item, start, end)) + + 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) + all_miss_embedding = data_embedding_func(miss_item_list) + + if isinstance(all_miss_embedding, list): + # Per-item embeddings: no split needed, and each cache entry owns + # its storage (a torch.split view would pin the whole concatenated + # buffer for as long as any single item stays cached). + assert len(all_miss_embedding) == len(miss_items), ( + f"per-item embedding count {len(all_miss_embedding)} != " + f"cache-miss item count {len(miss_items)}" + ) + split_embeddings = [ + emb.reshape(-1, emb.shape[-1]) for emb in all_miss_embedding + ] + else: + all_miss_embedding = all_miss_embedding.reshape( + -1, all_miss_embedding.shape[-1] + ) + # Split output by per-item token count + token_counts = [end - start + 1 for _, _, start, end in miss_items] + split_embeddings = torch.split(all_miss_embedding, token_counts, dim=0) + + for (idx, item, _, _), emb in zip(miss_items, split_embeddings): + cached_embeddings[idx] = emb + embedding_cache.set(item.hash, EmbeddingResult(embedding=emb)) + + chunk_slices = [] + for idx, _, start, end in overlapping: + emb = cached_embeddings[idx] + overlap_start = max(start, chunk_start) + overlap_end = min(end, chunk_end - 1) # inclusive + local_start = overlap_start - start + local_end = overlap_end - start + 1 # exclusive for slicing + chunk_slices.append(emb[local_start:local_end]) + + return torch.cat(chunk_slices, dim=0) + + +def _assemble_per_image_chunk( + overlapping: List[Tuple[int, MultimodalDataItem, int, int]], + hash_to_embedding: Dict[int, torch.Tensor], + extend_prefix_len: int, + extend_seq_len: int, +) -> Optional[torch.Tensor]: + """ + Assemble the chunk embedding for one request from pre-computed embeddings. + All overlapping items must already have their embeddings in hash_to_embedding. + """ + if not overlapping: + return None + + chunk_start = extend_prefix_len + chunk_end = extend_prefix_len + extend_seq_len # exclusive + + chunk_slices = [] + for _idx, item, start, end in overlapping: + emb = hash_to_embedding[item.hash] # shape: (end - start + 1, hidden) + overlap_start = max(start, chunk_start) + overlap_end = min(end, chunk_end - 1) # inclusive + local_start = overlap_start - start + local_end = overlap_end - start + 1 # exclusive for slicing + chunk_slices.append(emb[local_start:local_end]) + + return torch.cat(chunk_slices, dim=0) + + +def _get_chunked_prefill_embedding( + data_embedding_func: DataEmbeddingFunc, + embedding_items: List[MultimodalDataItem], + items_size: List[int], + prefix_length: List[int], + extend_length: List[int], + items_offset_list: List[List[Tuple[int, int]]], + input_ids: torch.Tensor, +) -> tuple[torch.Tensor | None, torch.Tensor]: + """ + Chunked prefill embedding: encode items across all requests and extract + per-request chunks. Images from all requests are batched into a single + ViT call for efficiency. + """ + device = input_ids.device + # FIXME(Xinyuan): temporary workaround for eagle3 + max_iterations = min(len(items_size) - 1, len(prefix_length)) + + # Phase 0: classify requests into per-image vs full/EVS path + per_image_requests = [] # batched ViT encoding + full_path_requests = [] # per-request encoding (EVS etc.) + all_chunks: List[Tuple[int, torch.Tensor]] = [] + + for i in range(max_iterations): + if items_size[i] == items_size[i + 1]: + continue + embedding_items_per_req = embedding_items[items_size[i] : items_size[i + 1]] + items_offset = items_offset_list[i] + assert items_offset is not None, items_offset + + extend_prefix_len = prefix_length[i] + extend_seq_len = extend_length[i] if i < len(extend_length) else 0 + if extend_seq_len <= 0: + continue + + # Skip if all items already prefilled. + if all(offset_end < prefix_length[i] for _, offset_end in items_offset): + continue + + req_info = PerImageRequestInfo( + req_idx=i, + items=embedding_items_per_req, + items_offset=items_offset, + extend_prefix_len=extend_prefix_len, + extend_seq_len=extend_seq_len, + ) + + is_per_image = all(len(item.offsets) == 1 for item in embedding_items_per_req) + if is_per_image: + if _is_hip or _is_npu: + # ROCm CI regressed with one large cross-request ViT batch; keep + # the previous per-request path on HIP while CUDA uses batching. + chunk = _get_chunked_embedding_by_item( + data_embedding_func, + embedding_items_per_req, + items_offset, + extend_prefix_len, + extend_seq_len, + device, + ) + if chunk is not None: + all_chunks.append((i, chunk)) + else: + per_image_requests.append(req_info) + else: + full_path_requests.append(req_info) + + # Phase 1: batch encode all per-image cache misses in ONE ViT call + hash_to_embedding: Dict[int, torch.Tensor] = {} + if per_image_requests: + hash_to_embedding = _batch_encode_per_image_misses( + data_embedding_func, per_image_requests, device + ) + + # Phase 2: assemble per-request chunks in original request order + for req_info in per_image_requests: + chunk = _assemble_per_image_chunk( + req_info.overlapping, + hash_to_embedding, + req_info.extend_prefix_len, + req_info.extend_seq_len, + ) + if chunk is not None: + all_chunks.append((req_info.req_idx, chunk)) + + for req_info in full_path_requests: + chunk_embedding, input_ids = _get_chunked_embedding_full( + data_embedding_func, + req_info.items, + req_info.items_offset, + req_info.extend_prefix_len, + req_info.extend_seq_len, + input_ids, + device, + ) + if chunk_embedding is not None: + all_chunks.append((req_info.req_idx, chunk_embedding)) + + # Sort by original request index to maintain correct output order + all_chunks.sort(key=lambda x: x[0]) + embedding_list = [chunk for _, chunk in all_chunks] + + if len(embedding_list) == 0: + return None, input_ids + return torch.concat(embedding_list, dim=0), input_ids + + +def _get_multimodal_mask( + input_ids: torch.Tensor, placeholder_tensor: torch.Tensor +) -> torch.Tensor: + return torch.isin(input_ids, placeholder_tensor).unsqueeze(-1) + + +def _adjust_embedding_length( + embedding: torch.Tensor, + mask: torch.Tensor, + logger, +) -> torch.Tensor: + num_mm_tokens_in_embedding = embedding.shape[0] + num_mm_tokens_in_input_ids = mask.sum().item() + if num_mm_tokens_in_input_ids != num_mm_tokens_in_embedding: + logger.warning( + f"Number of tokens in multimodal embedding does not match those in the input text. " + f"Got {num_mm_tokens_in_input_ids} tokens in the text but {num_mm_tokens_in_embedding} " + f"tokens from multimodal embeddings." + ) + if num_mm_tokens_in_input_ids < num_mm_tokens_in_embedding: + chunked_prefill_size = get_schedule().chunked_prefill_size + if chunked_prefill_size != -1: + logger.warning( + "You may want to avoid this issue by raising `chunked_prefill_size`, or disabling chunked prefill" + ) + # extract from the end: this is a compromise + if embedding.dim() == 2: + embedding = embedding[-num_mm_tokens_in_input_ids:, :] + else: + num_multimodal = num_mm_tokens_in_input_ids // embedding.shape[0] + embedding = embedding[-num_multimodal:, :] + else: + raise RuntimeError( + f"Insufficient multimodal embedding length: {num_mm_tokens_in_input_ids=} vs {num_mm_tokens_in_embedding=}. This is an internal error" + ) + return embedding + + +def get_embedding_and_mask( + data_embedding_func: DataEmbeddingFunc, + embedding_items: List[MultimodalDataItem], + placeholder_tensor: torch.Tensor, + input_ids: torch.Tensor, + items_size: List[int], + prefix_length: List[int], + extend_length: List[int], + items_offset_list: List[List[Tuple[int, int]]], +) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]: + """ + Generate multimodal embeddings and create a mask for identifying their positions in the input sequence. + + Args: + data_embedding_func: Function that generates embeddings for multimodal items + embedding_items: List of multimodal items to embed + placeholder_tensor: Tensor containing token IDs that serve as placeholders for multimodal content + input_ids: The input token IDs tensor + items_size: Cumulative sizes of multimodal items per request + prefix_length: Prefix lengths for each request + extend_length: Sequence lengths for each request + items_offset_list: List of offset ranges for multimodal items in each request + + Returns: + A tuple containing: + - The generated embeddings tensor + - A boolean mask tensor indicating where these embeddings should be placed + - If EVS is used, the pruned input ids tensor; otherwise, the original input ids tensor + """ + # 1. Get embedding + embedding = _get_precomputed_embedding( + embedding_items, items_size, prefix_length, extend_length, items_offset_list + ) + if embedding is None: + embedding, input_ids = _get_chunked_prefill_embedding( + data_embedding_func, + embedding_items, + items_size, + prefix_length, + extend_length, + items_offset_list, + input_ids, + ) + if embedding is None: + return None, None, input_ids + # 2. Get mask + if _is_npu: + torch.npu.current_stream().synchronize() + special_multimodal_mask = _get_multimodal_mask(input_ids, placeholder_tensor) + # 3. Adjust embedding length if needed + embedding = _adjust_embedding_length(embedding, special_multimodal_mask, logger) + return embedding, special_multimodal_mask, input_ids diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 97e96c06d..aebb16253 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -9,9 +9,8 @@ import pickle import sys from abc import abstractmethod from collections import defaultdict -from dataclasses import dataclass, field from multiprocessing import shared_memory -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple +from typing import Any, Dict, List, Literal, Optional, Tuple import numpy as np import torch @@ -24,29 +23,30 @@ from sglang.srt.managers.io_struct import ( TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, ) + +# Preserve the existing initialization import for downstream callers. +from sglang.srt.managers.mm_schedule import ( + DataEmbeddingFunc, + get_embedding_and_mask, +) +from sglang.srt.managers.mm_schedule import ( + init_mm_embedding_cache as init_mm_embedding_cache, +) from sglang.srt.managers.schedule_batch import ( CudaIpcTensorTransportProxy, Modality, - MultimodalDataItem, MultimodalInputs, ) -from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.multimodal.evs import EVSEmbeddingResult from sglang.srt.runtime_context import ( get_disagg, - get_parallel, - get_schedule, get_server_args, get_serving, ) -from sglang.srt.utils import flatten_nested_list, is_hip, is_npu, print_warning_once +from sglang.srt.utils import flatten_nested_list, print_warning_once from sglang.srt.utils.stale_shm_cleanup import make_shm_name from sglang.utils import logger -_is_hip = is_hip() -_is_npu = is_npu() - # NOTE: Using the shared logger from sglang.utils instead of creating a module-specific logger # to ensure consistent logging behavior across the codebase. This prevents issues with log # propagation that can cause some log messages (like 'server is fired up') to not appear @@ -365,649 +365,6 @@ class MultiModalityDataPaddingPatternMultimodalTokens(MultiModalityDataPaddingPa return ret_input_ids -embedding_cache: Optional[MultiModalStaticCache] = None - - -def init_mm_embedding_cache(max_size: int = 0): - global embedding_cache - embedding_cache = MultiModalStaticCache(max_size) - - -def get_embedding_chunk( - embedding: torch.Tensor, - extend_prefix_len: int, - extend_seq_len: int, - items_offset: List[Tuple[int, int]], -) -> Tuple[torch.Tensor, int, int]: - """ - Extract a chunk of embeddings based on the specified prefix length, sequence length, and offset ranges. - - Args: - embedding: The full embedding tensor to extract a chunk from - extend_prefix_len: The starting position (prefix length) for extraction - extend_seq_len: The number of tokens to extract - items_offset: List of [start, end] offset ranges for multimodal items in the input sequence - - Returns: - A tuple containing: - - The extracted embedding chunk as a tensor - - The start index used for extraction - - The end index used for extraction - - Note: - If there's no overlap between the requested range and the offset ranges, - an empty tensor is returned with zeros for start and end indices. - """ - start_index, end_index = 0, 0 - extend_start_index = extend_prefix_len - extend_end_index = extend_prefix_len + extend_seq_len - 1 - - for start, end in items_offset: - if extend_start_index >= start and extend_start_index <= end: - start_index += extend_start_index - start - elif extend_start_index > end: - start_index += end - start + 1 - - if extend_end_index >= start and extend_end_index <= end: - end_index += extend_end_index - start + 1 - elif extend_end_index > end: - end_index += end - start + 1 - # some models' embedding is 3-dim, reshape it to 2-dim - embedding = embedding.reshape(-1, embedding.shape[-1]) - embedding_chunk = embedding[start_index:end_index] - return embedding_chunk, start_index, end_index - - -def _get_precomputed_embedding( - items: List[MultimodalDataItem], - items_size: List[int], - prefix_length: List[int], - extend_length: List[int], - items_offset_list: List[List[Tuple[int, int]]], -) -> Optional[torch.Tensor]: - """ - If all items have precomputed_embeddings, return their concatenation. - If some but not all have precomputed_embeddings, raise NotImplementedError. - If none have precomputed_embeddings, return None. - """ - precomputed_embeddings = [] - max_iterations = min(len(items_size) - 1, len(prefix_length)) - - for i in range(max_iterations): - if items_size[i] == items_size[i + 1]: - continue - - items_per_req = items[items_size[i] : items_size[i + 1]] - extend_len = extend_length[i] if i < len(extend_length) else 0 - items_offset = items_offset_list[i] - - if any(item.precomputed_embeddings is None for item in items_per_req): - chunk = None - else: - req_embeddings = torch.concat( - [item.precomputed_embeddings for item in items_per_req] - ) - chunk, _, _ = get_embedding_chunk( - embedding=req_embeddings, - extend_prefix_len=prefix_length[i], - extend_seq_len=extend_len, - items_offset=items_offset, - ) - - if chunk is None and len(items_per_req) > 1: - return None - precomputed_embeddings.append(chunk) - - if any(feature is not None for feature in precomputed_embeddings): - if not all(feature is not None for feature in precomputed_embeddings): - raise NotImplementedError( - "MM inputs where only some items are precomputed." - ) - - # Normalize device across chunks before concat. - target_device = next( - (t.device for t in precomputed_embeddings if t.is_cuda), - precomputed_embeddings[0].device, - ) - precomputed_embeddings = [ - t if t.device == target_device else t.to(target_device, non_blocking=True) - for t in precomputed_embeddings - ] - result = torch.concat(precomputed_embeddings) - # some models embedding is 3-dim, reshape it to 2-dim (similar to get_embedding_chunk) - result = result.reshape(-1, result.shape[-1]) - return result - return None - - -# A modality's embedding function. May return the combined [tokens, hidden] -# tensor, an EVSEmbeddingResult, or one tensor per input item. The per-item -# form lets encoders that naturally produce per-item outputs (e.g. a wav -# AutoEncoder looping over clips) skip an encoder-side torch.cat that -# per-item consumers (_get_chunked_embedding_by_item) would immediately -# split back apart — and each cached entry then owns its storage instead of -# being a view pinning the concatenated buffer. -DataEmbeddingFunc = Callable[ - [List[MultimodalDataItem]], - torch.Tensor | List[torch.Tensor] | EVSEmbeddingResult, -] - - -def _flatten_embedding_result( - embedding: torch.Tensor | List[torch.Tensor], -) -> torch.Tensor: - """Normalize a DataEmbeddingFunc result to one [tokens, hidden] tensor.""" - if isinstance(embedding, list): - if not embedding: - raise ValueError( - "DataEmbeddingFunc returned an empty per-item list; expected " - "one entry per input item" - ) - flat = [e.reshape(-1, e.shape[-1]) for e in embedding] - return flat[0] if len(flat) == 1 else torch.cat(flat, dim=0) - return embedding - - -def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) -> bool: - """Models that materialize and batch visual features inside their encoder. - - instead of performing multiple H2D for each mm feature from all mm_items (followed by concatenation on device), - for some models which internally performs H2D on concated mm feature, these small H2D calls could be replaced with a single big H2D - """ - owner = getattr(data_embedding_func, "__self__", None) - if owner is None: - return False - if getattr(data_embedding_func, "__name__", None) not in ( - "get_image_feature", - "get_video_feature", - ): - return False - return owner.__class__.__name__ in { - "Qwen3VLForConditionalGeneration", - "Qwen3VLMoeForConditionalGeneration", - "Qwen3_5ForConditionalGeneration", - "Qwen3_5MoeForConditionalGeneration", - "KimiK25ForConditionalGeneration", - } - - -def _move_items_to_device( - items: List[MultimodalDataItem], device: torch.device -) -> None: - """Move item features to the target device (in-place, non-blocking).""" - for item in items: - if isinstance(item.feature, torch.Tensor) and item.feature.device != device: - item.feature = item.feature.to(device, non_blocking=True) - - -def _acknowledge_deferred_cuda_ipc_cache_hits( - items: List[MultimodalDataItem], -) -> None: - """Release lazy Kimi IPC slices when a cached embedding skips ViT. - - On an encoder-DP miss, exactly one rank copies an image and acknowledges - the full TP group. On a cache hit no rank copies it, so rank zero performs - the equivalent single acknowledgement. This preserves the fixed-pool - lifecycle without reintroducing an unnecessary GPU-to-GPU copy. - """ - parallel = get_parallel() - if parallel.attn_tp_rank != 0: - return - # The pool's recycler counts the whole TP group, so the acknowledgement must - # match that count even when an attention subgroup is smaller. - consumer_count = max(parallel.tp_size, 1) - for item in items: - item.acknowledge_deferred_cuda_ipc_feature(consumer_count) - - -def _get_chunked_embedding_full( - data_embedding_func: DataEmbeddingFunc, - embedding_items_per_req: List[MultimodalDataItem], - items_offset: List[Tuple[int, int]], - extend_prefix_len: int, - extend_seq_len: int, - input_ids: torch.Tensor, - device: torch.device, -) -> Tuple[Optional[torch.Tensor], torch.Tensor]: - """ - Fallback: encode all items at once, cache combined result, extract chunk. - Used for non-bundled items or EVS results. - """ - item_hashes = [item.hash for item in embedding_items_per_req] - embedding_items_hash = MultiModalStaticCache.combine_hashes(item_hashes) - embedding_per_req = embedding_cache.get(item_hashes) - - 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) - embedding = data_embedding_func(embedding_items_per_req) - if isinstance(embedding, list): - # This path caches the combined per-request embedding, so the - # per-item form is flattened here. - embedding = _flatten_embedding_result(embedding) - embedding_per_req = ( - EmbeddingResult(embedding=embedding) - if isinstance(embedding, torch.Tensor) - else embedding - ) - embedding_cache.set(embedding_items_hash, embedding_per_req) - else: - _acknowledge_deferred_cuda_ipc_cache_hits(embedding_items_per_req) - - if isinstance(embedding_per_req, EVSEmbeddingResult): - item = embedding_items_per_req[0] - input_ids, items_offset = ( - embedding_per_req.redistribute_pruned_frames_placeholders( - input_ids, - items_offset, - item=item, - extend_prefix_len=extend_prefix_len, - extend_seq_len=extend_seq_len, - ) - ) - - embedding_per_req_chunk, _, _ = get_embedding_chunk( - embedding=embedding_per_req.embedding, - extend_prefix_len=extend_prefix_len, - extend_seq_len=extend_seq_len, - items_offset=items_offset, - ) - return embedding_per_req_chunk, input_ids - - -@dataclass -class PerImageRequestInfo: - """Metadata for a single request using the per-image encoding path.""" - - req_idx: int - items: List[MultimodalDataItem] - items_offset: List[Tuple[int, int]] - extend_prefix_len: int - extend_seq_len: int - overlapping: List[Tuple[int, MultimodalDataItem, int, int]] = field( - default_factory=list - ) - - -def _batch_encode_per_image_misses( - data_embedding_func: DataEmbeddingFunc, - per_image_requests: List[PerImageRequestInfo], - device: torch.device, -) -> Dict[int, torch.Tensor]: - """ - Collect cache misses across ALL per-image requests, deduplicate by hash, - encode in a single ViT call, and populate the cache. - - Returns: - hash_to_embedding: mapping from item.hash to its full embedding tensor. - """ - unique_misses: Dict[int, Tuple[MultimodalDataItem, int]] = {} - hash_to_embedding: Dict[int, torch.Tensor] = {} - - # Phase 1a: find overlapping items per request and collect cache misses - for req_info in per_image_requests: - chunk_start = req_info.extend_prefix_len - chunk_end = chunk_start + req_info.extend_seq_len # exclusive - overlapping = [] - if req_info.extend_seq_len > 0: - for idx, (item, (start, end)) in enumerate( - zip(req_info.items, req_info.items_offset) - ): - if end >= chunk_start and start < chunk_end: - overlapping.append((idx, item, start, end)) - req_info.overlapping = overlapping - - for _idx, item, start, end in overlapping: - if item.hash in hash_to_embedding: - continue - cached = embedding_cache.get_single(item.hash) - if cached is not None: - hash_to_embedding[item.hash] = cached.embedding - elif item.hash not in unique_misses: - token_count = end - start + 1 - unique_misses[item.hash] = (item, token_count) - - # Phase 1b: single ViT call for all unique cache misses - if unique_misses: - ordered_hashes = list(unique_misses.keys()) - miss_items = [unique_misses[h][0] for h in ordered_hashes] - token_counts = [unique_misses[h][1] for h in ordered_hashes] - - if not _can_skip_pre_embed_feature_move(data_embedding_func): - _move_items_to_device(miss_items, device) - all_miss_embedding = data_embedding_func(miss_items) - - if isinstance(all_miss_embedding, list): - # Per-item embeddings: no split needed, and each cache entry owns - # its storage (a torch.split view would pin the whole concatenated - # buffer for as long as any single item stays cached). Mirrors - # _get_chunked_embedding_by_item. - assert len(all_miss_embedding) == len(miss_items), ( - f"per-item embedding count {len(all_miss_embedding)} != " - f"cache-miss item count {len(miss_items)}" - ) - split_embeddings = [ - emb.reshape(-1, emb.shape[-1]) for emb in all_miss_embedding - ] - else: - all_miss_embedding = all_miss_embedding.reshape( - -1, all_miss_embedding.shape[-1] - ) - split_embeddings = torch.split(all_miss_embedding, token_counts, dim=0) - for h, emb in zip(ordered_hashes, split_embeddings): - embedding_cache.set(h, EmbeddingResult(embedding=emb)) - # Keep a local ref (no extra GPU memory) so assembly never fails due to LRU eviction. - hash_to_embedding[h] = emb - - return hash_to_embedding - - -def _get_chunked_embedding_by_item( - data_embedding_func: DataEmbeddingFunc, - embedding_items_per_req: List[MultimodalDataItem], - items_offset: List[Tuple[int, int]], - extend_prefix_len: int, - extend_seq_len: int, - device: torch.device, -) -> Optional[torch.Tensor]: - """ - Per-image chunk-aware encoding for one request. - Items must already be split per-image (each item has exactly one offset). - """ - chunk_start = extend_prefix_len - chunk_end = extend_prefix_len + extend_seq_len # exclusive - - if extend_seq_len <= 0: - return None - - overlapping = [] - for idx, (item, (start, end)) in enumerate( - zip(embedding_items_per_req, items_offset) - ): - if end >= chunk_start and start < chunk_end: - overlapping.append((idx, item, start, end)) - - if not overlapping: - return None - - cached_embeddings = {} - miss_items = [] - for idx, item, start, end in overlapping: - cached = embedding_cache.get_single(item.hash) - if cached is not None: - cached_embeddings[idx] = cached.embedding - _acknowledge_deferred_cuda_ipc_cache_hits([item]) - else: - miss_items.append((idx, item, start, end)) - - 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) - all_miss_embedding = data_embedding_func(miss_item_list) - - if isinstance(all_miss_embedding, list): - # Per-item embeddings: no split needed, and each cache entry owns - # its storage (a torch.split view would pin the whole concatenated - # buffer for as long as any single item stays cached). - assert len(all_miss_embedding) == len(miss_items), ( - f"per-item embedding count {len(all_miss_embedding)} != " - f"cache-miss item count {len(miss_items)}" - ) - split_embeddings = [ - emb.reshape(-1, emb.shape[-1]) for emb in all_miss_embedding - ] - else: - all_miss_embedding = all_miss_embedding.reshape( - -1, all_miss_embedding.shape[-1] - ) - # Split output by per-item token count - token_counts = [end - start + 1 for _, _, start, end in miss_items] - split_embeddings = torch.split(all_miss_embedding, token_counts, dim=0) - - for (idx, item, _, _), emb in zip(miss_items, split_embeddings): - cached_embeddings[idx] = emb - embedding_cache.set(item.hash, EmbeddingResult(embedding=emb)) - - chunk_slices = [] - for idx, _, start, end in overlapping: - emb = cached_embeddings[idx] - overlap_start = max(start, chunk_start) - overlap_end = min(end, chunk_end - 1) # inclusive - local_start = overlap_start - start - local_end = overlap_end - start + 1 # exclusive for slicing - chunk_slices.append(emb[local_start:local_end]) - - return torch.cat(chunk_slices, dim=0) - - -def _assemble_per_image_chunk( - overlapping: List[Tuple[int, MultimodalDataItem, int, int]], - hash_to_embedding: Dict[int, torch.Tensor], - extend_prefix_len: int, - extend_seq_len: int, -) -> Optional[torch.Tensor]: - """ - Assemble the chunk embedding for one request from pre-computed embeddings. - All overlapping items must already have their embeddings in hash_to_embedding. - """ - if not overlapping: - return None - - chunk_start = extend_prefix_len - chunk_end = extend_prefix_len + extend_seq_len # exclusive - - chunk_slices = [] - for _idx, item, start, end in overlapping: - emb = hash_to_embedding[item.hash] # shape: (end - start + 1, hidden) - overlap_start = max(start, chunk_start) - overlap_end = min(end, chunk_end - 1) # inclusive - local_start = overlap_start - start - local_end = overlap_end - start + 1 # exclusive for slicing - chunk_slices.append(emb[local_start:local_end]) - - return torch.cat(chunk_slices, dim=0) - - -def _get_chunked_prefill_embedding( - data_embedding_func: DataEmbeddingFunc, - embedding_items: List[MultimodalDataItem], - items_size: List[int], - prefix_length: List[int], - extend_length: List[int], - items_offset_list: List[List[Tuple[int, int]]], - input_ids: torch.Tensor, -) -> tuple[torch.Tensor | None, torch.Tensor]: - """ - Chunked prefill embedding: encode items across all requests and extract - per-request chunks. Images from all requests are batched into a single - ViT call for efficiency. - """ - device = input_ids.device - # FIXME(Xinyuan): temporary workaround for eagle3 - max_iterations = min(len(items_size) - 1, len(prefix_length)) - - # Phase 0: classify requests into per-image vs full/EVS path - per_image_requests = [] # batched ViT encoding - full_path_requests = [] # per-request encoding (EVS etc.) - all_chunks: List[Tuple[int, torch.Tensor]] = [] - - for i in range(max_iterations): - if items_size[i] == items_size[i + 1]: - continue - embedding_items_per_req = embedding_items[items_size[i] : items_size[i + 1]] - items_offset = items_offset_list[i] - assert items_offset is not None, items_offset - - extend_prefix_len = prefix_length[i] - extend_seq_len = extend_length[i] if i < len(extend_length) else 0 - if extend_seq_len <= 0: - continue - - # Skip if all items already prefilled. - if all(offset_end < prefix_length[i] for _, offset_end in items_offset): - continue - - req_info = PerImageRequestInfo( - req_idx=i, - items=embedding_items_per_req, - items_offset=items_offset, - extend_prefix_len=extend_prefix_len, - extend_seq_len=extend_seq_len, - ) - - is_per_image = all(len(item.offsets) == 1 for item in embedding_items_per_req) - if is_per_image: - if _is_hip or _is_npu: - # ROCm CI regressed with one large cross-request ViT batch; keep - # the previous per-request path on HIP while CUDA uses batching. - chunk = _get_chunked_embedding_by_item( - data_embedding_func, - embedding_items_per_req, - items_offset, - extend_prefix_len, - extend_seq_len, - device, - ) - if chunk is not None: - all_chunks.append((i, chunk)) - else: - per_image_requests.append(req_info) - else: - full_path_requests.append(req_info) - - # Phase 1: batch encode all per-image cache misses in ONE ViT call - hash_to_embedding: Dict[int, torch.Tensor] = {} - if per_image_requests: - hash_to_embedding = _batch_encode_per_image_misses( - data_embedding_func, per_image_requests, device - ) - - # Phase 2: assemble per-request chunks in original request order - for req_info in per_image_requests: - chunk = _assemble_per_image_chunk( - req_info.overlapping, - hash_to_embedding, - req_info.extend_prefix_len, - req_info.extend_seq_len, - ) - if chunk is not None: - all_chunks.append((req_info.req_idx, chunk)) - - for req_info in full_path_requests: - chunk_embedding, input_ids = _get_chunked_embedding_full( - data_embedding_func, - req_info.items, - req_info.items_offset, - req_info.extend_prefix_len, - req_info.extend_seq_len, - input_ids, - device, - ) - if chunk_embedding is not None: - all_chunks.append((req_info.req_idx, chunk_embedding)) - - # Sort by original request index to maintain correct output order - all_chunks.sort(key=lambda x: x[0]) - embedding_list = [chunk for _, chunk in all_chunks] - - if len(embedding_list) == 0: - return None, input_ids - return torch.concat(embedding_list, dim=0), input_ids - - -def _get_multimodal_mask( - input_ids: torch.Tensor, placeholder_tensor: torch.Tensor -) -> torch.Tensor: - return torch.isin(input_ids, placeholder_tensor).unsqueeze(-1) - - -def _adjust_embedding_length( - embedding: torch.Tensor, - mask: torch.Tensor, - logger, -) -> torch.Tensor: - num_mm_tokens_in_embedding = embedding.shape[0] - num_mm_tokens_in_input_ids = mask.sum().item() - if num_mm_tokens_in_input_ids != num_mm_tokens_in_embedding: - logger.warning( - f"Number of tokens in multimodal embedding does not match those in the input text. " - f"Got {num_mm_tokens_in_input_ids} tokens in the text but {num_mm_tokens_in_embedding} " - f"tokens from multimodal embeddings." - ) - if num_mm_tokens_in_input_ids < num_mm_tokens_in_embedding: - chunked_prefill_size = get_schedule().chunked_prefill_size - if chunked_prefill_size != -1: - logger.warning( - "You may want to avoid this issue by raising `chunked_prefill_size`, or disabling chunked prefill" - ) - # extract from the end: this is a compromise - if embedding.dim() == 2: - embedding = embedding[-num_mm_tokens_in_input_ids:, :] - else: - num_multimodal = num_mm_tokens_in_input_ids // embedding.shape[0] - embedding = embedding[-num_multimodal:, :] - else: - raise RuntimeError( - f"Insufficient multimodal embedding length: {num_mm_tokens_in_input_ids=} vs {num_mm_tokens_in_embedding=}. This is an internal error" - ) - return embedding - - -def get_embedding_and_mask( - data_embedding_func: DataEmbeddingFunc, - embedding_items: List[MultimodalDataItem], - placeholder_tensor: torch.Tensor, - input_ids: torch.Tensor, - items_size: List[int], - prefix_length: List[int], - extend_length: List[int], - items_offset_list: List[List[Tuple[int, int]]], -) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]: - """ - Generate multimodal embeddings and create a mask for identifying their positions in the input sequence. - - Args: - data_embedding_func: Function that generates embeddings for multimodal items - embedding_items: List of multimodal items to embed - placeholder_tensor: Tensor containing token IDs that serve as placeholders for multimodal content - input_ids: The input token IDs tensor - items_size: Cumulative sizes of multimodal items per request - prefix_length: Prefix lengths for each request - extend_length: Sequence lengths for each request - items_offset_list: List of offset ranges for multimodal items in each request - - Returns: - A tuple containing: - - The generated embeddings tensor - - A boolean mask tensor indicating where these embeddings should be placed - - If EVS is used, the pruned input ids tensor; otherwise, the original input ids tensor - """ - # 1. Get embedding - embedding = _get_precomputed_embedding( - embedding_items, items_size, prefix_length, extend_length, items_offset_list - ) - if embedding is None: - embedding, input_ids = _get_chunked_prefill_embedding( - data_embedding_func, - embedding_items, - items_size, - prefix_length, - extend_length, - items_offset_list, - input_ids, - ) - if embedding is None: - return None, None, input_ids - # 2. Get mask - if _is_npu: - torch.npu.current_stream().synchronize() - special_multimodal_mask = _get_multimodal_mask(input_ids, placeholder_tensor) - # 3. Adjust embedding length if needed - embedding = _adjust_embedding_length(embedding, special_multimodal_mask, logger) - return embedding, special_multimodal_mask, input_ids - - def embed_mm_inputs( mm_inputs_list: List[MultimodalInputs], extend_prefix_lens: List[int], diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 50f797101..2a396aed8 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -32,7 +32,7 @@ from sglang.srt.configs.hybrid_arch import ( ) from sglang.srt.configs.model_config import ModelImpl, is_deepseek_dsa from sglang.srt.environ import envs -from sglang.srt.managers.mm_utils import init_mm_embedding_cache +from sglang.srt.managers.mm_schedule import init_mm_embedding_cache from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache from sglang.srt.model_loader.utils import get_resolved_model_impl diff --git a/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py b/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py index e92588071..e286ad565 100644 --- a/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py +++ b/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py @@ -1,18 +1,18 @@ """Unit tests for per-item DataEmbeddingFunc results in the chunked mm path. A DataEmbeddingFunc may return either one combined [tokens, hidden] tensor or -one tensor per item (see mm_utils.DataEmbeddingFunc). These tests assert the +one tensor per item (see mm_schedule.DataEmbeddingFunc). These tests assert the two forms produce bitwise-identical chunked-prefill embeddings, and that the per-item form yields cache entries that own their storage (a torch.split view of the combined tensor pins the whole concatenated buffer). -CPU-only: exercises mm_utils internals directly, no engine or GPU. +CPU-only: exercises mm_schedule internals directly, no engine or GPU. """ import pytest import torch -from sglang.srt.managers import mm_utils +from sglang.srt.managers import mm_schedule from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.runtime_context import get_context, get_parallel from sglang.test.ci.ci_register import register_cpu_ci @@ -66,6 +66,14 @@ CHUNKS = [(0, 8), (8, 8), (16, 8), (24, 6)] _CPU = torch.device("cpu") +@pytest.fixture(autouse=True) +def _skip_cuda_ipc_acknowledgement(monkeypatch): + """Keep CPU embedding tests independent of tensor-parallel runtime state.""" + monkeypatch.setattr( + mm_schedule, "_acknowledge_deferred_cuda_ipc_cache_hits", lambda _items: None + ) + + def _num_tokens(item: MultimodalDataItem) -> int: start, end = item.offsets[0] return end - start + 1 @@ -97,10 +105,10 @@ def _make_items(): def _run_by_item_chunks(encoder): - mm_utils.init_mm_embedding_cache(1 << 30) + mm_schedule.init_mm_embedding_cache(1 << 30) items = _make_items() return [ - mm_utils._get_chunked_embedding_by_item( + mm_schedule._get_chunked_embedding_by_item( encoder, items, ITEM_OFFSETS, prefix_len, extend_len, _CPU ) for prefix_len, extend_len in CHUNKS @@ -108,12 +116,12 @@ def _run_by_item_chunks(encoder): def _run_full_chunks(encoder): - mm_utils.init_mm_embedding_cache(1 << 30) + mm_schedule.init_mm_embedding_cache(1 << 30) items = _make_items() input_ids = torch.zeros(TOTAL_LEN, dtype=torch.long) outs = [] for prefix_len, extend_len in CHUNKS: - chunk, _ = mm_utils._get_chunked_embedding_full( + chunk, _ = mm_schedule._get_chunked_embedding_full( encoder, items, ITEM_OFFSETS, prefix_len, extend_len, input_ids, _CPU ) outs.append(chunk) @@ -150,13 +158,13 @@ def test_full_matches_by_item(): def test_list_cache_entries_own_storage(): - mm_utils.init_mm_embedding_cache(1 << 30) + mm_schedule.init_mm_embedding_cache(1 << 30) items = _make_items() - mm_utils._get_chunked_embedding_by_item( + mm_schedule._get_chunked_embedding_by_item( _encoder_list, items, ITEM_OFFSETS, 0, TOTAL_LEN, _CPU ) for item in items: - emb = mm_utils.embedding_cache.get_single(item.hash).embedding + emb = mm_schedule.embedding_cache.get_single(item.hash).embedding own_bytes = emb.numel() * emb.element_size() assert emb.untyped_storage().nbytes() == own_bytes @@ -164,14 +172,14 @@ def test_list_cache_entries_own_storage(): def test_tensor_cache_entries_share_storage(): # Documents the motivation for the per-item form: split views of the # combined tensor keep the whole concatenated buffer alive. - mm_utils.init_mm_embedding_cache(1 << 30) + mm_schedule.init_mm_embedding_cache(1 << 30) items = _make_items() - mm_utils._get_chunked_embedding_by_item( + mm_schedule._get_chunked_embedding_by_item( _encoder_tensor, items, ITEM_OFFSETS, 0, TOTAL_LEN, _CPU ) total_tokens = sum(_num_tokens(item) for item in items) for item in items: - emb = mm_utils.embedding_cache.get_single(item.hash).embedding + emb = mm_schedule.embedding_cache.get_single(item.hash).embedding assert ( emb.untyped_storage().nbytes() == total_tokens * HIDDEN * emb.element_size() )