diff --git a/python/sglang/srt/mem_cache/multimodal_cache.py b/python/sglang/srt/mem_cache/multimodal_cache.py index 0f4ee734f..de5bc7a75 100644 --- a/python/sglang/srt/mem_cache/multimodal_cache.py +++ b/python/sglang/srt/mem_cache/multimodal_cache.py @@ -1,6 +1,6 @@ import abc from collections import OrderedDict -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import List, Optional import torch @@ -65,7 +65,7 @@ class MultimodalCache(abc.ABC): def _get_tensor_size(embedding: torch.Tensor): - return embedding.element_size() * embedding.numel() + return embedding.untyped_storage().nbytes() @dataclass(kw_only=True) @@ -109,13 +109,19 @@ class MultiModalStaticCache(MultimodalCache): if mm_hash in self.mm_cache: self.mm_cache.move_to_end(mm_hash) return True - data_size = _get_tensor_size(embedding.embedding) + tensor = embedding.embedding + storage_size = _get_tensor_size(tensor) + data_size = min(storage_size, tensor.element_size() * tensor.numel()) while self.current_size + data_size > self.max_size: if not self.mm_cache: return False lru_hash, lru_embedding = self.mm_cache.popitem(last=False) self.current_size -= _get_tensor_size(lru_embedding.embedding) + if storage_size > data_size: + # Clone admitted slices so one cached item cannot retain a whole batch. + embedding = replace(embedding, embedding=tensor.clone()) + data_size = _get_tensor_size(embedding.embedding) self.mm_cache[mm_hash] = embedding self.current_size += data_size return True 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 ee53a3da0..fa1063506 100644 --- a/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py +++ b/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py @@ -177,20 +177,15 @@ def test_list_cache_entries_own_storage(): assert emb.untyped_storage().nbytes() == own_bytes -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. +def test_tensor_cache_entries_own_storage(): mm_schedule.init_mm_embedding_cache(1 << 30) items = _make_items() 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_schedule.embedding_cache.get_single(item.hash).embedding - assert ( - emb.untyped_storage().nbytes() == total_tokens * HIDDEN * emb.element_size() - ) + assert emb.untyped_storage().nbytes() == emb.numel() * emb.element_size() def test_by_item_mismatched_cache_entry_is_reencoded():