[mm] Accept per-item embedding lists from DataEmbeddingFunc (#31826)

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

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"]))