perf(multimodal): optimize Qwen vision runtime

python/tokenspeed/runtime/models/qwen3_5.py

@@ -1234,23 +1234,22 @@ def pre_encode(
         video, ``image_grid_thw`` otherwise) so a single shared encoder cudagraph
         wrapper can serve both image and video batches.
         """
-        pixel_values = torch.cat([item.feature for item in items], dim=0).type(
-            self.visual.dtype
-        )
-        grid = torch.concat(
-            [
-                getattr(
-                    item,
-                    (
-                        "video_grid_thw"
-                        if item.modality == Modality.VIDEO
-                        else "image_grid_thw"
-                    ),
-                )
-                for item in items
-            ],
-            dim=0,
-        )
+        features = [item.feature for item in items]
+        grids = [
+            getattr(
+                item,
+                (
+                    "video_grid_thw"
+                    if item.modality == Modality.VIDEO
+                    else "image_grid_thw"
+                ),
+            )
+            for item in items
+        ]
+        pixel_values = (
+            features[0] if len(features) == 1 else torch.cat(features, dim=0)
+        ).type(self.visual.dtype)
+        grid = grids[0] if len(grids) == 1 else torch.concat(grids, dim=0)
         if pixel_values.dim() != 2:
             raise ValueError(f"pixel_values must be 2D, got {pixel_values.dim()}D.")
         if grid.dim() != 2:
@@ -1262,6 +1261,8 @@ def post_encode(
         self, encoder_outs: list[torch.Tensor], grid: torch.Tensor
     ) -> torch.Tensor:
         """Eager step after the captured region; returns features."""
+        if len(encoder_outs) == 1:
+            return encoder_outs[0]
         return torch.cat(encoder_outs, dim=0)
 
     def _build_encoder_cudagraph_wrapper(
@@ -1350,8 +1351,14 @@ def forward(
             text_embedding=self.model.get_input_embeddings(),
             ctx=multimodal_context,
             encoders={
-                Modality.IMAGE: EncoderSpec(self.image_encoder, deepstack=True),
-                Modality.VIDEO: EncoderSpec(self.video_encoder, deepstack=True),
+                Modality.IMAGE: EncoderSpec(
+                    self.image_encoder,
+                    deepstack=self.num_deepstack_embeddings > 0,
+                ),
+                Modality.VIDEO: EncoderSpec(
+                    self.video_encoder,
+                    deepstack=self.num_deepstack_embeddings > 0,
+                ),
             },
             multimodal_model=self,
             is_decode_or_idle=ctx.forward_mode.is_decode_or_idle(),

test/runtime/test_qwen35_vision_assembly.py

@@ -0,0 +1,69 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tokenspeed.runtime.models.qwen3_5 import Qwen3_5ForConditionalGeneration
+from tokenspeed.runtime.multimodal.inputs import Modality
+
+
+class _Visual:
+    dtype = torch.float32
+
+    def prepare_patch_embed(self, pixel_values, grid):
+        self.pixel_values = pixel_values
+        self.grid = grid
+        return pixel_values
+
+
+@pytest.mark.parametrize(
+    ("modality", "grid_field"),
+    [
+        (Modality.IMAGE, "image_grid_thw"),
+        (Modality.VIDEO, "video_grid_thw"),
+    ],
+)
+def test_single_item_pre_encode_reuses_feature_and_grid(modality, grid_field):
+    visual = _Visual()
+    model = SimpleNamespace(visual=visual)
+    feature = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    grid = torch.tensor([[1, 2, 2]], dtype=torch.int64)
+    item = SimpleNamespace(feature=feature, modality=modality, **{grid_field: grid})
+
+    tokens, result_grid = Qwen3_5ForConditionalGeneration.pre_encode(model, [item])
+
+    assert tokens.data_ptr() == feature.data_ptr()
+    assert visual.pixel_values.data_ptr() == feature.data_ptr()
+    assert result_grid.data_ptr() == grid.data_ptr()
+    assert visual.grid.data_ptr() == grid.data_ptr()
+
+
+def test_multi_item_pre_encode_preserves_concatenation():
+    visual = _Visual()
+    model = SimpleNamespace(visual=visual)
+    first = SimpleNamespace(
+        feature=torch.ones((2, 4), dtype=torch.float32),
+        modality=Modality.IMAGE,
+        image_grid_thw=torch.tensor([[1, 2, 2]], dtype=torch.int64),
+    )
+    second = SimpleNamespace(
+        feature=torch.full((3, 4), 2.0, dtype=torch.float32),
+        modality=Modality.IMAGE,
+        image_grid_thw=torch.tensor([[1, 3, 2]], dtype=torch.int64),
+    )
+
+    tokens, grid = Qwen3_5ForConditionalGeneration.pre_encode(model, [first, second])
+
+    assert torch.equal(tokens, torch.cat([first.feature, second.feature], dim=0))
+    assert torch.equal(
+        grid,
+        torch.cat([first.image_grid_thw, second.image_grid_thw], dim=0),
+    )
+
+
+def test_single_encoder_output_is_returned_without_concatenation():
+    output = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+
+    result = Qwen3_5ForConditionalGeneration.post_encode(None, [output], None)
+
+    assert result is output