prefix caching for MM

sendnn_inference/envs.py

@@ -30,6 +30,7 @@
     SENDNN_INFERENCE_TP_MM_SHARING: bool = True
     SENDNN_INFERENCE_LONG_OUT_PRIO: bool = False
     SENDNN_INFERENCE_PAUSING_ENABLED: bool = True
+    SENDNN_INFERENCE_MM_ENCODER_CACHE_MB: int = 512
 
 logger = init_logger(__name__)
 
@@ -210,6 +211,14 @@ def clear_env_cache():
     "SENDNN_INFERENCE_PAUSING_ENABLED": lambda: bool(
         int(os.getenv("SENDNN_INFERENCE_PAUSING_ENABLED", "1"))
     ),
+    # Byte budget (in MiB) for the cross-request vision encoder-output cache in the
+    # MM encoder subprocess (and the inline fallback path), keyed by multimodal
+    # content hash (mm_hash). When an image is already cached, the vision tower is
+    # skipped: only the (cheap) text embeddings are recomputed and the cached image
+    # features are merged back in. Set to 0 to disable the cache entirely.
+    "SENDNN_INFERENCE_MM_ENCODER_CACHE_MB": lambda: int(
+        os.getenv("SENDNN_INFERENCE_MM_ENCODER_CACHE_MB", "512")
+    ),
 }
 # --8<-- [end:env-vars-definition]
 

sendnn_inference/model_executor/model_loader/spyre.py

@@ -579,6 +579,22 @@ def get_maybe_mm_embeddings(self, input_ids, mm_features, is_decode):
             self.mm_device,
         )
 
