From 96d4c41944abc2c62168842cc24090361e45182c Mon Sep 17 00:00:00 2001 From: Vorapol Assavasangthong Date: Fri, 23 Oct 2026 16:21:33 +0300 Subject: [PATCH 1/2] [AMD] vision: Batch ViT across requests on ROCm when the DP encoder is enabled ROCm keeps one ViT call per request (#32415). Under --mm-enable-dp-encoder that leaves a single-image request on one attention-TP rank while the others wait, so the DP encoder has nothing to split. Use the cross-request batched path on ROCm when the DP encoder is enabled and attention-TP > 1; NPU, XPU, and ROCm without the DP encoder or at attention-TP 1 are unchanged. Under that same condition, Qwen3-VL warms its vision tower on every rank in the post-load hook (precompile_kernels_after_loading), so a non-zero rank does not pay its first ViT call (rotary torch.compile, lazy kernel loads, the vectorized pos-embed path) inside the first served image burst. --- python/sglang/srt/managers/mm_schedule.py | 31 ++- python/sglang/srt/models/qwen3_vl.py | 35 ++++ python/sglang/srt/models/qwen4_exp.py | 1 + python/sglang/srt/multimodal/mm_utils.py | 19 +- .../test_mm_chunked_embedding_unit.py | 23 +- .../test_hip_dp_encoder_vit_batching.py | 198 ++++++++++++++++++ 6 files changed, 288 insertions(+), 19 deletions(-) create mode 100644 test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py diff --git a/python/sglang/srt/managers/mm_schedule.py b/python/sglang/srt/managers/mm_schedule.py index c5818e07b0a5..253a3b6ba43f 100644 --- a/python/sglang/srt/managers/mm_schedule.py +++ b/python/sglang/srt/managers/mm_schedule.py @@ -8,15 +8,14 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.multimodal.evs import EVSEmbeddingResult +from sglang.srt.multimodal.mm_utils import hip_dp_encoder_batches_vit from sglang.srt.multimodal.transport.cuda_ipc import BORROW_CUDA_IPC_FEATURE_KEY -from sglang.srt.runtime_context import get_schedule -from sglang.srt.utils import is_hip, is_npu, is_xpu +from sglang.srt.runtime_context import get_platform, get_schedule +from sglang.srt.utils import is_npu from sglang.srt.utils.async_probe import maybe_assert_sum from sglang.utils import logger -_is_hip = is_hip() _is_npu = is_npu() -_is_xpu = is_xpu() embedding_cache: Optional[MultiModalStaticCache] = None host_offload_event: Optional[torch.cuda.Event] = None @@ -538,6 +537,21 @@ def _assemble_per_image_chunk( ] +def _batch_vit_across_requests() -> bool: + """Whether chunked prefill encodes the images of all requests in one ViT call. + + CUDA batches across requests. NPU and XPU keep one ViT call per request, and + so does ROCm, where one large cross-request batch regressed CI (#32415), + except when the DP encoder has more than one rank to split the batch over. + """ + platform = get_platform() + if platform.is_npu or platform.is_xpu: + return False + if platform.is_hip: + return hip_dp_encoder_batches_vit() + return True + + def _get_chunked_prefill_embedding( data_embedding_func: DataEmbeddingFunc, embedding_items: List[MultimodalDataItem], @@ -560,6 +574,7 @@ def _get_chunked_prefill_embedding( per_image_requests = [] # batched ViT encoding full_path_requests = [] # per-request encoding (EVS etc.) all_chunks: List[Tuple[int, torch.Tensor]] = [] + batch_vit = _batch_vit_across_requests() for i in range(max_iterations): if items_size[i] == items_size[i + 1]: @@ -595,9 +610,9 @@ def _get_chunked_prefill_embedding( for item in embedding_items_per_req ) if is_per_image: - if _is_hip or _is_npu or _is_xpu: - # ROCm CI regressed with one large cross-request ViT batch; keep - # the previous per-request path on HIP/NPU/XPU while CUDA uses batching. + if batch_vit: + per_image_requests.append(req_info) + else: chunk_slices = _get_chunked_embedding_by_item( data_embedding_func, embedding_items_per_req, @@ -606,8 +621,6 @@ def _get_chunked_prefill_embedding( device, ) all_chunks.extend((i, chunk) for chunk in chunk_slices) - else: - per_image_requests.append(req_info) else: full_path_requests.append(req_info) diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 171777a77e87..1bdc490c81cb 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -77,6 +77,7 @@ compute_cu_seqlens_from_grid_numpy, ) from sglang.srt.multimodal.mm_utils import ( + hip_dp_encoder_batches_vit, materialize_multimodal_features, run_dp_sharded_mrope_vision_model, ) @@ -1439,6 +1440,40 @@ def pad_input_ids(self, input_ids: array, mm_inputs: MultimodalInputs) -> array: _require_vision(self) return pad_mm_tokens(input_ids, mm_inputs) + def precompile_kernels_after_loading(self) -> None: + if self.visual is not None and hip_dp_encoder_batches_vit(): + self._warm_up_vision_tower() + + def _warm_up_vision_tower(self) -> None: + """Run the vision tower on dummy images on this rank and discard the output. + + With batched ViT under the DP encoder, a non-zero rank otherwise runs its + first ViT call inside the first served image burst, paying the + torch.compile of the rotary kernel and lazy kernel loads there. Two grid + sizes cover the static and the dynamic-shape compile; the last call has + enough images to take the vectorized position-embedding path. + """ + patch_embed = self.visual.patch_embed + patch_dim = ( + patch_embed.in_channels + * patch_embed.temporal_patch_size + * patch_embed.patch_size**2 + ) + grids = ( + [[1, 32, 32]], + [[1, 48, 64]], + [[1, 32, 32]] * _VECTORIZED_VL_POS_EMBED_MIN_IMAGES, + ) + with torch.no_grad(): + for grid in grids: + pixel_values = torch.zeros( + sum(t * h * w for t, h, w in grid), + patch_dim, + device=self.visual.device, + dtype=self.visual.dtype, + ) + self.visual(pixel_values, grid_thw=torch.tensor(grid)) + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: _require_vision(self) image_grid_thw = torch.concat([item.image_grid_thw for item in items], dim=0) diff --git a/python/sglang/srt/models/qwen4_exp.py b/python/sglang/srt/models/qwen4_exp.py index f4cb98bc1110..8ece6a5c1da9 100644 --- a/python/sglang/srt/models/qwen4_exp.py +++ b/python/sglang/srt/models/qwen4_exp.py @@ -2674,6 +2674,7 @@ def load_qwen4_exp_ple_shard(name: str, loaded_weight: torch.Tensor) -> bool: return loaded_params def precompile_kernels_after_loading(self) -> None: + super().precompile_kernels_after_loading() from sglang.srt.layers.quantization.unquant import precompile_splitk_tactics if precompile_splitk_tactics(): diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index 075d742f759b..e11b35dafa3d 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/python/sglang/srt/multimodal/mm_utils.py @@ -41,7 +41,7 @@ from PIL import Image from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_mm, get_parallel, get_platform from sglang.srt.utils import flatten_nested_list @@ -417,6 +417,23 @@ def process_images(images, image_processor, model_cfg): return new_images +def hip_dp_encoder_batches_vit() -> bool: + """Whether ROCm encodes the images of all requests in one ViT call. + + ROCm otherwise encodes one request at a time (see + ``mm_schedule._batch_vit_across_requests``). Under the DP encoder, + ``run_dp_sharded_mrope_vision_model`` splits the images of one ViT call + across the attention-TP ranks, so a single-image request runs on one rank + while the others wait. A cross-request batch is what gives the DP encoder + work to split. With one attention-TP rank there is nothing to split. + """ + return ( + get_platform().is_hip + and get_mm().mm_enable_dp_encoder + and get_parallel().attn_tp_size > 1 + ) + + # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/vision.py def get_dp_encoder_lb_assignment( sizes: list[int], diff --git a/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py b/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py index 0bc0a4e278e3..0578cd8f7b8b 100644 --- a/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py +++ b/test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py @@ -22,7 +22,7 @@ DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, CudaIpcTensorTransportProxy, ) -from sglang.srt.runtime_context import get_context, get_parallel +from sglang.srt.runtime_context import get_context, get_parallel, override_platform from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -102,6 +102,13 @@ def _make_items(): ] +def _encode_per_request(request, per_request: bool) -> None: + # ROCm without the DP encoder keeps one ViT call per request. + platform_override = override_platform(is_hip=per_request) + platform_override.install() + request.addfinalizer(platform_override.restore) + + def _run_by_item_chunks(encoder): mm_schedule.init_mm_embedding_cache(1 << 30) items = _make_items() @@ -199,9 +206,9 @@ def test_tensor_cache_entries_own_storage(per_request): @pytest.mark.parametrize("per_request", [False, True]) @pytest.mark.parametrize("cache_bytes", [0, 6 * HIDDEN * 4]) def test_tensor_encoder_chunks_survive_cache_pressure( - monkeypatch, per_request, cache_bytes + request, per_request, cache_bytes ): - monkeypatch.setattr(mm_schedule, "_is_hip", per_request) + _encode_per_request(request, per_request) mm_schedule.init_mm_embedding_cache(cache_bytes) items = _make_items() chunk, _ = mm_schedule._get_chunked_prefill_embedding( @@ -341,10 +348,8 @@ def test_full_mismatched_cache_entry_is_reencoded(caplog): @pytest.mark.parametrize("per_request", [False, True]) -def test_multi_offset_audio_item_is_chunked_and_cached_per_item( - monkeypatch, per_request -): - monkeypatch.setattr(mm_schedule, "_is_hip", per_request) +def test_multi_offset_audio_item_is_chunked_and_cached_per_item(request, per_request): + _encode_per_request(request, per_request) mm_schedule.init_mm_embedding_cache(1 << 30) history_offsets = [(2, 3), (6, 7), (10, 11)] new_offsets = [(14, 15), (18, 19)] @@ -396,9 +401,9 @@ def run(items, prefix_len, extend_len): @pytest.mark.parametrize("num_requests", [1, 2]) @pytest.mark.parametrize("num_items", [1, 2]) def test_cached_audio_items_concatenate_once_in_request_order( - monkeypatch, per_request, num_requests, num_items + monkeypatch, request, per_request, num_requests, num_items ): - monkeypatch.setattr(mm_schedule, "_is_hip", per_request) + _encode_per_request(request, per_request) mm_schedule.init_mm_embedding_cache(1 << 30) items = [] expected_slices = [] diff --git a/test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py b/test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py new file mode 100644 index 000000000000..19718e0d97fa --- /dev/null +++ b/test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py @@ -0,0 +1,198 @@ +"""ROCm batches ViT across requests only when the DP encoder has ranks to split over. + +CPU-only: the platform, the `mm` bag and the attention-TP topology are stated +through the runtime context; the vision tower and the encoder are fakes. +""" + +import unittest +from contextlib import ExitStack +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from sglang.srt.managers import mm_schedule +from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem +from sglang.srt.model_executor.model_runner_components.load_model_utils import ( + maybe_precompile_model_kernels_after_loading, +) +from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration +from sglang.srt.multimodal.mm_utils import hip_dp_encoder_batches_vit +from sglang.srt.runtime_context import ( + get_context, + get_parallel, + override_platform, + publish, + reset_context, +) +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +_PLATFORMS = ("cuda", "hip", "npu", "xpu") + + +class _RecordingVisionTower: + device = torch.device("cpu") + dtype = torch.float32 + patch_embed = SimpleNamespace(in_channels=3, temporal_patch_size=2, patch_size=16) + + def __init__(self): + self.calls = [] + + def __call__(self, pixel_values, *, grid_thw): + self.calls.append( + (pixel_values.shape[0], grid_thw.tolist(), torch.is_grad_enabled()) + ) + + +class _RuntimeCase(CustomTestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="test") + + def _runtime(self, *, platform, dp_encoder, attn_tp_size): + stack = ExitStack() + stack.enter_context( + override_platform(**{f"is_{name}": name == platform for name in _PLATFORMS}) + ) + stack.enter_context( + get_context().override_server_args(mm_enable_dp_encoder=dp_encoder) + ) + stack.enter_context( + get_parallel().override( + tp_size=attn_tp_size, + attn_tp_size=attn_tp_size, + moe_tp_size=attn_tp_size, + attn_tp_rank=0, + ) + ) + return stack + + +# (platform, mm_enable_dp_encoder, attn_tp_size, batch across requests) +_GATE_CASES = ( + ("hip", True, 2, True), + ("hip", True, 1, False), + ("hip", False, 2, False), + ("hip", False, 1, False), + ("npu", True, 2, False), + ("npu", False, 1, False), + ("xpu", True, 2, False), + ("xpu", False, 1, False), + ("cuda", True, 2, True), + ("cuda", False, 1, True), +) + + +class TestVitBatchingGate(_RuntimeCase): + def test_gate(self): + for platform, dp_encoder, attn_tp_size, batched in _GATE_CASES: + with ( + self.subTest( + platform=platform, dp_encoder=dp_encoder, attn_tp=attn_tp_size + ), + self._runtime( + platform=platform, + dp_encoder=dp_encoder, + attn_tp_size=attn_tp_size, + ), + ): + self.assertEqual(mm_schedule._batch_vit_across_requests(), batched) + self.assertEqual( + hip_dp_encoder_batches_vit(), platform == "hip" and batched + ) + + def test_chunked_prefill_encodes_requests_in_one_call_only_under_gate(self): + offsets = [(2, 5)] + seq_len = 8 + for dp_encoder, attn_tp_size, expected_calls in ( + (True, 2, [2]), + (True, 1, [1, 1]), + (False, 2, [1, 1]), + ): + with ( + self.subTest(dp_encoder=dp_encoder, attn_tp=attn_tp_size), + self._runtime( + platform="hip", dp_encoder=dp_encoder, attn_tp_size=attn_tp_size + ), + ): + mm_schedule.init_mm_embedding_cache(1 << 30) + items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + hash=2000 + req, + feature=torch.zeros(1), + offsets=offsets, + ) + for req in range(2) + ] + encoder = Mock(side_effect=lambda batch: torch.ones(4 * len(batch), 8)) + + chunk, _ = mm_schedule._get_chunked_prefill_embedding( + encoder, + items, + [0, 1, 2], + [0, 0], + [seq_len, seq_len], + [offsets, offsets], + torch.zeros(2 * seq_len, dtype=torch.long), + ) + + self.assertEqual( + [len(call.args[0]) for call in encoder.call_args_list], + expected_calls, + ) + self.assertEqual(tuple(chunk.shape), (8, 8)) + + +class TestQwen3VLVisionWarmUp(_RuntimeCase): + @staticmethod + def _model(visual): + model = Qwen3VLForConditionalGeneration.__new__(Qwen3VLForConditionalGeneration) + torch.nn.Module.__init__(model) + model.visual = visual + return model + + def test_post_load_hook_warms_the_tower_once_under_gate(self): + tower = _RecordingVisionTower() + with self._runtime(platform="hip", dp_encoder=True, attn_tp_size=2): + maybe_precompile_model_kernels_after_loading(self._model(tower), "cpu") + + self.assertEqual( + tower.calls, + [ + (32 * 32, [[1, 32, 32]], False), + (48 * 64, [[1, 48, 64]], False), + (6 * 32 * 32, [[1, 32, 32]] * 6, False), + ], + ) + + def test_no_warm_up_off_gate(self): + for platform, dp_encoder, attn_tp_size, _ in _GATE_CASES: + if platform == "hip" and dp_encoder and attn_tp_size > 1: + continue + tower = _RecordingVisionTower() + with ( + self.subTest( + platform=platform, dp_encoder=dp_encoder, attn_tp=attn_tp_size + ), + self._runtime( + platform=platform, + dp_encoder=dp_encoder, + attn_tp_size=attn_tp_size, + ), + ): + maybe_precompile_model_kernels_after_loading(self._model(tower), "cpu") + self.assertEqual(tower.calls, []) + + def test_language_model_only_skips_warm_up(self): + with self._runtime(platform="hip", dp_encoder=True, attn_tp_size=2): + maybe_precompile_model_kernels_after_loading(self._model(None), "cpu") + + +if __name__ == "__main__": + unittest.main() From f9b6c0f23322165034c2f6f4d3c851703cec6b73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Mon, 26 Oct 2026 21:48:05 +0000 Subject: [PATCH 2/2] [AMD] vision: warm up the GLM-5-Next vision tower under the batched DP encoder With ViT batching across requests under the DP encoder on ROCm, a rank other than 0 runs its first vision call inside the first served image burst and compiles swiglu_clamped there, as Qwen3-VL does for its rotary kernel. - mm_utils.warm_up_vision_tower(visual, grids): the dummy-image loop from Qwen3-VL's warm-up, shared. - Qwen3-VL calls it with its grids (unchanged). - Glm5NextForConditionalGeneration.precompile_kernels_after_loading warms its tower under hip_dp_encoder_batches_vit() with two grids: the first compiles swiglu_clamped's 3-D (blocks) and 2-D (merger) calls, the second's new size recompiles the 3-D call with a dynamic shape. - Tests: the GLM hook's calls under the gate, none off the gate, none without a vision tower. --- python/sglang/srt/models/glm5_next.py | 10 +++++++ python/sglang/srt/models/qwen3_vl.py | 27 ++++++------------- python/sglang/srt/multimodal/mm_utils.py | 26 ++++++++++++++++++ .../test_hip_dp_encoder_vit_batching.py | 25 +++++++++++++++++ 4 files changed, 69 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/models/glm5_next.py b/python/sglang/srt/models/glm5_next.py index 8df16198c2a4..bde53b4e316c 100644 --- a/python/sglang/srt/models/glm5_next.py +++ b/python/sglang/srt/models/glm5_next.py @@ -111,8 +111,10 @@ ) from sglang.srt.models.utils import WeightsMapper from sglang.srt.multimodal.mm_utils import ( + hip_dp_encoder_batches_vit, run_dp_presharded_mrope_vision_model, run_dp_sharded_mrope_vision_model, + warm_up_vision_tower, ) from sglang.srt.runtime_context import ( get_lora, @@ -1556,6 +1558,14 @@ def set_dflash_layers_to_capture(self, layer_ids: List[int]): # Capturing before layer k + 1 gives the completed output of layer k. self.model.layers_to_capture = [val + 1 for val in layer_ids] + def precompile_kernels_after_loading(self) -> None: + # Under the batched DP encoder every rank's first ViT call would otherwise + # compile swiglu_clamped inside the first image burst. The first grid + # compiles its 3-D (blocks) and 2-D (merger) calls; the second grid's new + # size recompiles the 3-D call with a dynamic shape. + if self.visual is not None and hip_dp_encoder_batches_vit(): + warm_up_vision_tower(self.visual, ([[1, 32, 32]], [[1, 48, 64]])) + def pad_input_ids(self, input_ids: array, mm_inputs: MultimodalInputs) -> array: return pad_mm_tokens(input_ids, mm_inputs) diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 1bdc490c81cb..ca6b9952a7ab 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -80,6 +80,7 @@ hip_dp_encoder_batches_vit, materialize_multimodal_features, run_dp_sharded_mrope_vision_model, + warm_up_vision_tower, ) from sglang.srt.multimodal.transport.cuda_ipc import ( BORROW_CUDA_IPC_FEATURE_KEY, @@ -1453,26 +1454,14 @@ def _warm_up_vision_tower(self) -> None: sizes cover the static and the dynamic-shape compile; the last call has enough images to take the vectorized position-embedding path. """ - patch_embed = self.visual.patch_embed - patch_dim = ( - patch_embed.in_channels - * patch_embed.temporal_patch_size - * patch_embed.patch_size**2 + warm_up_vision_tower( + self.visual, + ( + [[1, 32, 32]], + [[1, 48, 64]], + [[1, 32, 32]] * _VECTORIZED_VL_POS_EMBED_MIN_IMAGES, + ), ) - grids = ( - [[1, 32, 32]], - [[1, 48, 64]], - [[1, 32, 32]] * _VECTORIZED_VL_POS_EMBED_MIN_IMAGES, - ) - with torch.no_grad(): - for grid in grids: - pixel_values = torch.zeros( - sum(t * h * w for t, h, w in grid), - patch_dim, - device=self.visual.device, - dtype=self.visual.dtype, - ) - self.visual(pixel_values, grid_thw=torch.tensor(grid)) def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: _require_vision(self) diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index e11b35dafa3d..ba549a767229 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/python/sglang/srt/multimodal/mm_utils.py @@ -434,6 +434,32 @@ def hip_dp_encoder_batches_vit() -> bool: ) +def warm_up_vision_tower( + visual: torch.nn.Module, grids: Sequence[list[list[int]]] +) -> None: + """Run the vision tower on zero images of each grid on this rank, without grad. + + Under ``hip_dp_encoder_batches_vit``, a non-zero rank otherwise runs its first + ViT call inside the first served image burst and pays its torch.compile + compiles there. ``grids`` holds one ``grid_thw`` list per call. + """ + patch_embed = visual.patch_embed + patch_dim = ( + patch_embed.in_channels + * patch_embed.temporal_patch_size + * patch_embed.patch_size**2 + ) + with torch.no_grad(): + for grid in grids: + pixel_values = torch.zeros( + sum(t * h * w for t, h, w in grid), + patch_dim, + device=visual.device, + dtype=visual.dtype, + ) + visual(pixel_values, grid_thw=torch.tensor(grid)) + + # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/vision.py def get_dp_encoder_lb_assignment( sizes: list[int], diff --git a/test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py b/test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py index 19718e0d97fa..686b48a421b3 100644 --- a/test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py +++ b/test/registered/unit/multimodal/test_hip_dp_encoder_vit_batching.py @@ -16,6 +16,7 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import ( maybe_precompile_model_kernels_after_loading, ) +from sglang.srt.models.glm5_next import Glm5NextForConditionalGeneration from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration from sglang.srt.multimodal.mm_utils import hip_dp_encoder_batches_vit from sglang.srt.runtime_context import ( @@ -194,5 +195,29 @@ def test_language_model_only_skips_warm_up(self): maybe_precompile_model_kernels_after_loading(self._model(None), "cpu") +class TestGlm5NextVisionWarmUp(TestQwen3VLVisionWarmUp): + @staticmethod + def _model(visual): + model = Glm5NextForConditionalGeneration.__new__( + Glm5NextForConditionalGeneration + ) + torch.nn.Module.__init__(model) + model.visual = visual + return model + + def test_post_load_hook_warms_the_tower_once_under_gate(self): + tower = _RecordingVisionTower() + with self._runtime(platform="hip", dp_encoder=True, attn_tp_size=2): + maybe_precompile_model_kernels_after_loading(self._model(tower), "cpu") + + self.assertEqual( + tower.calls, + [ + (32 * 32, [[1, 32, 32]], False), + (48 * 64, [[1, 48, 64]], False), + ], + ) + + if __name__ == "__main__": unittest.main()