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

[VLM] Introduce Cache for positional embedding ids for Qwen-VL family (#14292)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2025-12-04 12:32:00 +08:00
committed by GitHub
parent 04df80a9a1
commit b2b09f5f24
3 changed files with 47 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
python/sglang/srt/models/qwen2_5_vl.py
View File

@@ -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)

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

@@ -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)

38
python/sglang/srt/models/utils.py
View File

@@ -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))