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

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