From b2b09f5f24b9f5d9e9ad4c7503c84fabb6c293c2 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Thu, 4 Dec 2025 12:32:00 +0800 Subject: [PATCH] [VLM] Introduce Cache for positional embedding ids for Qwen-VL family (#14292) Co-authored-by: luoyuan.luo --- python/sglang/srt/models/qwen2_5_vl.py | 30 ++++---------------- python/sglang/srt/models/qwen3_vl.py | 25 +++-------------- python/sglang/srt/models/utils.py | 38 ++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 46 deletions(-) diff --git a/python/sglang/srt/models/qwen2_5_vl.py b/python/sglang/srt/models/qwen2_5_vl.py index fdc0c7650..32efa19d8 100644 --- a/python/sglang/srt/models/qwen2_5_vl.py +++ b/python/sglang/srt/models/qwen2_5_vl.py @@ -69,7 +69,7 @@ from sglang.srt.managers.schedule_batch import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.qwen2 import Qwen2Model -from sglang.srt.models.utils import permute_inv +from sglang.srt.models.utils import RotaryPosMixin, permute_inv from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix @@ -246,7 +246,7 @@ class Qwen2_5_VisionPatchMerger(nn.Module): return out -class Qwen2_5_VisionTransformer(nn.Module): +class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin): def __init__( self, @@ -362,30 +362,10 @@ class Qwen2_5_VisionTransformer(nn.Module): def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: pos_ids = [] - for i in range(grid_thw.size(0)): - t, h, w = grid_thw[i].tolist() - hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + for t, h, w in grid_thw: + base = self.rot_pos_ids(h, w, self.spatial_merge_size) + pos_ids.append(base if t == 1 else base.repeat(t, 1)) - hpos_ids = hpos_ids.reshape( - h // self.spatial_merge_size, - self.spatial_merge_size, - w // self.spatial_merge_size, - self.spatial_merge_size, - ) - hpos_ids = hpos_ids.permute(0, 2, 1, 3) - hpos_ids = hpos_ids.flatten() - - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) - wpos_ids = wpos_ids.reshape( - h // self.spatial_merge_size, - self.spatial_merge_size, - w // self.spatial_merge_size, - self.spatial_merge_size, - ) - wpos_ids = wpos_ids.permute(0, 2, 1, 3) - wpos_ids = wpos_ids.flatten() - - pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) pos_ids = torch.cat(pos_ids, dim=0) max_grid_size = grid_thw[:, 1:].max() rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 593bfceca..ed52f7ff4 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -50,7 +50,7 @@ from sglang.srt.managers.schedule_batch import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.qwen3 import Qwen3Model -from sglang.srt.models.utils import compute_cu_seqlens_from_grid_numpy +from sglang.srt.models.utils import RotaryPosMixin, compute_cu_seqlens_from_grid_numpy from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix @@ -257,7 +257,7 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module): return out -class Qwen3VLMoeVisionModel(nn.Module): +class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): def __init__( self, @@ -339,26 +339,9 @@ class Qwen3VLMoeVisionModel(nn.Module): def rot_pos_emb(self, grid_thw): pos_ids = [] for t, h, w in grid_thw: - hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) - hpos_ids = hpos_ids.reshape( - h // self.spatial_merge_size, - self.spatial_merge_size, - w // self.spatial_merge_size, - self.spatial_merge_size, - ) - hpos_ids = hpos_ids.permute(0, 2, 1, 3) - hpos_ids = hpos_ids.flatten() + base = self.rot_pos_ids(h, w, self.spatial_merge_size) + pos_ids.append(base if t == 1 else base.repeat(t, 1)) - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) - wpos_ids = wpos_ids.reshape( - h // self.spatial_merge_size, - self.spatial_merge_size, - w // self.spatial_merge_size, - self.spatial_merge_size, - ) - wpos_ids = wpos_ids.permute(0, 2, 1, 3) - wpos_ids = wpos_ids.flatten() - pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) pos_ids = torch.cat(pos_ids, dim=0) max_grid_size = grid_thw[:, 1:].max() rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 15c50e8a7..20854cc18 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -12,6 +12,8 @@ # limitations under the License. # ============================================================================== +from functools import lru_cache + import numpy as np import torch @@ -82,3 +84,39 @@ def compute_cu_seqlens_from_grid_numpy(grid_thw: torch.Tensor) -> torch.Tensor: cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens]) cu_seqlens = torch.from_numpy(cu_seqlens) return cu_seqlens + + +class RotaryPosMixin: + + @staticmethod + @lru_cache(maxsize=1024) + def rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor: + if isinstance(h, torch.Tensor): + h = int(h.item()) + if isinstance(w, torch.Tensor): + w = int(w.item()) + if isinstance(spatial_merge_size, torch.Tensor): + spatial_merge_size = int(spatial_merge_size.item()) + hpos_ids = np.broadcast_to(np.arange(h).reshape(h, 1), (h, w)) + h_div = h // spatial_merge_size + w_div = w // spatial_merge_size + hpos_ids = hpos_ids.reshape( + h_div, + spatial_merge_size, + w_div, + spatial_merge_size, + ) + hpos_ids = hpos_ids.transpose(0, 2, 1, 3) + hpos_ids = hpos_ids.flatten() + + wpos_ids = np.broadcast_to(np.arange(w).reshape(1, w), (h, w)) + wpos_ids = wpos_ids.reshape( + h_div, + spatial_merge_size, + w_div, + spatial_merge_size, + ) + wpos_ids = wpos_ids.transpose(0, 2, 1, 3) + wpos_ids = wpos_ids.flatten() + + return torch.from_numpy(np.stack([hpos_ids, wpos_ids], axis=-1))