[VLM] Qwen3-VL / Moss-VL ViT preprocessing optimizations (#28940)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/environ.py

@@ -663,6 +663,10 @@ class Envs:
     SGLANG_MM_BUFFER_SIZE_MB = EnvInt(0)
     SGLANG_MM_PRECOMPUTE_HASH = EnvBool(False)
     SGLANG_VIT_ENABLE_CUDA_GRAPH = EnvBool(False)
+    # Use the fully-vectorized ViT position-embedding interpolation (no per-image
+    # Python loop / CPU<->GPU sync). Bit-exact with the legacy implementation;
+    # set False to fall back to the per-image loop.
+    SGLANG_VIT_ENABLE_VECTORIZED_POS_EMBED = EnvBool(True)
     SGLANG_MM_SKIP_COMPUTE_HASH = EnvBool(False)
     # For pre-tokenized (list[int]) multimodal prompts,
     # preserve the user's original tokens to avoid retokenization drift.

python/sglang/srt/models/moss_vl.py

@@ -15,6 +15,7 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
     Qwen2_5_VisionRotaryEmbedding,
 )
 
+from sglang.srt.environ import envs
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.attention.vision import VisionAttention
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
@@ -49,6 +50,10 @@ from sglang.srt.utils import add_prefix
 
 logger = logging.getLogger(__name__)
 
+# Below this image count the per-image loop beats the vectorized path (which has a
+# fixed setup cost); both give the same result.
+_VECTORIZED_VL_POS_EMBED_MIN_IMAGES = 6
+
 
 # ==================== Vision Components ====================
 
@@ -391,6 +396,120 @@ class MossVLVisionModel(nn.Module):
 
         return torch.cat(patch_pos_embeds_permute)
 
