Separate mm encoder scheduling

.github/ci_model_cache.yaml

@@ -2,7 +2,7 @@
 #
 # Adding/removing/changing an entry here changes the cache key and will trigger
 # a fresh download + cache save on the next push to main. Total uncompressed
-# size of all entries must stay under 8 GiB (GHA cache limit is 10 GiB).
+# size of all entries must stay under 10 GiB (GHA cache hard limit).
 
 models:
   - repo: ibm-ai-platform/micro-g3.3-8b-instruct-1b
@@ -13,3 +13,8 @@ models:
     revision: cf74d8acd4f198de950bf004b262e6accfed5d2c
   - repo: cross-encoder/stsb-roberta-large
     revision: 2b12c2c0088918e76151fd5937b7bba986ef1f98
+  # Random-init CI fixture (~10 MB) that stands in for granite-vision-3.2-2b
+  # in the async-MM-encoder e2e tests. See tests/e2e/conftest.py — the
+  # revision must match NANO_GV_REVISION there.
+  - repo: joerunde/nano-gv
+    revision: c9470d9e54b023dd9ab8a8a98057489fdb18ba03

docs/contributing/parallel_vision_encode_plan.md

@@ -0,0 +1,107 @@
+# Parallel Vision Encoder Execution In Single Instance vLLM
+
+## Background
+
+Multimodal models on Spyre compute the vision encoder on CPU (rank 0 only) and broadcast embeddings to other ranks via POSIX shared memory. Today this encoding runs serially, once per request, at the start of that request's first prefill step. MM encoding is expensive operation and in current implementation its blocking, so no other operation like prefill and decode of other requests can run in parallel affecting overall performance.
+
+## Goal
+
+Overlap CPU / NNPA vision encoding with AIU prefill/decode by running the encoder in a separate subprocess. Embeddings are written to POSIX shared memory and all TP workers read them independently — no rank-0 broadcast of large tensors. The scheduler gates MM request prefill on encoding readiness, so a request only enters prefill once its embedding is available.
+
+## Evolution Path
+
+**Phase 1 and 2:** Combined in current implementation. Vision encoding runs in a dedicated non-daemon subprocess (`mm-encoder`) managed by `SpyreMultiprocExecutor`. The encoder subprocess loads only the vision model via `get_model(..., vision_only=True)`.
+The scheduler submits MM requests for encoding on every step and gates prefill on encoding.
+
+**Phase 3 (future):** Enable vision encoder batching within the encoder subprocess. This will further improve the performance by handling all pending MM requests in single batch. This requires FMS changes to stack same-resolution images instead of
+concatenating. See [Phase 3](#phase-3-add-vision-encoder-batching) below.
+
+### Current Flow
+
+```text
+Scheduler picks 1 MM request
+  → execute_model():
+      encode(request)          # CPU, rank 0, single request
+      broadcast embeddings     # SHM + dist.broadcast
+      prefill chunk on Spyre
+      (repeat for each chunk)
+  → decode steps
+```
+
+### Implemented Flow
+
+```text
+Encoder subprocess starts AFTER warmup completes
+  (SpyreMultiprocExecutor hooks on collective_rpc("compile_or_warm_up_model"))
+
+Scheduler emits unsubmitted waiting MM requests on EVERY step (prefill and decode).
+Scheduler gates MM prefill on _mm_encoding_ready
+  (only applies when SENDNN_INFERENCE_ASYNC_MM_ENCODER=1).
+
+SpyreMultiprocExecutor.execute_model() on every step:
+  1. Submit new _spyre_mm_encode_requests → job_queue   # non-blocking put_nowait
+  2. Drain result_queue (non-blocking)                   # collect completed encodings
+     if results:
+       collective_rpc("store_mm_embeddings")             # all TP workers read SHM
+       cleanup SHM blocks
+       set scheduler_output._spyre_newly_encoded_req_ids
+  3. super().execute_model() → workers run AIU forward   # concurrent with encoder subprocess encoder process runs in parallel
+
+scheduler.update_from_output():
+  _mm_encoding_ready.update(_spyre_newly_encoded_req_ids)
+
+Next schedule() call: request now in _mm_encoding_ready → scheduled for prefill
+  add_new_request(): cached_mm_embeddings = pending_mm_embeddings.pop(req_id)
+  _prepare_chunked_prefill(): uses cached embeddings, skips inline encoding
+```
+
+---
+
+## Changes Summary
+
+| File | Change |
+|---|---|
+| `sendnn_inference/platform.py` | Register `SpyreMultiprocExecutor` when `SENDNN_INFERENCE_ASYNC_MM_ENCODER=1` and TP > 1 |
+| `sendnn_inference/v1/executor/spyre_executor.py` | `SpyreMultiprocExecutor`: override `execute_model` to submit encode jobs, collect results, call `store_mm_embeddings` on workers |
+| `sendnn_inference/v1/worker/mm_encoder_process.py` | `VisionEncoderRunner` + `encoder_process_main`: load vision-only model, serve encode jobs, write embeddings to SHM |
+| `sendnn_inference/v1/worker/spyre_worker.py` | Add `store_mm_embeddings` — delegates to model runner |
+| `sendnn_inference/v1/worker/spyre_model_runner.py` | Add `pending_mm_embeddings` dict, `store_mm_embeddings` (reads from SHM), `_compute_and_cache_mm_embeddings` as inline fallback for warmup; consume in `add_new_request` |
+| `sendnn_inference/v1/core/scheduler.py` | Add `MMEncodeRequest` dataclass; emit encode jobs every step; track `_mm_encoding_submitted` / `_mm_encoding_ready`; gate MM prefill on encoding readiness (async mode only); update state in `update_from_output` and `finish_requests` |
+| `sendnn_inference/model_executor/model_loader/spyre.py` | Extract `cast_params_for_spyre` as module-level function reusable by encoder subprocess |
+| `sendnn_inference/envs.py` | Add `SENDNN_INFERENCE_ASYNC_MM_ENCODER` env var (default 0) |
+
+Non-MM requests, the warmup path, chunked prefill logic, and TP broadcast are unaffected.
+
+## Alternatives considered
+
+### Threading (abandoned)
+
+**What we tried:** Start a `threading.Thread` in the worker model runner. The thread uses the already-loaded `fms_model` directly (no copy) and encodes waiting MM requests in the background while the AIU runs.
+
+**Why it failed:** Spyre operations and vision encoding both are blocking operations. The background thread cannot make any progress during AIU execution. Encoding only runs in tiny Python gaps between AIU calls and prefill / decode gets impacted by encoding operations.
+
+**Verdict:** No benefit. Reverted.
+
+### Subprocess from worker (abandoned)
+
+**What we tried:** Start a `multiprocessing.Process` from inside the worker's `load_model` or
+`complete_warmup`.
+
+**Why it failed:** vLLM spawns worker processes as **daemon processes**
+(`multiprocessing.Process(daemon=True)`). Python forbids daemon processes from spawning children (`AssertionError: daemonic processes are not allowed to have children`).
+
+**Verdict:** Architecturally impossible from a worker process.
+
+### Subprocess from MultiprocExecutor (**implemented**)
+
+**The idea:** vLLM's `MultiprocExecutor` runs in the **main (non-daemon) process**. Any process it spawns is also non-daemon. By subclassing `MultiprocExecutor` as `SpyreMultiprocExecutor`, we can start the encoder process at the executor level.
+
+**Model weight loading:** FMS now supports `get_model(..., vision_only=True)`, which loads only vision tower + projector + text embedding from the checkpoint, skipping the LLM decoder. The encoder subprocess calls this directly.
+
+**SHM-based result delivery:** The encoder process writes completed embeddings to POSIX SHM and puts only `(req_id, shape, dtype)` metadata on the result queue (no large tensors in the queue). The executor calls `collective_rpc("store_mm_embeddings", metadata)` so all TP workers read from SHM independently — no rank-0 to others tensor broadcast.
+
+**Scheduler-level encoding readiness gate:** The scheduler tracks `_mm_encoding_submitted` and `_mm_encoding_ready` sets. MM requests are only eligible for prefill when their encoding is confirmed complete. Text-only requests are completely unaffected. The scheduler submits encoding jobs on every step (prefill AND decode) so the encoder stays ahead of the prefill queue.
+
+## Phase 3: Add Vision Encoder Batching
+
+For N same-resolution images, the vision transformer can runs once on `[N, P, D]` instead of N times on `[1, P, D]`. CPU / NNPA matmul efficiently scales with batch size, so the single batched call should be significantly faster than N sequential calls — particularly for large images where the `P²` self-attention dominates.

sendnn_inference/envs.py

@@ -26,6 +26,7 @@
     SENDNN_INFERENCE_MODEL_CONFIG_FILE: str | None = None
     SENDNN_INFERENCE_CPU_MM_DTYPE: torch.dtype = torch.float16
     SENDNN_INFERENCE_MM_DEVICE: str = "auto"
+    SENDNN_INFERENCE_ASYNC_MM_ENCODER: bool = True
     SENDNN_INFERENCE_TP_MM_SHARING: bool = True
     SENDNN_INFERENCE_LONG_OUT_PRIO: bool = False
     SENDNN_INFERENCE_PAUSING_ENABLED: bool = True
@@ -175,6 +176,16 @@ def clear_env_cache():
     "SENDNN_INFERENCE_MM_DEVICE": lambda: parse_mm_device(
         os.getenv("SENDNN_INFERENCE_MM_DEVICE", "auto")
     ),
+    # Enable the async vision encoder subprocess.
+    # When set to 1, SpyreMultiprocExecutor spawns a separate process that loads
+    # only the vision model via get_model(..., vision_only=True) and pre-encodes
+    # MM requests in parallel with AIU prefill/decode.  The scheduler gates MM
+    # request prefill on encoding readiness so a request only starts prefill once
+    # its embedding is available.  Only effective for decoder models with TP > 1.
+    # Defaults to 0 (disabled) — uses the Phase 1 blocking encode path.
+    "SENDNN_INFERENCE_ASYNC_MM_ENCODER": lambda: bool(
+        int(os.getenv("SENDNN_INFERENCE_ASYNC_MM_ENCODER", "1"))
+    ),
     # When "1" (default), rank 0 runs the vision encoder and shares the result
     # with other TP ranks via POSIX shared memory (one encoder call instead of
     # world_size calls).  Set to "0" to fall back to every TP rank running the

sendnn_inference/model_executor/model_loader/spyre.py

@@ -44,6 +44,67 @@ class SpyreAttentionMetadata:
     is_prefill: bool
 
 
+def cast_params_for_spyre(
+    fms_model: nn.Module,
+    mm_parameter_prefixes: tuple[str, ...],
+    is_fp8_model: bool = False,
+) -> str:
+    """Cast model params for Spyre execution and return the resolved mm_device.
+
+    Places MM submodules (vision_tower / multi_modal_projector) onto
+    SENDNN_INFERENCE_MM_DEVICE with SENDNN_INFERENCE_CPU_MM_DTYPE, and casts
+    all other submodules to fp16 for Spyre.  For FP8 models only bf16
+    params/buffers are converted to fp16 — fp8 weights and fp32 scales are
+    left untouched.
+
+    Callable from both SpyreCausalLM._cast_params_for_spyre (full model) and
+    VisionEncoderRunner (vision-only encoder subprocess).
+    """
+    cpu_mm_dtype = envs_spyre.SENDNN_INFERENCE_CPU_MM_DTYPE
+    mm_device = envs_spyre.SENDNN_INFERENCE_MM_DEVICE
+    mm_module_names = {p.rstrip(".") for p in mm_parameter_prefixes}
+
+    if mm_device == "nnpa" and mm_module_names and not utils_spyre.ensure_nnpa_registered():
+        raise RuntimeError(
+            "SENDNN_INFERENCE_MM_DEVICE resolved to nnpa (torch_nnpa is installed) "
+            "but the nnpa device could not be initialized. Refusing to fall back to "
+            "CPU; a broken nnpa backend must not be silently masked. Set "
+            "SENDNN_INFERENCE_MM_DEVICE=cpu to run the vision tower on CPU."
+        )
+
+    # Cast the (non-mm) model to fp16 for Spyre, and place the multimodal
+    # submodules on their configured device/dtype.  The mm submodules are
+    # moved wholesale with Module.to(...) (required because
+    # nn.Parameter.set_data can't swap the CPU->nnpa backend; Module._apply
+    # rebuilds the Parameters, and buffers move too).  Their descendants must
+    # be skipped so the non-mm branch doesn't re-cast them back to fp16
+    # after placement (named_modules yields parents before children).
+    mm_prefixes_tuple = tuple(mm_parameter_prefixes)
+    for module_name, module in fms_model.named_modules():
+        if module_name in mm_module_names:
+            logger.debug(
+                "Placing %s submodule on device=%s dtype=%s.",
+                module_name,
+                mm_device,
+                cpu_mm_dtype,
+            )
+            module.to(device=mm_device, dtype=cpu_mm_dtype)
+        elif module_name.startswith(mm_prefixes_tuple):
+            # Descendant of an mm submodule; already placed with its ancestor.
+            continue
+        elif is_fp8_model:
+            # Per-param cast restricted to bf16: leaves fp8 weights and
+            # fp32 scales alone.  recurse=False so each param is visited
+            # once via the outer named_modules() walk.
+            for param in module.parameters(recurse=False):
+                if param.dtype == torch.bfloat16:
+                    param.data = param.data.to(dtype=torch.float16)
+        else:
+            module.to(dtype=torch.float16)
+
+    return mm_device
+
+
 class SpyreCausalLM(nn.Module):
     def __init__(
         self,
@@ -295,49 +356,10 @@ def _cast_params_for_spyre(self):
         For quantized (e.g. FP8) models we only convert bf16 params/buffers to
         fp16 — fp8 weights and fp32 scales must be left untouched.
         """
-        cpu_mm_dtype = envs_spyre.SENDNN_INFERENCE_CPU_MM_DTYPE
-        mm_device = envs_spyre.SENDNN_INFERENCE_MM_DEVICE
         mm_prefixes = self.mm_model_utils.mm_parameter_prefixes if self.mm_model_utils else ()
-        mm_module_names = {p.rstrip(".") for p in mm_prefixes}
-
-        if mm_device == "nnpa" and mm_module_names and not utils_spyre.ensure_nnpa_registered():
-            raise RuntimeError(
-                "SENDNN_INFERENCE_MM_DEVICE resolved to nnpa (torch_nnpa is installed) "
-                "but the nnpa device could not be initialized. Refusing to fall back to "
-                "CPU; a broken nnpa backend must not be silently masked. Set "
-                "SENDNN_INFERENCE_MM_DEVICE=cpu to run the vision tower on CPU."
-            )
-        self.mm_device = mm_device
-
-        # Cast the (non-mm) model to fp16 for Spyre, and place the multimodal
-        # submodules on their configured device/dtype. The mm submodules are
-        # moved wholesale with Module.to(...) (required because
-        # nn.Parameter.set_data can't swap the CPU->nnpa backend; Module._apply
-        # rebuilds the Parameters, and buffers move too). Their descendants must
-        # be skipped so the non-mm branch doesn't re-cast them back to fp16
-        # after placement (named_modules yields parents before children).
-        mm_prefixes_tuple = tuple(mm_prefixes)
-        for module_name, module in self.fms_model.named_modules():
-            if module_name in mm_module_names:
-                logger.debug(
-                    "Placing %s submodule on device=%s dtype=%s.",
-                    module_name,
-                    mm_device,
-                    cpu_mm_dtype,
-                )
-                module.to(device=mm_device, dtype=cpu_mm_dtype)
-            elif module_name.startswith(mm_prefixes_tuple):
-                # Descendant of an mm submodule; already placed with its ancestor.
-                continue
-            elif self.is_fp8_model:
-                # Per-param cast restricted to bf16: leaves fp8 weights and
-                # fp32 scales alone. recurse=False so each param is visited
-                # once via the outer named_modules() walk.
-                for param in module.parameters(recurse=False):
-                    if param.dtype == torch.bfloat16:
-                        param.data = param.data.to(dtype=torch.float16)
-            else:
-                module.to(dtype=torch.float16)
+        self.mm_device = cast_params_for_spyre(
+            self.fms_model, mm_prefixes, is_fp8_model=self.is_fp8_model
+        )
 
     def _cast_to_f32(self):
         """Cast model parameters to f32."""

sendnn_inference/multimodal/mm_mappings/llava_next.py

@@ -53,13 +53,30 @@ def unwrap_mm_kv_cache_opts(self):
     @staticmethod
     def get_mm_specific_load_overrides(hf_config: PretrainedConfig):
         """Get any overrides needed for initializing the FMS model from the
-        transformers config. For this model, we need to fix the head_dim, which
-        currently surfaces as a problem for all 2b variants of granite 3.x LLMs
-        when running through FMS.
-
-        TODO: If additional variants of granite vision are added, or broader
-        llava next support is added in FMS, handle it properly here.
+        transformers config. For granite-vision-3.2-2b in HF format, the
+        text_config doesn't declare `head_dim` explicitly, and the default
+        fallback (`hidden_size / num_attention_heads` = 2048/32 = 64) is
+        wrong for the FMS Granite implementation, which expects 128. We
+        patch it in here.
+
+        A config that already sets `head_dim` correctly (whether to 128 or
+        any other legal value) needs no override — returning a partial
+        `text_config` dict here with `override_hf_pretrained_config=True`
+        would obliterate every other field FMS built from the HF config
+        (emb_dim, nheads, etc.) because FMS's override merges are
+        whole-key replacements, not deep merges.
         """
+        text_cfg = hf_config.text_config
+        explicit_head_dim = getattr(text_cfg, "head_dim", None)
+        if explicit_head_dim is not None:
+            # HF config knows its own head_dim — trust it.
+            return {}
+
+        implicit_head_dim = text_cfg.hidden_size // text_cfg.num_attention_heads
+        if implicit_head_dim == 128:
+            # Implicit head_dim already matches what FMS wants; no override.
+            return {}
+
         return {
             "override_hf_pretrained_config": True,
             "text_config": {"head_dim": 128},

sendnn_inference/platform.py

@@ -258,6 +258,28 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         if parallel_config.worker_cls == "auto":
             parallel_config.worker_cls = "sendnn_inference.v1.worker.spyre_worker.SpyreWorker"
 
+        # Use SpyreMultiprocExecutor when async MM encoding is enabled via
+        # SENDNN_INFERENCE_ASYNC_MM_ENCODER=1.  The executor manages a separate
+        # vision encoder subprocess that runs in parallel with AIU inference.
+        # Pass the class object directly — Executor.get_class handles
+        # isinstance(backend, type) before string-based dispatch, which avoids
+        # Pydantic's Literal validator silently dropping a string class path.
+        if (
+            is_decoder
+            and model_config.is_multimodal_model
+            and parallel_config.world_size > 1
+            and envs_spyre.SENDNN_INFERENCE_ASYNC_MM_ENCODER
+        ):
+            from sendnn_inference.v1.executor.spyre_executor import SpyreMultiprocExecutor
+
+            parallel_config.distributed_executor_backend = SpyreMultiprocExecutor
+            logger.info(
+                "Using SpyreMultiprocExecutor with async MM encoder subprocess "
+                "(world_size=%d, model=%s)",
+                parallel_config.world_size,
+                model_config.model,
+            )
+
         cls._check_threading_config(parallel_config.world_size)
 
         # set env vars based on the model
@@ -670,7 +692,7 @@ def _check_threading_config(cls, worker_count: int):
 
         # NOTE: math.ceil can output a number for each worker that sums
         # to a total greater than cpu_count.
-        if is_multimodal:
+        if is_multimodal and not envs_spyre.SENDNN_INFERENCE_ASYNC_MM_ENCODER:
             if cpu_count is None:
                 cpus_per_worker = None
             elif platform.machine() == "ppc64le":