optimization: shard kimi dp image feature transport and misc optimizations (#31227)

test/registered/unit/layers/attention/test_vision_max_seqlen.py

@@ -4,7 +4,7 @@ import torch
 from torch import nn
 
 from sglang.srt.layers.attention import vision
-from sglang.srt.models.kimi_k25 import MoonViTEncoderLayer
+from sglang.srt.models.kimi_k25 import MoonViT3dEncoder, MoonViTEncoderLayer
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=1, suite="base-a-test-cpu")
@@ -44,6 +44,70 @@ def test_vision_flash3_uses_precomputed_max_seqlen(monkeypatch):
     assert recorded["max_seqlen_k"] == 17
 
 
+def test_vision_triton_uses_precomputed_max_seqlen(monkeypatch):
+    """Triton vision attention must share the encoder-level host scalar."""
+
+    recorded = {}
+
+    def fake_context_attention(q, k, v, output, *args, **kwargs):
+        recorded["max_seqlen"] = args[2]
+        recorded["sequence_lengths"] = args[1]
+        output.copy_(q)
+
+    monkeypatch.setattr(vision, "context_attention_fwd", fake_context_attention)
+
+    attention = vision.VisionTritonAttention(use_data_parallel=True)
+    q = torch.zeros(3, 1, 8)
+    cu_seqlens = torch.tensor([0, 1, 3], dtype=torch.int32)
+    sequence_lengths = torch.tensor([1, 2], dtype=torch.int32)
+    output = attention(
+        q,
+        q,
+        q,
+        cu_seqlens=cu_seqlens,
+        bsz=1,
+        seq_len=3,
+        max_seqlen=17,
+        sequence_lengths=sequence_lengths,
+    )
+
+    assert torch.equal(output, q)
+    assert recorded["max_seqlen"] == 17
+    assert recorded["sequence_lengths"] is sequence_lengths
+
+
+def test_vision_flash4_uses_precomputed_max_seqlen(monkeypatch):
+    """FA4 must not re-synchronize for every vision transformer layer."""
+
+    recorded = {}
+
+    def fake_flash_attn(q, k, v, **kwargs):
+        recorded.update(kwargs)
+        return q
+
+    monkeypatch.setattr(vision, "_is_cuda", True)
+    monkeypatch.setattr(
+        vision, "flash_attn_varlen_func", fake_flash_attn, raising=False
+    )
+
+    attention = vision.VisionFlash4Attention(use_data_parallel=True)
+    q = torch.zeros(3, 1, 8)
+    cu_seqlens = torch.tensor([0, 1, 3], dtype=torch.int32)
+    output = attention(
+        q,
+        q,
+        q,
+        cu_seqlens=cu_seqlens,
+        bsz=1,
+        seq_len=3,
+        max_seqlen=17,
+    )
+
+    assert output is q
+    assert recorded["max_seqlen_q"] == 17
+    assert recorded["max_seqlen_k"] == 17
+
+
 def test_kimi_moonvit_forwards_one_precomputed_max_seqlen():
     """MoonViT must share its encoder-level scalar with each attention block."""
 
@@ -73,6 +137,45 @@ def test_kimi_moonvit_forwards_one_precomputed_max_seqlen():
     assert recorded["max_seqlen"] == 19
 
 
+def test_kimi_moonvit_precomputes_sequence_lengths_once():
+    """MoonViT shares packed sequence metadata across all attention blocks."""
+
+    recorded = {}
+
+    class CapturingRope:
+        def get_freqs_cis(self, grid_thws, device):
+            return torch.ones(7, 2, dtype=torch.complex64, device=device)
+
+    class CapturingBlock(nn.Module):
+        def forward(
+            self,
+            hidden_states,
+            cu_seqlens,
+            max_seqlen,
+            rope_freqs_cis,
+            sequence_lengths,
+        ):
+            recorded["cu_seqlens"] = cu_seqlens
+            recorded["max_seqlen"] = max_seqlen
+            recorded["sequence_lengths"] = sequence_lengths
+            return hidden_states
+
+    encoder = MoonViT3dEncoder.__new__(MoonViT3dEncoder)
+    nn.Module.__init__(encoder)
+    encoder.rope_2d = CapturingRope()
+    encoder.blocks = nn.ModuleList([CapturingBlock()])
+    encoder.final_layernorm = nn.Identity()
+
+    hidden_states = torch.ones(7, 4)
+    grid_thws = torch.tensor([[1, 1, 3], [1, 2, 2]], dtype=torch.int32)
+    output = encoder(hidden_states, grid_thws)
+
+    assert torch.equal(output, hidden_states)
+    assert torch.equal(recorded["sequence_lengths"], torch.tensor([3, 4]))
+    assert torch.equal(recorded["cu_seqlens"], torch.tensor([0, 3, 7]))
+    assert recorded["max_seqlen"] == 4
+
+
 if __name__ == "__main__":
     import pytest
 

