diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 7e2611587..3a957ae7f 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -364,6 +364,27 @@ class MultiModalityDataPaddingPatternMultimodalTokens(MultiModalityDataPaddingPa return ret_input_ids +# masked_scatter_ materializes the expanded [num_tokens, hidden] bool mask plus +# an int64 prefix-sum over it (~9 B per num_tokens x hidden element); the +# cumsum-derived row indices keep the transients O(num_tokens) and sync-free. +def _scatter_mm_embedding( + dest: torch.Tensor, mask: torch.Tensor, src: torch.Tensor +) -> None: + # mask: [num_tokens, 1] bool; src: [num_mm_tokens, width] in sequence order. + src = src.to(dest.device, dest.dtype) + num_src_rows = src.size(0) + flat_mask = mask.view(-1) + ranks = torch.cumsum(flat_mask, dim=0) - 1 + # False rows collapse into the discard slot num_src_rows; a mask/src + # row-count mismatch device-asserts in scatter_/index_copy_ (poison init). + ranks = ranks.masked_fill(~flat_mask, num_src_rows) + rows = torch.full( + (num_src_rows + 1,), dest.size(0), dtype=torch.long, device=dest.device + ) + rows.scatter_(0, ranks, torch.arange(flat_mask.numel(), device=dest.device)) + dest.index_copy_(0, rows[:num_src_rows], src) + + def embed_mm_inputs( mm_inputs_list: List[MultimodalInputs], extend_prefix_lens: List[int], @@ -485,18 +506,16 @@ def embed_mm_inputs( other_info["input_deepstack_embeds"] = input_deepstack_embeds # 4. scatter embeddings into input embedding - # masked_scatter_ avoids the cudaStreamSynchronize that torch.where triggers. - def _scatter(dest, mask, src): - dest.masked_scatter_(mask.expand_as(dest), src.to(dest.device, dest.dtype)) - for i, modality, embedding, mask in zip( range(len(embeddings)), modalities, embeddings, masks ): if embedding is None or mask is None: continue - _scatter(input_embeds, mask, embedding) + _scatter_mm_embedding(dest=input_embeds, mask=mask, src=embedding) if use_deepstack.get(modality, None): - _scatter(input_deepstack_embeds, mask, deepstack_embeddings[i]) + _scatter_mm_embedding( + dest=input_deepstack_embeds, mask=mask, src=deepstack_embeddings[i] + ) return input_embeds, other_info diff --git a/test/registered/unit/managers/test_mm_embed_scatter.py b/test/registered/unit/managers/test_mm_embed_scatter.py new file mode 100644 index 000000000..bd7810af2 --- /dev/null +++ b/test/registered/unit/managers/test_mm_embed_scatter.py @@ -0,0 +1,61 @@ +import pytest +import torch + +from sglang.srt.managers.mm_utils import _scatter_mm_embedding +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-b-test-cpu") + +NUM_TOKENS = 64 + + +def _make_mask(pattern: str) -> torch.Tensor: + mask = torch.zeros(NUM_TOKENS, dtype=torch.bool) + if pattern == "interleaved": + mask[::3] = True + elif pattern == "blocks": + mask[5:20] = True + mask[40:41] = True + elif pattern == "all_true": + mask[:] = True + return mask.unsqueeze(-1) + + +@pytest.mark.parametrize("width", [8, 24]) +@pytest.mark.parametrize("src_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "mask_pattern", ["interleaved", "blocks", "all_true", "all_false"] +) +def test_scatter_matches_masked_scatter_bitwise(width, src_dtype, mask_pattern): + """The row-index mm embedding merge must stay bitwise identical to + masked_scatter_ semantics, whose internal transients it avoids.""" + torch.manual_seed(0) + mask = _make_mask(mask_pattern) + dest = torch.randn(NUM_TOKENS, width).to(torch.bfloat16) + src = torch.randn(int(mask.sum()), width, dtype=src_dtype) + + expected = dest.clone() + expected.masked_scatter_(mask.expand_as(expected), src.to(expected.dtype)) + + actual = dest.clone() + _scatter_mm_embedding(dest=actual, mask=mask, src=src) + assert torch.equal(actual, expected) + + +def test_scatter_row_count_mismatch_fails_loud(): + """A mask/src row-count mismatch must raise, not silently corrupt rows.""" + dest = torch.zeros(8, 4) + src_short_mask = _make_mask("all_false")[:8] + src_short_mask[1] = True + with pytest.raises((RuntimeError, IndexError)): + _scatter_mm_embedding(dest=dest, mask=src_short_mask, src=torch.ones(3, 4)) + mask_heavy = src_short_mask.clone() + mask_heavy[2:6] = True + with pytest.raises((RuntimeError, IndexError)): + _scatter_mm_embedding(dest=dest, mask=mask_heavy, src=torch.ones(1, 4)) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))