+    def encode_images(self, mm_features):
+        """Run the vision encoder for the request's image(s), returning packed
+        image features [num_image_tokens, emb_dim]. See MMUtilsBase.encode_images."""
+        return self.mm_model_utils.encode_images(self.fms_model, mm_features, self.mm_device)
+
+    def embed_text(self, input_ids):
+        """Token-embedding lookup only (no vision tower)."""
+        return self.mm_model_utils.embed_text(self.fms_model, input_ids)
+
+    def merge_embeddings(self, input_ids, text_embeds, image_features):
+        """Scatter precomputed image features into text embeddings at the image
+        placeholder positions."""
+        return self.mm_model_utils.merge_embeddings(
+            self.fms_model, input_ids, text_embeds, image_features
+        )
+
     def sample(
         self,
         logits: torch.Tensor,

sendnn_inference/multimodal/mm_mappings/base.py

@@ -109,6 +109,45 @@ def get_maybe_mm_embeddings(
         """
         pass
 
+    # The three methods below decompose ``get_maybe_mm_embeddings`` into its
+    # independently-reusable stages so the (expensive) image encoding can be cached
+    # by mm_hash separately from the text embeddings. Composing
+    # merge(embed_text, encode_images) must reproduce the merged output of
+    # ``get_maybe_mm_embeddings`` (verified by a decomposition-equivalence test).
+
+    @staticmethod
+    @abstractmethod
+    def encode_images(
+        fms_model: torch.nn.Module,
+        mm_features: list[MultiModalFeatureSpec],
+        mm_device: str,
+    ) -> torch.Tensor:
+        """Run the vision tower + projector on the request's image(s).
+
+        Returns the packed image features, shape [num_image_tokens, emb_dim].
+        Depends only on the image content (not the prompt text), so the result is
+        cacheable by mm_hash. ``mm_device`` is where the vision_tower weights live.
+        """
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def embed_text(fms_model: torch.nn.Module, input_ids: torch.Tensor) -> torch.Tensor:
+        """Token-embedding lookup only (no vision tower). Shape [bsz, seq_len, emb_dim]."""
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def merge_embeddings(
+        fms_model: torch.nn.Module,
+        input_ids: torch.Tensor,
+        text_embeds: torch.Tensor,
+        image_features: torch.Tensor,
+    ) -> torch.Tensor:
+        """Scatter ``image_features`` into ``text_embeds`` at the image placeholder
+        positions (``input_ids == image_token_index``) and return the merged tensor."""
+        pass
+
     @abstractmethod
     def get_warmup_inputs(self, req_count: int) -> MMWarmupInputs:
         pass

sendnn_inference/multimodal/mm_mappings/llava_next.py

@@ -94,48 +94,84 @@ def get_maybe_mm_embeddings(
         the (potentially compiled) FMS model.
         """
         fms_kwargs = {"use_cache": True}
-        mm_spec_keys = ["pixel_values", "image_sizes"]
 
         # Only merge multimodal features in prefill; nothing mm in decode
         if mm_features:
             assert not is_decode  # We never pass features in decode
-            if len(mm_features) != 1:
-                raise ValueError("Currently we assume we only embed one mm request at a time")
-            mm_spec = mm_features[0].data
-            if mm_spec is not None:
-                # NOTE: This should be pretty safe as it's dependent on the
-                # vLLM/HF processor objects, but we check it anyway to be safe
-                # for now, since transformers 5.0 is just around the corner.
-                if any(k not in mm_spec for k in mm_spec_keys):
-                    raise KeyError(f"Llava Next requires kwargs: {mm_spec_keys}")
-
-                pixel_values = mm_spec["pixel_values"].data
-                # Place pixel_values on the same device/dtype as the
-                # vision_tower so the encoder forward can run on NNPA when the
-                # vision_tower weights ended up on nnpa (CPU otherwise).
-                mm_dtype = envs_spyre.SENDNN_INFERENCE_CPU_MM_DTYPE
-                if pixel_values.device.type != mm_device or pixel_values.dtype != mm_dtype:
-                    pixel_values = pixel_values.to(device=mm_device, dtype=mm_dtype)
-                fms_kwargs["pixel_values"] = pixel_values
-
-                image_sizes = mm_spec["image_sizes"].data
-
-                # Careful about this; if it's 1D, we'll a tensor of shape
-                # [x, y], which will break in a weird way in image packing,
-                # since it assumes it's 2D and will get sad about getting
-                # an int instead of an iterable
-                if image_sizes.ndim == 1:
-                    image_sizes = image_sizes.unsqueeze(0)
-                # image_sizes is an integer index tensor; keep it on CPU
-                # (NNPA dispatch for int tensors would just fall back anyway).
-                fms_kwargs["image_sizes"] = image_sizes
+            # Shared prep with encode_images (keeps the fused and decomposed
+            # paths from drifting). image_sizes stays on CPU (integer index tensor).
+            pixel_values, image_sizes = LlavaNextMMUtils._prepare_vision_inputs(
+                mm_features, mm_device
+            )
+            fms_kwargs["pixel_values"] = pixel_values
+            fms_kwargs["image_sizes"] = image_sizes
 
         # The value of iteration does not matter for decode as long as it's > 0
         input_embeds, _ = fms_model.prepare_inputs_for_generation(
             iteration=0 if not is_decode else 1, input_ids=input_ids, kwargs=fms_kwargs
         )  # ty: ignore[call-non-callable]
         return input_embeds
 
+    @staticmethod
+    def _prepare_vision_inputs(
+        mm_features: list[MultiModalFeatureSpec], mm_device: str
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Extract and prepare ``(pixel_values, image_sizes)`` for the vision tower.
+
+        Shared by ``get_maybe_mm_embeddings`` and ``encode_images`` so the two
+        paths cannot drift. ``pixel_values`` is placed on the vision_tower's
+        device/dtype; ``image_sizes`` is promoted to 2D (kept on CPU — it is an
+        integer index tensor).
+        """
+        if len(mm_features) != 1:
+            raise ValueError("Currently we assume we only embed one mm request at a time")
+        mm_spec = mm_features[0].data
+        mm_spec_keys = ["pixel_values", "image_sizes"]
+        if mm_spec is None or any(k not in mm_spec for k in mm_spec_keys):
+            raise KeyError(f"Llava Next requires kwargs: {mm_spec_keys}")
+
+        pixel_values = mm_spec["pixel_values"].data
+        mm_dtype = envs_spyre.SENDNN_INFERENCE_CPU_MM_DTYPE
+        if pixel_values.device.type != mm_device or pixel_values.dtype != mm_dtype:
+            pixel_values = pixel_values.to(device=mm_device, dtype=mm_dtype)
+
+        image_sizes = mm_spec["image_sizes"].data
+        if image_sizes.ndim == 1:
+            image_sizes = image_sizes.unsqueeze(0)
+        return pixel_values, image_sizes  # ty: ignore[invalid-return-type]
+
+    @staticmethod
+    def encode_images(
+        fms_model: torch.nn.Module,
+        mm_features: list[MultiModalFeatureSpec],
+        mm_device: str,
+    ) -> torch.Tensor:
+        """Run the SiglipVision tower + projector for Llava Next and return the
+        packed image features [num_image_tokens, emb_dim]."""
+        pixel_values, image_sizes = LlavaNextMMUtils._prepare_vision_inputs(mm_features, mm_device)
+        image_features = fms_model.get_image_features(pixel_values, image_sizes)  # ty: ignore[call-non-callable]
+        return fms_model.pack_image_features(  # ty: ignore[call-non-callable]
+            image_features, image_sizes, image_newline=fms_model.image_newline
+        )
+
+    @staticmethod
+    def embed_text(fms_model: torch.nn.Module, input_ids: torch.Tensor) -> torch.Tensor:
+        return fms_model._get_text_embeddings(input_ids)  # ty: ignore[call-non-callable]
+
+    @staticmethod
+    def merge_embeddings(
+        fms_model: torch.nn.Module,
+        input_ids: torch.Tensor,
+        text_embeds: torch.Tensor,
+        image_features: torch.Tensor,
+    ) -> torch.Tensor:
+        image_features = image_features.to(text_embeds.device, text_embeds.dtype)
+        image_positions = (input_ids[0] == fms_model.config.image_token_index).nonzero(
+            as_tuple=True
+        )[0]
+        text_embeds[0, image_positions] = image_features
+        return text_embeds
+
     def get_warmup_inputs(self, req_count: int) -> MMWarmupInputs:
         """Get the inputs to the huggingface processor to create the warmup
         features or feature shapes.

sendnn_inference/multimodal/mm_mappings/mistral3.py

@@ -55,60 +55,99 @@ def get_maybe_mm_embeddings(
 
         # Only merge multimodal features in prefill; nothing mm in decode
         if mm_features:
-            # Looks for ["pixel_values", "image_sizes"] in mm_features
-            if len(mm_features) != 1:
-                raise ValueError("Currently we assume we only embed one mm request at a time")
-            mm_spec = mm_features[0].data
-
-            # when using config and tokenizer are set to `mistral` we don't get
-            # pixel_values in mm_spec. So we are mapping these back here
-            if isinstance(mm_spec, MultiModalKwargsItem) and "images" in mm_spec:
-                mm_spec["pixel_values"] = mm_spec.pop("images")
-
-            if mm_spec is not None:
-                if "pixel_values" not in mm_spec:
-                    raise KeyError("Mistral3 requires pixel_values")
-
-                pixel_values = mm_spec["pixel_values"].data
-                # FMS vision tower expects pixel_values with batch dimension
-                # If squeezed during spec building, add it back
-                if pixel_values.ndim == 3:
-                    pixel_values = pixel_values.unsqueeze(0)
-                # Move pixel_values onto the same device/dtype as the
-                # vision_tower params so the encoder forward can run there
-                # (NNPA / privateuse1 backend, or CPU). mm_device is the device
-                # the vision_tower weights actually ended up on. The merge with
-                # text_embeds happens later inside FMS, where the merged
-                # output ends up on text_embeds.device (CPU) automatically.
-                mm_dtype = envs_spyre.SENDNN_INFERENCE_CPU_MM_DTYPE
-                if pixel_values.device.type != mm_device or pixel_values.dtype != mm_dtype:
-                    pixel_values = pixel_values.to(device=mm_device, dtype=mm_dtype)
-                fms_kwargs["pixel_values"] = pixel_values
-
-                if "image_sizes" in mm_spec:
-                    # Use the processor's image_sizes which tracks the logical image dimensions
-                    # This is used by the projector to correctly split/merge patches
-                    image_sizes_tensor = mm_spec["image_sizes"].data
-                    if image_sizes_tensor.ndim == 1:
-                        # Single image: convert to list of tuples
-                        image_sizes = [(image_sizes_tensor[0].item(), image_sizes_tensor[1].item())]
-                    else:
-                        # Multiple images
-                        image_sizes = [(h.item(), w.item()) for h, w in image_sizes_tensor]
-                else:
-                    # Mistral image input in vLLM doesn't contain image_sizes as attribute, so we
-                    # are calculating based on pixel_values
-                    # Ref: https://github.com/vllm-project/vllm/blob/f97ca671766c5201404e9fc812e35bf2c4e95a01/vllm/model_executor/models/mistral3.py#L516C9-L518C10
-                    image_sizes = [(img.shape[-2], img.shape[-1]) for img in pixel_values]
-
-                fms_kwargs["image_sizes"] = image_sizes
+            # Shared prep with encode_images (keeps the fused and decomposed
+            # paths from drifting). The merge with text_embeds happens later
+            # inside FMS, where the merged output lands on text_embeds.device.
+            pixel_values, image_sizes = Mistral3MMUtils._prepare_vision_inputs(
+                mm_features, mm_device
+            )
+            fms_kwargs["pixel_values"] = pixel_values
+            fms_kwargs["image_sizes"] = image_sizes
 
         # The value of iteration does not matter for decode as long as it's > 0
         input_embeds, _ = fms_model.prepare_inputs_for_generation(
             iteration=0 if not is_decode else 1, input_ids=input_ids, kwargs=fms_kwargs
         )  # ty: ignore[call-non-callable]
         return input_embeds
 
+    @staticmethod
+    def _prepare_vision_inputs(
+        mm_features: list[MultiModalFeatureSpec], mm_device: str
+    ) -> tuple[torch.Tensor, list[tuple[int, int]]]:
+        """Extract and prepare ``(pixel_values, image_sizes)`` for the vision tower.
+
+        Shared by ``get_maybe_mm_embeddings`` and ``encode_images`` so the two
+        paths cannot drift. Handles the mistral tokenizer emitting ``images``
+        instead of ``pixel_values``, restores the batch dim, places pixel_values
+        on the vision_tower's device/dtype, and derives logical ``image_sizes``
+        (falling back to the pixel_values shape when absent).
+        """
+        if len(mm_features) != 1:
+            raise ValueError("Currently we assume we only embed one mm request at a time")
+        mm_spec = mm_features[0].data
+
+        # when config and tokenizer are set to `mistral` we don't get
+        # pixel_values in mm_spec, so map them back here.
+        if isinstance(mm_spec, MultiModalKwargsItem) and "images" in mm_spec:
+            mm_spec["pixel_values"] = mm_spec.pop("images")
+        if mm_spec is None or "pixel_values" not in mm_spec:
+            raise KeyError("Mistral3 requires pixel_values")
+
+        pixel_values = mm_spec["pixel_values"].data
+        # FMS vision tower expects pixel_values with a batch dimension; if it was
+        # squeezed during spec building, add it back.
+        if pixel_values.ndim == 3:
+            pixel_values = pixel_values.unsqueeze(0)
+        # Move pixel_values onto the same device/dtype as the vision_tower params
+        # so the encoder forward can run there (NNPA / privateuse1, or CPU).
+        mm_dtype = envs_spyre.SENDNN_INFERENCE_CPU_MM_DTYPE
+        if pixel_values.device.type != mm_device or pixel_values.dtype != mm_dtype:
+            pixel_values = pixel_values.to(device=mm_device, dtype=mm_dtype)
+
+        if "image_sizes" in mm_spec:
+            # Use the processor's image_sizes (logical image dimensions); the
+            # projector uses these to correctly split/merge patches.
+            image_sizes_tensor = mm_spec["image_sizes"].data
+            if image_sizes_tensor.ndim == 1:
+                image_sizes = [(image_sizes_tensor[0].item(), image_sizes_tensor[1].item())]
+            else:
+                image_sizes = [(h.item(), w.item()) for h, w in image_sizes_tensor]
+        else:
+            # Mistral image input in vLLM has no image_sizes attribute, so derive
+            # it from pixel_values.
+            image_sizes = [(img.shape[-2], img.shape[-1]) for img in pixel_values]
+        return pixel_values, image_sizes  # ty: ignore[invalid-return-type]
+
+    @staticmethod
+    def encode_images(
+        fms_model: torch.nn.Module,
+        mm_features: list[MultiModalFeatureSpec],
+        mm_device: str,
+    ) -> torch.Tensor:
+        """Run the PixtralVision tower + projector for mistral3 and return the
+        packed image features [num_image_tokens, emb_dim]."""
+        pixel_values, image_sizes = Mistral3MMUtils._prepare_vision_inputs(mm_features, mm_device)
+        return fms_model._get_image_features(pixel_values, image_sizes)  # ty: ignore[call-non-callable]
+
+    @staticmethod
+    def embed_text(fms_model: torch.nn.Module, input_ids: torch.Tensor) -> torch.Tensor:
+        return fms_model._get_text_embeddings(input_ids, None)  # ty: ignore[call-non-callable]
+
+    @staticmethod
+    def merge_embeddings(
+        fms_model: torch.nn.Module,
+        input_ids: torch.Tensor,
+        text_embeds: torch.Tensor,
+        image_features: torch.Tensor,
+    ) -> torch.Tensor:
+        return fms_model._merge_multimodal_embeddings(  # ty: ignore[call-non-callable]
+            input_ids,
+            text_embeds,
+            image_features,
+            device=text_embeds.device,
+            dtype=text_embeds.dtype,
+        )
+
     def get_warmup_inputs(self, req_count: int) -> MMWarmupInputs:
         """Generate input for warmup using using dummy image."""