From 26d7b08b366fefbc32de39414d256103740ab3a0 Mon Sep 17 00:00:00 2001 From: nikheal2 Date: Wed, 15 Jul 2026 12:20:24 +0530 Subject: [PATCH 1/7] prefix caching for MM Signed-off-by: nikheal2 changing logger from debug to info Signed-off-by: nikheal2 changing logger to vllm logger Signed-off-by: nikheal2 refcatoring cache count logic Signed-off-by: nikheal2 fixing metrics Signed-off-by: nikheal2 adding test cases Signed-off-by: nikheal2 adding debug log Signed-off-by: nikheal2 adding debug log Signed-off-by: nikheal2 fix image identifier Signed-off-by: nikheal2 fix logging Signed-off-by: nikheal2 --- sendnn_inference/envs.py | 9 ++ .../model_executor/model_loader/spyre.py | 16 +++ .../multimodal/mm_mappings/base.py | 39 ++++++ .../multimodal/mm_mappings/llava_next.py | 47 +++++++ .../multimodal/mm_mappings/mistral3.py | 55 +++++++++ sendnn_inference/v1/core/scheduler.py | 48 +++++++- .../v1/executor/spyre_executor.py | 14 ++- .../v1/worker/mm_encoder_cache.py | 116 ++++++++++++++++++ .../v1/worker/mm_encoder_process.py | 91 ++++++++++++-- .../v1/worker/spyre_model_runner.py | 80 ++++++++++-- tests/e2e/test_spyre_mm.py | 95 ++++++++++++++ tests/v1/executor/test_spyre_executor.py | 10 +- tests/v1/worker/test_mm_encoder_process.py | 46 ++++--- tests/v1/worker/test_scheduler_tkv_limits.py | 47 +------ 14 files changed, 624 insertions(+), 89 deletions(-) create mode 100644 sendnn_inference/v1/worker/mm_encoder_cache.py diff --git a/sendnn_inference/envs.py b/sendnn_inference/envs.py index 21865df35..7528ce5cb 100644 --- a/sendnn_inference/envs.py +++ b/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] diff --git a/sendnn_inference/model_executor/model_loader/spyre.py b/sendnn_inference/model_executor/model_loader/spyre.py index 7b6ecb5d0..00f2c39d8 100644 --- a/sendnn_inference/model_executor/model_loader/spyre.py +++ b/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, diff --git a/sendnn_inference/multimodal/mm_mappings/base.py b/sendnn_inference/multimodal/mm_mappings/base.py index 9629f27d9..f19fde242 100644 --- a/sendnn_inference/multimodal/mm_mappings/base.py +++ b/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 diff --git a/sendnn_inference/multimodal/mm_mappings/llava_next.py b/sendnn_inference/multimodal/mm_mappings/llava_next.py index 85c881ee8..6a4c46713 100644 --- a/sendnn_inference/multimodal/mm_mappings/llava_next.py +++ b/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. diff --git a/sendnn_inference/multimodal/mm_mappings/mistral3.py b/sendnn_inference/multimodal/mm_mappings/mistral3.py index 0c47e7222..fb5cf26ac 100644 --- a/sendnn_inference/multimodal/mm_mappings/mistral3.py +++ b/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.""" diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 945de035c..ed5bbc232 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/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 diff --git a/sendnn_inference/v1/executor/spyre_executor.py b/sendnn_inference/v1/executor/spyre_executor.py index 9bb9754fc..3f74a6161 100644 --- a/sendnn_inference/v1/executor/spyre_executor.py +++ b/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 diff --git a/sendnn_inference/v1/worker/mm_encoder_cache.py b/sendnn_inference/v1/worker/mm_encoder_cache.py new file mode 100644 index 000000000..5f36ebddd --- /dev/null +++ b/sendnn_inference/v1/worker/mm_encoder_cache.py @@ -0,0 +1,116 @@ +"""Per-rank, cross-request cache of vision encoder outputs. + +The (expensive) vision tower + projector turn an image into packed feature vectors +(shape ``[num_image_tokens, emb_dim]``). Those features depend only on the image, +so they are cached here keyed by the multimodal content hash +(``MultiModalFeatureSpec.identifier``, a.k.a. mm_hash). + +On a later request containing the same image, the caller reuses the cached features +and merges them into freshly-computed text embeddings, skipping the vision tower. +See ``spyre_model_runner._compute_and_cache_mm_embeddings``. + +The cache is a byte-bounded LRU. It lives on each TP rank independently; because +every rank processes the identical request stream and stores identically-sized +tensors, the ranks' caches stay in lock-step (same contents, same evictions) with +no cross-rank coordination. Cached tensors are kept on CPU and cloned on insert so +they are detached from any request-scoped buffer. +""" + +from collections import OrderedDict +from typing import Any + +import torch + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Identifiers used for warmup features must never be cached: they are dummy +# images and would poison real lookups (and are reused across models). +_WARMUP_IDENTIFIER_PREFIX = "MM-warmup" + + +def placeholder_slices(mm_features: Any) -> list[tuple[str, int, int]]: + """Return ``(identifier, offset, length)`` for each cacheable image. + + ``offset``/``length`` come from each feature's ``mm_position``. ``length`` is + the reserved placeholder span, which may differ from the packed feature-row + count, so it must NOT be used to validate a cache entry — the ``identifier`` + (mm_hash) alone keys the cache. Only ``identifier`` is used by callers today. + """ + slices: list[tuple[str, int, int]] = [] + for feat in mm_features or []: + position = getattr(feat, "mm_position", None) + identifier = getattr(feat, "identifier", None) + if position is None or not MMEncoderCache.is_cacheable(identifier): + continue + slices.append((identifier, position.offset, position.length)) + return slices + + +class MMEncoderCache: + """Byte-bounded LRU mapping mm_hash -> packed image features (CPU).""" + + def __init__(self, capacity_bytes: int): + self.capacity_bytes = max(0, capacity_bytes) + self._store: OrderedDict[str, torch.Tensor] = OrderedDict() + self._nbytes = 0 + self.hits = 0 + self.misses = 0 + + @property + def enabled(self) -> bool: + return self.capacity_bytes > 0 + + @staticmethod + def is_cacheable(identifier: str | None) -> bool: + return bool(identifier) and not identifier.startswith(_WARMUP_IDENTIFIER_PREFIX) + + def get(self, identifier: str) -> torch.Tensor | None: + """Return the cached features for *identifier*, marking most-recently-used. + + Does not update hit/miss counters — call :meth:`record_lookup` once per + request after deciding hit vs. miss. + """ + tensor = self._store.get(identifier) + if tensor is not None: + self._store.move_to_end(identifier) + return tensor + + def put(self, identifier: str, tensor: torch.Tensor) -> None: + if not self.enabled or not self.is_cacheable(identifier): + return + tensor = tensor.detach().to("cpu").contiguous() + nbytes = tensor.numel() * tensor.element_size() + # A single entry larger than the whole budget is simply not cached. + over = nbytes > self.capacity_bytes + logger.info( + "MM encoder cache: entry '%s' size=%.2f MiB, budget=%.2f MiB " + "(SENDNN_INFERENCE_MM_ENCODER_CACHE_MB) — %s", + identifier, + nbytes / 1024 / 1024, + self.capacity_bytes / 1024 / 1024, + "OVER budget → NOT cached" if over else "under budget → cached", + ) + if over: + return + if identifier in self._store: + self._nbytes -= self._store[identifier].numel() * self._store[identifier].element_size() + self._store.pop(identifier) + self._store[identifier] = tensor + self._nbytes += nbytes + self._evict_to_fit() + + def _evict_to_fit(self) -> None: + while self._nbytes > self.capacity_bytes and self._store: + _, evicted = self._store.popitem(last=False) + self._nbytes -= evicted.numel() * evicted.element_size() + + def record_lookup(self, hit: bool) -> None: + if hit: + self.hits += 1 + else: + self.misses += 1 + + def __contains__(self, identifier: str) -> bool: + return identifier in self._store diff --git a/sendnn_inference/v1/worker/mm_encoder_process.py b/sendnn_inference/v1/worker/mm_encoder_process.py index b9e887a5b..155312832 100644 --- a/sendnn_inference/v1/worker/mm_encoder_process.py +++ b/sendnn_inference/v1/worker/mm_encoder_process.py @@ -12,7 +12,6 @@ of the full tensor. """ -import logging import math import os import platform @@ -21,13 +20,15 @@ import torch from vllm.config import VllmConfig +from vllm.logger import init_logger import sendnn_inference.envs as envs_spyre from sendnn_inference.model_executor.model_loader.spyre import SpyreCausalLM, cast_params_for_spyre from sendnn_inference.platform import SpyrePlatform, THREADING_ENVS +from sendnn_inference.v1.worker.mm_encoder_cache import MMEncoderCache, placeholder_slices from sendnn_inference.v1.worker.mm_shared_memory import write_embeddings -logger = logging.getLogger(__name__) +logger = init_logger(__name__) def _resolve_mm_utils_cls(hf_config): @@ -122,23 +123,80 @@ def __init__(self, vllm_config: VllmConfig) -> None: self.mm_utils_cls.mm_parameter_prefixes, is_fp8_model=False, ) + # Cross-request cache of pure image features, keyed by mm_hash. A repeat + # image reuses its features and skips the vision tower; text embedding and + # merge still run per request so the emitted (merged) embedding is correct. + # Process-local — there is a single encoder subprocess, so no coordination. + self.mm_encoder_cache = MMEncoderCache( + capacity_bytes=envs_spyre.SENDNN_INFERENCE_MM_ENCODER_CACHE_MB * 1024 * 1024 + ) logger.info("encoder_process: mm_utils=%s", self.mm_utils_cls.__name__) torch.set_grad_enabled(False) logger.info("encoder_process: vision model loaded in %.2fs", time.time() - t0) def execute_model(self, request) -> torch.Tensor: - """Encode a single MMEncodeRequest and return a CPU-contiguous tensor.""" + """Encode a single MMEncodeRequest and return a CPU-contiguous tensor. + + Three stages (mirroring FMS's ``prepare_inputs_for_generation``): reuse or + compute the image features (cached by mm_hash), look up the text embeddings, + and merge. Only the vision tower is skipped on a cache hit; the emitted + tensor is the full merged embedding, so the downstream SHM/worker path is + unchanged. + """ input_ids = torch.tensor(request.prompt_token_ids, dtype=torch.int64).unsqueeze(0) with torch.inference_mode(): - embeds = self.mm_utils_cls.get_maybe_mm_embeddings( - self.fms_model, - input_ids, - request.mm_features, - is_decode=False, - mm_device=self.mm_device, + image_features = self._get_or_encode_image_features(request) + text_embeds = self.mm_utils_cls.embed_text(self.fms_model, input_ids) + embeds = self.mm_utils_cls.merge_embeddings( + self.fms_model, input_ids, text_embeds, image_features ) return embeds.to(dtype=self._decoder_dtype).cpu().contiguous() + def _get_or_encode_image_features(self, request) -> torch.Tensor: + """Return packed image features for the request, using the mm_hash cache. + + One image per request (matching the mm_mapping asserts). Warmup identifiers + are never cached (handled by ``MMEncoderCache.is_cacheable``). + """ + mm_features = request.mm_features + slices = placeholder_slices(mm_features) + cacheable = self.mm_encoder_cache.enabled and len(slices) == 1 == len(mm_features) + identifier = slices[0][0] if cacheable else None + + if cacheable: + # The mm_hash identifier (sha256 of the image) uniquely identifies the + # image, so a stored entry is exactly what encode_images would produce. + # (Do NOT gate on mm_position.length: the reserved placeholder span can + # differ from the packed feature-row count, e.g. 704 vs 682.) + cached = self.mm_encoder_cache.get(identifier) + if cached is not None: + self.mm_encoder_cache.record_lookup(hit=True) + logger.info( + "encoder_process: vision-encoder-cache hit rate: %.1f%% " + "(req '%s', HIT — skipped tower)", + self._hit_rate() * 100, + request.request_id, + ) + return cached + + image_features = self.mm_utils_cls.encode_images( + self.fms_model, mm_features, self.mm_device + ) + if cacheable: + self.mm_encoder_cache.put(identifier, image_features) + self.mm_encoder_cache.record_lookup(hit=False) + logger.info( + "encoder_process: vision-encoder-cache hit rate: %.1f%% (req '%s', MISS)", + self._hit_rate() * 100, + request.request_id, + ) + return image_features + + def _hit_rate(self) -> float: + cache = self.mm_encoder_cache + total = cache.hits + cache.misses + return (cache.hits / total) if total else 0.0 + # ── Process entry point ─────────────────────────────────────────────────────── @@ -297,7 +355,10 @@ def encoder_process_main( # can clean up promptly instead of waiting for a timeout. if req_id in skip_ids: skip_ids.discard(req_id) - result_queue.put((req_id, None, None)) + # Carry the current cumulative cache counters even on skip so the + # scheduler always sees the latest hit/miss totals. + cache = runner.mm_encoder_cache + result_queue.put((req_id, None, None, cache.hits, cache.misses)) logger.debug("encoder_process: skipped encode for cancelled req '%s'", req_id) continue @@ -312,11 +373,17 @@ def encoder_process_main( shm.close() t_elapsed = time.time() - t0 - result_queue.put((req_id, tuple(embeds.shape), embeds.dtype)) + # Cumulative encoder-cache counters ride along so the scheduler can + # surface the real MM cache hit rate (see scheduler.make_stats). + cache = runner.mm_encoder_cache + result_queue.put( + (req_id, tuple(embeds.shape), embeds.dtype, cache.hits, cache.misses) + ) # Tombstone: a late cancel may still arrive on cancel_queue for this req_id. processed_ids.add(req_id) logger.info("maybe_mm_embedding processing time: %.2fms", t_elapsed * 1000) except Exception as exc: logger.exception("encoder_process: failed to execute_model '%s': %s", req_id, exc) - result_queue.put((req_id, None, None)) + cache = runner.mm_encoder_cache + result_queue.put((req_id, None, None, cache.hits, cache.misses)) processed_ids.add(req_id) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 78c710101..0958cd50d 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -42,6 +42,10 @@ from sendnn_inference.platform import SpyrePlatform from sendnn_inference.utils import exact_div from sendnn_inference.v1.sample.spyre_logits_processor import build_logitsprocs_for_cb +from sendnn_inference.v1.worker.mm_encoder_cache import ( + MMEncoderCache, + placeholder_slices, +) from sendnn_inference.v1.worker.mm_shared_memory import ( cleanup_embeddings, dtype_to_idx, @@ -113,6 +117,11 @@ class SpyreModelRunnerOutput(ModelRunnerOutput): # available than the number of scheduled tokens. In that case, the scheduler # needs to update its state to reflect the correct number of computed tokens prefix_cache_hit_len: dict[str, int] = field(default_factory=dict) + # Cumulative MM encoder-cache counters from the inline encode path + # (TP=1 / async encoder off / warmup). Used by the scheduler to report the + # real MM cache hit rate. Zero when the async encoder subprocess owns encoding. + mm_cache_hits: int = 0 + mm_cache_queries: int = 0 @dataclass @@ -791,6 +800,14 @@ def __init__( # Initialize performance metric logger for tracking embedding times self.perf_logger = create_perf_metric_logger(rank=rank) + # Cross-request cache of pure image features (keyed by mm_hash), used only + # by the inline fallback encode path (TP=1 / async encoder off / warmup). + # The async encoder subprocess has its own cache. Per-worker and + # deterministic across ranks (identical request stream), so no coordination. + self.mm_encoder_cache = MMEncoderCache( + capacity_bytes=envs_spyre.SENDNN_INFERENCE_MM_ENCODER_CACHE_MB * 1024 * 1024 + ) + # Pre-computed MM embeddings for waiting requests (keyed by request_id). # Populated by store_mm_embeddings() (async path via executor). # Consumed and removed by add_new_request() when the request begins prefill. @@ -1006,10 +1023,8 @@ def _compute_and_cache_mm_embeddings( # the other ranks that are waiting on broadcast A. t0 = time.time() with torch.inference_mode(): - full_embeds = self.model.get_maybe_mm_embeddings( - full_input_tokens, - mm_features=mm_features, - is_decode=False, + full_embeds = self._encode_and_merge_cached( + full_input_tokens, mm_features, req_id ) t_elapsed = time.time() - t0 logger.info("maybe_mm_embedding processing time: %.2fms", (t_elapsed * 1000)) @@ -1071,11 +1086,7 @@ def _compute_and_cache_mm_embeddings( # original behaviour — no SHM, no coordination overhead. t0 = time.time() with torch.inference_mode(): - full_embeds = self.model.get_maybe_mm_embeddings( - full_input_tokens, - mm_features=mm_features, - is_decode=False, - ) + full_embeds = self._encode_and_merge_cached(full_input_tokens, mm_features, req_id) # Ensure embeddings are on CPU before passing to the Spyre decoder. # When the vision tower runs on NNPA, get_maybe_mm_embeddings returns # an NNPA tensor; slicing it into a CPU input_embeds in @@ -1098,6 +1109,51 @@ def _compute_and_cache_mm_embeddings( self.rank, ) + def _encode_and_merge_cached( + self, full_input_tokens: torch.Tensor, mm_features: Any, req_id: str + ) -> torch.Tensor: + """Build merged MM embeddings, reusing cached image features by mm_hash. + + Three stages (mirroring FMS's ``prepare_inputs_for_generation``): reuse or + compute the image features (cached; only the vision tower is skipped on a + hit), look up the text embeddings, and merge. Must run inside + ``torch.inference_mode`` (merge updates the text-embedding tensor in place). + Equivalent to ``get_maybe_mm_embeddings(mm_features, is_decode=False)``. + """ + slices = placeholder_slices(mm_features) + cacheable = self.mm_encoder_cache.enabled and len(slices) == 1 == len(mm_features) + identifier = slices[0][0] if cacheable else None + + image_features = None + if cacheable: + # Keyed by the mm_hash identifier only; do NOT gate on + # mm_position.length (the reserved placeholder span can differ from the + # packed feature-row count, so a length check rejects valid entries). + cached = self.mm_encoder_cache.get(identifier) + if cached is not None: + image_features = cached + self.mm_encoder_cache.record_lookup(hit=True) + cache = self.mm_encoder_cache + logger.info( + "vision-encoder-cache hit rate: %.1f%% (req '%s', HIT — skipped tower)", + 100 * cache.hits / (cache.hits + cache.misses), + req_id, + ) + if image_features is None: + image_features = self.model.encode_images(mm_features) + if cacheable: + self.mm_encoder_cache.put(identifier, image_features) + self.mm_encoder_cache.record_lookup(hit=False) + cache = self.mm_encoder_cache + logger.info( + "vision-encoder-cache hit rate: %.1f%% (req '%s', MISS)", + 100 * cache.hits / (cache.hits + cache.misses), + req_id, + ) + + text_embeds = self.model.embed_text(full_input_tokens) + return self.model.merge_embeddings(full_input_tokens, text_embeds, image_features) + def _prepare_chunked_prefill(self, req_id: str) -> SamplingForwardInputs: """ Cases / Scenarios for the chunked prefill with right padding. @@ -1653,6 +1709,8 @@ def get_empty_output(self) -> SpyreModelRunnerOutput: num_nans_in_logits=None, tkv=0, left_padding={}, + mm_cache_hits=self.mm_encoder_cache.hits, + mm_cache_queries=self.mm_encoder_cache.hits + self.mm_encoder_cache.misses, ) def check_incomplete_prefill(self, scheduler_output: SchedulerOutput): @@ -1924,6 +1982,8 @@ def prefill_output(self) -> SpyreModelRunnerOutput: tkv=self.tkv, left_padding=left_padding, prefix_cache_hit_len=self.get_prefix_cache_len(), + mm_cache_hits=self.mm_encoder_cache.hits, + mm_cache_queries=self.mm_encoder_cache.hits + self.mm_encoder_cache.misses, ) def sample_tokens(self, grammar_output: "GrammarOutput | None") -> ModelRunnerOutput | None: @@ -2000,6 +2060,8 @@ def sampled_output(self, output: SamplerOutput, is_prefill: bool) -> SpyreModelR pooler_output=[], tkv=self.tkv, left_padding=left_padding, + mm_cache_hits=self.mm_encoder_cache.hits, + mm_cache_queries=self.mm_encoder_cache.hits + self.mm_encoder_cache.misses, ) def get_prefix_cache_len(self) -> dict[str, int]: diff --git a/tests/e2e/test_spyre_mm.py b/tests/e2e/test_spyre_mm.py index 5d9d499f1..21981e107 100644 --- a/tests/e2e/test_spyre_mm.py +++ b/tests/e2e/test_spyre_mm.py @@ -18,6 +18,7 @@ # the FMS serialization utilities are patched at import time, # and the patching is currently NOT idempotent. import sendnn_inference.multimodal.mm_mappings.llava_next # noqa: F401 +from sendnn_inference.multimodal.mm_mappings.llava_next import LlavaNextMMUtils # We should not use a very large value here, because # we do not have tiny multimodal models at the moment. @@ -72,6 +73,49 @@ def generate_fms_results(processor, model_path, prompts): return generated_texts +@pytest.mark.skip("Multimodal E2E tests are currently disabled; no tiny model") +@pytest.mark.cpu +@pytest.mark.parametrize("model", get_spyre_model_list(isMultimodal=True)) +def test_mm_embedding_decomposition_equivalence(model, monkeypatch): + """Gate for the vision-encoder cache: the decomposed embedding path + (encode images + embed text + merge) must reproduce the fused + ``get_maybe_mm_embeddings`` output bit-for-bit. This is what lets the encoder + subprocess cache and reuse pure image features independently of the prompt text. + """ + processor = AutoProcessor.from_pretrained(model.name) + hf_config = AutoConfig.from_pretrained(model.name) + image_token = processor.decode(hf_config.image_token_index) + + prompts = get_single_image_prompts(1, image_token, tile_size=hf_config.vision_config.image_size) + proc_res = processor( + text=prompts[0]["prompt"], + images=prompts[0]["multi_modal_data"]["image"], + return_tensors="pt", + ) + input_ids = proc_res.input_ids + mm_features = LlavaNextMMUtils._build_multimodal_spec(proc_res) + + fms_model = get_model( + "hf_pretrained", + model.name, + data_type=torch.bfloat16, + fused_weights=False, + override_hf_pretrained_config=True, + text_config={"head_dim": 128}, + ) + + fused = LlavaNextMMUtils.get_maybe_mm_embeddings( + fms_model, input_ids, mm_features, is_decode=False, mm_device="cpu" + ) + text_embeds = LlavaNextMMUtils.embed_text(fms_model, input_ids) + image_features = LlavaNextMMUtils.encode_images(fms_model, mm_features, "cpu") + decomposed = LlavaNextMMUtils.merge_embeddings( + fms_model, input_ids, text_embeds, image_features + ) + + assert torch.equal(fused, decomposed) + + @pytest.mark.skip("Multimodal E2E tests are currently disabled; no tiny model") @pytest.mark.cpu @pytest.mark.parametrize("model", get_spyre_model_list(isMultimodal=True)) @@ -119,3 +163,54 @@ def test_alignment_with_fms(model, mode, monkeypatch): # and sendnn_inference running with the eager backend. for fms_text, vllm_result in zip(fms_texts, vllm_results): assert vllm_result["text"] == fms_text + + +@pytest.mark.skip("Multimodal E2E tests are currently disabled; no tiny model") +@pytest.mark.cpu +@pytest.mark.parametrize("model", get_spyre_model_list(isMultimodal=True)) +def test_mm_encoder_cache_repeat_consistency(model, monkeypatch): + """E2E coverage for prefix caching for MM. + + Repeating the same (image + prompt) must yield an identical result: the + vision-encoder cache hit path and the KV-block prefix reuse must produce the + same output as the freshly-computed path (no corruption / stale features). + The encoder cache is on by default (SENDNN_INFERENCE_MM_ENCODER_CACHE_MB > 0), + so the second occurrence is a cache hit. + """ + processor = AutoProcessor.from_pretrained(model.name) + hf_config = AutoConfig.from_pretrained(model.name) + image_token = processor.decode(hf_config.image_token_index) + + single = list( + get_single_image_prompts( + 1, + image_token, + tile_size=hf_config.vision_config.image_size, + ) + ) + # Same image + prompt twice → the 2nd request is a vision-encoder cache hit. + prompts = single + single + + sampling_params = SamplingParams( + max_tokens=MAX_TOKENS, + temperature=0.0, + ignore_eos=True, + logprobs=0, + ) + + results = generate_spyre_vllm_output( + model=model, + prompts=prompts, + sampling_params=sampling_params, + tensor_parallel_size=1, + backend="eager", + max_num_seqs=2, + monkeypatch=monkeypatch, + max_model_len=2048, + max_num_batched_tokens=1024, + ) + + assert results[0]["text"] == results[1]["text"], ( + "Repeated identical image+prompt produced different output — the MM " + "encoder/prefix cache changed the result" + ) diff --git a/tests/v1/executor/test_spyre_executor.py b/tests/v1/executor/test_spyre_executor.py index 589c07f6a..fba58e3da 100644 --- a/tests/v1/executor/test_spyre_executor.py +++ b/tests/v1/executor/test_spyre_executor.py @@ -117,7 +117,7 @@ def test_successful_result_triggers_store_and_cleanup(self, executor): """When a result is drained, collective_rpc + cleanup must fire.""" shape = (1, 4, 8) dtype = torch.float16 - _install_queues(executor, result_items=[("req-done", shape, dtype)]) + _install_queues(executor, result_items=[("req-done", shape, dtype, 3, 1)]) executor._mm_in_flight = 1 sched = _make_scheduler_output() @@ -132,16 +132,20 @@ def test_successful_result_triggers_store_and_cleanup(self, executor): rpc_call = executor._parent_collective_rpc.call_args assert rpc_call[0][0] == "store_mm_embeddings" # args is passed as a tuple wrapping the metadata list: args=([...],) + # The cumulative cache counters are stripped before the RPC (3-tuples). assert rpc_call[1]["args"][0] == [("req-done", shape, dtype)] mock_cleanup.assert_called_once_with("req-done") assert sched._spyre_newly_encoded_req_ids == ["req-done"] + # Cumulative MM cache counters are stamped for the scheduler. + assert sched._spyre_mm_cache_hits == 3 + assert sched._spyre_mm_cache_queries == 4 assert executor._mm_in_flight == 0 def test_error_result_sets_failed_req_ids_for_scheduler_retry(self, executor): """(req_id, None, None) must be collected into _spyre_failed_encode_req_ids so the scheduler can clear _mm_encoding_submitted and allow retry.""" - _install_queues(executor, result_items=[("req-err", None, None)]) + _install_queues(executor, result_items=[("req-err", None, None, 0, 0)]) executor._mm_in_flight = 1 sched = _make_scheduler_output() @@ -212,7 +216,7 @@ def test_in_flight_zero_skips_result_drain(self, executor): """When _mm_in_flight == 0, the result queue must not be polled.""" # Even though the queue conceptually has an item, _mm_in_flight==0 # should prevent any get_nowait() call. - _install_queues(executor, result_items=[("req-sneaky", (1, 4, 8), torch.float16)]) + _install_queues(executor, result_items=[("req-sneaky", (1, 4, 8), torch.float16, 0, 0)]) executor._mm_in_flight = 0 sched = _make_scheduler_output() diff --git a/tests/v1/worker/test_mm_encoder_process.py b/tests/v1/worker/test_mm_encoder_process.py index 0db4c3dd0..7a46fdae2 100644 --- a/tests/v1/worker/test_mm_encoder_process.py +++ b/tests/v1/worker/test_mm_encoder_process.py @@ -152,8 +152,12 @@ def test_vision_only_and_fused_weights_always_set(self, tmp_path): class TestVisionEncoderRunnerExecuteModel: - def _make_runner_direct(self): - """Build a VisionEncoderRunner instance bypassing __init__.""" + def _make_runner_direct(self, cache_mb: int = 0): + """Build a VisionEncoderRunner instance bypassing __init__. + + cache_mb=0 disables the encoder cache (every request encodes). + """ + from sendnn_inference.v1.worker.mm_encoder_cache import MMEncoderCache from sendnn_inference.v1.worker.mm_encoder_process import VisionEncoderRunner runner = VisionEncoderRunner.__new__(VisionEncoderRunner) @@ -161,14 +165,16 @@ def _make_runner_direct(self): runner.mm_device = "cpu" runner.mm_utils_cls = MagicMock() runner.fms_model = MagicMock() + runner.mm_encoder_cache = MMEncoderCache(capacity_bytes=cache_mb * 1024 * 1024) return runner def test_output_is_float16_cpu_contiguous(self): """execute_model must cast to _decoder_dtype and return a CPU tensor.""" runner = self._make_runner_direct() - # Simulate vision encoder returning float32 - raw_embeds = torch.ones(1, 8, 16, dtype=torch.float32) - runner.mm_utils_cls.get_maybe_mm_embeddings.return_value = raw_embeds + # Simulate the merge stage returning float32 merged embeddings. + runner.mm_utils_cls.merge_embeddings.return_value = torch.ones( + 1, 8, 16, dtype=torch.float32 + ) job = _make_mm_encode_request() result = runner.execute_model(job) @@ -178,9 +184,9 @@ def test_output_is_float16_cpu_contiguous(self): assert result.is_contiguous() def test_input_ids_built_from_prompt_token_ids(self): - """execute_model must pass the job's prompt_token_ids as input_ids.""" + """execute_model must pass the job's prompt_token_ids as input_ids to embed_text.""" runner = self._make_runner_direct() - runner.mm_utils_cls.get_maybe_mm_embeddings.return_value = torch.zeros( + runner.mm_utils_cls.merge_embeddings.return_value = torch.zeros( 1, 4, 8, dtype=torch.float16 ) job = _make_mm_encode_request() @@ -188,8 +194,8 @@ def test_input_ids_built_from_prompt_token_ids(self): runner.execute_model(job) - call_kwargs = runner.mm_utils_cls.get_maybe_mm_embeddings.call_args - input_ids = call_kwargs[0][1] # second positional arg + # embed_text(fms_model, input_ids) — input_ids is the second positional arg. + input_ids = runner.mm_utils_cls.embed_text.call_args[0][1] assert input_ids.shape == (1, 3) assert input_ids.tolist() == [[10, 20, 30]] @@ -253,6 +259,9 @@ def test_job_processed_result_on_queue(self): mock_runner = MagicMock() mock_runner.execute_model.return_value = fake_embeds + # Real ints so the (widened) result tuple stays picklable across the queue. + mock_runner.mm_encoder_cache.hits = 0 + mock_runner.mm_encoder_cache.misses = 1 fake_shm = MagicMock() with ( @@ -268,10 +277,11 @@ def test_job_processed_result_on_queue(self): encoder_process_main(_make_vllm_config(), jq, rq, stop) assert rq.get(timeout=2) == "READY" - req_id, shape, dtype = rq.get(timeout=2) + req_id, shape, dtype, hits, misses = rq.get(timeout=2) assert req_id == "req-job" assert shape == tuple(fake_embeds.shape) assert dtype == fake_embeds.dtype + assert (hits, misses) == (0, 1) def test_encode_failure_puts_none_metadata(self): """When execute_model raises, (req_id, None, None) must be put on result queue.""" @@ -287,6 +297,8 @@ def test_encode_failure_puts_none_metadata(self): mock_runner = MagicMock() mock_runner.execute_model.side_effect = RuntimeError("encode error") + mock_runner.mm_encoder_cache.hits = 0 + mock_runner.mm_encoder_cache.misses = 0 with patch( "sendnn_inference.v1.worker.mm_encoder_process.VisionEncoderRunner", @@ -295,7 +307,7 @@ def test_encode_failure_puts_none_metadata(self): encoder_process_main(_make_vllm_config(), jq, rq, stop) assert rq.get(timeout=2) == "READY" - req_id, shape, dtype = rq.get(timeout=2) + req_id, shape, dtype, hits, misses = rq.get(timeout=2) assert req_id == "req-fail" assert shape is None assert dtype is None @@ -316,6 +328,8 @@ def test_cancel_queue_skips_job_before_encode(self): jq.put(None) mock_runner = MagicMock() + mock_runner.mm_encoder_cache.hits = 0 + mock_runner.mm_encoder_cache.misses = 0 with ( patch( @@ -327,7 +341,7 @@ def test_cancel_queue_skips_job_before_encode(self): encoder_process_main(_make_vllm_config(), jq, rq, stop, cq) assert rq.get(timeout=2) == "READY" - assert rq.get(timeout=2) == ("req-cancel", None, None) + assert rq.get(timeout=2) == ("req-cancel", None, None, 0, 0) assert not mock_runner.execute_model.called def test_resubmitted_request_encodes_after_cancel_consumed(self): @@ -359,6 +373,8 @@ def test_resubmitted_request_encodes_after_cancel_consumed(self): fake_embeds = torch.zeros(1, 4, 8, dtype=torch.float16) mock_runner = MagicMock() mock_runner.execute_model.return_value = fake_embeds + mock_runner.mm_encoder_cache.hits = 0 + mock_runner.mm_encoder_cache.misses = 0 mock_shm = MagicMock() with ( @@ -374,10 +390,10 @@ def test_resubmitted_request_encodes_after_cancel_consumed(self): encoder_process_main(_make_vllm_config(), jq, rq, stop, cq) assert rq.get(timeout=2) == "READY" - # First job: cancelled → abort result - assert rq.get(timeout=2) == ("req-1", None, None) + # First job: cancelled → abort result (with cumulative cache counters) + assert rq.get(timeout=2) == ("req-1", None, None, 0, 0) # Re-request: skip_ids cleared → encoded normally - req_id, shape, dtype = rq.get(timeout=2) + req_id, shape, dtype, hits, misses = rq.get(timeout=2) assert req_id == "req-1" assert shape is not None mock_runner.execute_model.assert_called_once() diff --git a/tests/v1/worker/test_scheduler_tkv_limits.py b/tests/v1/worker/test_scheduler_tkv_limits.py index 8dd05feaf..51471ad8f 100644 --- a/tests/v1/worker/test_scheduler_tkv_limits.py +++ b/tests/v1/worker/test_scheduler_tkv_limits.py @@ -187,47 +187,6 @@ def test_scheduler_tkv_limits_ongoing_batch(monkeypatch: pytest.MonkeyPatch): break -@pytest.mark.cpu -@pytest.mark.chunked_prefill -def test_chunked_prefill_make_stats_zeros_mm_cache_hits( - monkeypatch: pytest.MonkeyPatch, -): - """ - Regression test: Spyre forces MM cache hit reporting to zero in make_stats(). - - Spyre does not support cross-request MM cache reuse today. This test - verifies that ChunkedPrefillSpyreScheduler.make_stats() forces - mm_cache_stats.hits to zero while still applying the existing - prefix-cache hit correction. - """ - model_runner = InstrumentedModelRunner.build( - monkeypatch=monkeypatch, - max_num_batched_tokens=512, - max_num_seqs=32, - max_model_len=32768, - available_blocks=32768, - ) - scheduler = model_runner.scheduler - - fake_stats = SimpleNamespace( - prefix_cache_stats=SimpleNamespace(queries=256, hits=128), - mm_cache_stats=SimpleNamespace(hits=5), - kv_connector_stats=None, - ) - - monkeypatch.setattr( - SpyreScheduler, - "make_stats", - lambda self, *args, **kwargs: fake_stats, - ) - - stats = scheduler.make_stats() - - assert stats is fake_stats - assert stats.mm_cache_stats.hits == 0 - assert stats.prefix_cache_stats.hits == scheduler.adjust_hit(256, 128) - - @pytest.mark.cpu @pytest.mark.chunked_prefill def test_chunked_prefill_make_stats_without_mm_cache_stats( @@ -236,9 +195,9 @@ def test_chunked_prefill_make_stats_without_mm_cache_stats( """ Regression test: make_stats() handles stats objects without mm_cache_stats. - This verifies that the defensive getattr() guard avoids attribute errors - when the returned stats object does not expose mm_cache_stats, while the - existing prefix-cache correction still applies. + make_stats() no longer reads or writes mm_cache_stats, so a base stats object + that omits the attribute must not raise, while the prefix-cache correction and + the dedicated sendnn vision-encoder metric still apply. """ model_runner = InstrumentedModelRunner.build( monkeypatch=monkeypatch, From af6527e0d4f7a6c923df69776cf4967d608e781c Mon Sep 17 00:00:00 2001 From: nikheal2 Date: Mon, 27 Jul 2026 15:06:07 +0530 Subject: [PATCH 2/7] fix pre-commit Signed-off-by: nikheal2 --- sendnn_inference/v1/worker/mm_encoder_process.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/sendnn_inference/v1/worker/mm_encoder_process.py b/sendnn_inference/v1/worker/mm_encoder_process.py index 155312832..5479e6812 100644 --- a/sendnn_inference/v1/worker/mm_encoder_process.py +++ b/sendnn_inference/v1/worker/mm_encoder_process.py @@ -376,9 +376,7 @@ def encoder_process_main( # Cumulative encoder-cache counters ride along so the scheduler can # surface the real MM cache hit rate (see scheduler.make_stats). cache = runner.mm_encoder_cache - result_queue.put( - (req_id, tuple(embeds.shape), embeds.dtype, cache.hits, cache.misses) - ) + result_queue.put((req_id, tuple(embeds.shape), embeds.dtype, cache.hits, cache.misses)) # Tombstone: a late cancel may still arrive on cancel_queue for this req_id. processed_ids.add(req_id) logger.info("maybe_mm_embedding processing time: %.2fms", t_elapsed * 1000) From f7aab4f15636c8f2287ee79f3bee414fabed080d Mon Sep 17 00:00:00 2001 From: nikheal2 Date: Tue, 4 Aug 2026 23:52:55 +0530 Subject: [PATCH 3/7] removed dup code Signed-off-by: nikheal2 --- .../multimodal/mm_mappings/llava_next.py | 67 ++++++------- .../multimodal/mm_mappings/mistral3.py | 96 ++++++++----------- sendnn_inference/v1/core/scheduler.py | 18 ++-- .../v1/worker/mm_encoder_cache.py | 22 ++--- .../v1/worker/mm_encoder_process.py | 12 +-- .../v1/worker/spyre_model_runner.py | 6 +- tests/v1/worker/test_scheduler_tkv_limits.py | 47 ++++++++- 7 files changed, 145 insertions(+), 123 deletions(-) diff --git a/sendnn_inference/multimodal/mm_mappings/llava_next.py b/sendnn_inference/multimodal/mm_mappings/llava_next.py index 6a4c46713..e4e0d91b4 100644 --- a/sendnn_inference/multimodal/mm_mappings/llava_next.py +++ b/sendnn_inference/multimodal/mm_mappings/llava_next.py @@ -94,41 +94,17 @@ 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( @@ -137,13 +113,16 @@ def get_maybe_mm_embeddings( 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].""" + 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 @@ -159,7 +138,19 @@ def encode_images( image_sizes = mm_spec["image_sizes"].data if image_sizes.ndim == 1: image_sizes = image_sizes.unsqueeze(0) + return pixel_values, image_sizes + @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) return fms_model.pack_image_features( image_features, image_sizes, image_newline=fms_model.image_newline diff --git a/sendnn_inference/multimodal/mm_mappings/mistral3.py b/sendnn_inference/multimodal/mm_mappings/mistral3.py index fb5cf26ac..3d7f49aa6 100644 --- a/sendnn_inference/multimodal/mm_mappings/mistral3.py +++ b/sendnn_inference/multimodal/mm_mappings/mistral3.py @@ -55,53 +55,14 @@ 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( @@ -110,39 +71,64 @@ def get_maybe_mm_embeddings( 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].""" + 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 - # As in get_maybe_mm_embeddings: mistral tokenizer emits "images" not "pixel_values". + # 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 + @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) @staticmethod diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index ed5bbc232..3d7f49dc5 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -249,10 +249,10 @@ def __init__(self, *args, **kwargs) -> None: self.pause_events = 0 self.resume_events = 0 - # Cross-request MM encoder-cache stats. The worker/encoder-subprocess + # Cross-request vision-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. + # delta into the dedicated sendnn ``vision_encoder_cache`` metric. self._mm_async_cum_hits = 0 self._mm_async_cum_queries = 0 self._mm_inline_cum_hits = 0 @@ -992,11 +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. - 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). + Spyre forces upstream ``mm_cache_stats`` (vLLM's multimodal processor/input + cache) hit reporting to 0.0% — pre-existing Spyre behavior, unrelated to our + cache. Our cross-request vision-encoder cache is reported *separately* as a + dedicated sendnn ``vision_encoder_cache`` metric, from cumulative counters + plumbed back via update_from_output. """ base_stats = super().make_stats(*args, **kwargs) @@ -1006,6 +1006,10 @@ 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. diff --git a/sendnn_inference/v1/worker/mm_encoder_cache.py b/sendnn_inference/v1/worker/mm_encoder_cache.py index 5f36ebddd..31cb0c738 100644 --- a/sendnn_inference/v1/worker/mm_encoder_cache.py +++ b/sendnn_inference/v1/worker/mm_encoder_cache.py @@ -30,22 +30,22 @@ _WARMUP_IDENTIFIER_PREFIX = "MM-warmup" -def placeholder_slices(mm_features: Any) -> list[tuple[str, int, int]]: - """Return ``(identifier, offset, length)`` for each cacheable image. +def placeholder_slices(mm_features: Any) -> list[str]: + """Return the mm_hash ``identifier`` of each cacheable image in the request. - ``offset``/``length`` come from each feature's ``mm_position``. ``length`` is - the reserved placeholder span, which may differ from the packed feature-row - count, so it must NOT be used to validate a cache entry — the ``identifier`` - (mm_hash) alone keys the cache. Only ``identifier`` is used by callers today. + The identifier (mm_hash) alone keys the cache. Features without an + ``mm_position`` (not a real image placeholder) and warmup/non-cacheable + identifiers are skipped (see ``MMEncoderCache.is_cacheable``). """ - slices: list[tuple[str, int, int]] = [] + identifiers: list[str] = [] for feat in mm_features or []: - position = getattr(feat, "mm_position", None) identifier = getattr(feat, "identifier", None) - if position is None or not MMEncoderCache.is_cacheable(identifier): + if getattr(feat, "mm_position", None) is None or not MMEncoderCache.is_cacheable( + identifier + ): continue - slices.append((identifier, position.offset, position.length)) - return slices + identifiers.append(identifier) + return identifiers class MMEncoderCache: diff --git a/sendnn_inference/v1/worker/mm_encoder_process.py b/sendnn_inference/v1/worker/mm_encoder_process.py index 5479e6812..2ce2bb1ea 100644 --- a/sendnn_inference/v1/worker/mm_encoder_process.py +++ b/sendnn_inference/v1/worker/mm_encoder_process.py @@ -159,9 +159,9 @@ def _get_or_encode_image_features(self, request) -> torch.Tensor: are never cached (handled by ``MMEncoderCache.is_cacheable``). """ mm_features = request.mm_features - slices = placeholder_slices(mm_features) - cacheable = self.mm_encoder_cache.enabled and len(slices) == 1 == len(mm_features) - identifier = slices[0][0] if cacheable else None + identifiers = placeholder_slices(mm_features) + cacheable = self.mm_encoder_cache.enabled and len(identifiers) == 1 == len(mm_features) + identifier = identifiers[0] if cacheable else None if cacheable: # The mm_hash identifier (sha256 of the image) uniquely identifies the @@ -179,18 +179,18 @@ def _get_or_encode_image_features(self, request) -> torch.Tensor: ) return cached - image_features = self.mm_utils_cls.encode_images( + image_embed = self.mm_utils_cls.encode_images( self.fms_model, mm_features, self.mm_device ) if cacheable: - self.mm_encoder_cache.put(identifier, image_features) + self.mm_encoder_cache.put(identifier, image_embed) self.mm_encoder_cache.record_lookup(hit=False) logger.info( "encoder_process: vision-encoder-cache hit rate: %.1f%% (req '%s', MISS)", self._hit_rate() * 100, request.request_id, ) - return image_features + return image_embed def _hit_rate(self) -> float: cache = self.mm_encoder_cache diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 0958cd50d..3dba0512d 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1120,9 +1120,9 @@ def _encode_and_merge_cached( ``torch.inference_mode`` (merge updates the text-embedding tensor in place). Equivalent to ``get_maybe_mm_embeddings(mm_features, is_decode=False)``. """ - slices = placeholder_slices(mm_features) - cacheable = self.mm_encoder_cache.enabled and len(slices) == 1 == len(mm_features) - identifier = slices[0][0] if cacheable else None + identifiers = placeholder_slices(mm_features) + cacheable = self.mm_encoder_cache.enabled and len(identifiers) == 1 == len(mm_features) + identifier = identifiers[0] if cacheable else None image_features = None if cacheable: diff --git a/tests/v1/worker/test_scheduler_tkv_limits.py b/tests/v1/worker/test_scheduler_tkv_limits.py index 51471ad8f..325231436 100644 --- a/tests/v1/worker/test_scheduler_tkv_limits.py +++ b/tests/v1/worker/test_scheduler_tkv_limits.py @@ -187,6 +187,47 @@ def test_scheduler_tkv_limits_ongoing_batch(monkeypatch: pytest.MonkeyPatch): break +@pytest.mark.cpu +@pytest.mark.chunked_prefill +def test_chunked_prefill_make_stats_zeros_mm_cache_hits( + monkeypatch: pytest.MonkeyPatch, +): + """ + Regression test: Spyre forces MM cache hit reporting to zero in make_stats(). + + Spyre does not support cross-request MM cache reuse today. This test + verifies that ChunkedPrefillSpyreScheduler.make_stats() forces + mm_cache_stats.hits to zero while still applying the existing + prefix-cache hit correction. + """ + model_runner = InstrumentedModelRunner.build( + monkeypatch=monkeypatch, + max_num_batched_tokens=512, + max_num_seqs=32, + max_model_len=32768, + available_blocks=32768, + ) + scheduler = model_runner.scheduler + + fake_stats = SimpleNamespace( + prefix_cache_stats=SimpleNamespace(queries=256, hits=128), + mm_cache_stats=SimpleNamespace(hits=5), + kv_connector_stats=None, + ) + + monkeypatch.setattr( + SpyreScheduler, + "make_stats", + lambda self, *args, **kwargs: fake_stats, + ) + + stats = scheduler.make_stats() + + assert stats is fake_stats + assert stats.mm_cache_stats.hits == 0 + assert stats.prefix_cache_stats.hits == scheduler.adjust_hit(256, 128) + + @pytest.mark.cpu @pytest.mark.chunked_prefill def test_chunked_prefill_make_stats_without_mm_cache_stats( @@ -195,9 +236,9 @@ def test_chunked_prefill_make_stats_without_mm_cache_stats( """ Regression test: make_stats() handles stats objects without mm_cache_stats. - make_stats() no longer reads or writes mm_cache_stats, so a base stats object - that omits the attribute must not raise, while the prefix-cache correction and - the dedicated sendnn vision-encoder metric still apply. + The mm_cache_stats override uses a defensive getattr() guard, so a base stats + object that omits the attribute must not raise, while the prefix-cache + correction still applies. """ model_runner = InstrumentedModelRunner.build( monkeypatch=monkeypatch, From 08f495d04c0c5221692eb05394caa70a1629179b Mon Sep 17 00:00:00 2001 From: nikheal2 Date: Tue, 4 Aug 2026 23:55:59 +0530 Subject: [PATCH 4/7] removed un-necesary code changes Signed-off-by: nikheal2 --- tests/v1/worker/test_scheduler_tkv_limits.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/v1/worker/test_scheduler_tkv_limits.py b/tests/v1/worker/test_scheduler_tkv_limits.py index 325231436..8dd05feaf 100644 --- a/tests/v1/worker/test_scheduler_tkv_limits.py +++ b/tests/v1/worker/test_scheduler_tkv_limits.py @@ -236,9 +236,9 @@ def test_chunked_prefill_make_stats_without_mm_cache_stats( """ Regression test: make_stats() handles stats objects without mm_cache_stats. - The mm_cache_stats override uses a defensive getattr() guard, so a base stats - object that omits the attribute must not raise, while the prefix-cache - correction still applies. + This verifies that the defensive getattr() guard avoids attribute errors + when the returned stats object does not expose mm_cache_stats, while the + existing prefix-cache correction still applies. """ model_runner = InstrumentedModelRunner.build( monkeypatch=monkeypatch, From 7524c3d2d0b7c4e61dfab626298d1fb36e56b55b Mon Sep 17 00:00:00 2001 From: nikheal2 Date: Wed, 5 Aug 2026 00:22:40 +0530 Subject: [PATCH 5/7] removed debug lines Signed-off-by: nikheal2 --- sendnn_inference/v1/worker/mm_encoder_process.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/sendnn_inference/v1/worker/mm_encoder_process.py b/sendnn_inference/v1/worker/mm_encoder_process.py index 2ce2bb1ea..9c444e9f3 100644 --- a/sendnn_inference/v1/worker/mm_encoder_process.py +++ b/sendnn_inference/v1/worker/mm_encoder_process.py @@ -164,10 +164,8 @@ def _get_or_encode_image_features(self, request) -> torch.Tensor: identifier = identifiers[0] if cacheable else None if cacheable: - # The mm_hash identifier (sha256 of the image) uniquely identifies the + # The mm_hash identifier (sßha256 of the image) uniquely identifies the # image, so a stored entry is exactly what encode_images would produce. - # (Do NOT gate on mm_position.length: the reserved placeholder span can - # differ from the packed feature-row count, e.g. 704 vs 682.) cached = self.mm_encoder_cache.get(identifier) if cached is not None: self.mm_encoder_cache.record_lookup(hit=True) From 823bf0ef130bc54233417610aae40df44984e8ad Mon Sep 17 00:00:00 2001 From: nikheal2 Date: Wed, 5 Aug 2026 23:10:11 +0530 Subject: [PATCH 6/7] renamed function Signed-off-by: nikheal2 --- sendnn_inference/v1/worker/mm_encoder_cache.py | 4 ++-- sendnn_inference/v1/worker/mm_encoder_process.py | 6 +++--- sendnn_inference/v1/worker/spyre_model_runner.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/sendnn_inference/v1/worker/mm_encoder_cache.py b/sendnn_inference/v1/worker/mm_encoder_cache.py index 31cb0c738..f1c7e2091 100644 --- a/sendnn_inference/v1/worker/mm_encoder_cache.py +++ b/sendnn_inference/v1/worker/mm_encoder_cache.py @@ -30,7 +30,7 @@ _WARMUP_IDENTIFIER_PREFIX = "MM-warmup" -def placeholder_slices(mm_features: Any) -> list[str]: +def cacheable_identifiers(mm_features: Any) -> list[str]: """Return the mm_hash ``identifier`` of each cacheable image in the request. The identifier (mm_hash) alone keys the cache. Features without an @@ -84,7 +84,7 @@ def put(self, identifier: str, tensor: torch.Tensor) -> None: nbytes = tensor.numel() * tensor.element_size() # A single entry larger than the whole budget is simply not cached. over = nbytes > self.capacity_bytes - logger.info( + logger.debug( "MM encoder cache: entry '%s' size=%.2f MiB, budget=%.2f MiB " "(SENDNN_INFERENCE_MM_ENCODER_CACHE_MB) — %s", identifier, diff --git a/sendnn_inference/v1/worker/mm_encoder_process.py b/sendnn_inference/v1/worker/mm_encoder_process.py index 9c444e9f3..941c74ede 100644 --- a/sendnn_inference/v1/worker/mm_encoder_process.py +++ b/sendnn_inference/v1/worker/mm_encoder_process.py @@ -25,7 +25,7 @@ import sendnn_inference.envs as envs_spyre from sendnn_inference.model_executor.model_loader.spyre import SpyreCausalLM, cast_params_for_spyre from sendnn_inference.platform import SpyrePlatform, THREADING_ENVS -from sendnn_inference.v1.worker.mm_encoder_cache import MMEncoderCache, placeholder_slices +from sendnn_inference.v1.worker.mm_encoder_cache import MMEncoderCache, cacheable_identifiers from sendnn_inference.v1.worker.mm_shared_memory import write_embeddings logger = init_logger(__name__) @@ -159,12 +159,12 @@ def _get_or_encode_image_features(self, request) -> torch.Tensor: are never cached (handled by ``MMEncoderCache.is_cacheable``). """ mm_features = request.mm_features - identifiers = placeholder_slices(mm_features) + identifiers = cacheable_identifiers(mm_features) cacheable = self.mm_encoder_cache.enabled and len(identifiers) == 1 == len(mm_features) identifier = identifiers[0] if cacheable else None if cacheable: - # The mm_hash identifier (sßha256 of the image) uniquely identifies the + # The mm_hash identifier (sha256 of the image) uniquely identifies the # image, so a stored entry is exactly what encode_images would produce. cached = self.mm_encoder_cache.get(identifier) if cached is not None: diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 3dba0512d..cc4dd1796 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -44,7 +44,7 @@ from sendnn_inference.v1.sample.spyre_logits_processor import build_logitsprocs_for_cb from sendnn_inference.v1.worker.mm_encoder_cache import ( MMEncoderCache, - placeholder_slices, + cacheable_identifiers, ) from sendnn_inference.v1.worker.mm_shared_memory import ( cleanup_embeddings, @@ -1120,7 +1120,7 @@ def _encode_and_merge_cached( ``torch.inference_mode`` (merge updates the text-embedding tensor in place). Equivalent to ``get_maybe_mm_embeddings(mm_features, is_decode=False)``. """ - identifiers = placeholder_slices(mm_features) + identifiers = cacheable_identifiers(mm_features) cacheable = self.mm_encoder_cache.enabled and len(identifiers) == 1 == len(mm_features) identifier = identifiers[0] if cacheable else None From 8407907df1ccc957d3acb416b2edd7afe69bac87 Mon Sep 17 00:00:00 2001 From: nikheal2 Date: Wed, 5 Aug 2026 23:36:31 +0530 Subject: [PATCH 7/7] fix pre-commit Signed-off-by: nikheal2 --- .../multimodal/mm_mappings/llava_next.py | 12 +++++------- sendnn_inference/multimodal/mm_mappings/mistral3.py | 12 +++++------- sendnn_inference/v1/worker/mm_encoder_cache.py | 2 +- sendnn_inference/v1/worker/mm_encoder_process.py | 8 +++----- sendnn_inference/v1/worker/spyre_model_runner.py | 4 ++-- 5 files changed, 16 insertions(+), 22 deletions(-) diff --git a/sendnn_inference/multimodal/mm_mappings/llava_next.py b/sendnn_inference/multimodal/mm_mappings/llava_next.py index e4e0d91b4..b0f5213b2 100644 --- a/sendnn_inference/multimodal/mm_mappings/llava_next.py +++ b/sendnn_inference/multimodal/mm_mappings/llava_next.py @@ -138,7 +138,7 @@ def _prepare_vision_inputs( image_sizes = mm_spec["image_sizes"].data if image_sizes.ndim == 1: image_sizes = image_sizes.unsqueeze(0) - return pixel_values, image_sizes + return pixel_values, image_sizes # ty: ignore[invalid-return-type] @staticmethod def encode_images( @@ -148,17 +148,15 @@ def encode_images( ) -> 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) - return fms_model.pack_image_features( + 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) + return fms_model._get_text_embeddings(input_ids) # ty: ignore[call-non-callable] @staticmethod def merge_embeddings( diff --git a/sendnn_inference/multimodal/mm_mappings/mistral3.py b/sendnn_inference/multimodal/mm_mappings/mistral3.py index 3d7f49aa6..a66ae6a9c 100644 --- a/sendnn_inference/multimodal/mm_mappings/mistral3.py +++ b/sendnn_inference/multimodal/mm_mappings/mistral3.py @@ -116,7 +116,7 @@ def _prepare_vision_inputs( # 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 + return pixel_values, image_sizes # ty: ignore[invalid-return-type] @staticmethod def encode_images( @@ -126,14 +126,12 @@ def encode_images( ) -> 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) + 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) + return fms_model._get_text_embeddings(input_ids, None) # ty: ignore[call-non-callable] @staticmethod def merge_embeddings( @@ -142,7 +140,7 @@ def merge_embeddings( text_embeds: torch.Tensor, image_features: torch.Tensor, ) -> torch.Tensor: - return fms_model._merge_multimodal_embeddings( + return fms_model._merge_multimodal_embeddings( # ty: ignore[call-non-callable] input_ids, text_embeds, image_features, diff --git a/sendnn_inference/v1/worker/mm_encoder_cache.py b/sendnn_inference/v1/worker/mm_encoder_cache.py index f1c7e2091..dc17b3eae 100644 --- a/sendnn_inference/v1/worker/mm_encoder_cache.py +++ b/sendnn_inference/v1/worker/mm_encoder_cache.py @@ -44,7 +44,7 @@ def cacheable_identifiers(mm_features: Any) -> list[str]: identifier ): continue - identifiers.append(identifier) + identifiers.append(identifier) # ty: ignore[invalid-argument-type] return identifiers diff --git a/sendnn_inference/v1/worker/mm_encoder_process.py b/sendnn_inference/v1/worker/mm_encoder_process.py index 941c74ede..61bbef29c 100644 --- a/sendnn_inference/v1/worker/mm_encoder_process.py +++ b/sendnn_inference/v1/worker/mm_encoder_process.py @@ -166,7 +166,7 @@ def _get_or_encode_image_features(self, request) -> torch.Tensor: if cacheable: # The mm_hash identifier (sha256 of the image) uniquely identifies the # image, so a stored entry is exactly what encode_images would produce. - cached = self.mm_encoder_cache.get(identifier) + cached = self.mm_encoder_cache.get(identifier) # ty: ignore[invalid-argument-type] if cached is not None: self.mm_encoder_cache.record_lookup(hit=True) logger.info( @@ -177,11 +177,9 @@ def _get_or_encode_image_features(self, request) -> torch.Tensor: ) return cached - image_embed = self.mm_utils_cls.encode_images( - self.fms_model, mm_features, self.mm_device - ) + image_embed = self.mm_utils_cls.encode_images(self.fms_model, mm_features, self.mm_device) if cacheable: - self.mm_encoder_cache.put(identifier, image_embed) + self.mm_encoder_cache.put(identifier, image_embed) # ty: ignore[invalid-argument-type] self.mm_encoder_cache.record_lookup(hit=False) logger.info( "encoder_process: vision-encoder-cache hit rate: %.1f%% (req '%s', MISS)", diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index cc4dd1796..7f56852a6 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1129,7 +1129,7 @@ def _encode_and_merge_cached( # Keyed by the mm_hash identifier only; do NOT gate on # mm_position.length (the reserved placeholder span can differ from the # packed feature-row count, so a length check rejects valid entries). - cached = self.mm_encoder_cache.get(identifier) + cached = self.mm_encoder_cache.get(identifier) # ty: ignore[invalid-argument-type] if cached is not None: image_features = cached self.mm_encoder_cache.record_lookup(hit=True) @@ -1142,7 +1142,7 @@ def _encode_and_merge_cached( if image_features is None: image_features = self.model.encode_images(mm_features) if cacheable: - self.mm_encoder_cache.put(identifier, image_features) + self.mm_encoder_cache.put(identifier, image_features) # ty: ignore[invalid-argument-type] self.mm_encoder_cache.record_lookup(hit=False) cache = self.mm_encoder_cache logger.info(