[VLM] Replace torch.repeat_interleave with faster np.repeat for Qwen-VL series (#13736)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/models/qwen2_vl.py

@@ -44,6 +44,7 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInp
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen2 import Qwen2Model
+from sglang.srt.models.utils import compute_cu_seqlens_from_grid_numpy
 from sglang.srt.utils import add_prefix
 from sglang.srt.utils.hf_transformers_utils import get_processor
 
@@ -387,10 +388,7 @@ class Qwen2VisionTransformer(nn.Module):
         emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
         position_embeddings = (emb.cos(), emb.sin())
         # compute cu_seqlens
-        cu_seqlens = torch.repeat_interleave(
-            grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
-        ).cumsum(dim=0, dtype=torch.int32)
-        cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
+        cu_seqlens = compute_cu_seqlens_from_grid_numpy(grid_thw)
 
         # transformers
         x = x.unsqueeze(1)

python/sglang/srt/models/qwen3_vl.py

@@ -46,6 +46,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.utils import add_prefix
 from sglang.srt.utils.hf_transformers_utils import get_processor
 
@@ -434,15 +435,7 @@ class Qwen3VLMoeVisionModel(nn.Module):
         position_embeddings = (emb.cos(), emb.sin())
 
         # compute cu_seqlens
-        cu_seqlens = torch.repeat_interleave(
-            grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
-        ).cumsum(dim=0)
-        cu_seqlens = torch.cat(
-            [
-                torch.zeros(1, dtype=torch.int32, device=cu_seqlens.device),
-                cu_seqlens.to(torch.int32),
-            ]
-        )
+        cu_seqlens = compute_cu_seqlens_from_grid_numpy(grid_thw)
 
         x = x.unsqueeze(1)
 

python/sglang/srt/models/utils.py

@@ -12,6 +12,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import numpy as np
 import torch
 
 from sglang.srt.layers.radix_attention import RadixAttention
@@ -59,3 +60,25 @@ def permute_inv(perm: torch.Tensor) -> torch.Tensor:
     inv_perm = torch.empty_like(perm)
     inv_perm[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype)
     return inv_perm
+
+
+def compute_cu_seqlens_from_grid_numpy(grid_thw: torch.Tensor) -> torch.Tensor:
+    """
+    Compute cu_seqlens from grid_thw using NumPy.
+
+    grid_thw: [T, 3] int tensor on CPU.
+              columns: [repeat_count, H, W]
+    Returns:
+        cu_seqlens: 1D int32 tensor on CPU, shape [N + 1]
+    """
+    assert (
+        grid_thw.device.type == "cpu"
+    ), "compute_cu_seqlens_from_grid_numpy expects a CPU tensor"
+    arr = grid_thw.numpy()
+
+    cu_seqlens = np.repeat(arr[:, 1] * arr[:, 2], arr[:, 0]).cumsum(
+        axis=0, dtype=np.int32
+    )
+    cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
+    cu_seqlens = torch.from_numpy(cu_seqlens)
+    return cu_seqlens