A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 22:15 UTC
[Model][QwenVL] Replace `torch.repeat_interleave` with faster `np.repeat` (#28964)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
Lukas Geiger committed
Nov 20, 2025 at 06:04 UTC
a9705a290af05ad71023714074ad8bf1a50c60a3
3 files changed
+18
-23
tests/models/multimodal/generation/test_qwen2_vl.py
+2
-12
@@ -128,12 +128,7 @@ def batch_make_image_embeddings(
128
visual = model.visual
129
130
pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype)
131
- image_grid_thw_on_device = image_grid_thw.to(
132
- visual.device, dtype=torch.int64
133
- )
134
- return visual(
135
- pixel_values_on_device, grid_thw=image_grid_thw_on_device
136
- ).cpu()
131
+ return visual(pixel_values_on_device, grid_thw=image_grid_thw).cpu()
132
133
image_embeds = torch.concat(llm.apply_model(get_image_embeds))
134
@@ -217,12 +212,7 @@ def batch_make_video_embeddings(
212
visual = model.visual
213
214
pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype)
220
- video_grid_thw_on_device = video_grid_thw.to(
221
- visual.device, dtype=torch.int64
222
- )
223
- return visual(
224
- pixel_values_on_device, grid_thw=video_grid_thw_on_device
225
- ).cpu()
215
+ return visual(pixel_values_on_device, grid_thw=video_grid_thw).cpu()
216
217
video_embeds = torch.concat(llm.apply_model(get_image_embeds))
218
vllm/model_executor/models/qwen2_vl.py
+9
-6
@@ -29,6 +29,7 @@ from collections.abc import Callable, Iterable, Mapping, Sequence
29
from functools import partial
30
from typing import Annotated, Any, Literal, TypeAlias
31
32
+import numpy as np
33
import torch
34
import torch.nn as nn
35
import torch.nn.functional as F
@@ -751,25 +752,27 @@ class Qwen2VisionTransformer(nn.Module):
752
753
if isinstance(grid_thw, list):
754
grid_thw_list = grid_thw
754
- grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
755
+ grid_thw = np.array(grid_thw, dtype=np.int32)
756
else:
757
grid_thw_list = grid_thw.tolist()
758
+ grid_thw = grid_thw.numpy()
759
760
# compute position embedding
761
rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
762
763
# compute cu_seqlens
762
- cu_seqlens = torch.repeat_interleave(
763
- grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
764
- ).cumsum(dim=0, dtype=torch.int32)
765
- cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
766
- cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
764
+ cu_seqlens = np.repeat(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
765
+ axis=0, dtype=np.int32
766
+ )
767
+ cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
768
+ cu_seqlens = torch.from_numpy(cu_seqlens)
769
770
# transformers
771
x = x.unsqueeze(1)
772
773
# pre-compute seqlens for attn mask to reduce cuMemcpy operations
774
max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
775
+ cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
776
for blk in self.blocks:
777
x = blk(
778
x,
vllm/model_executor/models/qwen3_vl.py
+7
-5
@@ -553,18 +553,20 @@ class Qwen3_VisionTransformer(nn.Module):
553
554
if isinstance(grid_thw, list):
555
grid_thw_list = grid_thw
556
- grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
556
+ grid_thw = np.array(grid_thw, dtype=np.int32)
557
else:
558
grid_thw_list = grid_thw.tolist()
559
+ grid_thw = grid_thw.numpy()
560
561
pos_embeds = self.fast_pos_embed_interpolate(grid_thw_list)
562
hidden_states = hidden_states + pos_embeds
563
rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
564
564
- cu_seqlens = torch.repeat_interleave(
565
- grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
566
- ).cumsum(dim=0, dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32)
567
- cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
565
+ cu_seqlens = np.repeat(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
566
+ axis=0, dtype=np.int32
567
+ )
568
+ cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
569
+ cu_seqlens = torch.from_numpy(cu_seqlens)
570
571
hidden_states = hidden_states.unsqueeze(1)
572
max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)