From 448662e85e325f830cb2d1d423fa23099677cee4 Mon Sep 17 00:00:00 2001 From: Lu Fang <30275821+houseroad@users.noreply.github.com> Date: Fri, 24 Jul 2026 08:27:24 -0700 Subject: [PATCH] [mm] Accept per-item embedding lists from DataEmbeddingFunc (#31826) --- python/sglang/srt/managers/mm_utils.py | 52 +++++- .../test_mm_chunked_embedding_unit.py | 149 ++++++++++++++++++ 2 files changed, 195 insertions(+), 6 deletions(-) create mode 100644 test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 891d8df6c..9d0e4ccb1 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -474,11 +474,34 @@ def _get_precomputed_embedding( 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 | EVSEmbeddingResult + [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. @@ -553,6 +576,10 @@ def _get_chunked_embedding_full( 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) @@ -700,12 +727,25 @@ def _get_chunked_embedding_by_item( 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) - all_miss_embedding = all_miss_embedding.reshape( - -1, all_miss_embedding.shape[-1] - ) - token_counts = [end - start + 1 for _, _, start, end in miss_items] - 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). + 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 diff --git a/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py b/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py new file mode 100644 index 000000000..35814459f --- /dev/null +++ b/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py @@ -0,0 +1,149 @@ +"""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 +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. +""" + +import pytest +import torch + +from sglang.srt.managers import mm_utils +from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-b-test-cpu") + +HIDDEN = 16 + +# Three items with text gaps between their placeholder runs; offsets are +# (start, end) inclusive, mirroring processor output. +ITEM_OFFSETS = [(2, 5), (9, 14), (20, 24)] +TOTAL_LEN = 30 + +# Chunk windows (prefix_len, extend_len) covering the sequence, sized so item +# boundaries fall both inside and across chunks. +CHUNKS = [(0, 8), (8, 8), (16, 8), (24, 6)] + +_CPU = torch.device("cpu") + + +def _num_tokens(item: MultimodalDataItem) -> int: + start, end = item.offsets[0] + return end - start + 1 + + +def _item_embedding(item: MultimodalDataItem) -> torch.Tensor: + gen = torch.Generator().manual_seed(item.hash) + return torch.randn(_num_tokens(item), HIDDEN, generator=gen) + + +def _encoder_tensor(items): + return torch.cat([_item_embedding(item) for item in items], dim=0) + + +def _encoder_list(items): + return [_item_embedding(item) for item in items] + + +def _make_items(): + return [ + MultimodalDataItem( + modality=Modality.IMAGE, + hash=1000 + i, + feature=torch.zeros(1), + offsets=[offset], + ) + for i, offset in enumerate(ITEM_OFFSETS) + ] + + +def _run_by_item_chunks(encoder): + mm_utils.init_mm_embedding_cache(1 << 30) + items = _make_items() + return [ + mm_utils._get_chunked_embedding_by_item( + encoder, items, ITEM_OFFSETS, prefix_len, extend_len, _CPU + ) + for prefix_len, extend_len in CHUNKS + ] + + +def _run_full_chunks(encoder): + mm_utils.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( + encoder, items, ITEM_OFFSETS, prefix_len, extend_len, input_ids, _CPU + ) + outs.append(chunk) + return outs + + +def _assert_chunks_equal(chunks_a, chunks_b): + assert len(chunks_a) == len(chunks_b) + for a, b in zip(chunks_a, chunks_b): + if a is None or b is None: + assert a is None and b is None + continue + assert a.shape == b.shape + torch.testing.assert_close(a, b, rtol=0, atol=0) + + +def test_by_item_list_matches_tensor(): + _assert_chunks_equal( + _run_by_item_chunks(_encoder_tensor), _run_by_item_chunks(_encoder_list) + ) + + +def test_full_list_matches_tensor(): + _assert_chunks_equal( + _run_full_chunks(_encoder_tensor), _run_full_chunks(_encoder_list) + ) + + +def test_full_matches_by_item(): + # The two chunked strategies agree with each other for single-offset items. + _assert_chunks_equal( + _run_full_chunks(_encoder_tensor), _run_by_item_chunks(_encoder_list) + ) + + +def test_list_cache_entries_own_storage(): + mm_utils.init_mm_embedding_cache(1 << 30) + items = _make_items() + mm_utils._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 + own_bytes = emb.numel() * emb.element_size() + 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. + mm_utils.init_mm_embedding_cache(1 << 30) + items = _make_items() + mm_utils._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 + assert ( + emb.untyped_storage().nbytes() == total_tokens * HIDDEN * emb.element_size() + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))