[VLM] Avoid synchronizing multimodal placeholder counts (#34995)

Co-authored-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>

python/sglang/srt/managers/mm_schedule.py

@@ -10,6 +10,7 @@ from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalSta
 from sglang.srt.multimodal.evs import EVSEmbeddingResult
 from sglang.srt.runtime_context import get_parallel, get_schedule
 from sglang.srt.utils import is_hip, is_npu
+from sglang.srt.utils.async_probe import maybe_assert_sum
 from sglang.utils import logger
 
 _is_hip = is_hip()
@@ -573,13 +574,31 @@ def _get_multimodal_mask(
     return torch.isin(input_ids, placeholder_tensor).unsqueeze(-1)
 
 
+def _count_mm_tokens_in_extend(
+    prefix_length: List[int],
+    extend_length: List[int],
+    items_offset_list: List[List[Tuple[int, int]]],
+) -> int:
+    """Count MM placeholders from host offsets without reading back the GPU mask."""
+    num_mm_tokens = 0
+    for i, (extend_start, items_offset) in enumerate(
+        zip(prefix_length, items_offset_list)
+    ):
+        extend_end = extend_start + (extend_length[i] if i < len(extend_length) else 0)
+        for item_start, item_end in items_offset:
+            overlap_start = max(item_start, extend_start)
+            overlap_end = min(item_end + 1, extend_end)
+            num_mm_tokens += max(overlap_end - overlap_start, 0)
+
+    return num_mm_tokens
+
+
 def _adjust_embedding_length(
     embedding: torch.Tensor,
-    mask: torch.Tensor,
+    num_mm_tokens_in_input_ids: int,
     logger,
 ) -> torch.Tensor:
     num_mm_tokens_in_embedding = embedding.shape[0]
-    num_mm_tokens_in_input_ids = mask.sum().item()
     if num_mm_tokens_in_input_ids != num_mm_tokens_in_embedding:
         logger.warning(
             f"Number of tokens in multimodal embedding does not match those in the input text. "
@@ -634,6 +653,13 @@ def get_embedding_and_mask(
         - A boolean mask tensor indicating where these embeddings should be placed
         - If EVS is used, the pruned input ids tensor; otherwise, the original input ids tensor
     """
+    original_input_ids = input_ids
+    num_mm_tokens_in_input_ids = _count_mm_tokens_in_extend(
+        prefix_length,
+        extend_length,
+        items_offset_list,
+    )
+
     # 1. Get embedding
     embedding = _get_precomputed_embedding(
         embedding_items, items_size, prefix_length, extend_length, items_offset_list
@@ -655,5 +681,14 @@ def get_embedding_and_mask(
         torch.npu.current_stream().synchronize()
     special_multimodal_mask = _get_multimodal_mask(input_ids, placeholder_tensor)
     # 3. Adjust embedding length if needed
-    embedding = _adjust_embedding_length(embedding, special_multimodal_mask, logger)
+    if input_ids is not original_input_ids:
+        # EVS rewrites placeholder spans after pruning, making the original offsets stale.
+        num_mm_tokens_in_input_ids = special_multimodal_mask.sum().item()
+    else:
+        maybe_assert_sum(
+            special_multimodal_mask,
+            num_mm_tokens_in_input_ids,
+            "MM placeholder count derived from offsets does not match input_ids",
+        )
+    embedding = _adjust_embedding_length(embedding, num_mm_tokens_in_input_ids, logger)
     return embedding, special_multimodal_mask, input_ids

python/sglang/srt/utils/async_probe.py

@@ -82,6 +82,12 @@ def maybe_assert_async(cond: torch.Tensor, msg: str = ""):
     torch._assert_async(cond, msg)
 
 
+def maybe_assert_sum(tensor: torch.Tensor, expected: int, msg: str = "") -> None:
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    torch._assert_async(tensor.sum() == expected, msg)
+
+
 def maybe_detect_nan(tensor: Optional[torch.Tensor], msg: str = ""):
     """Async NaN check — no GPU-CPU sync, error surfaces at next sync point."""
     if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():

test/registered/unit/managers/test_mm_embedding_length.py

@@ -0,0 +1,168 @@
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.managers import mm_schedule as mm_utils
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-b-test-cpu")
+
+
+@pytest.mark.parametrize(
+    (
+        "prefix_length",
+        "extend_length",
+        "items_offset_list",
+        "expected",
+    ),
+    [
+        ([8], [16], [[(2, 5), (9, 14), (20, 24)]], 10),
+        ([30], [0], [[(2, 5), (9, 14), (20, 24)]], 0),
+        (
+            [4, 0, 10],
+            [4, 10, 10],
+            [[(2, 5)], [], [(5, 12), (18, 25)]],
+            7,
+        ),
+    ],
+)
+def test_count_mm_tokens_in_extend(
+    prefix_length, extend_length, items_offset_list, expected
+):
+    input_ids = []
+    for prefix, extend, item_offsets in zip(
+        prefix_length, extend_length, items_offset_list
+    ):
+        seq_len = max(
+            prefix + extend,
+            max((item_end + 1 for _, item_end in item_offsets), default=0),
+        )
+        req_input_ids = torch.zeros(seq_len, dtype=torch.long)
+        for item_start, item_end in item_offsets:
+            req_input_ids[item_start : item_end + 1] = 1
+        input_ids.append(req_input_ids[prefix : prefix + extend])
+
+    actual = torch.isin(torch.cat(input_ids), torch.tensor([1])).sum().item()
+    derived = mm_utils._count_mm_tokens_in_extend(
+        prefix_length=prefix_length,
+        extend_length=extend_length,
+        items_offset_list=items_offset_list,
+    )
+    assert actual == derived == expected
+
+
+def test_get_embedding_and_mask_uses_offset_count_without_readback():
+    input_ids = torch.zeros(8, dtype=torch.long)
+    input_ids[2:5] = 1
+    embedding = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    mask = Mock()
+    mask.sum.side_effect = AssertionError("mask count must stay on device")
+
+    with (
+        envs.SGLANG_ENABLE_ASYNC_ASSERT.override(False),
+        patch.object(mm_utils, "_get_precomputed_embedding", return_value=embedding),
+        patch.object(mm_utils, "_get_multimodal_mask", return_value=mask),
+    ):
+        result, result_mask, result_input_ids = mm_utils.get_embedding_and_mask(
+            data_embedding_func=Mock(),
+            embedding_items=[],
+            placeholder_tensor=torch.tensor([1]),
+            input_ids=input_ids,
+            items_size=[0, 1],
+            prefix_length=[0],
+            extend_length=[8],
+            items_offset_list=[[(2, 4)]],
+        )
+
+    mask.sum.assert_not_called()
+    assert result is embedding
+    assert result_mask is mask
+    assert result_input_ids is input_ids
+
+
+def test_get_embedding_and_mask_async_asserts_offset_count():
+    input_ids = torch.zeros(8, dtype=torch.long)
+    input_ids[2:5] = 1
+    embedding = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+
+    with (
+        envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True),
+        patch.object(mm_utils, "_get_precomputed_embedding", return_value=embedding),
+        patch.object(mm_utils.torch, "_assert_async") as assert_async,
+    ):
+        mm_utils.get_embedding_and_mask(
+            data_embedding_func=Mock(),
+            embedding_items=[],
+            placeholder_tensor=torch.tensor([1]),
+            input_ids=input_ids,
+            items_size=[0, 1],
+            prefix_length=[0],
+            extend_length=[8],
+            items_offset_list=[[(2, 4)]],
+        )
+
+    assert_async.assert_called_once()
+    condition, message = assert_async.call_args.args
+    assert condition.item()
+    assert "derived from offsets" in message
+
+
+def test_adjust_embedding_length_crops_overlong_embedding():
+    embedding = torch.arange(20, dtype=torch.float32).reshape(5, 4)
+    server_args = Mock(chunked_prefill_size=-1)
+
+    with patch.object(mm_utils, "get_schedule", return_value=server_args):
+        result = mm_utils._adjust_embedding_length(embedding, 3, Mock())
+
+    torch.testing.assert_close(result, embedding[-3:], rtol=0, atol=0)
+
+
+def test_adjust_embedding_length_rejects_short_embedding():
+    embedding = torch.zeros(2, 4)
+
+    with pytest.raises(RuntimeError, match="Insufficient multimodal embedding length"):
+        mm_utils._adjust_embedding_length(embedding, 3, Mock())
+
+
+def test_get_embedding_and_mask_falls_back_after_input_ids_rewrite():
+    input_ids = torch.zeros(8, dtype=torch.long)
+    rewritten_input_ids = input_ids.clone()
+    embedding = torch.zeros(2, 4)
+    mask_sum = Mock()
+    mask_sum.item.return_value = 2
+    mask = Mock()
+    mask.sum.return_value = mask_sum
+
+    with (
+        patch.object(mm_utils, "_get_precomputed_embedding", return_value=None),
+        patch.object(
+            mm_utils,
+            "_get_chunked_prefill_embedding",
+            return_value=(embedding, rewritten_input_ids),
+        ),
+        patch.object(mm_utils, "_get_multimodal_mask", return_value=mask),
+    ):
+        result, result_mask, result_input_ids = mm_utils.get_embedding_and_mask(
+            data_embedding_func=Mock(),
+            embedding_items=[],
+            placeholder_tensor=torch.tensor([1]),
+            input_ids=input_ids,
+            items_size=[0, 1],
+            prefix_length=[0],
+            extend_length=[8],
+            items_offset_list=[[(2, 4)]],
+        )
+
+    mask.sum.assert_called_once_with()
+    mask_sum.item.assert_called_once_with()
+    assert result is embedding
+    assert result_mask is mask
+    assert result_input_ids is rewritten_input_ids
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))