[Model] Add vision encoder and input embeddings merger warmup for Qwen2.5 VL model (vllm-project#972)

Signed-off-by: Kewei Wang <[email protected]>

Author: kwang3939

tests/models/jax/test_qwen2_5_vl.py

@@ -16,12 +16,11 @@
 # Import the module itself to allow patching
 # Corrected imports for the code under test
 from tpu_inference.models.jax.qwen2_5_vl import (
-    AttentionMetadata, MultiModalEmbeddings, Qwen2_5_VisionAttention,
-    Qwen2_5_VisionBlock, Qwen2_5_VisionMLP, Qwen2_5_VisionPatchEmbed,
-    Qwen2_5_VisionPatchMerger, Qwen2_5_VisionRotaryEmbedding,
-    Qwen2_5_VisionTransformer, Qwen2_5_VLForConditionalGeneration,
-    Qwen2_5_VLImagePixelInputs, SegmentIds, apply_rotary_pos_emb_vision,
-    generate_window_segment_ids)
+    AttentionMetadata, Qwen2_5_VisionAttention, Qwen2_5_VisionBlock,
+    Qwen2_5_VisionMLP, Qwen2_5_VisionPatchEmbed, Qwen2_5_VisionPatchMerger,
+    Qwen2_5_VisionRotaryEmbedding, Qwen2_5_VisionTransformer,
+    Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLImagePixelInputs, SegmentIds,
+    apply_rotary_pos_emb_vision, generate_window_segment_ids)
 
 
 # --- Configuration Mocking ---
