[Perf] [AMD] vision: Batch ViT across requests on ROCm when the DP encoder is enabled

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)
 

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)
 

python/sglang/srt/models/qwen3_vl.py

@@ -77,8 +77,10 @@
     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,
+    warm_up_vision_tower,
 )
 from sglang.srt.multimodal.transport.cuda_ipc import (
     BORROW_CUDA_IPC_FEATURE_KEY,
@@ -1439,6 +1441,28 @@ 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.
+        """
+        warm_up_vision_tower(
+            self.visual,
+            (
+                [[1, 32, 32]],
+                [[1, 48, 64]],
+                [[1, 32, 32]] * _VECTORIZED_VL_POS_EMBED_MIN_IMAGES,
+            ),
+        )
+
     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)

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():

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,49 @@ 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
+    )
+
+
+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],

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 = []