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

@@ -136,6 +136,53 @@ def get_maybe_mm_embeddings(
         )  # ty: ignore[call-non-callable]
         return input_embeds
 
+    @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]."""
+        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)
+
+        image_features = fms_model.get_image_features(pixel_values, image_sizes)
+        return fms_model.pack_image_features(
+            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)
+
+    @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

@@ -109,6 +109,61 @@ def get_maybe_mm_embeddings(
         )  # ty: ignore[call-non-callable]
         return input_embeds
 
+    @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]."""
+        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
+
+        # As in get_maybe_mm_embeddings: mistral tokenizer emits "images" not "pixel_values".
+        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
+        if pixel_values.ndim == 3:
+            pixel_values = pixel_values.unsqueeze(0)
+        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:
+            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:
+            image_sizes = [(img.shape[-2], img.shape[-1]) for img in pixel_values]
+
+        return fms_model._get_image_features(pixel_values, image_sizes)
+
+    @staticmethod
+    def embed_text(fms_model: torch.nn.Module, input_ids: torch.Tensor) -> torch.Tensor:
+        return fms_model._get_text_embeddings(input_ids, None)
+
+    @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(
+            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."""
 

sendnn_inference/v1/core/scheduler.py

@@ -160,6 +160,12 @@ class ChunkedPrefillSpyreSchedulerStats:
     num_paused_reqs: int = 0
     pause_events: int = 0
     resume_events: int = 0
+    # Per-interval delta of the cross-request vision-encoder cache: how many
+    # images reused cached vision features (hits) out of those looked up
+    # (queries). Whole-image granularity, keyed by mm_hash — distinct from
+    # upstream mm_cache_stats and from text prefix/block caching.
+    vision_encoder_cache_hits: int = 0
+    vision_encoder_cache_queries: int = 0
 
 
 class ChunkedPrefillSpyreScheduler(SpyreScheduler):
@@ -243,6 +249,17 @@ def __init__(self, *args, **kwargs) -> None:
         self.pause_events = 0
         self.resume_events = 0
 
+        # Cross-request MM encoder-cache stats. The worker/encoder-subprocess
+        # caches report cumulative counters (only one path is active per run, so
+        # they sum without double counting); make_stats emits the per-interval
+        # delta into base_stats.mm_cache_stats.
+        self._mm_async_cum_hits = 0
+        self._mm_async_cum_queries = 0
+        self._mm_inline_cum_hits = 0
+        self._mm_inline_cum_queries = 0
+        self._mm_reported_hits = 0
+        self._mm_reported_queries = 0
+
         self.request_last_decode_step = defaultdict(int)
         self.long_output_prio = envs_spyre.SENDNN_INFERENCE_LONG_OUT_PRIO
 
@@ -269,6 +286,16 @@ def update_from_output(self, scheduler_output, model_runner_output):
             logger.error("MM encode failed for req '%s' — aborting request", req_id)
             self.finish_requests([req_id], RequestStatus.FINISHED_ABORTED)
 
+        # Track the latest cumulative MM encoder-cache counters. The async encoder
+        # subprocess reports via scheduler_output (only when it encoded this step);
+        # the inline fallback path reports via model_runner_output every step.
+        async_hits = getattr(scheduler_output, "_spyre_mm_cache_hits", None)
+        if async_hits is not None:
+            self._mm_async_cum_hits = async_hits
+            self._mm_async_cum_queries = getattr(scheduler_output, "_spyre_mm_cache_queries", 0)
+        self._mm_inline_cum_hits = getattr(model_runner_output, "mm_cache_hits", 0)
+        self._mm_inline_cum_queries = getattr(model_runner_output, "mm_cache_queries", 0)
+
         # Remove completed prefills
         self.ongoing_prefills = [
             req for req in self.ongoing_prefills if req.num_computed_tokens < req.num_prompt_tokens
@@ -965,8 +992,11 @@ def make_stats(self, *args, **kwargs) -> SchedulerStats | None:
         """Update the scheduler stats from the base scheduler.
         In sendnn-inference the last chunk is always recomputed, even though
         the space is not duplicated.
-        Spyre does not support cross-request MM cache reuse today, so MM cache
-        hit reporting is forced to 0.0%.
+        The cross-request vision-encoder cache lives worker-side; its cumulative
+        counters are plumbed back via update_from_output, and the per-interval
+        delta is reported as a dedicated sendnn ``vision_encoder_cache`` metric
+        (NOT folded into upstream's mm_cache_stats, which tracks a different
+        cache — vLLM's multimodal processor/input cache).
         """
         base_stats = super().make_stats(*args, **kwargs)
 
@@ -976,9 +1006,15 @@ def make_stats(self, *args, **kwargs) -> SchedulerStats | None:
                     base_stats.prefix_cache_stats.queries, base_stats.prefix_cache_stats.hits
                 )
 
-                mm_cache_stats = getattr(base_stats, "mm_cache_stats", None)
-                if mm_cache_stats is not None:
-                    mm_cache_stats.hits = 0
+            # Per-interval delta of the real vision-encoder cache hit/query counts.
+            # Sum the two sources (async subprocess + inline fallback); only one is
+            # active per run, so the inactive one stays 0.
+            cum_hits = self._mm_async_cum_hits + self._mm_inline_cum_hits
+            cum_queries = self._mm_async_cum_queries + self._mm_inline_cum_queries
+            ve_cache_hits = max(0, cum_hits - self._mm_reported_hits)
+            ve_cache_queries = max(0, cum_queries - self._mm_reported_queries)
+            self._mm_reported_hits = cum_hits
+            self._mm_reported_queries = cum_queries
 
             decode_batch_size = sum(1 for r in self.running if r not in self.ongoing_prefills)
             num_paused_reqs = len(self.paused_decoding_requests)
@@ -993,6 +1029,8 @@ def make_stats(self, *args, **kwargs) -> SchedulerStats | None:
                 num_paused_reqs=num_paused_reqs,
                 pause_events=self.pause_events,
                 resume_events=self.resume_events,
+                vision_encoder_cache_hits=ve_cache_hits,
+                vision_encoder_cache_queries=ve_cache_queries,
             )
             self.pause_events = 0
             self.resume_events = 0

sendnn_inference/v1/executor/spyre_executor.py

@@ -131,13 +131,19 @@ def execute_model(self, scheduler_output: Any, non_block: bool = False) -> Any:
                     )
                     failed_encode_req_ids.append(req.request_id)
 
+        # Latest cumulative encoder-cache counters seen this drain (None if the
+        # encoder reported none this step); stamped onto scheduler_output below.
+        mm_cache_hits: int | None = None
+        mm_cache_misses: int | None = None
         if self._mm_result_queue is not None and self._mm_in_flight > 0:
             # Collect completed results (non-blocking drain).
             newly_encoded_metadata: list[tuple] = []
             while True:
                 try:
-                    req_id, shape, dtype = self._mm_result_queue.get_nowait()
+                    req_id, shape, dtype, c_hits, c_misses = self._mm_result_queue.get_nowait()
                     self._mm_in_flight -= 1
+                    # Cumulative snapshots grow monotonically; keep the newest.
+                    mm_cache_hits, mm_cache_misses = c_hits, c_misses
                     if shape is not None and dtype is not None:
                         newly_encoded_metadata.append((req_id, shape, dtype))
                     else:
@@ -176,6 +182,12 @@ def execute_model(self, scheduler_output: Any, non_block: bool = False) -> Any:
             scheduler_output._spyre_newly_encoded_req_ids = newly_encoded_req_ids
         if failed_encode_req_ids:
             scheduler_output._spyre_failed_encode_req_ids = failed_encode_req_ids
+        # Surface the async encoder subprocess's cumulative cache counters so the
+        # scheduler can report the real MM cache hit rate.  Only stamp when the
+        # encoder reported this step; otherwise the scheduler keeps its last value.
+        if mm_cache_hits is not None:
+            scheduler_output._spyre_mm_cache_hits = mm_cache_hits
+            scheduler_output._spyre_mm_cache_queries = mm_cache_hits + mm_cache_misses
 
         # Clear _spyre_mm_encode_requests before dispatching to workers.
         # The async encoder owns all MM encoding jobs