[Bugfix] Merge MM embeddings by index instead of token IDs

vllm/v1/spec_decode/eagle.py

@@ -18,6 +18,7 @@
 from vllm.model_executor.model_loader import get_model
 from vllm.model_executor.models import supports_multimodal
 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
+from vllm.model_executor.models.utils import _merge_multimodal_embeddings
 from vllm.platforms import current_platform
 from vllm.utils import is_pin_memory_available
 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
@@ -157,6 +158,7 @@ def propose(
         next_token_ids: torch.Tensor,
         common_attn_metadata: CommonAttentionMetadata,
         sampling_metadata: SamplingMetadata,
+        is_mm_embed: Optional[torch.Tensor] = None,
         mm_embeds: Optional[list[torch.Tensor]] = None,
     ) -> torch.Tensor:
         num_tokens = target_token_ids.shape[0]
@@ -196,18 +198,22 @@ def propose(
         # copy inputs to buffer for cudagraph
         self.positions[:num_tokens] = target_positions
         self.hidden_states[:num_tokens] = target_hidden_states
-        if self.is_multimodal_model:
-            input_ids = self.input_ids[:num_tokens]
-            inputs_embeds = self.model.get_input_embeddings(
-                input_ids,
-                multimodal_embeddings=mm_embeds or None,
+
+        if mm_embeds:
+            assert is_mm_embed is not None
+
+            inputs_embeds_scheduled = _merge_multimodal_embeddings(
+                self.input_ids[:num_tokens],
+                is_mm_embed,
+                multimodal_embeddings=mm_embeds,
             )
-            self.inputs_embeds[:num_tokens] = inputs_embeds
-            inputs_embeds = self.inputs_embeds[:num_input_tokens]
+            self.inputs_embeds[:num_tokens] = inputs_embeds_scheduled
+
             input_ids = None
+            inputs_embeds = self.inputs_embeds[:num_input_tokens]
         else:
-            inputs_embeds = None
             input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = None
 
         with set_forward_context(per_layer_attn_metadata,
                                  self.vllm_config,

vllm/v1/worker/gpu_model_runner.py

@@ -44,6 +44,7 @@
                                                    supports_transcription)
 from vllm.model_executor.models.interfaces_base import (
     VllmModelForPooling, is_pooling_model, is_text_generation_model)
+from vllm.model_executor.models.utils import _merge_multimodal_embeddings
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargsItem,
                                     PlaceholderRange)
@@ -257,6 +258,10 @@ def __init__(
             dtype=self.dtype,
             device=self.device)
 
+        # Only relevant for multimodal models
+        self.is_mm_embed = self._make_buffer(self.max_num_tokens,
+                                             dtype=torch.bool)
+
         # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
         if self.uses_mrope:
             # NOTE: `mrope_positions` is implemented with one additional dummy
@@ -1185,8 +1190,11 @@ def _gather_mm_embeddings(
         self,
         scheduler_output: "SchedulerOutput",
         shift_computed_tokens: int = 0,
-    ) -> list[torch.Tensor]:
-        mm_embeds: list[torch.Tensor] = []
+    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        is_mm_embed = self.is_mm_embed.cpu
+        mm_embeds = list[torch.Tensor]()
+
+        req_start_idx = 0
         for req_id in self.input_batch.req_ids:
             num_scheduled_tokens = scheduler_output.num_scheduled_tokens[
                 req_id]
@@ -1195,6 +1203,7 @@ def _gather_mm_embeddings(
                 req_state.num_computed_tokens + shift_computed_tokens
             mm_positions = req_state.mm_positions
             mm_hashes = req_state.mm_hashes
+
             for i, pos_info in enumerate(mm_positions):
                 start_pos = pos_info.offset
                 num_encoder_tokens = pos_info.length
@@ -1211,6 +1220,10 @@ def _gather_mm_embeddings(
                     # in the decoder's KV cache.
                     continue
 
+                req_start_pos = req_start_idx + start_pos
+                is_mm_embed[req_start_pos:req_start_pos + num_encoder_tokens] \
+                    = True if pos_info.is_embed is None else pos_info.is_embed
+
                 start_idx = max(num_computed_tokens - start_pos, 0)
                 end_idx = min(
                     num_computed_tokens - start_pos + num_scheduled_tokens,
@@ -1231,7 +1244,13 @@ def _gather_mm_embeddings(
                     is_embed=is_embed,
                 )
                 mm_embeds.append(mm_embeds_item)
-        return mm_embeds
+
+            req_start_idx += num_scheduled_tokens
+
+        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        self.is_mm_embed.copy_to_gpu(total_num_scheduled_tokens)
+
+        return self.is_mm_embed.gpu[:total_num_scheduled_tokens], mm_embeds
 
     def get_model(self) -> nn.Module:
         # get raw model out of the cudagraph wrapper.
@@ -1514,18 +1533,24 @@ def execute_model(
         if self.supports_mm_inputs:
             # Run the multimodal encoder if any.
             self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+            is_mm_embed, mm_embeds = self._gather_mm_embeddings(
+                scheduler_output)
         else:
-            mm_embeds = []
+            is_mm_embed, mm_embeds = torch.tensor(False), []
 
         if self.supports_mm_inputs and get_pp_group().is_first_rank:
             # NOTE(woosuk): To unify token ids and soft tokens (vision
             # embeddings), we always use embeddings (rather than token ids)
             # as input to the multimodal model, even when the input is text.
             inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids.gpu[:num_scheduled_tokens],
-                multimodal_embeddings=mm_embeds or None,
-            )
+                self.input_ids.gpu[:num_scheduled_tokens])
+
+            if mm_embeds:
+                inputs_embeds_scheduled = _merge_multimodal_embeddings(
+                    inputs_embeds_scheduled,
+                    is_mm_embed,
+                    multimodal_embeddings=mm_embeds,
+                )
 
             # TODO(woosuk): Avoid the copy. Optimize.
             self.inputs_embeds[:num_scheduled_tokens].copy_(
@@ -1860,10 +1885,14 @@ def propose_draft_token_ids(
                         [h[token_indices] for h in aux_hidden_states], dim=-1)
                 else:
                     target_hidden_states = hidden_states[token_indices]
-            mm_embeds = None
+
             if self.supports_mm_inputs:
-                mm_embeds = self._gather_mm_embeddings(scheduler_output,
-                                                       shift_computed_tokens=1)
+                is_mm_embed, mm_embeds = self._gather_mm_embeddings(
+                    scheduler_output,
+                    shift_computed_tokens=1,
+                )
+            else:
+                is_mm_embed, mm_embeds = torch.tensor(False), []
 
             draft_token_ids = self.drafter.propose(
                 target_token_ids=target_token_ids,
@@ -1872,6 +1901,7 @@ def propose_draft_token_ids(
                 next_token_ids=next_token_ids,
                 sampling_metadata=sampling_metadata,
                 common_attn_metadata=common_attn_metadata,
+                is_mm_embed=is_mm_embed,
                 mm_embeds=mm_embeds,
             )
         return draft_token_ids

vllm/v1/worker/tpu_model_runner.py

@@ -31,6 +31,7 @@
 from vllm.model_executor.models.interfaces import supports_transcription
 from vllm.model_executor.models.interfaces_base import (
     is_pooling_model, is_text_generation_model)
+from vllm.model_executor.models.utils import _merge_multimodal_embeddings
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargsItem,
                                     PlaceholderRange)
@@ -261,6 +262,12 @@ def __init__(
                                         pin_memory=self.pin_memory)
         self.seq_lens_np = self.seq_lens_cpu.numpy()
 
+        # Only relevant for multimodal models
+        self.is_mm_embed_cpu = torch.zeros(self.max_num_tokens,
+                                           dtype=torch.bool,
+                                           device="cpu",
+                                           pin_memory=self.pin_memory)
+
         # Range tensor with values [0 .. self.max_num_tokens - 1].
         # Used to initialize positions / context_lens / seq_lens
         # Keep in int64 to avoid overflow with long context
@@ -809,31 +816,6 @@ def _prepare_inputs(self, scheduler_output: "SchedulerOutput",
         return per_layer_attn_metadata, logits_indices, padded_num_reqs,\
             num_reqs, end_index
 
-    def _scatter_placeholders(
-        self,
-        embeds: torch.Tensor,
-        is_embed: Optional[torch.Tensor],
-    ) -> torch.Tensor:
-        if is_embed is None:
-            return embeds
-
-        placeholders = embeds.new_full(
-            (is_embed.shape[0], embeds.shape[-1]),
-            fill_value=torch.nan,
-        )
-        placeholders[is_embed] = embeds
-        return placeholders
-
-    def _gather_placeholders(
-        self,
-        placeholders: torch.Tensor,
-        is_embed: Optional[torch.Tensor],
-    ) -> torch.Tensor:
-        if is_embed is None:
-            return placeholders
-
-        return placeholders[is_embed]
-
     def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
         scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
         if not scheduled_encoder_inputs:
@@ -906,8 +888,11 @@ def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
     def _gather_mm_embeddings(
         self,
         scheduler_output: "SchedulerOutput",
-    ) -> list[torch.Tensor]:
-        mm_embeds: list[torch.Tensor] = []
+    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        is_mm_embed = self.is_mm_embed_cpu
+        mm_embeds = list[torch.Tensor]()
+
+        req_start_idx = 0
         for req_id in self.input_batch.req_ids:
             num_scheduled_tokens = scheduler_output.num_scheduled_tokens[
                 req_id]
@@ -935,6 +920,10 @@ def _gather_mm_embeddings(
                     # in the decoder's KV cache.
                     continue
 
+                req_start_pos = req_start_idx + start_pos
+                is_mm_embed[req_start_pos:req_start_pos + num_encoder_tokens] \
+                    = True
+
                 mm_hash = mm_hashes[i]
                 encoder_output = self.encoder_cache.get(mm_hash, None)
                 assert encoder_output is not None,\
@@ -943,18 +932,33 @@ def _gather_mm_embeddings(
                 " be contiguous and embeddings."
                 encoder_output = self.encoder_cache[mm_hash]
                 mm_embeds.append(encoder_output)
-        return mm_embeds
 
-    def _get_model_inputs(self, input_ids: torch.Tensor,
-                          mm_embeds: list[torch.Tensor]):
+            req_start_idx += num_scheduled_tokens
+
+        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        is_mm_embed = is_mm_embed[:total_num_scheduled_tokens].to(self.device)
+
+        return is_mm_embed, mm_embeds
+
+    def _get_model_inputs(
+        self,
+        input_ids: torch.Tensor,
+        is_mm_embed: torch.Tensor,
+        mm_embeds: list[torch.Tensor],
+    ):
         if self.supports_mm_inputs:
             # NOTE(woosuk): To unify token ids and soft tokens (vision
             # embeddings), we always use embeddings (rather than token ids)
             # as input to the multimodal model, even when the input is text.
-            inputs_embeds = self.model.get_input_embeddings(
-                input_ids=input_ids,
-                multimodal_embeddings=mm_embeds,
-            )
+            inputs_embeds = self.model.get_input_embeddings(input_ids)
+
+            if mm_embeds:
+                inputs_embeds = _merge_multimodal_embeddings(
+                    inputs_embeds,
+                    is_mm_embed,
+                    multimodal_embeddings=mm_embeds,
+                )
+
             return None, inputs_embeds
         else:
             # For text-only models, we use token ids as input.
@@ -982,9 +986,11 @@ def execute_model(
         if self.supports_mm_inputs:
             # Run the multimodal encoder if any.
             self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+            is_mm_embed, mm_embeds = self._gather_mm_embeddings(
+                scheduler_output)
         else:
-            mm_embeds = []
+            is_mm_embed, mm_embeds = torch.tensor(False), []
+
         xm.mark_step()
         # Prepare inputs, the requests might be split into multiple
         # executions, combine the result of each execution.
@@ -1001,7 +1007,7 @@ def execute_model(
             attn_metadata, logits_indices, padded_num_reqs, num_reqs,\
                 end_index = self._prepare_inputs(scheduler_output, start_index)
             input_ids, inputs_embeds = self._get_model_inputs(
-                self.input_ids, mm_embeds)
+                self.input_ids, is_mm_embed, mm_embeds)
             xm.mark_step()
             # Run the decoder
             with set_forward_context(
@@ -1358,6 +1364,7 @@ def _precompile_mm_encoder(self) -> None:
                         placeholders_ids = placeholders_ids.to(self.device)
                         # Assign outputs or the graph will be cut short.
                         a, b = self._get_model_inputs(placeholders_ids,
+                                                      torch.tensor(True),
                                                       [mm_embeds])
                         assert a is None
                         xm.mark_step()
@@ -1369,7 +1376,8 @@ def _precompile_mm_encoder(self) -> None:
                                                dtype=torch.int32,
                                                device="cpu")
                 placeholders_ids = placeholders_ids.to(self.device)
-                a, b = self._get_model_inputs(placeholders_ids, [])
+                a, b = self._get_model_inputs(placeholders_ids,
+                                              torch.tensor(False), [])
                 assert a is None
                 xm.mark_step()