diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9cc2f7710..61d8391d9 100644 --- a/python/sglang/srt/environ.py +++ b/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. diff --git a/python/sglang/srt/models/moss_vl.py b/python/sglang/srt/models/moss_vl.py index 8ea3be829..8d4007162 100644 --- a/python/sglang/srt/models/moss_vl.py +++ b/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,7 +518,13 @@ class MossVLVisionModel(nn.Module): x = x.to(device=self.device, dtype=self.dtype) x = self.patch_embed(x) - pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + 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) diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 76d6922ac..6a89e803a 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/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,7 +896,10 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): grid_thw_list = grid_thw.tolist() grid_thw = grid_thw.cpu().numpy() - pos_embeds = self.fast_pos_embed_interpolate_from_list(grid_thw_list) + 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 rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list) @@ -948,7 +1080,14 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): else: grid_thw_list = grid_thw.tolist() - pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + 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 # rotary embedding -> (cos, sin) diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 52951b226..cc2b1002e 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/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 ( - discard_alpha_channel - and not isinstance(img, torch.Tensor) - and img.mode != "RGB" - ): - # Needed only when `img` is a PIL image - img = img.convert("RGB") + if isinstance(img, torch.Tensor): + return img # JPEG already decoded on GPU by nvJPEG + # PIL decodes lazily; do it here in the io worker so the decode + # doesn't run later on the event-loop thread. + if discard_alpha_channel and img.mode != "RGB": + return img.convert("RGB") + img.load() return img elif modality == Modality.VIDEO: return load_video(data, frame_count_limit) diff --git a/test/registered/models/test_vit_pos_embed_interpolate.py b/test/registered/models/test_vit_pos_embed_interpolate.py new file mode 100644 index 000000000..ef9313797 --- /dev/null +++ b/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() diff --git a/test/registered/unit/multimodal/test_base_processor_image_decode.py b/test/registered/unit/multimodal/test_base_processor_image_decode.py new file mode 100644 index 000000000..171abdda5 --- /dev/null +++ b/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()