Skip to content

Commit 124f649

Browse files
[None][refactor] Drop batched pos-embed kernel and prune vLLM references
Stripped layer of commit 7368e5b after measurement showed the batched single-launch kernel did not move the conc=64 throughput needle and added non-trivial setup cost (5 small H->D copies for per-image metadata) on the low-concurrency path. * Vision pos-embed interpolation goes back to the per-image Triton kernel ``_triton_pos_embed_interpolate`` + ``torch.cat`` join in ``Qwen3VisionModel.fast_pos_embed_interpolate``. The kernel takes ``(h, w, h_scale, w_scale)`` as scalar args, so no H<->D metadata transfers happen on the hot path. The batched ``_bilinear_pos_embed_batched_kernel`` / ``_triton_pos_embed_interpolate_batched`` pair is removed. * Vision rotary cache no longer hard-codes ``max_rope_seqlen = 8192``; ``Qwen3VisionModel.__init__`` now constructs a standard ``RotaryEmbedding`` whose ``rotary_cos_sin`` buffer is sized to ``text_config.max_position_embeddings`` via ``RopeParams.from_config``. ``_rotary_pos_emb_thw`` slices ``self.rotary_pos_emb.rotary_cos_sin [:max(h, w)]`` then indexes with ``pos_ids``, mirroring upstream/main. * Removed vLLM cross-reference comments in modeling_qwen3vl.py, modeling_qwen2vl.py, and modules/rotary_embedding.py while keeping the local technical descriptions intact. * Updated unit tests to mock ``rotary_pos_emb.rotary_cos_sin`` (the new buffer location) and dropped the batched-vs-per-image comparison test together with the batched kernel. 3-run aiperf sweep (conc=64, ISL=1000, OSL=100, warmup=20, single 512x512 image, Qwen3-VL-8B): throughput 18.01 +/- 0.94 req/s vs main 17.25 +/- 0.72 (+4.4%), TTFT 1082 +/- 107 ms vs 1240 +/- 137 (-12.7%), latency 3528 +/- 178 ms vs 3692 +/- 149 (-4.4%). Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
1 parent 7368e5b commit 124f649

4 files changed

Lines changed: 99 additions & 365 deletions

File tree

tensorrt_llm/_torch/models/modeling_qwen2vl.py

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -773,12 +773,11 @@ def forward(
773773
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
774774
**kwargs,
775775
) -> torch.Tensor:
776-
# vLLM-style residual fusion: collapse the post-attn residual add
777-
# into ``norm2``'s residual path (``RMSNorm.forward`` with ``residual=``
778-
# uses the FlashInfer ``fused_add_rmsnorm`` kernel). Saves one
779-
# ``add<c10::BFloat16>`` launch per vision block (= 32 launches per
780-
# executor iter at full-batch). Mirrors
781-
# ``vllm/model_executor/models/qwen2_5_vl.py::Qwen2_5_VisionBlock``.
776+
# Collapse the post-attn residual add into ``norm2``'s residual path
777+
# (``RMSNorm.forward`` with ``residual=`` uses the FlashInfer
778+
# ``fused_add_rmsnorm`` kernel). Saves one ``add<c10::BFloat16>``
779+
# launch per vision block (= 32 launches per executor iter at
780+
# full-batch).
782781
x_attn = self.attn(
783782
hidden_states=self.norm1(hidden_states),
784783
attn_metadata=attn_metadata,
@@ -951,9 +950,9 @@ def get_rope_and_window_index_by_thw(
951950
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Tuple[int, ...]]:
952951
"""GPU (cos_thw, sin_thw); CPU (window_index_thw, seq_lens_thw).
953952
954-
Cached per ``(t, h, w)``. Mirrors vLLM's ``get_rope_by_thw`` pattern:
955-
pos_ids and window_index are built on CPU, then moved to device once
956-
per unique tile. The gather (``cos[pos_ids]``) and window reorder
953+
Cached per ``(t, h, w)``. pos_ids and window_index are built on CPU,
954+
then moved to device once per unique tile. The gather
955+
(``cos[pos_ids]``) and window reorder
957956
(``cos_thw[window_index_thw_dev]``) run on device, so the cached
958957
cos/sin tensors are already-on-device -- ``forward``'s ``torch.cat``
959958
on them is a device-side cat with no H->D transfer. ``window_index``
@@ -978,13 +977,6 @@ def get_rope_and_window_index_by_thw(
978977
bytes fp32). Typical production VLM serving has 10-30 unique tile
979978
shapes, so the cache settles around 10-20 MB. ``maxsize=1024`` is
980979
a safety cap; reaching it requires >1024 distinct (t, h, w).
981-
982-
Note: vLLM's equivalent ``get_rope_by_thw`` casts the rotary
983-
cos_sin_cache to the model dtype (bf16) via ``get_rope(dtype=torch.
984-
get_default_dtype())`` -- so the equivalent vLLM cache is half the
985-
size (320 B/token, ~5-10 MB production). ``RopeParams.create_rope_
986-
const_params`` hardcodes fp32; switching to model dtype would halve
987-
this footprint but is a separate change.
988980
"""
989981
hpos_ids = torch.arange(h, dtype=torch.long).unsqueeze(1).expand(-1, w)
990982
wpos_ids = torch.arange(w, dtype=torch.long).unsqueeze(0).expand(h, -1)
@@ -1031,9 +1023,8 @@ def get_rope_and_window_index_by_thw(
10311023
-1, sin_thw.shape[-1])
10321024

10331025
# Cast once (per cached (t, h, w)) to the vision-tower dtype so the
1034-
# per-block ``cos.to(dtype=q.dtype)`` cast in ``apply_rope`` becomes a
1035-
# no-op. Mirrors vLLM's ``get_rope_by_thw`` which builds the rotary
1036-
# cache directly in the model dtype (see note above).
1026+
# per-block ``cos.to(dtype=q.dtype)`` cast in ``apply_rope`` becomes
1027+
# a no-op on the hot path.
10371028
target_dtype = (
10381029
self.model_config.pretrained_config.vision_config.torch_dtype
10391030
if hasattr(self.model_config.pretrained_config.vision_config,

0 commit comments

Comments
 (0)