diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 89de28da6..57c6a3973 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -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, diff --git a/test/srt/test_embed_interpolate_unittest.py b/test/srt/test_embed_interpolate_unittest.py index cb09935bc..d18f71091 100644 --- a/test/srt/test_embed_interpolate_unittest.py +++ b/test/srt/test_embed_interpolate_unittest.py @@ -71,9 +71,10 @@ class TestEmbedInterpolate(unittest.TestCase): norm_eps=1e-6, prefix="visual", ) - embeddings = model.fast_pos_embed_interpolate( - [(t, s, s) for t, s in zip(t_dim, s_dim)] + grid_thw = torch.tensor( + [(t, s, s) for t, s in zip(t_dim, s_dim)], dtype=torch.int32 ) + embeddings = model.fast_pos_embed_interpolate(grid_thw) embeddings_s0 = embeddings[: s_dim[0] * s_dim[0], :] embeddings_s1 = embeddings[s_dim[0] * s_dim[0] : 2 * s_dim[0] * s_dim[0], :]