[Model][QwenVL] Replace torch.repeat_interleave with faster np.repeat
#28964
Conversation
`np.repeat` Signed-off-by: Lukas Geiger <[email protected]>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| cu_seqlens = np.repeat(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 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) |
There was a problem hiding this comment.
The new cu_seqlens path now runs entirely in numpy (np.repeat, np.concatenate, torch.from_numpy), so when this forward is traced/exported the tensor is computed on the concrete sample input and baked into the graph as a constant. Subsequent executions of the traced model with different grid_thw values will therefore reuse the wrong sequence layout, breaking attention masking for any shape that differs from the trace input. The previous Torch-only code (with a torch.jit.is_tracing() dtype guard) kept this computation traceable and input-dependent.
Useful? React with 👍 / 👎.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request replaces torch.repeat_interleave with np.repeat to improve performance in qwen2_vl.py and qwen3_vl.py. The changes are logical and correctly implemented. By performing the cu_seqlens computation on the CPU with NumPy and then moving the tensor to the GPU, the code avoids unnecessary device transfers and leverages faster CPU operations, which should result in the claimed performance gain. The code modifications are clean and I have no concerns. Good work.
github-actions
bot
added
the
ready
2 participants
Purpose
torch.repeat_interleaveis surprisingly slow so this PR replaces it with numpy ops. This slightly speeds up TTFT by ~1.5%.Part of #23884
Before:
After: