[None][fix] Restore Qwen3-VL vision-block RoPE shape after EXAONE-4.5 rebase

The qwen3vl_opt branch's ``1c4e86892d "Further optimize Qwen3-VL"``
removed ``rope_position_ids`` from the per-block call and added a
``rotary_pos_emb.flatten(1).repeat(1, 2).cos()/.sin()`` chain in the
vision tower forward. At authoring time that matched the HF
``apply_rotary_pos_emb_vision`` ABI (cos/sin shape
``(seq, head_dim)``), which ``Qwen2_5_VLVisionAttention.forward`` used.

After upstream ``7279d6322d "EXAONE-4.5 Support"`` refactored that
attention class to route through the new ``apply_rope`` helper
(``RotaryEmbedding.apply_rotary_pos_emb`` / FlashInfer), the expected
``cos.shape[-1]`` became ``head_dim // 2``: ``apply_rotary_pos_emb``
computes ``rot_dim = cos.shape[-1] * 2`` and chunks q/k into halves
of size ``cos.shape[-1]``. With Qwen3-VL's ``head_dim = 72``
(``72 % 64 != 0``, so the FlashInfer gate misses and the PyTorch
fallback runs), the prior ``.repeat(1, 2)`` made cos/sin
``(total, 72)`` and the elementwise multiply against
``q.chunk(2)[i]`` (shape ``(..., 36)``) raised
``RuntimeError: The size of tensor a (36) must match the size of
tensor b (72) at non-singleton dimension 2`` -- silently broken at
rebase time, missed by every unit test on this branch because they
verified isolated cos/sin construction (pos_ids, lru_cache identity,
gather equivalence) but never piped them through the actual attention
block. CI's model-level test would have caught it.

Fix:

* ``Qwen3VisionModel.forward``: drop ``cos.repeat(1, 2)`` /
  ``sin.repeat(1, 2)``. The cos/sin tuple returned by ``rot_pos_emb``
  is already at the post-EXAONE shape ``(total_tokens, 2 * freq_dim)
  = (total_tokens, head_dim // 2)``.

* ``Qwen3VisionModel.forward``: restore the pre-1c4e86892d
  ``rope_position_ids = torch.arange(seq_len, dtype=int32,
  pin_memory=prefer_pinned()).to(device, non_blocking=True)`` and
  pass ``position_ids=rope_position_ids`` into each block. This
  keeps the FlashInfer ``position_ids is not None`` gate satisfied
  for any future config where ``head_dim % 64 == 0``; for Qwen3-VL
  itself it still falls through to the PyTorch path, which now has
  the correct broadcast.

All prior optimizations on this branch are preserved: cos/sin
pre-computed init buffers (``rope_cos_cache`` / ``rope_sin_cache``),
L2 per-(t, h, w) cache, batched Triton pos-embed kernel,
rot_pos_ids lru_cache, async_tensor_h2d helper, argsort inverse
permutation, deepstack dict lookup. The fix is a layout correction
on the output side of ``rot_pos_emb``, not a revert of the
optimizations.

Test: ``test_qwen3vl_vision_block_forward_end_to_end`` builds one
``Qwen3VLVisionBlock`` with random weights and runs forward with the
post-EXAONE cos/sin shape and ``position_ids``. Without this fix it
raises the same broadcasting error; with it the block returns the
expected ``(seq_len, hidden)`` tensor. Any future regression that
re-introduces a stray ``.repeat(1, 2)`` on cos/sin (or drops
``position_ids``) will surface here, not only at CI time.

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Author: yechank-nvidia

tensorrt_llm/_torch/models/modeling_qwen3vl.py

@@ -1132,22 +1132,34 @@ def forward(
         grid_rows = grid_thw.tolist()
         # ``rot_pos_emb`` returns (cos, sin) gathered from the pre-computed
         # cos/sin buffers -- no ``.cos()``/``.sin()`` kernels in forward.
-        # Each half has shape (total_tokens, 2*freq_dim); ``.repeat(1, 2)``
-        # below expands to head_dim so the per-token layout matches the
-        # prior ``freqs.flatten(1).repeat(1, 2).cos()`` chain bit-for-bit.
+        # Each half has shape (total_tokens, 2*freq_dim) = (total_tokens,
+        # rotary_dim), which is what ``Qwen2_5_VLVisionAttention.apply_rope``
+        # / ``RotaryEmbedding.apply_rotary_pos_emb`` expects (it computes
+        # ``rot_dim = cos.shape[-1] * 2 = head_dim`` and chunks q/k into
+        # halves of size ``cos.shape[-1]``).
         cos, sin = self.rot_pos_emb(grid_rows)
         pos_embeds = self.fast_pos_embed_interpolate(grid_rows)
 
         hidden_states = self.patch_embed(pixel_values)
         hidden_states += pos_embeds
         hidden_states = hidden_states.flatten(1)
 
-        position_embeddings = (cos.repeat(1, 2), sin.repeat(1, 2))
+        # Vision RoPE backend (FlashInfer path) gates on ``position_ids is
+        # not None``; supply trivial 0..seq_len-1 positions on device so
+        # the gate clears when ``head_dim % 64 == 0``. For Qwen3-VL
+        # (head_dim=72) the gate misses and we fall through to the
+        # PyTorch path, which broadcasts cos/sin over the chunked q/k.
+        seq_len = hidden_states.shape[0]
+        rope_position_ids = torch.arange(seq_len, dtype=torch.int32, pin_memory=prefer_pinned()).to(
+            device=self.device, non_blocking=True
+        )
+        position_embeddings = (cos, sin)
 
         deepstack_feature_lists = []
         for layer_num, block in enumerate(self.blocks):
             hidden_states = block(
                 hidden_states,
+                position_ids=rope_position_ids,
                 attn_metadata=attn_metadata,
                 position_embeddings=position_embeddings,
             )

tests/unittest/_torch/modeling/test_modeling_qwen3vl.py

@@ -951,3 +951,104 @@ def test_argsort_matches_scatter_inverse(n):
     torch.testing.assert_close(
         identity, torch.arange(n, device="cuda", dtype=window_index.dtype), atol=0, rtol=0
     )
+
+
+# ---------------------------------------------------------------------------
+# End-to-end Qwen3VLVisionBlock forward path.
+#
+# The pre-EXAONE-4.5 ``apply_rotary_pos_emb_vision`` (HF) accepted
+# ``cos / sin`` of shape ``(seq, head_dim)``; the post-EXAONE-4.5
+# ``RotaryEmbedding.apply_rotary_pos_emb`` (TRT-LLM) expects
+# ``(seq, head_dim // 2)`` and broadcasts. This test pushes one block
+# through with the post-EXAONE shape so any future regression (e.g.
+# re-introducing a stray ``.repeat(1, 2)`` on cos/sin) surfaces here
+# rather than only at full-model CI.
+# ---------------------------------------------------------------------------
+
+
+def test_qwen3vl_vision_block_forward_end_to_end():
+    """One ``Qwen3VLVisionBlock.forward`` call with random weights and the
+    correct (post-EXAONE) ``cos/sin`` shape must not raise a broadcasting
+    error in ``apply_rope``."""
+    from transformers import Qwen3VLConfig
+
+    from tensorrt_llm._torch.attention_backend.utils import get_attention_backend
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLVisionBlock
+
+    hf_config = Qwen3VLConfig(
+        text_config={
+            "hidden_size": 256,
+            "num_attention_heads": 4,
+            "num_hidden_layers": 1,
+            "intermediate_size": 256,
+            "head_dim": 64,
+            "rope_scaling": {
+                "mrope_section": [8, 8, 8],
+                "mrope_interleaved": True,
+                "rope_type": "default",
+            },
+            "max_position_embeddings": 1024,
+            "rms_norm_eps": 1e-6,
+            "num_key_value_heads": 4,
+            "vocab_size": 1024,
+            "dtype": "bfloat16",
+            "rope_theta": 10000.0,
+            "tie_word_embeddings": False,
+            "model_type": "qwen3_vl_text",
+        },
+        vision_config={
+            "hidden_size": 1152,
+            "num_heads": 16,  # head_dim = 72
+            "depth": 1,
+            "intermediate_size": 4304,
+            "patch_size": 16,
+            "spatial_merge_size": 2,
+            "temporal_patch_size": 2,
+            "in_channels": 3,
+            "num_position_embeddings": 2304,
+            "out_hidden_size": 256,
+            "deepstack_visual_indexes": [],
+            "hidden_act": "gelu_pytorch_tanh",
+            "model_type": "qwen3_vl",
+            "initializer_range": 0.02,
+        },
+        architectures=["Qwen3VLForConditionalGeneration"],
+        image_token_id=151655,
+        video_token_id=151656,
+        vision_start_token_id=151652,
+        vision_end_token_id=151653,
+    )
+
+    mc = ModelConfig(pretrained_config=hf_config, attn_backend="TRTLLM")
+    block = Qwen3VLVisionBlock(mc, layer_idx=0).to("cuda", dtype=torch.bfloat16)
+
+    seq_len = 256
+    head_dim = 72
+    freq_dim = head_dim // 2  # 36 -- the (post-EXAONE) cos/sin last-dim
+
+    hidden = torch.randn(seq_len, 1152, device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn(seq_len, freq_dim, device="cuda", dtype=torch.float32)
+    sin = torch.randn(seq_len, freq_dim, device="cuda", dtype=torch.float32)
+    position_ids = torch.arange(seq_len, dtype=torch.int32, device="cuda")
+
+    metadata_cls = get_attention_backend("TRTLLM").Metadata
+    attn_metadata = metadata_cls(
+        max_num_requests=64,
+        max_num_tokens=seq_len,
+        kv_cache_manager=None,
+    )
+    attn_metadata.num_contexts = 1
+    attn_metadata.request_ids = [1]
+    attn_metadata.prompt_lens = [seq_len]
+    attn_metadata.seq_lens = torch.tensor([seq_len], dtype=torch.int)
+    attn_metadata.max_seq_len = seq_len
+    attn_metadata.prepare()
+
+    out = block(
+        hidden,
+        position_ids=position_ids,
+        position_embeddings=(cos, sin),
+        attn_metadata=attn_metadata,
+    )
+    assert out.shape == hidden.shape