@@ -508,12 +507,12 @@ def test_get_input_embeddings(self, mock_merge_embeddings: MagicMock,
         np.testing.assert_array_equal(embeds, mock_text_embeds)
         mock_merge_embeddings.assert_not_called()
 
-        embeds_empty_mm = model.get_input_embeddings(input_ids, tuple())
+        empty_mm = jnp.ones((0, model.config.hidden_size), )
+        embeds_empty_mm = model.get_input_embeddings(input_ids, empty_mm)
         np.testing.assert_array_equal(embeds_empty_mm, mock_text_embeds)
         mock_merge_embeddings.assert_not_called()
 
-        mm_embeds: MultiModalEmbeddings = (jnp.ones(
-            (5, model.config.hidden_size)), )
+        mm_embeds = jnp.ones((5, model.config.hidden_size))
         mock_merged = jnp.ones((1, 15, model.config.hidden_size))
         mock_merge_embeddings.return_value = mock_merged
 

tests/models/jax/utils/test_multi_modal_utils.py

@@ -4,7 +4,7 @@
 import pytest
 
 from tpu_inference.models.jax.utils.multi_modal_utils import (
-    MultiModalEmbeddings, NestedTensors, _flatten_embeddings,
+    MultiModalEmbeddings, NestedTensors, flatten_embeddings,
     merge_multimodal_embeddings, sanity_check_mm_encoder_outputs)
 
 # --- Tests for sanity_check_mm_encoder_outputs ---
@@ -65,45 +65,45 @@ def test_sanity_check_wrong_dimensions_in_list():
         sanity_check_mm_encoder_outputs(embeddings, 1)
 
 
-# --- Tests for _flatten_embeddings ---
+# --- Tests for flatten_embeddings ---
 
 
 def test_flatten_single_array():
-    """Tests _flatten_embeddings with a single 2D array."""
+    """Tests flatten_embeddings with a single 2D array."""
     emb: NestedTensors = jnp.arange(12).reshape((3, 4))
-    result = _flatten_embeddings(emb)
+    result = flatten_embeddings(emb)
     np.testing.assert_array_equal(result, emb)
 
 
 def test_flatten_single_3d_array():
-    """Tests _flatten_embeddings with a single 3D array."""
+    """Tests flatten_embeddings with a single 3D array."""
     emb: NestedTensors = jnp.arange(24).reshape((2, 3, 4))
-    result = _flatten_embeddings(emb)
+    result = flatten_embeddings(emb)
     expected = jnp.arange(24).reshape((6, 4))
     np.testing.assert_array_equal(result, expected)
 
 
 def test_flatten_list_of_arrays():
-    """Tests _flatten_embeddings with a list of 2D arrays."""
+    """Tests flatten_embeddings with a list of 2D arrays."""
     emb: NestedTensors = [
         jnp.arange(12).reshape((3, 4)),
         jnp.arange(12, 20).reshape((2, 4))
     ]
-    result = _flatten_embeddings(emb)
+    result = flatten_embeddings(emb)
     expected = jnp.arange(20).reshape((5, 4))
     np.testing.assert_array_equal(result, expected)
 
 
 def test_flatten_nested_list():
-    """Tests _flatten_embeddings with a nested list of arrays."""
+    """Tests flatten_embeddings with a nested list of arrays."""
     emb: NestedTensors = [
         jnp.arange(6).reshape((2, 3)),
         [
             jnp.arange(6, 12).reshape((2, 3)),
             jnp.arange(12, 15).reshape((1, 3))
         ]
     ]
-    result = _flatten_embeddings(emb)
+    result = flatten_embeddings(emb)
     expected = jnp.arange(15).reshape((5, 3))
     np.testing.assert_array_equal(result, expected)
 
@@ -191,7 +191,7 @@ def test_merge_mm_embeds_count_too_many_no_raise(placeholder_id, base_embeds):
         # Check that the first 2 embeddings from mm_embeds_too_many were used.
         expected = np.array(inputs_embeds)
         is_mm = np.isin(input_ids, np.array(placeholder_id))
-        expected[is_mm] = _flatten_embeddings(mm_embeds_too_many)[:2]
+        expected[is_mm] = flatten_embeddings(mm_embeds_too_many)[:2]
         np.testing.assert_array_equal(result, expected)
     except Exception as e:
         pytest.fail(

tpu_inference/models/common/model_loader.py

@@ -260,6 +260,9 @@ def combine_hidden_states(graphdef, state, hidden_states):
         model = nnx.merge(graphdef, state)
         return model.combine_hidden_states(hidden_states)
 
+    model = nnx.merge(graphdef, state)
+    precompile_vision_encoder_fn = getattr(model, "precompile_vision_encoder",
+                                           None)
     model_fn = functools.partial(run_model, graphdef)
     compute_logits_fn = functools.partial(run_compute_logits, graphdef)
     get_multimodal_embeddings_fn = functools.partial(
@@ -274,7 +277,14 @@ def combine_hidden_states(graphdef, state, hidden_states):
         jit_model,
         "get_mrope_input_positions") else jit_model.get_mrope_input_positions
 
-    return model_fn, compute_logits_fn, combine_hidden_states_fn, get_multimodal_embeddings_fn, get_input_embeddings_fn, get_mrope_input_positions_fn, state, lora_manager, model
+    multimodal_fns = {
+        "precompile_vision_encoder_fn": precompile_vision_encoder_fn,
+        "get_multimodal_embeddings_fn": get_multimodal_embeddings_fn,
+        "get_input_embeddings_fn": get_input_embeddings_fn,
+        "get_mrope_input_positions_fn": get_mrope_input_positions_fn,
+    }
+
+    return model_fn, compute_logits_fn, combine_hidden_states_fn, multimodal_fns, state, lora_manager, model
 
 
 def get_vllm_model(
@@ -295,7 +305,7 @@ def get_vllm_model(
     compute_logits_fn = model.jit_compute_logits_func()
     # the model needs to be returned because lora weights are neither torch.nn.parameter nor torch.nn.buffer. After we load the lora weights and set it to the torch.nn.Module, we can shard it and move it to TPU.
     combine_hidden_states_fn = None
-    return jit_model, compute_logits_fn, combine_hidden_states_fn, None, None, None, params, lora_manager, model
+    return jit_model, compute_logits_fn, combine_hidden_states_fn, None, params, lora_manager, model
 
 
 def get_model(

tpu_inference/models/jax/qwen2_5_vl.py

@@ -959,14 +959,13 @@ def get_multimodal_embeddings(self, image_grid_thw: tuple[tuple[int, int,
 
     def get_input_embeddings(
             self, input_ids: jax.Array,
-            multimodal_embeddings: Optional[MultiModalEmbeddings]
-    ) -> jax.Array:
+            multimodal_embeddings: Optional[jax.Array]) -> jax.Array:
 
         inputs_embeds = self.language_model.model.embed(input_ids)
 
 
         if multimodal_embeddings is not None \
-            and len(multimodal_embeddings) != 0:
+            and multimodal_embeddings.shape[0] != 0:
             inputs_embeds = merge_multimodal_embeddings(
                 input_ids, inputs_embeds, multimodal_embeddings,
                 [self.config.image_token_id, self.config.video_token_id])
@@ -1067,3 +1066,35 @@ def load_weights(self, rng_key: jax.Array) -> None:
                         model=self,
                         metadata_map=metadata_map,
                         mesh=self.mesh)
+
+    def precompile_vision_encoder(
+        self,
+        run_compilation_fn: Callable,
+    ) -> None:
+        image_shapes = []
+        if (warmup_config := self.vllm_config.additional_config.get(
+                "vision_warmup_config")):
+            image_shapes = warmup_config.get("image_shapes")
+
+        vc = self.vllm_config.model_config.hf_config.vision_config
+        for input_hw in image_shapes:
+            if not isinstance(input_hw, list) or len(input_hw) != 2:
+                logger.warning(f"Skipping invalid shape {input_hw}.")
+                continue
+            h_input, w_input = input_hw
+            t, h, w = 1, h_input // vc.patch_size, w_input // vc.patch_size
+            grid_thw = (t, h, w)
+            num_patches = t * h * w
+            patch_input_dim = vc.in_channels * vc.temporal_patch_size * vc.patch_size * vc.patch_size
+
+            dummy_pixel_values = jnp.ones(
+                (num_patches, patch_input_dim),
+                self.vllm_config.model_config.dtype,
+            )
+            dummy_grid_thw = grid_thw
+
+            run_compilation_fn("single_image_encoder",
+                               self.get_single_image_embedding,
+                               dummy_pixel_values,
+                               dummy_grid_thw,
+                               image_shape=input_hw)

tpu_inference/models/jax/utils/multi_modal_utils.py

@@ -50,7 +50,7 @@ def sanity_check_mm_encoder_outputs(
         "of the model's `get_multimodal_embeddings` method.")
 
 
-def _flatten_embeddings(embeddings: NestedTensors) -> jax.Array:
+def flatten_embeddings(embeddings: NestedTensors) -> jax.Array:
     """
     Recursively flattens and concatenates NestedTensors on all but the last
     dimension.
@@ -59,8 +59,7 @@ def _flatten_embeddings(embeddings: NestedTensors) -> jax.Array:
     if isinstance(embeddings, jax.Array):
         return embeddings.reshape(-1, embeddings.shape[-1])
 
-    return jnp.concatenate([_flatten_embeddings(t) for t in embeddings],
-                           axis=0)
+    return jnp.concatenate([flatten_embeddings(t) for t in embeddings], axis=0)
 
 
 def _embedding_count_expression(embeddings: NestedTensors) -> str:
@@ -79,7 +78,7 @@ def _embedding_count_expression(embeddings: NestedTensors) -> str:
 def _merge_multimodal_embeddings(
     inputs_embeds: jax.Array,
     is_multimodal: jax.Array,
-    multimodal_embeddings: NestedTensors,
+    multimodal_embeddings: jax.Array,
 ) -> jax.Array:
     """
     Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
@@ -89,7 +88,6 @@ def _merge_multimodal_embeddings(
     Note:
         This returns a new array with the updated values.
     """
-    flattened = _flatten_embeddings(multimodal_embeddings)
     # The check for matching number of tokens is removed as it is not
     # JIT-compatible. If the shapes mismatch, JAX will raise an error
     # during execution anyway. The user-friendly error message is
@@ -99,10 +97,11 @@ def _merge_multimodal_embeddings(
     # NonConcreteBooleanIndexError.
     # Create a dummy row to handle indices for non-multimodal tokens.
     # The content of the dummy row does not matter as it will be masked out.
-    dummy_row = jnp.zeros_like(flattened[0:1])
+    dummy_row = jnp.zeros_like(multimodal_embeddings[0:1])
 
     # Prepend the dummy row to the flattened embeddings.
-    flattened_padded = jnp.concatenate([dummy_row, flattened], axis=0)
+    flattened_padded = jnp.concatenate([dummy_row, multimodal_embeddings],
+                                       axis=0)
 
     # Create gather indices. For each token in the input sequence, this gives
     # the index into `flattened_padded`.
@@ -121,7 +120,7 @@ def _merge_multimodal_embeddings(
 def merge_multimodal_embeddings(
     input_ids: jax.Array,
     inputs_embeds: jax.Array,
-    multimodal_embeddings: NestedTensors,
+    multimodal_embeddings: jax.Array,
     placeholder_token_id: Union[int, list[int]],
 ) -> jax.Array:
     """

tpu_inference/runner/compilation_manager.py

@@ -74,6 +74,9 @@ def capture_model(self) -> None:
         with self.runner.maybe_setup_dummy_loras(self.runner.lora_config):
             self._precompile_backbone_text_only()
             if self.runner.is_multimodal_model:
+                self.runner.precompile_vision_encoder_fn(
+                    self._run_compilation, )
+                self._precompile_input_embeddings_merger()
                 self._precompile_backbone_with_inputs_embeds()
             if self.runner.scheduler_config.async_scheduling:
                 self._precompile_substitute_placeholder_token()
@@ -86,6 +89,36 @@ def capture_model(self) -> None:
             if self.runner.speculative_config:
                 self._precompile_speculative_decoding()
 
+    def _precompile_input_embeddings_merger(self) -> None:
+        for num_tokens in self.runner.num_tokens_paddings:
+            hidden_size = self.runner.vllm_config.model_config.get_hidden_size(
+            )
+            sharding = NamedSharding(self.runner.mesh, PartitionSpec())
+            dummy_multimodal_embeddings = self._create_dummy_tensor(
+                (num_tokens, hidden_size),
+                self.runner.vllm_config.model_config.dtype,
+                sharding=sharding)
+            dummy_input_ids = self._create_dummy_tensor((num_tokens, ),
+                                                        jnp.int32)
+
+            self._run_compilation(
+                "input_embeddings_merger",
+                self.runner.get_input_embeddings_fn,
+                self.runner.state,
+                dummy_input_ids,
+                dummy_multimodal_embeddings,
+                num_tokens=num_tokens,
+            )
+
+            self._run_compilation(
+                "input_embeddings_merger_text_only",
+                self.runner.get_input_embeddings_fn,
+                self.runner.state,
+                dummy_input_ids,
+                None,
+                num_tokens=num_tokens,
+            )
+
     def _precompile_backbone_helper(self, name, *, input_ids, positions,
                                     inputs_embeds) -> None:
         num_tokens = None

tpu_inference/runner/multimodal_manager.py

@@ -9,8 +9,8 @@
 from vllm.v1.worker.utils import (gather_mm_placeholders,
                                   scatter_mm_placeholders)
 
-from tpu_inference.models.jax.utils.multi_modal_utils import \
-    sanity_check_mm_encoder_outputs
+from tpu_inference.models.jax.utils.multi_modal_utils import (
+    flatten_embeddings, sanity_check_mm_encoder_outputs)
 
 if TYPE_CHECKING:
     from tpu_inference.runner.tpu_jax_runner import TPUModelRunner
@@ -158,10 +158,8 @@ def execute_mm_encoder(self, scheduler_output: "VllmSchedulerOutput"):
                 is_embed=pos_info.is_embed,
             )
 
-    def gather_mm_embeddings(
-        self,
-        scheduler_output: "VllmSchedulerOutput",
-    ) -> list[jax.Array]:
+    def gather_mm_embeddings(self, scheduler_output: "VllmSchedulerOutput",
+                             target_pad_len: int) -> list[jax.Array]:
         mm_embeds: list[jax.Array] = []
         for req_id in self.runner.input_batch.req_ids:
             num_scheduled_tokens = scheduler_output.num_scheduled_tokens[
@@ -205,4 +203,15 @@ def gather_mm_embeddings(
                     is_embed=is_embed,
                 )
                 mm_embeds.append(mm_embeds_item)
-        return mm_embeds
+        if not mm_embeds:
+            return None
+        flattened_embeds = flatten_embeddings(mm_embeds)
+        if flattened_embeds.shape[0] == 0:
+            return None
+
+        padding = jnp.zeros((target_pad_len - flattened_embeds.shape[0],
+                             flattened_embeds.shape[1]),
+                            dtype=flattened_embeds.dtype)
+        flattened_embeds = jnp.concatenate([flattened_embeds, padding], axis=0)
+
+        return flattened_embeds

tpu_inference/runner/tpu_jax_runner.py

@@ -381,12 +381,22 @@ def _init_inputs(self) -> None:
                                             dtype=np.int64)
 
     def load_model(self):
-        self.model_fn, self.compute_logits_fn, self.combine_hidden_states_fn, self.get_multimodal_embeddings_fn, self.get_input_embeddings_fn, self.get_mrope_input_positions_fn, self.state, self.lora_manager, self.model = get_model(
+        self.model_fn, self.compute_logits_fn, self.combine_hidden_states_fn, multimodal_fns, self.state, self.lora_manager, self.model = get_model(
             self.vllm_config,
             self.rng_key,
             self.mesh,
         )
 
+        multimodal_fns = multimodal_fns or {}
+        self.precompile_vision_encoder_fn = multimodal_fns.get(
+            "precompile_vision_encoder_fn", None)
+        self.get_multimodal_embeddings_fn = multimodal_fns.get(
+            "get_multimodal_embeddings_fn", None)
+        self.get_input_embeddings_fn = multimodal_fns.get(
+            "get_input_embeddings_fn", None)
+        self.get_mrope_input_positions_fn = multimodal_fns.get(
+            "get_mrope_input_positions_fn", None)
+
         if self.drafter is not None:
             logger.info("Loading drafter model...")
             self.drafter.load_model(self.state)
@@ -529,7 +539,8 @@ def _execute_model(
             # Run the multimodal encoder if any.
             # We have the modality embeds at this time.
             self.mm_manager.execute_mm_encoder(scheduler_output)
-            mm_embeds = self.mm_manager.gather_mm_embeddings(scheduler_output)
+            mm_embeds = self.mm_manager.gather_mm_embeddings(
+                scheduler_output, input_ids.shape[0])
         else:
             mm_embeds = []
 
@@ -970,8 +981,8 @@ def _get_input_ids_embeds(self, input_ids: jax.Array,
         if self.is_multimodal_model:
             inputs_embeds = self.get_input_embeddings_fn(
                 self.state,
-                input_ids=input_ids,
-                multimodal_embeddings=mm_embeds,
+                input_ids,
+                mm_embeds,
             )
             return None, inputs_embeds
         else:

tpu_inference/spec_decode/jax/eagle3.py

@@ -49,7 +49,7 @@ def __init__(
 
     def load_model(self, target_model: Any) -> None:
         """Loads the draft model."""
-        self.model_fn, self.compute_logits_fn, self.combine_hidden_states_fn, _, _, _, self.state, _, _ = get_model(
+        self.model_fn, self.compute_logits_fn, self.combine_hidden_states_fn, _, self.state, _, _ = get_model(
             self.vllm_config, self.rng_key, self.mesh, is_draft_model=True)
         del self.state.model['embed_tokens']
         self.state.model.embed_tokens = target_model.model.embed