diff --git a/python/sglang/srt/managers/mm_schedule.py b/python/sglang/srt/managers/mm_schedule.py index 9e91f6366..ed291525f 100644 --- a/python/sglang/srt/managers/mm_schedule.py +++ b/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 diff --git a/python/sglang/srt/utils/async_probe.py b/python/sglang/srt/utils/async_probe.py index 92d98f87b..abe25cd9a 100644 --- a/python/sglang/srt/utils/async_probe.py +++ b/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(): diff --git a/test/registered/unit/managers/test_mm_embedding_length.py b/test/registered/unit/managers/test_mm_embedding_length.py new file mode 100644 index 000000000..ff5e0f40a --- /dev/null +++ b/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"]))