diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 9d0e4ccb1..d8cf47e1f 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -671,11 +671,24 @@ def _batch_encode_per_image_misses( 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) - all_miss_embedding = all_miss_embedding.reshape( - -1, all_miss_embedding.shape[-1] - ) - split_embeddings = torch.split(all_miss_embedding, token_counts, dim=0) + 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.