feat: basic support for server-level multimodal cache (#10775)
This commit is contained in:
@@ -18,7 +18,7 @@ from sglang.srt.managers.schedule_batch import (
|
||||
MultimodalDataItem,
|
||||
MultimodalInputs,
|
||||
)
|
||||
from sglang.srt.mem_cache.multimodal_cache import MultiModalCache
|
||||
from sglang.srt.mem_cache.multimodal_cache import MultiModalStaticCache
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.utils import flatten_nested_list, is_npu, print_warning_once
|
||||
@@ -285,17 +285,12 @@ class MultiModalityDataPaddingPatternMultimodalTokens(MultiModalityDataPaddingPa
|
||||
return ret_input_ids
|
||||
|
||||
|
||||
embedding_cache: Optional[MultiModalCache] = None
|
||||
embedding_cache: Optional[MultiModalStaticCache] = None
|
||||
|
||||
|
||||
def init_embedding_cache(max_size: int = 0):
|
||||
def init_mm_embedding_cache(max_size: int = 0):
|
||||
global embedding_cache
|
||||
embedding_cache = MultiModalCache(max_size)
|
||||
|
||||
|
||||
def get_embedding_hash(embedding_items: List[MultimodalDataItem]) -> int:
|
||||
hash_list = [item.hash for item in embedding_items]
|
||||
return hash(tuple(hash_list))
|
||||
embedding_cache = MultiModalStaticCache(max_size)
|
||||
|
||||
|
||||
def get_embedding_chunk(
|
||||
@@ -382,14 +377,15 @@ def _get_chunked_prefill_embedding(
|
||||
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
|
||||
embedding_items_hash = get_embedding_hash(embedding_items_per_req)
|
||||
# if all items has been prefixed, we do not need to calculate embedding
|
||||
if all([offset_end < prefix_length[i] for _, offset_end in items_offset]):
|
||||
continue
|
||||
embedding_per_req = embedding_cache.get(embedding_items_hash)
|
||||
item_hashes = [item.hash for item in embedding_items]
|
||||
embedding_items_hash = MultiModalStaticCache.combine_hashes(item_hashes)
|
||||
embedding_per_req = embedding_cache.get(item_hashes)
|
||||
if embedding_per_req is None:
|
||||
embedding_per_req = data_embedding_func(embedding_items_per_req)
|
||||
if not embedding_cache.put(embedding_items_hash, embedding_per_req):
|
||||
if not embedding_cache.set(embedding_items_hash, embedding_per_req):
|
||||
print_warning_once(
|
||||
"Multimodal embedding cache is full. This typically occurs when a single "
|
||||
"embedding exceeds the cache size limit. Consider increasing the "
|
||||
|
||||
Reference in New Issue
Block a user