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)