test/registered/unit/models/test_kimi_k25.py

@@ -0,0 +1,241 @@
+"""CPU coverage for Kimi-K2.5/K2.7 encoder-DP wiring."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.srt.managers.schedule_batch import (
+    Modality,
+    MultimodalDataItem,
+    MultimodalInputs,
+    MultimodalProcessorOutput,
+)
+from sglang.srt.models.kimi_k25 import KimiK25ForConditionalGeneration
+from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
+from sglang.srt.multimodal.processors.kimi_k25 import (
+    _resize_images_by_source_shape,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils.cuda_ipc_transport_utils import (
+    DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
+    CudaIpcTensorTransportProxy,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class _MoonViT3dTower:
+    device = torch.device("cpu")
+    merge_kernel_size = (2, 2)
+
+    def __init__(self):
+        self.config = SimpleNamespace(hidden_size=2)
+        self.patch_embed = SimpleNamespace(
+            proj=SimpleNamespace(weight=torch.empty(1, dtype=torch.float32))
+        )
+        self.grid_thws = None
+
+    def __call__(self, pixel_values, grid_thws):
+        self.grid_thws = grid_thws
+        # MoonViT3d returns a list of [tokens, merge_area, hidden] tensors.
+        return [pixel_values.reshape(-1, 4, pixel_values.shape[-1])]
+
+
+class _Projector:
+    def __call__(self, image_embeds):
+        return image_embeds
+
+
+def _image_item(feature, grid_thw):
+    return MultimodalDataItem(
+        modality=Modality.IMAGE,
+        offsets=[(0, 1)],
+        feature=feature,
+        model_specific_data={"image_grid_thw": torch.tensor(grid_thw)},
+    )
+
+
+def test_kimi_gpu_preprocess_batches_only_source_compatible_images():
+    torch.manual_seed(0)
+    indexed_images = [
+        (0, torch.randn(3, 32, 24)),
+        (1, torch.randn(3, 32, 24)),
+        (2, torch.randn(3, 28, 20)),
+    ]
+    expected = [
+        F.interpolate(
+            image.unsqueeze(0), size=(16, 12), mode="bicubic", align_corners=False
+        )
+        for _, image in indexed_images
+    ]
+    real_interpolate = F.interpolate
+    input_shapes = []
+
+    def record_interpolate(image, *args, **kwargs):
+        input_shapes.append(tuple(image.shape))
+        return real_interpolate(image, *args, **kwargs)
+
+    with patch(
+        "sglang.srt.multimodal.processors.kimi_k25.F.interpolate",
+        side_effect=record_interpolate,
+    ):
+        actual = _resize_images_by_source_shape(indexed_images, 16, 12)
+
+    assert input_shapes == [(2, 3, 32, 24), (1, 3, 28, 20)]
+    assert len(actual) == len(expected)
+    for result, reference in zip(actual, expected):
+        torch.testing.assert_close(result, reference)
+
+
+def test_dp_helper_supports_moonvit3d_packed_embeddings_on_tp1():
+    tower = _MoonViT3dTower()
+    pixel_values = torch.randn(4, 2)
+
+    with get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0):
+        output = run_dp_sharded_mrope_vision_model(
+            tower, pixel_values, [[1, 2, 2]], rope_type="rope_2d_packed"
+        )
+
+    assert torch.equal(output, pixel_values.reshape(1, 4, 2))
+    assert torch.equal(tower.grid_thws, torch.tensor([[1, 2, 2]]))
+
+
+def test_dp_helper_can_lazily_load_kimi_features_on_tp1():
+    tower = _MoonViT3dTower()
+    pixel_values = torch.randn(4, 2)
+    loader = Mock(return_value=pixel_values)
+
+    with get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0):
+        output = run_dp_sharded_mrope_vision_model(
+            tower,
+            None,
+            [[1, 2, 2]],
+            rope_type="rope_2d_packed",
+            load_local_pixel_values=loader,
+            pixel_values_device=pixel_values.device,
+            pixel_values_dtype=pixel_values.dtype,
+        )
+
+    assert torch.equal(output, pixel_values.reshape(1, 4, 2))
+    loader.assert_called_once_with([0])
+
+
+def test_dp_helper_uses_config_hidden_size_for_empty_moonvit3d_rank():
+    class _GatherGroup:
+        def all_gather(self, tensor, dim):
+            return torch.cat([torch.ones_like(tensor), tensor], dim=dim)
+
+    tower = _MoonViT3dTower()
+    parallel = SimpleNamespace(
+        attn_tp_size=2,
+        attn_tp_rank=1,
+        attn_tp_group=_GatherGroup(),
+    )
+
+    with patch("sglang.srt.multimodal.mm_utils.get_parallel", return_value=parallel):
+        output = run_dp_sharded_mrope_vision_model(
+            tower,
+            torch.randn(4, 2),
+            [[1, 2, 2]],
+            rope_type="rope_2d_packed",
+        )
+
+    assert output.shape == (1, 4, 2)
+    assert tower.grid_thws is None
+
+
+def test_dp_helper_lazily_loads_only_its_local_image_shard():
+    class _GatherGroup:
+        def all_gather(self, tensor, dim):
+            # Rank one's embedding is irrelevant to this rank's loader call;
+            # retain the expected gathered shape for output reconstruction.
+            return torch.cat([tensor, torch.zeros_like(tensor)], dim=dim)
+
+    tower = _MoonViT3dTower()
+    features = [torch.full((4, 2), 1.0), torch.full((4, 2), 2.0)]
+    loader = Mock(side_effect=lambda indices: torch.cat([features[i] for i in indices]))
+    parallel = SimpleNamespace(
+        attn_tp_size=2,
+        attn_tp_rank=0,
+        attn_tp_group=_GatherGroup(),
+    )
+
+    with patch("sglang.srt.multimodal.mm_utils.get_parallel", return_value=parallel):
+        output = run_dp_sharded_mrope_vision_model(
+            tower,
+            None,
+            [[1, 2, 2], [1, 2, 2]],
+            rope_type="rope_2d_packed",
+            load_local_pixel_values=loader,
+            pixel_values_device=torch.device("cpu"),
+            pixel_values_dtype=torch.float32,
+        )
+
+    loader.assert_called_once_with([0])
+    assert output.shape == (2, 4, 2)
+
+
+def test_kimi_k25_encoder_dp_selects_packed_moonvit_contract():
+    model = KimiK25ForConditionalGeneration.__new__(KimiK25ForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.use_data_parallel = True
+    model.vision_tower = _MoonViT3dTower()
+    model.mm_projector = _Projector()
+    items = [_image_item(torch.randn(4, 2), [[1, 2, 2]])]
+    sharded_embeddings = torch.randn(1, 2)
+
+    with patch(
+        "sglang.srt.models.kimi_k25.run_dp_sharded_mrope_vision_model",
+        return_value=sharded_embeddings,
+    ) as run_dp:
+        output = model.get_image_feature(items)
+
+    assert output is sharded_embeddings
+    tower, pixel_values, grid_thws = run_dp.call_args.args
+    assert tower is model.vision_tower
+    assert pixel_values is None
+    assert grid_thws == [[1, 2, 2]]
+    assert run_dp.call_args.kwargs["rope_type"] == "rope_2d_packed"
+    assert callable(run_dp.call_args.kwargs["load_local_pixel_values"])
+
+
+def test_kimi_lazy_ipc_feature_skips_scheduler_reconstruction():
+    proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy)
+    proxy.reconstruct_on_target_device = Mock()
+    item = MultimodalDataItem(
+        modality=Modality.IMAGE,
+        hash=123,
+        pad_value=456,
+        offsets=[(0, 1)],
+        feature=proxy,
+        model_specific_data={DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY: True},
+    )
+
+    with patch(
+        "sglang.srt.managers.schedule_batch.torch.cuda.current_device", return_value=0
+    ):
+        mm_inputs = MultimodalInputs.from_processor_output(
+            MultimodalProcessorOutput(mm_items=[item])
+        )
+
+    assert mm_inputs.mm_items[0].feature is proxy
+    proxy.reconstruct_on_target_device.assert_not_called()
+
+
+def test_kimi_lazy_ipc_feature_acknowledges_all_tp_consumers():
+    proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy)
+    proxy.reconstruct_on_target_device = Mock(return_value=torch.randn(1, 2))
+    item = MultimodalDataItem(modality=Modality.IMAGE, feature=proxy)
+
+    item.reconstruct(0, ipc_consumer_count=8)
+
+    proxy.reconstruct_on_target_device.assert_called_once_with(0, consumer_count=8)
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/multimodal/test_cuda_ipc_pool_budget.py

@@ -0,0 +1,35 @@
+"""CPU-only regression tests for CUDA IPC multimodal pool budgeting."""
+
+import unittest
+
+from sglang.srt.utils.cuda_ipc_transport_utils import (
+    get_mm_feature_pool_size_per_worker,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+class TestCudaIpcPoolBudget(unittest.TestCase):
+    def test_budget_is_not_multiplied_by_tokenizer_workers(self):
+        budget = 1_024 * 1024 * 1024
+        worker_num = 16
+
+        per_worker = get_mm_feature_pool_size_per_worker(budget, worker_num)
+
+        self.assertEqual(per_worker, 64 * 1024 * 1024)
+        self.assertLessEqual(per_worker * worker_num, budget)
+
+    def test_remainder_is_not_overallocated(self):
+        self.assertEqual(get_mm_feature_pool_size_per_worker(1_001, 8), 125)
+        self.assertLessEqual(get_mm_feature_pool_size_per_worker(1_001, 8) * 8, 1_001)
+
+    def test_rejects_invalid_budget_or_worker_count(self):
+        with self.assertRaisesRegex(ValueError, "total_pool_size"):
+            get_mm_feature_pool_size_per_worker(0, 1)
+        with self.assertRaisesRegex(ValueError, "tokenizer_worker_num"):
+            get_mm_feature_pool_size_per_worker(1, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/multimodal/test_cuda_ipc_transport.py

@@ -0,0 +1,124 @@
+"""CUDA IPC multimodal feature transport regression tests.
+
+This covers the production path where a tokenizer worker places a feature in
+the bounded pool and the scheduler process opens the shared CUDA allocation.
+CPU-only policy tests intentionally cannot exercise this cross-process handle.
+"""
+
+import gc
+import multiprocessing as mp
+import queue
+import unittest
+
+import torch
+
+from sglang.srt.utils.cuda_ipc_transport_utils import (
+    CudaIpcTensorTransportProxy,
+    MmItemMemoryPool,
+    _pool_handle_cache_clear,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
+
+
+def _produce_pooled_tensor(proxy_queue, consumer_done, result_queue):
+    """Create a tokenizer-worker-like CUDA IPC pool in a spawned producer."""
+    pool = source = pool_slice = proxy = None
+    try:
+        torch.cuda.set_device(0)
+        pool = MmItemMemoryPool(
+            memory_size=1 << 20,
+            recycle_interval=60,
+            base_gpu_id=0,
+        )
+        source = torch.arange(35, dtype=torch.float32, device="cuda").reshape(5, 7)
+        expected = source.cpu().tolist()
+        sync_meta, pool_slice, byte_offset = pool.return_a_slice_tensor_with_flag(
+            source
+        )
+        if pool_slice is None:
+            raise RuntimeError("test tensor did not fit in the CUDA IPC pool")
+        pool_slice.copy_(source.view(torch.int8).view(-1), non_blocking=True)
+        torch.cuda.synchronize()
+        proxy = CudaIpcTensorTransportProxy(
+            data=pool_slice,
+            info_data=source,
+            sync_buffer_meta=sync_meta,
+            pool_ipc_handle=pool._pool_ipc_handle,
+            pool_byte_offset=byte_offset,
+            pool_device_index=pool._pool_device_index,
+        )
+        proxy_queue.put((proxy, expected))
+        if not consumer_done.wait(timeout=60):
+            raise TimeoutError("consumer did not release the CUDA IPC tensor")
+    except Exception as exc:  # pragma: no cover - returned to the parent
+        result_queue.put(("error", repr(exc)))
+        return
+    finally:
+        del proxy, pool_slice, source
+        if pool is not None:
+            pool.shutdown()
+            del pool
+        gc.collect()
+        torch.cuda.ipc_collect()
+    result_queue.put(("ok", None))
+
+
+class TestCudaIpcTransport(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is required")
+
+    def test_pooled_tensor_reconstructs_in_spawned_process(self):
+        """Consumer releases the pool mapping before the producer tears down."""
+        ctx = mp.get_context("spawn")
+        proxy_queue = ctx.Queue()
+        producer_results = ctx.Queue()
+        consumer_done = ctx.Event()
+        producer = ctx.Process(
+            target=_produce_pooled_tensor,
+            args=(proxy_queue, consumer_done, producer_results),
+        )
+        producer.start()
+        proxy = reconstructed = None
+        producer_result = None
+        try:
+            try:
+                proxy, expected = proxy_queue.get(timeout=60)
+            except queue.Empty:
+                producer_result = producer_results.get(timeout=5)
+                _status, payload = producer_result
+                self.fail(
+                    f"CUDA IPC producer failed before sending its proxy: {payload}"
+                )
+
+            reconstructed = proxy.reconstruct_on_target_device(0)
+            torch.cuda.synchronize()
+            self.assertEqual(reconstructed.cpu().tolist(), expected)
+        finally:
+            # The scheduler retains this cache for its lifetime. The test's
+            # consumer exits quickly, so it must close the mapping before the
+            # producer destroys the shared allocation.
+            del reconstructed, proxy
+            _pool_handle_cache_clear()
+            gc.collect()
+            torch.cuda.ipc_collect()
+            consumer_done.set()
+            producer.join(timeout=60)
+            try:
+                if producer_result is None:
+                    producer_result = producer_results.get(timeout=5)
+                status, payload = producer_result
+                self.assertEqual(status, "ok", payload)
+            finally:
+                if producer.is_alive():
+                    producer.terminate()
+                    producer.join(timeout=10)
+            self.assertEqual(producer.exitcode, 0)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

test/registered/unit/multimodal/test_feature_materialization.py

@@ -0,0 +1,39 @@
+"""Tests for the shared multimodal feature materialization helper."""
+
+import unittest
+
+import torch
+
+from sglang.srt.multimodal.mm_utils import materialize_multimodal_features
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestFeatureMaterialization(CustomTestCase):
+    def test_packs_variable_length_features_and_converts_dtype(self):
+        features = [
+            torch.arange(6, dtype=torch.float32).view(2, 3),
+            torch.arange(9, dtype=torch.float32).view(3, 3) + 10,
+        ]
+
+        result = materialize_multimodal_features(
+            features, device=torch.device("cpu"), dtype=torch.bfloat16
+        )
+
+        self.assertEqual(result.shape, (5, 3))
+        self.assertEqual(result.dtype, torch.bfloat16)
+        torch.testing.assert_close(result.float(), torch.cat(features, dim=0))
+
+    def test_rejects_incompatible_trailing_shapes(self):
+        with self.assertRaisesRegex(ValueError, "matching trailing shapes"):
+            materialize_multimodal_features(
+                [torch.empty(2, 3), torch.empty(1, 4)],
+                device=torch.device("cpu"),
+                dtype=torch.float32,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)