qwen3_omni: prefix caching for the Thinker

docs/adding_models.rst

@@ -440,13 +440,18 @@ A model with a keyed label is checked at load time. Loading fails in two more ca
   reuse; either drop the stream or give the resource the prefix hooks``.
 
 Some requests stay uncached even when the model keys a label. The engine checks each
-request's inputs on the keyed walk and cuts a prefix only from inputs that the default
-``split_inputs`` can slice. A request whose inputs carry ``tensor_inputs``, ``kwargs`` or
+request's inputs on the keyed walk and cuts a prefix only from inputs that
+``split_inputs`` can slice. A request whose inputs carry ``kwargs`` or
 ``resource_step_info`` is never probed; BAGEL's guided requests carry ``requires_cfg``
-there. BAGEL keys an understanding prompt's text and images, and a generation's first
-span when it is text: an edit's images are resized before their walks see them, and a
-file's digest doesn't cover the resize. A label under a sliding window files pages only
-up to its protected prefix, which the window never releases.
+there. The default ``split_inputs`` slices only the sequence, so a request whose inputs
+carry ``tensor_inputs`` is probed only when the submodule overrides it to cut those
+tensors. For a label keyed over several walks, where a keyed node streams into a
+partition other than its own, a request that runs that partition is not cached either,
+since a hit would skip the rows the node streams for the tokens the cache holds. BAGEL
+keys an understanding prompt's text and images, and a generation's first span when it
+is text: an edit's images are resized before their walks see them, and a file's digest
+doesn't cover the resize. A label under a sliding window files pages only up to its
+protected prefix, which the window never releases.
 
 .. _Step 2a — Declare your resources:
 
@@ -930,10 +935,13 @@ methods are:
    The engine applies ``torch.compile`` to both ``forward`` and ``forward_batched``, for
    every submodule. It also captures CUDA graphs for them when you declare capture
    configs (see :ref:`Step 5 <step-5>`). To disable compilation for a submodule, set the
-   class attribute ``disable_torch_compile = True``. Keep the compiled paths
-   compile-friendly. If a helper must not be traced, because it uses data-dependent
-   Python control flow or forces a host synchronization, exclude it explicitly with
-   ``@torch.compiler.disable``.
+   class attribute ``disable_torch_compile = True``. An encoder whose output feeds a
+   prefix cache layout walk compiles with dynamic shapes instead. One that marks every
+   dimension that varies with ``torch._dynamo.mark_dynamic`` in ``prepare_inputs`` sets
+   ``marks_dynamic_dims = True`` and compiles as any other submodule. Keep the compiled
+   paths compile-friendly. If a helper must not be traced, because it uses
+   data-dependent Python control flow or forces a host synchronization, exclude it
+   explicitly with ``@torch.compiler.disable``.
 
    Some submodules must run in their own parameter dtype. One example is an fp32 vocoder
    that is numerically sensitive. Such a submodule sets ``disable_autocast = True``. The

docs/installation.rst

@@ -121,7 +121,8 @@ Model families and some output formats need extra packages, exposed as pip *extr
        Install the matching XPU build of PyTorch.
    * - ``.[qwen3_omni]``
      - Qwen3-Omni runtime: the BAGEL set plus ``qwen-omni-utils``, ``datasets``, and
-       ``ninja`` (speeds up the JIT build of the vendored MoE align kernel).
+       ``ninja`` (speeds up the JIT build of the vendored MoE align kernel). It needs
+       ``transformers>=5.18``, whose ``transformers.vision_utils`` the vision encoder uses.
        **Also needs** ``flash-attn``, which is installed separately —
        see `flash-attn (Qwen3-Omni, LingBot-World)`_.
    * - ``.[orpheus]``
@@ -270,8 +271,9 @@ its dependencies:
 
 ``qwen-tts`` pins ``transformers==4.57.3`` and ``accelerate==1.12.0``. Installed with
 them, it can't share an environment with ``diffusers>=0.40`` (Wan 2.2, Cosmos3,