+    def fast_pos_embed_interpolate_vectorized(
+        self, grid_thw: torch.Tensor
+    ) -> torch.Tensor:
+        """Vectorized fast_pos_embed_interpolate (no per-image loop).
+
+        Same result as the loop version; the cost no longer scales with the number
+        of images.
+        """
+        num_grid_per_side = int(self.num_position_embeddings**0.5)
+        m = self.spatial_merge_size
+        device = self.pos_embed.weight.device
+        dtype = self.pos_embed.weight.dtype
+
+        grid_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist()
+        ts = [int(g[0]) for g in grid_list]
+        hs = [int(g[1]) for g in grid_list]
+        ws = [int(g[2]) for g in grid_list]
+        num_images = len(grid_list)
+
+        hw_list = [h * w for h, w in zip(hs, ws)]
+        thw_list = [t * s for t, s in zip(ts, hw_list)]
+        total_hw = sum(hw_list)
+        total_out = sum(thw_list)
+
+        def _exclusive_prefix(sizes):
+            out, acc = [], 0
+            for s in sizes:
+                out.append(acc)
+                acc += s
+            return torch.tensor(out, device=device, dtype=torch.long)
+
+        hw_off = _exclusive_prefix(hw_list)
+        thw_off = _exclusive_prefix(thw_list)
+        image_arange = torch.arange(num_images, device=device)
+
+        base_image_id = torch.repeat_interleave(
+            image_arange, torch.tensor(hw_list, device=device)
+        )
+        base_local = torch.arange(total_hw, device=device) - hw_off[base_image_id]
+        w_of = torch.tensor(ws, device=device)[base_image_id]
+        row = base_local // w_of
+        col = base_local % w_of
+
+        uniq_h, inv_h = torch.unique(
+            torch.tensor(hs, device=device), return_inverse=True
+        )
+        uniq_w, inv_w = torch.unique(
+            torch.tensor(ws, device=device), return_inverse=True
+        )
+        h_luts = [
+            torch.linspace(0, num_grid_per_side - 1, int(h), device=device)
+            for h in uniq_h.tolist()
+        ]
+        w_luts = [
+            torch.linspace(0, num_grid_per_side - 1, int(w), device=device)
+            for w in uniq_w.tolist()
+        ]
+        h_lut_off = _exclusive_prefix([len(x) for x in h_luts])
+        w_lut_off = _exclusive_prefix([len(x) for x in w_luts])
+        h_idxs = torch.cat(h_luts)[h_lut_off[inv_h[base_image_id]] + row]
+        w_idxs = torch.cat(w_luts)[w_lut_off[inv_w[base_image_id]] + col]
+
+        h_floor = h_idxs.int()
+        w_floor = w_idxs.int()
+        h_ceil = (h_idxs.int() + 1).clip(max=num_grid_per_side - 1)
+        w_ceil = (w_idxs.int() + 1).clip(max=num_grid_per_side - 1)
+        dh = h_idxs - h_floor
+        dw = w_idxs - w_floor
+
+        base_h = h_floor * num_grid_per_side
+        base_h_ceil = h_ceil * num_grid_per_side
+        indices = torch.stack(
+            [
+                base_h + w_floor,
+                base_h + w_ceil,
+                base_h_ceil + w_floor,
+                base_h_ceil + w_ceil,
+            ],
+            dim=0,
+        ).to(dtype=torch.long)
+        weights = torch.stack(
+            [
+                (1 - dh) * (1 - dw),
+                (1 - dh) * dw,
+                dh * (1 - dw),
+                dh * dw,
+            ],
+            dim=0,
+        ).to(dtype=dtype)
+        pe = self.pos_embed(indices) * weights[:, :, None]
+        base_embeds = pe[0] + pe[1] + pe[2] + pe[3]  # [total_hw, C]
+
+        out_image_id = torch.repeat_interleave(
+            image_arange, torch.tensor(thw_list, device=device)
+        )
+        pos_in_image = torch.arange(total_out, device=device) - thw_off[out_image_id]
+        hw_of_out = torch.tensor(hw_list, device=device)[out_image_id]
+        frame_idx = pos_in_image // hw_of_out
+        local_idx = pos_in_image % hw_of_out
+        patch = base_embeds[hw_off[out_image_id] + local_idx]
+
+        all_w = torch.tensor(ws, device=device)[out_image_id]
+        rows = local_idx // all_w
+        cols = local_idx % all_w
+        out_within = (
+            frame_idx * hw_of_out
+            + ((rows // m) * (all_w // m) + (cols // m)) * m * m
+            + (rows % m) * m
+            + (cols % m)
+        )
+        merged = torch.empty_like(patch)
+        merged[out_within + thw_off[out_image_id]] = patch
+        return merged
+
     def forward(
         self,
         x: torch.Tensor,
@@ -399,6 +518,12 @@ class MossVLVisionModel(nn.Module):
         x = x.to(device=self.device, dtype=self.dtype)
         x = self.patch_embed(x)
 
+        if (
+            envs.SGLANG_VIT_ENABLE_VECTORIZED_POS_EMBED.get()
+            and grid_thw.shape[0] >= _VECTORIZED_VL_POS_EMBED_MIN_IMAGES
+        ):
+            pos_embeds = self.fast_pos_embed_interpolate_vectorized(grid_thw)
+        else:
             pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
         x = x + pos_embeds
         rotary_pos_emb = self.rot_pos_emb(grid_thw)

python/sglang/srt/models/qwen3_vl.py

@@ -94,6 +94,10 @@ logger = logging.getLogger(__name__)
 _is_cpu_amx_available = cpu_has_amx_support()
 _is_cpu = is_cpu()
 
+# Below this image count the per-image loop beats the vectorized path (which has a
+# fixed setup cost; measured crossover ~6 on H20); both give the same result.
+_VECTORIZED_VL_POS_EMBED_MIN_IMAGES = 6
+
 
 class Qwen3_VisionMLP(nn.Module):
 
@@ -597,6 +601,131 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
 
         return torch.cat(outputs, dim=0)
 
+    def _use_vectorized_pos_embed(self, num_images: int) -> bool:
+        """Use the vectorized path only past a few images.
+
+        It drops the per-image loop but has a fixed setup cost, so the loop is
+        faster for a handful of images. Both give the same result.
+        """
+        return (
+            envs.SGLANG_VIT_ENABLE_VECTORIZED_POS_EMBED.get()
+            and num_images >= _VECTORIZED_VL_POS_EMBED_MIN_IMAGES
+        )
+
+    def fast_pos_embed_interpolate_vectorized(self, grid_thw):
+        """Vectorized fast_pos_embed_interpolate_from_list (no per-image loop).
+
+        Same result as the loop version; the cost no longer scales with the number
+        of images.
+        """
+        num_grid_per_side = self.num_grid_per_side
+        m = self.spatial_merge_size
+        dtype = self.dtype
+        device = self.device
+
+        grid_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist()
+        ts = [int(g[0]) for g in grid_list]
+        hs = [int(g[1]) for g in grid_list]
+        ws = [int(g[2]) for g in grid_list]
+        num_images = len(grid_list)
+
+        hw_list = [h * w for h, w in zip(hs, ws)]  # base tokens / frame / image
+        thw_list = [t * s for t, s in zip(ts, hw_list)]  # output tokens / image
+        total_hw = sum(hw_list)
+        total_out = sum(thw_list)
+
+        def _exclusive_prefix(sizes):
+            out, acc = [], 0
+            for s in sizes:
+                out.append(acc)
+                acc += s
+            return torch.tensor(out, device=device, dtype=torch.long)
+
+        hw_off = _exclusive_prefix(hw_list)  # image offset in the base layout
+        thw_off = _exclusive_prefix(thw_list)  # image offset in the output layout
+        image_arange = torch.arange(num_images, device=device)
+
+        # --- 1. per base-token image id + local (row, col) (single frame) ---
+        base_image_id = torch.repeat_interleave(
+            image_arange, torch.tensor(hw_list, device=device)
+        )
+        base_local = torch.arange(total_hw, device=device) - hw_off[base_image_id]
+        w_of = torch.tensor(ws, device=device)[base_image_id]
+        row = base_local // w_of
+        col = base_local % w_of
+
+        # per-size linspace LUT (one entry per unique h/w), so images of the same
+        # size share coords without the per-image loop
+        uniq_h, inv_h = torch.unique(
+            torch.tensor(hs, device=device), return_inverse=True
+        )
+        uniq_w, inv_w = torch.unique(
+            torch.tensor(ws, device=device), return_inverse=True
+        )
+        h_luts = [
+            torch.linspace(0, num_grid_per_side - 1, int(h), device=device)
+            for h in uniq_h.tolist()
+        ]
+        w_luts = [
+            torch.linspace(0, num_grid_per_side - 1, int(w), device=device)
+            for w in uniq_w.tolist()
+        ]
+        h_lut_off = _exclusive_prefix([len(x) for x in h_luts])
+        w_lut_off = _exclusive_prefix([len(x) for x in w_luts])
+        h_idxs = torch.cat(h_luts)[h_lut_off[inv_h[base_image_id]] + row]
+        w_idxs = torch.cat(w_luts)[w_lut_off[inv_w[base_image_id]] + col]
+
+        h_floor = h_idxs.to(torch.long)
+        w_floor = w_idxs.to(torch.long)
+        h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)
+        w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)
+        dh = h_idxs - h_floor
+        dw = w_idxs - w_floor
+        # bilinear weights (same form as ..._from_list)
+        w11 = dh * dw
+        w10 = dh - w11
+        w01 = dw - w11
+        w00 = 1 - dh - w01
+
+        base_h = h_floor * num_grid_per_side
+        base_h_ceil = h_ceil * num_grid_per_side
+        indices = torch.stack(
+            [
+                base_h + w_floor,
+                base_h + w_ceil,
+                base_h_ceil + w_floor,
+                base_h_ceil + w_ceil,
+            ],
+            dim=0,
+        )
+        weights = torch.stack([w00, w01, w10, w11], dim=0).to(dtype=dtype)
+        embeds = self.pos_embed(indices) * weights[:, :, None]
+        base_embeds = embeds.sum(dim=0)  # [total_hw, C]
+
+        # --- 2. temporal repeat (gather) ---
+        out_image_id = torch.repeat_interleave(
+            image_arange, torch.tensor(thw_list, device=device)
+        )
+        pos_in_image = torch.arange(total_out, device=device) - thw_off[out_image_id]
+        hw_of_out = torch.tensor(hw_list, device=device)[out_image_id]
+        frame_idx = pos_in_image // hw_of_out
+        local_idx = pos_in_image % hw_of_out
+        patch = base_embeds[hw_off[out_image_id] + local_idx]  # [total_out, C]
+
+        # --- 3. spatial-merge reorder (scatter) ---
+        all_w = torch.tensor(ws, device=device)[out_image_id]
+        rows = local_idx // all_w
+        cols = local_idx % all_w
+        out_within = (
+            frame_idx * hw_of_out
+            + ((rows // m) * (all_w // m) + (cols // m)) * m * m
+            + (rows % m) * m
+            + (cols % m)
+        )
+        merged = torch.empty_like(patch)
+        merged[out_within + thw_off[out_image_id]] = patch
+        return merged
+
     def add_padding_to_fi_seqlens(
         self, seq: np.ndarray, batch_size: int, padding_value: int
     ) -> np.ndarray:
@@ -767,6 +896,9 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
             grid_thw_list = grid_thw.tolist()
             grid_thw = grid_thw.cpu().numpy()
 
+        if self._use_vectorized_pos_embed(len(grid_thw_list)):
+            pos_embeds = self.fast_pos_embed_interpolate_vectorized(grid_thw_list)
+        else:
             pos_embeds = self.fast_pos_embed_interpolate_from_list(grid_thw_list)
         x += pos_embeds
 
@@ -948,6 +1080,13 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
         else:
             grid_thw_list = grid_thw.tolist()
 
+        if self.align_corners and self._use_vectorized_pos_embed(len(grid_thw_list)):
+            # The vectorized implementation uses linspace coordinates. In graph mode
+            # the legacy fallback honors enable_precise_embedding_interpolation, so
+            # only use the vectorized path when the active graph interpolation mode
+            # is also linspace; otherwise image count would change the output.
+            pos_embeds = self.fast_pos_embed_interpolate_vectorized(grid_thw_list)
+        else:
             pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
         x += pos_embeds
 

python/sglang/srt/multimodal/processors/base_processor.py

@@ -537,13 +537,13 @@ class BaseMultimodalProcessor(ABC):
         try:
             if modality == Modality.IMAGE:
                 img, _ = load_image(data, cls.gpu_image_decode)
-                if (
+                if isinstance(img, torch.Tensor):
-                    discard_alpha_channel
+                    return img  # JPEG already decoded on GPU by nvJPEG
-                    and not isinstance(img, torch.Tensor)
+                # PIL decodes lazily; do it here in the io worker so the decode
-                    and img.mode != "RGB"
+                # doesn't run later on the event-loop thread.
-                ):
+                if discard_alpha_channel and img.mode != "RGB":
-                    # Needed only when `img` is a PIL image
+                    return img.convert("RGB")
-                    img = img.convert("RGB")
+                img.load()
                 return img
             elif modality == Modality.VIDEO:
                 return load_video(data, frame_count_limit)

test/registered/models/test_vit_pos_embed_interpolate.py

@@ -0,0 +1,118 @@
+"""Bit-exact unit test for the vectorized ViT position-embedding interpolation.
+
+The vectorized path (``fast_pos_embed_interpolate_vectorized``) removes the
+per-image Python loop / CPU<->GPU sync of the legacy implementations. It is meant
+to be a pure speedup, so it must be numerically *identical* (bit-exact, rtol=0
+atol=0) to the loop version it replaces -- for single images, many images, video
+(t>1), and mixed-size batches, in both bf16 and fp32.
+
+The interpolation is a sequence of embedding lookups + arithmetic, so it runs and
+is bit-exact on CPU; the test exercises CUDA too when available. It calls the real
+model methods on a lightweight stub holding a real ``nn.Embedding`` (no model
+weights / distributed init needed).
+
+    python -m pytest test/registered/models/test_vit_pos_embed_interpolate.py -v
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+import torch.nn as nn
+
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
+register_cuda_ci(est_time=20, stage="base-a", runner_config="1-gpu-small")
+
+NUM_POS = 2304  # Qwen3-VL num_position_embeddings -> 48x48 grid
+HIDDEN = 64  # small hidden dim keeps the unit test fast
+MERGE = 2
+
+# t, h, w grids (h, w are multiples of MERGE). Covers single / large-upsample /
+# multi-mixed / video / video+image / many-duplicate.
+GRID_CASES = {
+    "single": [[1, 16, 16]],
+    "single_large": [[1, 64, 98]],  # h, w may exceed grid side (upsample)
+    "multi_mixed": [[1, 16, 24], [1, 32, 12], [1, 8, 40]],
+    "video": [[4, 16, 20]],
+    "video_plus_image": [[3, 12, 16], [1, 20, 28], [2, 8, 8]],
+    "many": [[1, 24, 24]] * 8,
+}
+
+
+def _devices():
+    devs = [torch.device("cpu")]
+    if torch.cuda.is_available():
+        devs.append(torch.device("cuda"))
+    return devs
+
+
+class TestViTPosEmbedInterpolate(CustomTestCase):
+    def _check(self, stub, legacy_fn, vectorized_fn, grid, label):
+        ref = legacy_fn(stub, grid)
+        out = vectorized_fn(stub, grid)
+        self.assertEqual(ref.shape, out.shape, f"{label}: shape mismatch")
+        self.assertTrue(
+            torch.equal(ref, out),
+            f"{label}: not bit-exact, max|diff|="
+            f"{(ref.float() - out.float()).abs().max().item():.3e}",
+        )
+
+    def test_qwen3_vl_vectorized_matches_loop(self):
+        try:
+            from sglang.srt.models.qwen3_vl import Qwen3VLMoeVisionModel as M
+        except Exception as e:  # heavy optional deps (flashinfer, ...) unavailable
+            self.skipTest(f"cannot import Qwen3VLMoeVisionModel: {e}")
+
+        for device in _devices():
+            for dtype in (torch.bfloat16, torch.float32):
+                stub = SimpleNamespace(
+                    num_grid_per_side=int(NUM_POS**0.5),
+                    spatial_merge_size=MERGE,
+                    num_position_embeddings=NUM_POS,
+                    pos_embed=nn.Embedding(NUM_POS, HIDDEN).to(
+                        device=device, dtype=dtype
+                    ),
+                    dtype=dtype,
+                    device=device,
+                )
+                for name, grid in GRID_CASES.items():
+                    self._check(
+                        stub,
+                        M.fast_pos_embed_interpolate_from_list,
+                        M.fast_pos_embed_interpolate_vectorized,
+                        grid,
+                        f"qwen3_vl/{name}/{dtype}/{device.type}",
+                    )
+
+    def test_moss_vl_vectorized_matches_loop(self):
+        try:
+            from sglang.srt.models.moss_vl import MossVLVisionModel as M
+        except Exception as e:
+            self.skipTest(f"cannot import MossVLVisionModel: {e}")
+
+        for device in _devices():
+            for dtype in (torch.bfloat16, torch.float32):
+                stub = SimpleNamespace(
+                    spatial_merge_size=MERGE,
+                    num_position_embeddings=NUM_POS,
+                    pos_embed=nn.Embedding(NUM_POS, HIDDEN).to(
+                        device=device, dtype=dtype
+                    ),
+                )
+                for name, grid in GRID_CASES.items():
+                    # the legacy moss method consumes a [num_images, 3] tensor
+                    grid_t = torch.tensor(grid, device=device)
+                    self._check(
+                        stub,
+                        M.fast_pos_embed_interpolate,
+                        M.fast_pos_embed_interpolate_vectorized,
+                        grid_t,
+                        f"moss_vl/{name}/{dtype}/{device.type}",
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/multimodal/test_base_processor_image_decode.py

@@ -0,0 +1,80 @@
+"""Unit tests for ``BaseMultimodalProcessor._load_single_item`` image decoding.
+
+Regression test for the change that forces the (otherwise lazy) PIL decode inside
+``_load_single_item`` — which runs in the ``io_executor`` worker thread — instead of
+letting it fire lazily on the main event-loop thread later (inside
+``pil_to_tensor``/``tobytes`` during processing). The behavior of the returned image
+(mode, pixels) must be unchanged; only *when/where* the decode happens differs.
+
+No server, no model loading — pure CPU.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+import io
+import unittest
+
+import numpy as np
+from PIL import Image
+
+from sglang.srt.managers.schedule_batch import Modality
+from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
+from sglang.test.test_utils import CustomTestCase
+
+
+class _StubProcessor(BaseMultimodalProcessor):
+    # gpu_image_decode=False forces the PIL (CPU) path so the test needs no GPU and
+    # exercises exactly the lazy-decode branch the fix targets. The abstract methods
+    # are never called: we only invoke the _load_single_item classmethod.
+    gpu_image_decode = False
+
+
+def _png_bytes(mode: str = "RGB", size=(8, 8)) -> bytes:
+    arr = (np.random.RandomState(0).rand(size[1], size[0], 3) * 255).astype("uint8")
+    img = Image.fromarray(arr, "RGB").convert(mode)
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def _is_decoded(img: Image.Image) -> bool:
+    """A lazily-opened PIL image has no decoded core yet; ``load()`` populates it.
+    PIL's ``.im`` property requires a completed load and raises otherwise."""
+    try:
+        return img.im is not None
+    except Exception:
+        return False
+
+
+class TestLoadSingleItemImageDecode(CustomTestCase):
+    def test_plain_open_is_lazy(self):
+        # Documents why the fix matters: a bare Image.open is not decoded yet, so
+        # without the fix the decode would land on the caller (main) thread.
+        lazy = Image.open(io.BytesIO(_png_bytes()))
+        self.assertFalse(_is_decoded(lazy))
+
+    def test_load_single_item_forces_decode(self):
+        img = _StubProcessor._load_single_item(_png_bytes("RGB"), Modality.IMAGE)
+        self.assertIsInstance(img, Image.Image)
+        self.assertEqual(img.mode, "RGB")
+        # The fix: decode is forced inside _load_single_item, not lazily later.
+        self.assertTrue(_is_decoded(img))
+
+    def test_rgba_converted_to_rgb_and_decoded(self):
+        img = _StubProcessor._load_single_item(_png_bytes("RGBA"), Modality.IMAGE)
+        # Existing alpha-discard behavior preserved.
+        self.assertEqual(img.mode, "RGB")
+        self.assertTrue(_is_decoded(img))
+
+    def test_pixels_match_reference(self):
+        # Output must be bit-identical to the pre-fix path (open -> [convert]).
+        data = _png_bytes("RGB")
+        img = _StubProcessor._load_single_item(data, Modality.IMAGE)
+        ref = Image.open(io.BytesIO(data)).convert("RGB")
+        np.testing.assert_array_equal(np.asarray(img), np.asarray(ref))
+
+
+if __name__ == "__main__":
+    unittest.main()