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

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

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-06-24 14:36:29 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent 534ac98eb2
commit 0df796473b
6 changed files with 476 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+4
View File
@@ -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
+126 -1
View File
@@ -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)
python/sglang/srt/models/qwen3_vl.py
+141 -2
View File
@@ -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)
python/sglang/srt/multimodal/processors/base_processor.py
+7 -7
View File
@@ -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)
test/registered/models/test_vit_pos_embed_interpolate.py
+118
View File
@@ -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
+80
View File
@@ -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()