minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(vlm): stream-order cuda-ipc feature pool lifecycle and streamline multimodal transport module (#33949)

Browse Source
This commit is contained in:
Mick
2026-08-10 18:47:19 +08:00
committed by GitHub
parent 0977b22431
commit 443b62db57
18 changed files with 1058 additions and 663 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_mm_process_config.py
+88
View File
@@ -174,6 +174,7 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
mm_io_worker_num=0,
tokenizer_worker_num=1,
base_gpu_id=2,
tp_size=8,
)
@staticmethod
@@ -263,6 +264,93 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
memory_pool.assert_not_called()
class TestStreamOrderedMmFeaturePool(CustomTestCase):
def test_consumer_slot_uses_global_tp_rank(self):
from sglang.srt.multimodal.transport.memory_pool import resolve_consumer_rank
parallel = SimpleNamespace(tp_rank=6, attn_tp_rank=2)
with patch("sglang.srt.runtime_context.get_parallel", return_value=parallel):
self.assertEqual(resolve_consumer_rank(8), 6)
def test_complete_group_acknowledges_each_consumer_slot(self):
from sglang.srt.multimodal.transport import memory_pool
consumer = memory_pool.StreamOrderedPoolConsumerMixin()
consumer._init_stream_ordered_consumer(
ready_byte_offset=64,
ack_byte_offset=68,
generation=3,
total_consumer_count=4,
transport_name="test",
)
with patch.object(memory_pool, "stream_write_value32") as write:
consumer._acknowledge_on_stream(1000, 0, consumer_count=4)
self.assertEqual(
[call.args[1] for call in write.call_args_list],
[1068, 1072, 1076, 1080],
)
self.assertEqual([call.args[2] for call in write.call_args_list], [3] * 4)
def test_reused_pool_slot_gets_new_generation(self):
from sglang.srt.multimodal.transport.memory_pool import (
StreamOrderedMmFeaturePool,
)
pool = object.__new__(StreamOrderedMmFeaturePool)
pool._available_ranges = [(256, 4096)]
pool._available_slots = [0]
pool._slot_generations = [0]
pool._occupied = {}
pool.control_words_per_slot = 2
pool.transport_name = "test"
first = pool._allocate_locked(512)
pool._release_locked(first)
pool._merge_ranges_locked()
second = pool._allocate_locked(512)
self.assertEqual(first.generation, 1)
self.assertEqual(second.generation, 2)
def test_pool_rejects_duplicate_release(self):
from sglang.srt.multimodal.transport.memory_pool import (
StreamOrderedMmFeaturePool,
)
pool = object.__new__(StreamOrderedMmFeaturePool)
pool._available_ranges = [(256, 4096)]
pool._available_slots = [0]
pool._slot_generations = [0]
pool._occupied = {}
pool.control_words_per_slot = 2
pool.transport_name = "test"
lease = pool._allocate_locked(512)
pool._release_locked(lease)
with self.assertRaisesRegex(RuntimeError, "inactive test pool lease"):
pool._release_locked(lease)
def test_pool_shutdown_wakes_recycler_before_returning(self):
from sglang.srt.multimodal.transport.memory_pool import (
StreamOrderedMmFeaturePool,
)
pool = object.__new__(StreamOrderedMmFeaturePool)
pool._recycler_stop_event = threading.Event()
pool._recycle_thread = threading.Thread(
target=pool._recycler_stop_event.wait,
args=(60,),
daemon=True,
)
pool._recycle_thread.start()
pool.shutdown()
self.assertFalse(pool._recycle_thread.is_alive())
class TestPrecomputeHashBeforeCpuTransfer(CustomTestCase):
@staticmethod
def _processor(enabled):
test/registered/unit/models/test_kimi_k25.py
+2 -2
View File
@@ -39,11 +39,11 @@ from sglang.srt.multimodal.processors.kimi_k25 import (
_resize_bicubic_if_needed,
_resize_images_by_source_shape,
)
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.srt.utils.cuda_ipc_transport_utils import (
from sglang.srt.multimodal.transport.cuda_ipc import (
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
CudaIpcTensorTransportProxy,
)
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
test/registered/unit/models/test_qwen3_vl_feature_materialization.py
+113 -4
View File
@@ -2,10 +2,16 @@
import unittest
from types import SimpleNamespace
from unittest.mock import Mock, patch
import torch
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
from sglang.srt.multimodal.processors.qwen_vl import QwenVLImageProcessor
from sglang.srt.multimodal.transport.cuda_ipc import (
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -27,9 +33,61 @@ class _RecordingVisual:
class TestQwen3VLFeatureMaterialization(CustomTestCase):
@staticmethod
def _model(visual, *, use_data_parallel):
model = Qwen3VLForConditionalGeneration.__new__(Qwen3VLForConditionalGeneration)
torch.nn.Module.__init__(model)
model.visual = visual
model.use_data_parallel = use_data_parallel
return model
def test_processor_defers_gpu_transport_for_encoder_dp(self):
for transport in ("cuda_ipc", "cuda_vmm"):
with self.subTest(transport=transport):
processor = QwenVLImageProcessor.__new__(QwenVLImageProcessor)
processor.mm_feature_transport = transport
processor.server_args = SimpleNamespace(mm_enable_dp_encoder=True)
processor.model_type = "qwen3_vl"
items = [
MultimodalDataItem(modality=Modality.IMAGE),
MultimodalDataItem(modality=Modality.VIDEO),
MultimodalDataItem(modality=Modality.AUDIO),
]
processor._mark_dp_encoder_features_for_deferred_reconstruction(items)
self.assertTrue(
items[0].model_specific_data[
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY
]
)
self.assertTrue(
items[1].model_specific_data[
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY
]
)
self.assertNotIn(
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
items[2].model_specific_data,
)
def test_processor_does_not_defer_cpu_transport(self):
processor = QwenVLImageProcessor.__new__(QwenVLImageProcessor)
processor.mm_feature_transport = "cpu"
processor.server_args = SimpleNamespace(mm_enable_dp_encoder=True)
processor.model_type = "qwen3_vl"
item = MultimodalDataItem(modality=Modality.IMAGE)
processor._mark_dp_encoder_features_for_deferred_reconstruction([item])
self.assertNotIn(
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
item.model_specific_data,
)
def test_image_features_are_packed_on_the_visual_device(self):
visual = _RecordingVisual()
model = SimpleNamespace(visual=visual, use_data_parallel=False)
model = self._model(visual, use_data_parallel=False)
items = [
SimpleNamespace(
feature=torch.ones(2, 3),
@@ -40,7 +98,7 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase):
image_grid_thw=torch.tensor([[1, 1, 1]]),
),
]
output = Qwen3VLForConditionalGeneration.get_image_feature(model, items)
output = model.get_image_feature(items)
self.assertIs(visual.pixel_values, output)
self.assertEqual(output.shape, (3, 3))
@@ -49,7 +107,7 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase):
def test_video_features_are_packed_on_the_visual_device(self):
visual = _RecordingVisual()
model = SimpleNamespace(visual=visual, use_data_parallel=False)
model = self._model(visual, use_data_parallel=False)
items = [
SimpleNamespace(
feature=torch.ones(3, 4),
@@ -60,7 +118,7 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase):
video_grid_thw=torch.tensor([[1, 1, 2]]),
),
]
output = Qwen3VLForConditionalGeneration.get_video_feature(model, items)
output = model.get_video_feature(items)
self.assertIs(visual.pixel_values, output)
self.assertEqual(output.shape, (5, 4))
@@ -70,6 +128,57 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase):
torch.equal(visual.grid_thw, torch.tensor([[1, 1, 3], [1, 1, 2]]))
)
def test_encoder_dp_materializes_only_locally_assigned_visual_items(self):
visual = SimpleNamespace(device=torch.device("cuda:0"), dtype=torch.bfloat16)
model = self._model(visual, use_data_parallel=True)
for modality, feature_method, grid_attribute in (
("image", model.get_image_feature, "image_grid_thw"),
("video", model.get_video_feature, "video_grid_thw"),
):
with self.subTest(modality=modality):
items = [
SimpleNamespace(
feature=torch.ones(2, 3),
reconstruct=Mock(),
**{grid_attribute: torch.tensor([[1, 1, 2]])},
),
SimpleNamespace(
feature=torch.ones(1, 3),
reconstruct=Mock(),
**{grid_attribute: torch.tensor([[1, 1, 1]])},
),
]
local_features = object()
encoded = object()
def run_dp(_visual, pixel_values, grid_thw, **kwargs):
self.assertIsNone(pixel_values)
self.assertEqual(grid_thw, [[1, 1, 2], [1, 1, 1]])
self.assertIs(
kwargs["load_local_pixel_values"]([1]), local_features
)
return encoded
with patch(
"sglang.srt.models.qwen3_vl.run_dp_sharded_mrope_vision_model",
side_effect=run_dp,
), patch(
"sglang.srt.models.qwen3_vl.materialize_multimodal_features",
return_value=local_features,
) as materialize, patch(
"sglang.srt.models.qwen3_vl.get_parallel",
return_value=SimpleNamespace(tp_size=8),
):
output = feature_method(items)
self.assertIs(output, encoded)
items[0].reconstruct.assert_not_called()
items[1].reconstruct.assert_called_once_with(0, ipc_consumer_count=8)
materialize.assert_called_once_with(
[items[1].feature], device=visual.device, dtype=visual.dtype
)
if __name__ == "__main__":
unittest.main(verbosity=2)
test/registered/unit/multimodal/test_cuda_ipc_pool_budget.py
+1 -1
View File
@@ -2,7 +2,7 @@
import unittest
from sglang.srt.utils.cuda_ipc_transport_utils import (
from sglang.srt.multimodal.transport.cuda_ipc import (
get_mm_feature_pool_size_per_worker,
)
from sglang.test.ci.ci_register import register_cpu_ci
test/registered/unit/multimodal/test_cuda_ipc_transport.py
+36 -18
View File
@@ -8,11 +8,13 @@ CPU-only policy tests intentionally cannot exercise this cross-process handle.
import gc
import multiprocessing as mp
import queue
import time
import unittest
from unittest.mock import Mock, patch
import torch
from sglang.srt.utils.cuda_ipc_transport_utils import (
from sglang.srt.multimodal.transport.cuda_ipc import (
CudaIpcTensorTransportProxy,
MmItemMemoryPool,
_pool_handle_cache_clear,
@@ -25,39 +27,40 @@ 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
pool = source = proxy = None
try:
torch.cuda.set_device(0)
pool = MmItemMemoryPool(
memory_size=1 << 20,
recycle_interval=60,
recycle_interval=0.01,
base_gpu_id=0,
consumer_count=1,
)
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
expected = torch.arange(35, dtype=torch.float32).reshape(5, 7).tolist()
proxy = pool.wrap_tensor(
source,
use_pool_handle_cache=True,
)
if pool_slice is None:
if proxy 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,
)
# Intentionally do not synchronize the producer. The consumer stream
# wait must order its copy after the producer-ready write.
proxy_queue.put((proxy, expected))
if not consumer_done.wait(timeout=60):
raise TimeoutError("consumer did not release the CUDA IPC tensor")
deadline = time.monotonic() + 5
while time.monotonic() < deadline:
if pool.active_lease_count == 0:
break
time.sleep(0.01)
else:
raise TimeoutError("pool did not observe the stream-ordered consumer ack")
except Exception as exc: # pragma: no cover - returned to the parent
result_queue.put(("error", repr(exc)))
return
finally:
del proxy, pool_slice, source
del proxy, source
if pool is not None:
pool.shutdown()
del pool
@@ -119,6 +122,21 @@ class TestCudaIpcTransport(CustomTestCase):
producer.join(timeout=10)
self.assertEqual(producer.exitcode, 0)
def test_uncached_mapping_waits_before_proxy_release(self):
proxy = object.__new__(CudaIpcTensorTransportProxy)
proxy.proxy_state = {"ipc_extra": {"use_pool_handle_cache": False}}
proxy._pool_storage = None
stream = Mock()
with patch(
"sglang.srt.multimodal.transport.cuda_ipc.torch.cuda.current_stream",
return_value=stream,
):
proxy._retain_storage_until_stream_completes(object(), 0)
stream.synchronize.assert_called_once_with()
self.assertIsNone(proxy._pool_storage)
if __name__ == "__main__":
unittest.main(verbosity=2)
test/registered/unit/test_global_config_read_ratchet.py
+3 -9
View File
@@ -29,8 +29,6 @@ Where the remaining reads live (``runtime_context.py``, exempt by module):
the short circuit is the point: with PP off the group is never touched, which
is what lets the ``Indexer`` be constructed before distributed init. The live
property would demand the group either way.
- ``cuda_ipc_transport_utils.tp_size`` runs in the tokenizer process, which has
no groups at all (the call site already guards for "not published yet").
- ``dp_attention.attn_cp_size`` / ``moe_dp_size``: the configuration the
predicate detects (``attn_cp_size > moe_dp_size``) is the one where
``initialize_model_parallel`` aliases ``_MOE_DP`` to ``_ATTN_CP``, so the live
@@ -102,14 +100,10 @@ _CONFIGURED_SIZE_CALL_SITES = {
"the same dict already carries the live moe_dp_size under 'dp'; this entry "
"is the configured intent"
),
("srt/utils/cuda_ipc_transport_utils.py", "configured_tp_size"): (
"runs in the tokenizer process, which has no parallel groups at all"
),
("srt/models/kimi_k25.py", "configured_tp_size"): (
"the IPC refcount has to name the same number the recycler waits on, and "
"that waiter (MmItemMemoryPool.try_to_recycle) reads configured_tp_size() "
"because it runs in the tokenizer process; a refcount taken from the live "
"attention subgroup would strand items in the bounded pool"
"the IPC refcount must match the configured TP consumer count captured "
"when the tokenizer creates MmItemMemoryPool; a live attention subgroup "
"size could strand leases in the bounded pool"
),
("srt/models/kimi_k3.py", "configured_tp_size"): (
"same as kimi_k25: the IPC refcount must agree with the recycler's waiter"