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

Refactor(qwen3-vl) optimize position encoding interpolation (#16781)

Browse Source 
Signed-off-by: chenzhenyang <andy271828@163.com>
Signed-off-by: chenzhenyang <chenzhenyang@moonshot.cn>
Co-authored-by: chenzhenyang <chenzhenyang@moonshot.cn>
Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>
This commit is contained in:
aaaandychen
2026-02-06 02:26:35 +08:00
committed by GitHub
parent d22163eb8c
commit 6a4b81e2d9
2 changed files with 126 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

146
python/sglang/srt/models/qwen3_vl.py
View File

@@ -19,6 +19,7 @@ import re
from functools import lru_cache, partial
from typing import Callable, Iterable, List, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
from einops import rearrange
@@ -396,31 +397,130 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
return cos_combined, sin_combined
def fast_pos_embed_interpolate(self, grid_thw):
patch_pos_embeds_permute = []
m_size = self.spatial_merge_size
def _get_interpolation_indices(self, dim_size: int) -> torch.Tensor:
"""
Compute continuous interpolation indices for a single dimension.
embeds = torch.arange(self.num_grid, device=self.pos_embed.weight.device)
embeds = (
self.pos_embed(embeds)
.permute(1, 0)
.reshape(1, -1, self.num_grid_per_side, self.num_grid_per_side)
Returns continuous indices.
"""
if self.align_corners:
indices = np.linspace(
0, self.num_grid_per_side - 1, dim_size, dtype=np.float32
)
else:
indices = (np.arange(dim_size, dtype=np.float32) + 0.5) * (
self.num_grid_per_side / dim_size
) - 0.5
indices = np.clip(indices, 0, self.num_grid_per_side - 1)
return indices
def _calculate_indices_and_weights(self, h_idxs, w_idxs):
"""
Compute bilinear interpolation indices and weights.
Returns tuple of (indices, weights), each as 4 numpy arrays for the 4 corner points.
"""
h_f = np.floor(h_idxs).astype(np.int64)
h_c = np.clip(h_f + 1, 0, self.num_grid_per_side - 1)
dh = h_idxs - h_f
w_f = np.floor(w_idxs).astype(np.int64)
w_c = np.clip(w_f + 1, 0, self.num_grid_per_side - 1)
dw = w_idxs - w_f
side = self.num_grid_per_side
indices = [
(h_f[:, None] * side + w_f).flatten(),
(h_f[:, None] * side + w_c).flatten(),
(h_c[:, None] * side + w_f).flatten(),
(h_c[:, None] * side + w_c).flatten(),
]
weights = [
((1 - dh)[:, None] * (1 - dw)).flatten(),
((1 - dh)[:, None] * dw).flatten(),
(dh[:, None] * (1 - dw)).flatten(),
(dh[:, None] * dw).flatten(),
]
return indices, weights
def _get_position_embedding(self, patch_pos_embeds, grid_ts, grid_hs, grid_ws):
"""
Tile and reorganize position embeddings to align with the token sequence.
"""
result_parts = []
merge_size = self.spatial_merge_size
for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
pos_embed = pos_embed.repeat(t, 1)
h_merge = h // merge_size
w_merge = w // merge_size
pos_embed = (
pos_embed.view(t, h_merge, merge_size, w_merge, merge_size, -1)
.permute(0, 1, 3, 2, 4, 5)
.flatten(0, 4)
)
result_parts.append(pos_embed)
return torch.cat(result_parts, dim=0)
def fast_pos_embed_interpolate(self, grid_thw):
"""Interpolate position embeddings for (batch, 3) size input dimensions.
Performs bilinear interpolation on spatial dimensions (height, width) and replicates
along temporal dimension. The result is reorganized according to spatial_merge_size.
Args:
grid_thw: Tensor of shape [batch_size, 3] with (temporal, height, width) dimensions
in patches for each sample.
Returns:
Interpolated position embeddings tensor.
"""
grid_thw_cpu = grid_thw.cpu().numpy()
# transfer data to CPU before loop
temporal_dims = grid_thw_cpu[:, 0].tolist()
height_dims = grid_thw_cpu[:, 1].tolist()
width_dims = grid_thw_cpu[:, 2].tolist()
device = self.pos_embed.weight.device
dtype = self.pos_embed.weight.dtype
patches_size = [h * w for h, w in zip(height_dims, width_dims)]
total_patches = sum(patches_size)
all_indices_np = np.zeros((4, total_patches), dtype=np.int64)
all_weights_np = np.zeros((4, total_patches), dtype=np.float32)
current_idx = 0
# calculate indices and weights on CPU
for t, h, w in zip(temporal_dims, height_dims, width_dims):
h_idxs = self._get_interpolation_indices(h)
w_idxs = self._get_interpolation_indices(w)
indices, weights = self._calculate_indices_and_weights(h_idxs, w_idxs)
end_idx = current_idx + h * w
for i in range(4):
all_indices_np[i, current_idx:end_idx] = indices[i]
all_weights_np[i, current_idx:end_idx] = weights[i]
current_idx = end_idx
idx_tensor = torch.from_numpy(all_indices_np).to(device)
weight_tensor = torch.from_numpy(all_weights_np).to(dtype=dtype, device=device)
# calculate interpolation
pos_embeds = self.pos_embed(idx_tensor.view(-1))
pos_embeds = pos_embeds.view(4, total_patches, -1)
patch_pos_embeds = (pos_embeds * weight_tensor.unsqueeze(-1)).sum(dim=0)
patch_pos_embeds = patch_pos_embeds.split(patches_size)
return self._get_position_embedding(
patch_pos_embeds, temporal_dims, height_dims, width_dims
)
for t, h, w in grid_thw:
pos_embed = torch.nn.functional.interpolate(
embeds, size=(h, w), mode="bilinear", align_corners=self.align_corners
)
pos_embed = pos_embed.reshape(
-1,
h // self.spatial_merge_size,
self.spatial_merge_size,
w // self.spatial_merge_size,
self.spatial_merge_size,
)
pos_embed = pos_embed.permute(1, 3, 2, 4, 0)
pos_embed = pos_embed.flatten(0, 3).repeat(t, 1)
patch_pos_embeds_permute.append(pos_embed)
return torch.cat(patch_pos_embeds_permute)
def forward(
self,