-FLUX.2 [klein]), which needs ``huggingface-hub>=1.23``, or with OmniVoice, which needs
-``transformers>=5.3``, and installing it after them downgrades their dependencies.
+FLUX.2 [klein]), which needs ``huggingface-hub>=1.23``, or with OmniVoice
+(``transformers>=5.3``) or Qwen3-Omni (``transformers>=5.18``). Installing it after them
+downgrades their dependencies.
 mstar reads only the speech-tokenizer source files and runs them under transformers
 4.57.3 or 5.x, so the package's own requirements are not needed. ``pip check`` still
 lists them (``transformers==4.57.3``, ``accelerate==1.12.0``, ``gradio``, ``sox`` and

mstar/conductor/conductor.py

@@ -61,6 +61,7 @@
 from mstar.utils.logging_config import quiet_noisy_loggers
 from mstar.utils.orphan import exit_when_orphaned
 from mstar.utils.profiler import nvtx_enabled, range_pop, range_push
+from mstar.worker.engine_manager import streamed_partitions
 
 logger = logging.getLogger(__name__)
 
@@ -613,6 +614,13 @@ def _resource_specs(self) -> dict:
             self._resource_specs_cache = specs
         return specs
 
+    def _streamed_partitions(self) -> frozenset[str]:
+        streamed = getattr(self, "_streamed_partitions_cache", None)
+        if streamed is None:
+            streamed = streamed_partitions(list(self._resource_specs().values()), self.model)
+            self._streamed_partitions_cache = streamed
+        return streamed
+
     def _fit_loop_caps(self, model_kwargs: dict, configs: dict, cap: int) -> tuple[int, dict[str, int]]:
         """The largest ``max_output_tokens`` up to ``cap``, and its loop caps, that keeps every loop
         within what the request's resources hold. A loop that ``max_output_tokens`` doesn't drive is
@@ -1256,6 +1264,8 @@ def _do_ingest_request(
         prefix_tail = kwargs.get("prefix_tail") or {}
         prefix_decode = kwargs.get("prefix_decode") or {}
         prefix_cache = kwargs.get("prefix_cache")
+        if not all(partition_states[name].is_done for name in self._streamed_partitions()):
+            prefix_cache = False
         kv_layout = kwargs.get("kv_layout") or {}
         refuse_undeclared_labels(type(self.model).__name__, self._kv_labels, kv_layout, request_data.resource_configs)
         for key, cfg in request_data.resource_configs.items():

mstar/engine/engine.py

@@ -49,6 +49,7 @@
 from mstar.model.submodule_base import (
     EMPTY_INPUT_METADATA,
     ARNodeInputs,
+    ARNodeSubmodule,
     BatchedModelOutput,
     HostRows,
     InputMetadata,
@@ -614,19 +615,22 @@ def _compile_submodules(self) -> None:
         A node in ``_item_encoders`` is compiled with dynamic=True instead. Under
         the default it recompiles on its second input shape and the new graph
         rounds differently, so the KV cached from its earlier outputs would no
-        longer match what it computes.
+        longer match what it computes. A node in ``_item_encoders`` whose submodule
+        sets ``marks_dynamic_dims`` keeps the default, since the submodule's
+        `prepare_inputs` marks the dimensions that vary.
         """
         if not self.device_capability.is_available:
             return
 
         for node_name, submodule_mgmt in self._submodules.items():
             submodule = submodule_mgmt.submodule
+            item_encoder = node_name in self._item_encoders
 
             if getattr(submodule, "disable_torch_compile", False):
                 logger.info("Engine: torch.compile disabled for %s (submodule opt-out)", node_name)
                 continue
 
-            dynamic = True if node_name in self._item_encoders else None
+            dynamic = True if item_encoder and not getattr(submodule, "marks_dynamic_dims", False) else None
             try:
                 # no token count, these serve eager batches of any size
                 submodule_mgmt.forward = with_inductor_overrides(torch.compile(
@@ -1184,7 +1188,10 @@ def _skip_cached_prefix(
             f"{self._prefix_model} keys the {walk!r} walk of {batch.node_name} "
             "for prefix reuse, but that walk's inputs are not a sequence to cut"
         )
-        if inputs.tensor_inputs or inputs.kwargs or inputs.resource_step_info:
+        submodule = self._submodules[batch.node_name].submodule
+        if inputs.kwargs or inputs.resource_step_info or (
+            inputs.tensor_inputs and type(submodule).split_inputs is ARNodeSubmodule.split_inputs
+        ):
             self._runner.apply_cached_prefix(rid, batch.node_name, walk, inputs, None)
             return inputs
         prefix = self._runner.resolve_cached_prefix(rid, batch.node_name, walk, inputs)
@@ -1197,7 +1204,7 @@ def _skip_cached_prefix(
         if matched >= inputs.input_seq_len:
             batch.cached_rids.add(rid)
             return None
-        return self._submodules[batch.node_name].submodule.split_inputs(
+        return submodule.split_inputs(
             walk, batch.per_request_info_wrapped[rid], inputs,
             matched, inputs.input_seq_len,
         )

mstar/model/qwen3_omni/config.py

@@ -156,6 +156,7 @@ class AudioEncoderConfig:
     encoder_attention_heads: int = 20
     output_dim: int = 3584
     max_source_positions: int = 1500
+    n_window: int = 50
 
     @classmethod
     def from_dict(cls, d: dict[str, Any]) -> AudioEncoderConfig:

mstar/model/qwen3_omni/qwen3_omni_model.py

@@ -30,6 +30,7 @@
 
 import logging
 from pathlib import Path
+from typing import NamedTuple
 
 import torch
 from jinja2 import TemplateError
@@ -61,6 +62,7 @@
     KVLabel,
     LabelLayout,
     Model,
+    PrefixKeys,
     ProcessPromptOutput,
     Span,
     TensorAndMetadata,
@@ -134,6 +136,19 @@ def _resolve_local_hf_snapshot(repo_id: str, cache_dir: str | None = None) -> st
     return str(Path(local_dir))
 
 
+class ImageSpan(NamedTuple):
+    length: int
+    advance: int
+
+
+def image_span(grid_thw: torch.Tensor, merge_size: int) -> ImageSpan:
+    """An image's tokens share one temporal position and sit on the merged
+    grid, so it advances the position by the grid's longer side."""
+    t, h, w = (int(n) for n in grid_thw)
+    h, w = h // merge_size, w // merge_size
+    return ImageSpan(t * h * w + 2, max(h, w) + 2)
+
+
 # ---------------------------------------------------------------------------
 # Qwen3OmniModel
 # ---------------------------------------------------------------------------
@@ -310,6 +325,14 @@ def get_request_resource_configs(
             )
         }
 
+    def checkpoint_path(self) -> str:
+        return self.local_dir
+
+    def _audio_frames(self, seqlens: torch.Tensor) -> torch.Tensor:
+        from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import _get_feat_extract_output_lengths
+
+        return _get_feat_extract_output_lengths(seqlens, n_window=self.config.audio_encoder.n_window)
+
     def _get_max_talker_output_tokens(self, **model_kwargs):
         return model_kwargs.get("talker_max_output_tokens", MAX_OUTPUT_TOKENS)
 
@@ -1281,12 +1304,9 @@ def _kv_layout(self, plan: list[PromptPart], tensors: NameToTensorList, audio_ou
         The Talker prefills every step's rows except those of a system or earlier
         assistant turn. It then prefills its assistant prefix, and writes a row per decode step.
         """
-        from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import _get_feat_extract_output_lengths
-
         from mstar.model.qwen3_omni.submodules import TalkerSubmodule, talker_text_mask
 
         merge = self.config.vision.spatial_merge_size
-        grids = {"image": tensors["image_grid_thw"], "video": tensors["video_grid_thw"]}
         thinker, talker, position = [], [], 0
         for step in plan:
             if step.modality == TEXT:
@@ -1295,27 +1315,35 @@ def _kv_layout(self, plan: list[PromptPart], tensors: NameToTensorList, audio_ou
                 kept = int(talker_text_mask(ids, self.config).sum())
             elif step.modality == "audio":
                 # the encoder's rows between two sentinels
-                rows = int(_get_feat_extract_output_lengths(tensors["audio_seqlens"][step.index])) + 2
+                rows = int(self._audio_frames(tensors["audio_seqlens"][step.index])) + 2
                 span = Span("digest", "prefill_audio", rows, rows, ("audio", step.index))
                 kept = rows
+            elif step.modality == "image":
+                image = image_span(tensors["image_grid_thw"][step.index], merge)
+                span = Span("digest", "prefill_vision", image.length, image.advance, ("image", step.index))
+                kept = image.length
             else:
-                grid = grids[step.modality][step.index]
-                per_grid = tensors["video_second_per_grid"][step.index].item() if step.modality == "video" else None
+                grid = tensors["video_grid_thw"][step.index]
+                per_grid = tensors["video_second_per_grid"][step.index].item()
                 rows = int(grid.prod()) // merge ** 2 + 2
                 # the walk moves the counter one past its end sentinel, which follows the grid's largest position
                 positions = get_rope_index_vision(
                     grid, position + 1, self.config.thinker.position_id_per_seconds,
                     spatial_merge_size=merge, seconds_per_grid=per_grid,
                 )
                 advance = int(positions.max()) + 2 - position
-                span = Span("digest", "prefill_vision", rows, advance, (step.modality, step.index))
+                span = Span("digest", "prefill_vision", rows, advance, ("video", step.index))
                 kept = rows
             thinker.append(span)
             talker.append(Span("rows", "talker_prefill", kept, kept, "thinker_states"))
             position += span.advance
         prefix = TalkerSubmodule.TALKER_LAST_PREFILL_TOKENS_PER_REQ
+        # A prompt with a video keys only its leading text, since keying past it has never been run or tested.
+        video = any(step.modality == "video" for step in plan)
         return {
-            THINKER_KV: {"main": LabelLayout(spans=tuple(thinker), growth=(("thinker_decode_loop", 1),))},
+            THINKER_KV: {"main": LabelLayout(
+                spans=tuple(thinker), growth=(("thinker_decode_loop", 1),), keyed_spans=1 if video else None,
+            )},
             TALKER_KV: {"main": LabelLayout(
                 spans=(*talker, Span("rows", "talker_last_prefill", prefix, prefix, "thinker_states")),
                 growth=(("talker_decode_loop", 1),),
@@ -1329,9 +1357,12 @@ def get_loop_caps(self, **model_kwargs) -> dict[str, int]:
         }
 
     def kv_labels(self) -> dict[str, dict[str, KVLabel]]:
+        """The Thinker keys its prompt by text ids, each image and audio clip by
+        file, and its replies as they decode."""
         return {
             THINKER_KV: {"main": KVLabel(
                 growth=(("thinker_decode", "thinker_decode_loop"),),
+                keys=PrefixKeys("text_inputs", ("prefill_text", "prefill_vision", "prefill_audio"), growth=True),
             )},
             TALKER_KV: {"main": KVLabel(
                 growth=(("talker_decode", "talker_decode_loop"),),