Scatter mm embeddings with row index_copy_ instead of masked_scatter_ to cut transient GPU memory (#37070)

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
 

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