diff --git a/.claude/skills/prefix-caching/SKILL.md b/.claude/skills/prefix-caching/SKILL.md index d870d21b27..0e706fed36 100644 --- a/.claude/skills/prefix-caching/SKILL.md +++ b/.claude/skills/prefix-caching/SKILL.md @@ -10,6 +10,6 @@ Prefix caching retains a finished request's KV pages so the next request with th - **Sharing is by owner count, and only sealed pages may be shared.** `PageArena` tracks `num_owners` per page and is the only path to the allocator; a page is freed when the count hits zero, not when a request ends. Appends land in the last page until it fills, so a partially filled page under two owners would have one owner mutating bytes another reads — the tail page therefore stays exclusive, and there is an assertion for it. Sealed pages are never written again; a stream that rewinds across one drops its pages and reacquires private ones. - **The key must name everything the stored bytes depend on.** SHA-256, as a chain (`resources/kv/keys.py`: `fingerprint`, `page_key`, `chain`) rooted in a process fingerprint covering weights, TP head slice, dtype, KV layout, page size, attention backend and position config. Change any of those and every old key becomes unreachable, so the cache empties itself. Media folds in a digest of raw bytes plus preprocessor params plus encoder fingerprint — placeholder token ids alone are identical across images, which is a collision both vLLM and SGLang have shipped. **Adding anything that changes hidden state without adding it to the chain is a correctness bug, not a cache-miss bug.** - **The matched length is the minimum across the node's resources** (`StepRunner.resolve_cached_prefix`; a resource with no opinion is not an answer of zero), resolved on the CPU before `prepare_inputs`. Bucket selection, attention planning, positions and commit all see a shorter request and never learn a hit happened. Keep it that way. -- **Eligibility is narrow deliberately:** a fresh stream, a keyed walk, a linear position scheme with no custom position ids (asserted), and no extra tensor inputs, kwargs or resource step info. Widening it means proving the stored bytes are valid at the positions they will be read at. +- **Eligibility is narrow deliberately:** a fresh stream, a keyed walk, a linear position scheme with no custom position ids (asserted), and no kwargs or resource step info, nor extra tensor inputs unless the submodule sets `cuts_tensor_inputs` and its `split_inputs` cuts them. Widening it means proving the stored bytes are valid at the positions they will be read at. - **A live request must never fail to allocate because of cached data.** Release runs where the shortfall surfaces, in `_alloc` under the arena lock, and removes **leaves only**, so parent pointers never go stale. - **Above one rank, rank 0 decides each hit.** With `prefix_cache_tp` on, rank 0 leases the hit when it broadcasts the step, and each follower adopts it and exits if its pages differ. Lookups above one rank don't stamp the LRU and hits are stamped at commit in step order, so every rank evicts alike. Any page move one rank makes and the others don't breaks that. Without the flag the cache is off above one rank. \ No newline at end of file diff --git a/docs/adding_models.rst b/docs/adding_models.rst index b55b9d78cf..5d37ceed03 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -444,13 +444,19 @@ 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 sets ``cuts_tensor_inputs`` +and overrides ``split_inputs`` 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: @@ -976,10 +982,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 `). 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 diff --git a/docs/installation.rst b/docs/installation.rst index c4cd8d668e..f2b544c132 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -129,7 +129,8 @@ Model families and some output formats need extra packages, exposed as pip *extr when installing both. * - ``.[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]`` @@ -286,8 +287,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 diff --git a/docs/serving.rst b/docs/serving.rst index 3d9434cada..b81e878801 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -733,8 +733,9 @@ or ``increase()``. - ``mstar_prefix_cache_query_tokens_total{model_name, component}``, ``mstar_prefix_cache_hit_tokens_total{model_name, component}``: counters, prompt tokens looked up in the prefix cache (once per request, its whole keyed prompt) and - tokens found there. Only models that key their prompts for prefix reuse (BAGEL and - Orpheus) count, and above one rank only with ``prefix_cache_tp`` set. + tokens found there. Only models that key their prompts for prefix reuse (BAGEL, + Orpheus and Qwen3-Omni's Thinker) count, and above one rank only with + ``prefix_cache_tp`` set. - ``mstar_component_pending_steps{model_name, component}``: gauge, requests whose next step on the node is ready (its inputs arrived) and not running: waiting for room in a step or for the worker to pick the node. diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index faec8e2b5a..6848f884df 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -69,6 +69,7 @@ from mstar.utils.orphan import exit_when_orphaned from mstar.utils.profiler import nvtx_enabled, range_pop, range_push from mstar.validation import validate_deployment +from mstar.validation.resources import streamed_partitions logger = logging.getLogger(__name__) @@ -650,6 +651,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 @@ -1339,6 +1347,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(): diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 72d259c652..aefac9f5ea 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -615,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( @@ -1187,7 +1190,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 not getattr(submodule, "cuts_tensor_inputs", False) + ): 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) @@ -1200,7 +1206,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, ) diff --git a/mstar/model/components/mrope.py b/mstar/model/components/mrope.py new file mode 100644 index 0000000000..08cf05220c --- /dev/null +++ b/mstar/model/components/mrope.py @@ -0,0 +1,37 @@ +"""3D positions for a packed MRoPE step (Qwen3.5, Qwen3-Omni's Thinker).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from mstar.model.submodule_base import ARNodeInputs + + +def mrope_position_ids(pos_ids: torch.Tensor | None, inputs: list[ARNodeInputs]) -> torch.Tensor: + """``[3, total_tokens]`` for the step, in packed request order. + + A pure-text step is ``pos_ids``, the position resource's planned 1D + positions, broadcast across the three grids (a view, no copy). An image's + grids do not advance together, so a request with `custom_pos_ids` supplies + its own and the step concatenates, using the slice of ``pos_ids`` for text. + """ + text_rows = [inp.custom_pos_ids is None for inp in inputs] + assert pos_ids is not None or not any(text_rows), ( + "position resource has no plan for this step; `plan` must run before `preprocess`" + ) + if all(text_rows): + return pos_ids[:sum(inp.input_seq_len for inp in inputs)].unsqueeze(0).expand(3, -1) + + parts: list[torch.Tensor] = [] + offset = 0 + for inp in inputs: + span = inp.input_seq_len + if inp.custom_pos_ids is not None: + parts.append(inp.custom_pos_ids.float()) + else: + parts.append(pos_ids[offset:offset + span].float().unsqueeze(0).expand(3, -1)) + offset += span + return torch.cat(parts, dim=1) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 55ec91e4e8..af61198f0d 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -36,6 +36,7 @@ from mstar.engine.resources.sampler.config import SamplerStep from mstar.engine.resources.sampler.resource import SamplerResource from mstar.engine.resources.step import Segment, SlotLease, SubmoduleStep +from mstar.model.components.mrope import mrope_position_ids from mstar.model.qwen3_5.components.language_model import Qwen3_5ForCausalLM from mstar.model.qwen3_5.components.rope import ( text_position_ids, @@ -376,35 +377,7 @@ def _decode_cos_sin( ) def _position_ids_3d(self, inputs: list[ARNodeInputs]) -> torch.Tensor: - """``[3, total_tokens]`` for the step, in packed request order. - - A pure-text step is the position resource's planned 1D positions - broadcast across the three grids (a view, no copy). An image's grids - do not advance together, so a request with `custom_pos_ids` supplies - its own and the step concatenates, using the resource's slice for text. - """ - total_tokens = sum(inp.input_seq_len for inp in inputs) - pos_ids = self.node_resources[ROPE].pos_ids("main") - assert pos_ids is not None, ( - "position resource has no plan for this step; `plan` must run " - "before `preprocess`" - ) - if not any(inp.custom_pos_ids is not None for inp in inputs): - return pos_ids[:total_tokens].unsqueeze(0).expand(3, -1) - - parts: list[torch.Tensor] = [] - offset = 0 - for inp in inputs: - span = inp.input_seq_len - if inp.custom_pos_ids is not None: - parts.append(inp.custom_pos_ids.float()) - else: - parts.append( - pos_ids[offset:offset + span].float() - .unsqueeze(0).expand(3, -1) - ) - offset += span - return torch.cat(parts, dim=1) + return mrope_position_ids(self.node_resources[ROPE].pos_ids("main"), inputs) def preprocess( self, diff --git a/mstar/model/qwen3_omni/config.py b/mstar/model/qwen3_omni/config.py index b8997a0d33..d03ca6dea2 100644 --- a/mstar/model/qwen3_omni/config.py +++ b/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: diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py index ffe6f2b664..ad72eb711c 100644 --- a/mstar/model/qwen3_omni/qwen3_omni_model.py +++ b/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, @@ -135,6 +137,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 # --------------------------------------------------------------------------- @@ -313,6 +328,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) @@ -1284,12 +1307,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: @@ -1298,12 +1318,16 @@ 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( @@ -1311,14 +1335,18 @@ def _kv_layout(self, plan: list[PromptPart], tensors: NameToTensorList, audio_ou 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),), @@ -1332,9 +1360,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"),), diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index 11540fcc2d..d8a0b8e798 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -16,10 +16,22 @@ import inspect import logging from collections.abc import Mapping -from typing import Any, Optional +from dataclasses import replace +from typing import Any import torch from torch import nn +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + _get_feat_extract_output_lengths, + chunk_and_pad_features, + get_audio_cu_seqlens, + get_valid_indices, +) +from transformers.vision_utils import ( + get_vision_cu_seqlens, + get_vision_interpolation_indices_and_weights, + get_vision_position_ids, +) from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo @@ -32,6 +44,7 @@ from mstar.engine.resources import AttentionStep, KVStep, PositionStep, SamplerStep, Segment, SlotLease, SubmoduleStep from mstar.engine.resources.attn.flashinfer import FlashInferManager from mstar.engine.resources.sampler.resource import SamplerResource +from mstar.model.components.mrope import mrope_position_ids from mstar.model.qwen3_omni.components.code2wav import Qwen3OmniMoeCode2Wav from mstar.model.qwen3_omni.components.rope import ( compute_3d_cos_sin, @@ -54,6 +67,7 @@ THINKER_SAMPLER, Qwen3OmniModelConfig, ) +from mstar.model.qwen3_omni.qwen3_omni_model import image_span from mstar.model.submodule_base import ( ARNodeInputs, ARNodeSubmodule, @@ -117,12 +131,19 @@ class AudioEncoderSubmodule(NodeSubmodule): Runs once per request (not batched across requests). """ + marks_dynamic_dims = True + def __init__(self, audio_encoder: nn.Module, config: Qwen3OmniModelConfig): super().__init__() self.audio_encoder = audio_encoder self.config = config # Older transformers take return_dict; newer ones dropped it. self._encoder_takes_return_dict = _accepts_kwarg(audio_encoder.forward, "return_dict") + n_window = audio_encoder.n_window + # the length of a full attention window, since no clip's window is longer + self._max_seqlen = _get_feat_extract_output_lengths(2 * n_window, n_window) * ( + audio_encoder.n_window_infer // (2 * n_window) + ) def prepare_inputs( self, @@ -133,21 +154,31 @@ def prepare_inputs( ) -> NodeInputs: # Edge name from graph walk is "audio_features" audio_features = inputs["audio_features"][0] - audio_seqlens = inputs.get("audio_seqlens", [None])[0] - - return NodeInputs( - tensor_inputs={ - "audio_features": audio_features, - "audio_seqlens": audio_seqlens, - } - ) + audio_seqlens = inputs["audio_seqlens"][0] + + # Built outside the compiled forward, since these helpers read the clip's length on the + # host, which would compile a graph for every length. + encoder = self.audio_encoder + padded_feature, chunk_lengths = chunk_and_pad_features(audio_features, audio_seqlens, encoder.n_window) + tensor_inputs = { + "padded_feature": padded_feature, + "chunk_lengths": chunk_lengths, + "valid_indices": get_valid_indices(chunk_lengths, encoder.n_window), + "cu_seqlens": get_audio_cu_seqlens(chunk_lengths, audio_seqlens, encoder.n_window_infer, encoder.n_window), + } + for name in ("chunk_lengths", "valid_indices", "cu_seqlens"): + torch._dynamo.mark_dynamic(tensor_inputs[name], 0) + # A clip shorter than one chunk is padded only to its own length, so the frame dimension of + # `padded_feature` varies too. Some clips give a dimension of size 1 inside the encoder, which torch + # specializes, so `mark_dynamic` would raise for them. + torch._dynamo.maybe_mark_dynamic(padded_feature, 0) + torch._dynamo.maybe_mark_dynamic(padded_feature, 2) + return NodeInputs(tensor_inputs=tensor_inputs) def forward( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - audio_features: torch.Tensor, - audio_seqlens: Optional[torch.Tensor] = None, **kwargs, ) -> NameToTensorList: """Run the audio encoder and return embeddings. @@ -155,14 +186,10 @@ def forward( Returns: {"audio_embeds": [tensor of shape (audio_tokens, hidden_size)]} """ - logger.debug( - "Running AudioEncoder with audio_features shape=%s", - audio_features.shape, - ) extra_kwargs = {"return_dict": True} if self._encoder_takes_return_dict else {} audio_out = self.audio_encoder( - audio_features, - feature_lens=audio_seqlens, + max_seqlen=self._max_seqlen, + **kwargs, **extra_kwargs, ) audio_embeds = getattr(audio_out, "last_hidden_state", audio_out) @@ -189,6 +216,8 @@ class VisionEncoderSubmodule(NodeSubmodule): Runs once per request (not batched across requests). """ + marks_dynamic_dims = True + def __init__(self, vision_encoder: nn.Module, config: Qwen3OmniModelConfig): super().__init__() self.vision_encoder = vision_encoder @@ -224,12 +253,24 @@ def prepare_inputs( if grid_thw.dim() == 1: grid_thw = grid_thw.unsqueeze(0) # (1, 3) - return NodeInputs( - tensor_inputs={ - "pixel_values": pixel_values, - "grid_thw": grid_thw, - } + # Built outside the compiled forward, since these helpers call `.tolist()` on the + # grid, which would compile a graph for every image size. + encoder = self.vision_encoder + interp_indices, interp_weights = get_vision_interpolation_indices_and_weights( + grid_thw, num_grid_per_side=encoder.num_grid_per_side, mode=encoder.interpolation_mode, + align_corners=encoder.interpolation_align_corners, spatial_merge_size=encoder.spatial_merge_size, ) + tensor_inputs = { + "pixel_values": pixel_values, + "grid_thw": grid_thw, + "interp_indices": interp_indices, + "interp_weights": interp_weights, + "position_ids": get_vision_position_ids(grid_thw, encoder.spatial_merge_size), + "cu_seqlens": get_vision_cu_seqlens(grid_thw), + } + for name in ("pixel_values", "interp_indices", "interp_weights", "position_ids", "cu_seqlens"): + torch._dynamo.mark_dynamic(tensor_inputs[name], 0) + return NodeInputs(tensor_inputs=tensor_inputs) def forward( self, @@ -256,6 +297,9 @@ def forward( encoder_output = self.vision_encoder( pixel_values, grid_thw=grid_thw, + # each video frame is its own segment, and a call encodes one item, so this is one frame's patches + max_seqlen=pixel_values.shape[0] // (kwargs["cu_seqlens"].shape[0] - 1), + **kwargs, ) if isinstance(encoder_output, tuple): @@ -292,6 +336,8 @@ class ThinkerSubmodule(ARNodeSubmodule): produces ``logits`` for text token sampling. """ + cuts_tensor_inputs = True + # Default MRoPE section for head_dim=128: [24, 20, 20] MROPE_SECTION = [24, 20, 20] @@ -427,24 +473,16 @@ def prepare_inputs( text_ids = inputs["text_inputs"][0].to(device) # (seq_len,) embeds = self.model.model.embed_tokens(text_ids) seq_len = text_ids.shape[0] - - # Compute 3D MRoPE position IDs for a pure-text span. Each - # prefill graph walk is single-modality so we use the simple - # per-modality helper instead of the full HF parser. - # - # ``start_pos`` is the next MRoPE position for this request, - # carried forward across walks by the stream's position counter - # (advanced by the runner's commit of each declared step). - pos_ids = get_rope_index_text(seq_len, start_pos, device) masks_for_talker = torch.stack([ torch.zeros(text_ids.shape, dtype=torch.bool, device=device), # multimodal talker_text_mask(text_ids, self.config) # text inclusion ]) + # A text row leaves its positions to the plan, which starts from the + # counter a cache hit has already moved. return ARNodeInputs( input_seq_len=seq_len, input_embeds=embeds, input_ids=text_ids, - custom_pos_ids=pos_ids, tensor_inputs={ "masks_for_talker": masks_for_talker } @@ -524,7 +562,11 @@ def prepare_inputs( # Sentinel token positions (text-like). start_pos_ids = get_rope_index_text(1, start_pos, device) - end_pos_base = float(vision_pos_ids.max().item()) + 1 + if seconds_per_grid is None: + # Use the prompt layout's helper, so the cache and this walk move the position by the same amount. + end_pos_base = start_pos + image_span(grid_thw, self.config.vision.spatial_merge_size).advance - 1 + else: + end_pos_base = float(vision_pos_ids.max().item()) + 1 end_pos_ids = get_rope_index_text(1, end_pos_base, device) pos_ids = torch.cat( @@ -554,6 +596,24 @@ def prepare_inputs( } ) + def split_inputs( + self, + graph_walk: str, + fwd_info: CurrentForwardPassInfo, + inputs: ARNodeInputs, + start: int, + end: int, + ) -> ARNodeInputs: + tensors = inputs.tensor_inputs + cut = super().split_inputs(graph_walk, fwd_info, replace(inputs, tensor_inputs={}), start, end) + cut.tensor_inputs = {"masks_for_talker": tensors["masks_for_talker"][:, start:end]} + if graph_walk == "prefill_vision": + cut.tensor_inputs["deepstack"] = [layer[start:end] for layer in tensors["deepstack"]] + cut.tensor_inputs["mrope_pos_advance"] = tensors["mrope_pos_advance"] + elif graph_walk == "prefill_audio": + cut.tensor_inputs["mrope_pos_advance"] = inputs.input_seq_len + return cut + def declare_step( self, graph_walk: str, request_ids: list[str], @@ -581,17 +641,17 @@ def declare_step( } pos_advance = None - if "prefill_vision" in row_walks: - # A vision row's 3D-grid MRoPE span is larger than its token count, - # so it advances the counter by its own; other rows by their span. + if graph_walk != "thinker_decode" and any("mrope_pos_advance" in inp.tensor_inputs for inp in inputs): + # A vision row's 3D-grid MRoPE span differs from its token count, and + # a cut audio row moves past its whole clip, so each advances the + # counter by its own `mrope_pos_advance`. pos_advance = tuple( - int(inp.tensor_inputs.get("mrope_pos_advance", 0)) - if walk == "prefill_vision" else inp.input_seq_len - for inp, walk in zip(inputs, row_walks, strict=True) + int(inp.tensor_inputs.get("mrope_pos_advance", inp.input_seq_len)) + for inp in inputs ) return SubmoduleStep( # must match `cg_key_info`, which picked the leased capture - cg_key_info=self.DEEPSTACK_KEY if pos_advance is not None else None, + cg_key_info=self.DEEPSTACK_KEY if "prefill_vision" in row_walks else None, segments=[ Segment( request_id=rid, @@ -626,9 +686,7 @@ def preprocess( input_embeds = torch.cat([ inp.input_embeds for inp in inputs ], dim=0) - position_ids_3d = torch.cat([ - inp.custom_pos_ids for inp in inputs - ], dim=1) # (3, total_tokens) + position_ids_3d = self._position_ids_3d(inputs) # (3, total_tokens) seq_lens = [ inp.input_seq_len for inp in inputs ] @@ -676,6 +734,9 @@ def preprocess( **extra_inputs } + def _position_ids_3d(self, inputs: list[ARNodeInputs]) -> torch.Tensor: + return mrope_position_ids(self.node_resources[THINKER_POS].pos_ids("main"), inputs) + def _deepstack_inputs( self, inputs: list[ARNodeInputs], dtype: torch.dtype, device: torch.device, ) -> dict[str, torch.Tensor]: diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 204e1f9176..5472e82c15 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -1027,6 +1027,10 @@ def cleanup_request(self, request_id: str): class ARNodeSubmodule(NodeSubmodule): + # Set True on a submodule whose `split_inputs` cuts its own ``tensor_inputs``. + # The engine probes the prefix cache for a walk that carries them only then. + cuts_tensor_inputs: bool = False + def split_inputs( self, graph_walk: str, diff --git a/mstar/validation/resources.py b/mstar/validation/resources.py index 4b5e160b6c..18f79c94d7 100644 --- a/mstar/validation/resources.py +++ b/mstar/validation/resources.py @@ -15,6 +15,7 @@ ) from mstar.graph.special_destinations import SPECIAL_DESTINATIONS from mstar.model.base import Model, keyed_labels +from mstar.streaming.topology import StreamingGraphEdge def check_resources( @@ -173,14 +174,15 @@ def _undeclared_pools(specs: Sequence[NodeResourceSpec], model: Model) -> Iterat def _unservable_walks(specs: Sequence[NodeResourceSpec], model: Model) -> Iterator[str]: """A keyed label over several walks if the cache could not serve one of - them whole: a walk where a keyed node feeds another node, or where no keyed - node runs. + them whole: a walk where a keyed node feeds another node in the walk's + partition, or where no keyed node runs. A walk the cache serves whole completes with no outputs, so a node that reads them would get nothing. A label keyed over one walk is not checked, because its one span always runs its last token. """ graphs = model.get_graph_walk_graphs() + partition_of = {walk: p.name for p in model.get_partitions() for walk in p.graph_walks} by_key = {spec.resource_key: spec for spec in specs} for key, by_label in keyed_labels(model).items(): for label, kv_label in by_label.items(): @@ -206,7 +208,11 @@ def _unservable_walks(specs: Sequence[NodeResourceSpec], model: Model) -> Iterat ) for name in ran: fed = sorted( - {edge.next_node for edge in nodes[name].outputs} + { + edge.next_node for edge in nodes[name].outputs + if not isinstance(edge, StreamingGraphEdge) + or edge.target_partition == partition_of[walk] + } - SPECIAL_DESTINATIONS ) if fed: @@ -215,3 +221,23 @@ def _unservable_walks(specs: Sequence[NodeResourceSpec], model: Model) -> Iterat f"{walk!r}, where {name} feeds {fed}, which would get " "nothing from a walk the cache serves whole" ) + + +def streamed_partitions(specs: Sequence[NodeResourceSpec], model: Model) -> frozenset[str]: + """The partitions a node streams into when its label is keyed over several walks. + + A walk the cache serves streams nothing for the tokens it skips, so the + conductor doesn't cache a request that runs one of these partitions. + """ + graphs = model.get_graph_walk_graphs() + partition_of = {walk: p.name for p in model.get_partitions() for walk in p.graph_walks} + nodes_of = {spec.resource_key: spec.nodes for spec in specs} + return frozenset( + edge.target_partition + for key, by_label in keyed_labels(model).items() + for kv_label in by_label.values() if len(kv_label.keys.walks) > 1 + for walk in kv_label.keys.walks + for name, node in graphs[walk].get_nodes().items() if name in nodes_of[key] + for edge in node.outputs + if isinstance(edge, StreamingGraphEdge) and edge.target_partition != partition_of[walk] + ) diff --git a/pyproject.toml b/pyproject.toml index 3890dc09b0..72bc41c49c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -203,7 +203,8 @@ qwen3_omni = [ "torchvision", "torchaudio", "torchcodec", - "transformers", + # the vision encoder builds its grid inputs with transformers.vision_utils + "transformers>=5.18", "datasets", "ninja", # speeds up the first-use JIT build of the vendored MoE align kernel ] @@ -397,8 +398,9 @@ all = [ "torchvision", "torchaudio", "torchcodec", - # Qwen3-TTS's speech tokenizer (qwen-tts, installed --no-deps) needs check_model_inputs - "transformers>=4.57.3", + # Qwen3-TTS's speech tokenizer (qwen-tts, installed --no-deps) needs check_model_inputs, + # and Qwen3-Omni's vision encoder needs transformers.vision_utils + "transformers>=5.18", "safetensors", "triton", "qwen-omni-utils", diff --git a/test/modular/test_kv_prefix_request_keys.py b/test/modular/test_kv_prefix_request_keys.py index b87e2200b3..8af7d6be03 100644 --- a/test/modular/test_kv_prefix_request_keys.py +++ b/test/modular/test_kv_prefix_request_keys.py @@ -18,14 +18,19 @@ import pytest import torch +from test_ingest_unpersist import _conductor, _ToyModel from mstar.api_server.data_worker import PreprocessWorkerThread from mstar.api_server.request_types import PreprocessInput from mstar.conductor.conductor import Conductor +from mstar.conductor.request_info import CurrentForwardConductorMetadata, PartitionDefinition from mstar.engine.resources import SamplingReqConfig from mstar.engine.resources.kv.config import KVReqConfig, KVSpec, PagedKVConfig from mstar.engine.resources.kv.keys import chain +from mstar.graph.base import GraphEdge, GraphNode, Sequential +from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.base import ( + ForwardPassArgs, KVLabel, LabelLayout, PrefixKeys, @@ -36,7 +41,8 @@ ) from mstar.model.orpheus.config import OrpheusModelConfig from mstar.model.orpheus.orpheus_model import OrpheusModel -from mstar.streaming.topology import PartitionTopology +from mstar.streaming.topology import PartitionTopology, StreamingGraphEdge +from mstar.utils.ipc_format import NewRequestConductor PAGE_SIZE = 16 PROMPT = list(range(100)) @@ -351,6 +357,70 @@ def test_a_request_that_says_nothing_is_still_cached(): ) +class _Talking(_ToyModel): + """Keys a layout on an LLM that streams every prefill walk to a talker, + which runs only when the request asks for audio.""" + + def __init__(self): + def llm(): + return GraphNode(name="LLM", input_names=["text_inputs"], outputs=[ + GraphEdge(next_node=EMIT_TO_CLIENT, name="new_token", output_modality="text"), + StreamingGraphEdge(next_node="talker", name="hidden", target_partition="Talker"), + ]) + + vit = GraphNode(name="vit", input_names=["pixels"], outputs=[GraphEdge(next_node="LLM", name="text_inputs")]) + talker = GraphNode(name="talker", input_names=["hidden"], outputs=[ + GraphEdge(next_node=EMIT_TO_CLIENT, name="audio", output_modality="audio"), + ]) + super().__init__( + {"prefill": llm(), "prefill_image": Sequential([vit, llm()]), "talk": talker}, None, + partitions=[ + PartitionDefinition("LLM", {"prefill", "prefill_image"}, "prefill"), + PartitionDefinition("Talker", {"talk"}, "talk", ["LLM"]), + ], + ) + + def get_node_resources(self): + return [KVSpec(resource_key="kv", nodes={"LLM"}, config=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=64, + ))] + + def kv_labels(self): + return {"kv": {"main": KVLabel(keys=PrefixKeys("text_inputs", ("prefill", "prefill_image")))}} + + def get_initial_forward_pass_args( + self, partition_name, input_modalities, output_modalities, input_signals, model_kwargs=None, + ): + walk = "prefill" if partition_name == "LLM" else "talk" + return ForwardPassArgs( + full_metadata=CurrentForwardConductorMetadata( + input_modalities=input_modalities, output_modalities=output_modalities, + graph_walk=walk, is_prefill=True, + ), + inputs=[], unpersist_tensors=[], + request_done=partition_name == "Talker" and "audio" not in output_modalities, + ) + + +def test_a_request_whose_talker_runs_is_not_cached(tmp_path, monkeypatch): + conductor = _conductor(_Talking(), [{"node_names": ["vit", "LLM", "talker"], "ranks": [0]}], tmp_path, monkeypatch) + + layout = LabelLayout(spans=(Span("ids", "prefill", 4, 4, "text_inputs"),)) + cached = {} + for out in ("audio", "text"): + conductor._do_ingest_request(NewRequestConductor( + request_id=out, initial_signals={}, initial_input_modalities=["text"], + initial_output_modalities=[out], input_metadata={}, + model_kwargs={"kv_layout": {"kv": {"main": layout.row()}}}, + )) + cached[out] = conductor.requests[out].resource_configs["kv"].prefix_cache + + assert cached == {"audio": False, "text": True}, ( + "a request whose prefill streams to a running talker was cached, so a " + "walk the cache serves would leave the talker short of rows" + ) + + # ── a config for every resource that declared a stream ────────────────── diff --git a/test/modular/test_kv_prefix_root.py b/test/modular/test_kv_prefix_root.py index 98147f21f8..e704318be9 100644 --- a/test/modular/test_kv_prefix_root.py +++ b/test/modular/test_kv_prefix_root.py @@ -28,6 +28,7 @@ import pytest import torch +from mstar.conductor.request_info import PartitionDefinition from mstar.engine.engine import Engine from mstar.engine.resources import Resource from mstar.engine.resources.kv import manager as manager_mod @@ -38,6 +39,7 @@ from mstar.graph.base import GraphEdge, GraphNode, Sequential from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.base import KVLabel, PrefixKeys +from mstar.streaming.topology import StreamingGraphEdge from mstar.validation.resources import _unservable_walks KV = "kv" @@ -436,8 +438,10 @@ def node(name, feeds): } walks = ("prefill", "prefill_image") label = KVLabel(keys=PrefixKeys("text_inputs", walks)) + partition = PartitionDefinition(name="LLM", graph_walks=set(graphs), initial_walk="prefill", producer_partitions=[]) return SimpleNamespace( kv_labels=lambda: {KV: {"main": label}}, get_graph_walk_graphs=lambda: graphs, + get_partitions=lambda: [partition], checkpoint_path=lambda: None, preprocess_fingerprint=lambda: "laid out", ) @@ -463,6 +467,23 @@ def test_a_walk_the_cache_cannot_serve_whole_is_refused_at_startup(node, model, assert problems and re.search(message, problems[0]), problems +def test_a_keyed_node_that_streams_into_another_partition_loads(): + model = _laid_out() + for graph in model.get_graph_walk_graphs().values(): + graph.get_nodes()["LLM"].outputs.append( + StreamingGraphEdge(next_node="talker", name="hidden", target_partition="Talker"), + ) + spec = KVSpec( + resource_key=KV, nodes={"LLM"}, + config=PagedKVConfig(num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=64), + ) + + assert not list(_unservable_walks([spec], model)), ( + "a keyed node that streams into another partition was refused, though a walk the cache " + "serves whole leaves its own partition nothing to read" + ) + + def test_only_a_node_whose_output_is_cached_under_a_digest_is_compiled_dynamic(monkeypatch): engine = Engine.__new__(Engine) engine._resources = {} @@ -483,3 +504,23 @@ def test_only_a_node_whose_output_is_cached_under_a_digest_is_compiled_dynamic(m "an image encoder left to recompile caches KV it no longer computes for the same " "file, or a node the cache never keys from lost the compile main gives it" ) + + +def test_an_encoder_that_marks_its_own_dims_compiles_with_the_default_while_its_output_is_cached(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch, "compile", lambda fn, dynamic, **kwargs: dynamic) + engine = Engine.__new__(Engine) + engine._resources = {} + engine._device = torch.device("cuda") + spec = KVSpec(resource_key=KV, nodes={"LLM"}, config=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=64, + )) + engine._open_prefix_caches({KV: spec}, _laid_out()) + vit = SimpleNamespace(forward="eager", forward_batched="eager", marks_dynamic_dims=True) + engine._submodules = {"vit": SimpleNamespace(submodule=vit, forward="eager")} + + engine._compile_submodules() + + assert engine._submodules["vit"].forward is None, ( + "an encoder that marks its own dynamic dimensions was compiled with dynamic=True while its KV is cached" + ) diff --git a/test/modular/test_kv_prefix_stage.py b/test/modular/test_kv_prefix_stage.py index 7e829f1dee..91fc9a1808 100644 --- a/test/modular/test_kv_prefix_stage.py +++ b/test/modular/test_kv_prefix_stage.py @@ -184,6 +184,24 @@ def test_a_walk_carrying_anything_opaque_never_probes(opaque): ) +class _Overriding(_Submodule): + """Overrides `split_inputs` without cutting its ``tensor_inputs``.""" + + def split_inputs(self, graph_walk, fwd_info, inputs, start, end): + return super().split_inputs(graph_walk, fwd_info, inputs, start, end) + + +def test_overriding_split_inputs_alone_never_probes_a_walk_carrying_tensor_inputs(): + resource = _Answering(96) + engine = _engine({"kv": resource}, {NODE: ["kv"]}) + engine._submodules[NODE].submodule = _Overriding() + inputs = _inputs(tensor_inputs={"extra": torch.zeros(PROMPT)}) + + assert _stage(engine, inputs) is inputs and resource.resolved == 0, ( + "a submodule that never set cuts_tensor_inputs was probed, and its cut raises mid-step" + ) + + def test_a_falsy_step_info_is_not_opaque(): resource = _Answering(96) engine = _engine({"kv": resource}, {NODE: ["kv"]}) diff --git a/test/modular/test_mrope_position_ids.py b/test/modular/test_mrope_position_ids.py new file mode 100644 index 0000000000..2d67860c08 --- /dev/null +++ b/test/modular/test_mrope_position_ids.py @@ -0,0 +1,34 @@ +"""Where each row of a packed MRoPE step reads its positions. + +Qwen3.5 and Qwen3-Omni's Thinker pack text, image and audio rows into one step. +A text row's positions come from the position resource's plan, which covers every +row of the step in packed order, so a text row after an image must read the plan +past the image's tokens, not where the previous text row ended. +""" + +from __future__ import annotations + +import sys + +sys.path.insert(0, ".") + +import torch + +from mstar.model.components.mrope import mrope_position_ids +from mstar.model.submodule_base import ARNodeInputs + + +def test_a_text_row_after_an_image_reads_the_plan_past_the_images_tokens(): + plan = torch.arange(10, 19) + image = torch.tensor([[5.0, 5.0], [5.0, 6.0], [5.0, 7.0]]) + rows = [ + ARNodeInputs(input_seq_len=3), + ARNodeInputs(input_seq_len=2, custom_pos_ids=image), + ARNodeInputs(input_seq_len=4), + ] + + positions = mrope_position_ids(plan, rows) + + assert positions.tolist() == [ + [10, 11, 12, *grid, 15, 16, 17, 18] for grid in image.tolist() + ], "a row took positions from the wrong slice of the plan, so its tokens rotate as if elsewhere in the prompt" diff --git a/test/modular/test_qwen3_omni_prefix_cache.py b/test/modular/test_qwen3_omni_prefix_cache.py new file mode 100644 index 0000000000..bb5e48fd6a --- /dev/null +++ b/test/modular/test_qwen3_omni_prefix_cache.py @@ -0,0 +1,392 @@ +"""Qwen3-Omni's Thinker keys its prompt past images and audio, and a hit cuts a walk short.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch._dynamo.utils import counters +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeAudioEncoderConfig, + Qwen3OmniMoeVisionEncoderConfig, +) +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeAudioEncoder, + Qwen3OmniMoeVisionEncoder, +) + +from mstar.engine.engine import Engine +from mstar.engine.resources import PositionStep, Resource, StepContext, StepRunner +from mstar.engine.resources.base import CachedPrefix +from mstar.engine.resources.kv.plan import KVPlanOutputs +from mstar.engine.resources.position.config import PositionConfig +from mstar.engine.resources.position.manager import RopeManager +from mstar.model.qwen3_omni.components.rope import get_rope_index_vision +from mstar.model.qwen3_omni.config import THINKER_KV, THINKER_MIXED, THINKER_POS, Qwen3OmniModelConfig +from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel, image_span +from mstar.model.qwen3_omni.submodules import AudioEncoderSubmodule, ThinkerSubmodule, VisionEncoderSubmodule + +RID = "r0" +NODE = "Thinker" +HIDDEN = 8 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="flash-attn runs on a GPU") + + +class _StubThinker(nn.Module): + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(16, HIDDEN) + + +class _Holding(Resource): + + supports_preplan = True + + def __init__(self, tokens: int, position: int): + self._prefix = CachedPrefix("main", tokens, position) + + @classmethod + def build(cls, spec, info): + raise NotImplementedError + + def resolve_cached_prefix(self, rid, node_name, graph_walk, inputs): + return self._prefix + + +@pytest.fixture +def thinker(): + config = SimpleNamespace( + im_start_token_id=1, system_token_id=2, assistant_token_id=3, thinker_hidden_size=HIDDEN, + thinker=SimpleNamespace( + audio_start_token_id=4, audio_end_token_id=5, vision_start_token_id=6, vision_end_token_id=7, + position_id_per_seconds=25, + ), + vision=SimpleNamespace(spatial_merge_size=2, deepstack_visual_indexes=[1, 2]), + ) + thinker = ThinkerSubmodule(thinker_model=_StubThinker(), config=config) + thinker.node_resources = {THINKER_POS: RopeManager( + config=PositionConfig(kv_cache=THINKER_KV), device=torch.device("cpu"), dtype=torch.float32, + )} + return thinker + + +def _prepare(thinker, walk: str, inputs: dict, held: CachedPrefix | None = None, rid: str = RID): + info = SimpleNamespace(rid_handle=rid) + prepared = thinker.prepare_inputs(walk, info, inputs) + if held is None: + return prepared + engine = Engine.__new__(Engine) + engine._runner = StepRunner( + {THINKER_KV: _Holding(held.tokens, held.position), THINKER_POS: thinker.node_resources[THINKER_POS]}, + node_resources={NODE: [THINKER_KV, THINKER_POS]}, + ) + engine._submodules = {NODE: SimpleNamespace(submodule=thinker)} + engine._keyed_walks = {NODE: {"prefill_text", "prefill_vision", "prefill_audio"}} + batch = SimpleNamespace( + node_name=NODE, cached_rids=set(), per_request_info_wrapped={rid: info}, + step_context=StepContext(request_ids=(rid,), graph_walk=walk, slot=0, capture=False), + ) + return engine._skip_cached_prefix(batch, rid, prepared) + + +def test_a_cut_text_walk_prepared_again_is_planned_where_the_cache_keyed_it(thinker): + rope = thinker.node_resources[THINKER_POS] + rope._counters[RID] = {"main": 10} + text = {"text_inputs": [torch.arange(8, 14).repeat(2)]} + + _prepare(thinker, "prefill_text", text, CachedPrefix("main", 6, 16)) + # a refused admit prepares the walk again, after the probe moved the counter + again = _prepare(thinker, "prefill_text", text, CachedPrefix("main", 6, 16)) + kv_plan = KVPlanOutputs({"main": SimpleNamespace(rows=[(RID, "main", again.input_seq_len)])}) + rope.plan(PositionStep(), StepContext((RID,), "prefill_text", 0, False, plan_results={THINKER_KV: kv_plan})) + + assert again.custom_pos_ids is None and thinker._position_ids_3d([again]).tolist() == [list(range(16, 22))] * 3, ( + "the retried walk's tail runs past the positions the cache keyed it at" + ) + + +def test_a_cut_vision_walk_cuts_its_masks_and_deepstack_with_the_embeds(thinker): + image = { + "vision_embeds": [torch.randn(6, HIDDEN)], "deepstack": [torch.randn(6, HIDDEN)] * 2, + "image_grid_thw": [torch.tensor([1, 4, 6])], + } + whole = _prepare(thinker, "prefill_vision", image) + + cut = thinker.split_inputs("prefill_vision", SimpleNamespace(rid_handle=RID), whole, 3, whole.input_seq_len) + + assert torch.equal(cut.tensor_inputs["masks_for_talker"], whole.tensor_inputs["masks_for_talker"][:, 3:]) and all( + torch.equal(layer, full[3:]) + for layer, full in zip(cut.tensor_inputs["deepstack"], whole.tensor_inputs["deepstack"], strict=True) + ), "the Talker masks or deepstack rows kept rows the embeds lost, so they would land on other tokens" + + +def test_each_cut_row_of_a_mixed_step_moves_its_own_counter_past_its_whole_walk(thinker): + # The cache holds 6 of 12 text tokens, 28 of a clip's 30 slots, and 3 of an image's 8 slots, + # which advance the counter by 5. A partly held item leaves the counter at its start. + rope = thinker.node_resources[THINKER_POS] + rows = { + "text": ("prefill_text", {"text_inputs": [torch.arange(8, 14).repeat(2)]}, 10, CachedPrefix("main", 6, 16)), + "clip": ("prefill_audio", {"audio_embeds": [torch.zeros(28, HIDDEN)]}, 20, CachedPrefix("main", 28, 20)), + "image": ("prefill_vision", { + "vision_embeds": [torch.randn(6, HIDDEN)], "deepstack": [torch.randn(6, HIDDEN)] * 2, + "image_grid_thw": [torch.tensor([1, 4, 6])], + }, 30, CachedPrefix("main", 3, 30)), + } + inputs = [] + for rid, (walk, tensors, counter, held) in rows.items(): + rope._counters[rid] = {"main": counter} + cut = _prepare(thinker, walk, tensors, held, rid=rid) + cut.graph_walk = walk + inputs.append(cut) + + step = thinker.declare_step(THINKER_MIXED, list(rows), inputs) + rope.commit(step.steps[THINKER_POS], StepContext(tuple(rows), THINKER_MIXED, slot=0, capture=False)) + + assert [cut.input_seq_len for cut in inputs] == [6, 2, 5], "a row ran slots the cache already holds" + assert [rope.position(rid, "main") for rid in rows] == [22, 50, 35], ( + "a row moved its counter by the slots it ran, so its later walks would run at wrong positions" + ) + + +def test_an_images_span_is_what_its_positions_place(): + grid, start = (1, 30, 40), 7 + positions = get_rope_index_vision(torch.tensor(grid), start + 1, position_id_per_seconds=25, spatial_merge_size=2) + + # the end sentinel sits one past the image's last position, and the next walk one past that + assert image_span(torch.tensor(grid), 2) == (positions.shape[1] + 2, int(positions.max()) + 2 - start), ( + "the layout's image span disagrees with the positions its walk places" + ) + + +_SPECIALS = { + "<|vision_start|>": 151652, "<|vision_end|>": 151653, "<|image_pad|>": 151655, "<|video_pad|>": 151656, + "<|audio_start|>": 151669, "<|audio_end|>": 151670, "<|audio_pad|>": 151675, +} +_PLACEHOLDERS = { + "image": "<|vision_start|><|image_pad|><|vision_end|>", + "audio": "<|audio_start|><|audio_pad|><|audio_end|>", + "video": "<|vision_start|><|video_pad|><|vision_end|>", +} + + +class _Tokenizer: + unk_token_id = 0 + + def convert_tokens_to_ids(self, token): + return _SPECIALS.get(token, 0) + + def __call__(self, text, return_tensors=None): + for token, id_ in _SPECIALS.items(): + text = text.replace(token, chr(id_)) + return {"input_ids": torch.tensor([[ord(c) for c in text]])} + + +class _VideoProcessor: + temporal_patch_size = 2 + + def __call__(self, videos, size): + return {"video_grid_thw": torch.tensor([[2, 4, 6]]), "pixel_values_videos": torch.zeros(1)} + + +class _Processor: + video_processor = _VideoProcessor() + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): + turns = "".join( + f"<|im_start|>{m['role']}\n" + (m["content"] if isinstance(m["content"], str) else "".join( + part["text"] if part["type"] == "text" else _PLACEHOLDERS[part["type"]] for part in m["content"] + )) + "<|im_end|>\n" + for m in messages + ) + return turns + "<|im_start|>assistant\n" + + @staticmethod + def image_processor(images, return_tensors): + height, width, _ = images[0].shape + return {"pixel_values": torch.zeros(1), "image_grid_thw": torch.tensor([[1, height // 16, width // 16]])} + + @staticmethod + def feature_extractor(audio, sampling_rate, **kwargs): + frames = len(audio) // 160 + return {"input_features": torch.zeros(1, 128, frames), "attention_mask": torch.ones(1, frames)} + + +@pytest.fixture +def model(): + model = Qwen3OmniModel.__new__(Qwen3OmniModel) + model.tokenizer = _Tokenizer() + model._processor = _Processor() + model.config = Qwen3OmniModelConfig() + return model + + +def _prompt(model, modalities: list[str], tensors: dict | None, **kwargs): + texts = ["Say what you see."] * modalities.count("text") + return model.process_prompt(texts, modalities, ["text"], tensors=tensors, **kwargs) + + +# A 64x96 image is a 4x6 patch grid, 2x3 after the merge, so it writes 8 slots with its sentinels and advances by 5. +_IMAGE = ("prefill_vision", 8, 5) +# A second of audio is 100 feature frames, which the encoder turns into 13, so it writes 15 slots with its sentinels. +_AUDIO = ("prefill_audio", 15, 15) +_PROMPTS = { + "text, image, text": (["text", "image", "text"], {"image_inputs": [torch.rand(3, 64, 96)]}, [_IMAGE]), + "audio then text": (["audio", "text"], {"audio_inputs": [torch.zeros(16000)]}, [_AUDIO]), +} + + +@pytest.mark.parametrize(("modalities", "tensors", "items"), _PROMPTS.values(), ids=_PROMPTS.keys()) +def test_the_layout_lists_each_walk_the_schedule_runs_with_the_slots_it_writes(model, modalities, tensors, items): + out = _prompt(model, modalities, tensors) + layout = out.metadata["kv_layout"][THINKER_KV]["main"].spans + signals = {name: [object()] * len(value) for name, value in out.new_input_tensors.items()} + + walks = [walk for walk, _ in model._build_thinker_prefill_schedule(modalities, signals)] + segments = iter(out.new_input_tensors["text_inputs"]) + media = iter(items) + expected = [ + ("prefill_text", len(next(segments)), None) if walk == "prefill_text" else next(media) for walk in walks + ] + assert [(span.walk, span.length, None if span.kind == "ids" else span.advance) for span in layout] == expected, ( + "the layout disagrees with the walks the schedule runs, so the cache would drop it partway through" + ) + + +def _vision(attn_implementation: str = "sdpa") -> VisionEncoderSubmodule: + # The tensors `prepare_inputs` builds depend only on the grid, the merge size and + # `num_grid_per_side`, so a small encoder is enough. + config = Qwen3OmniMoeVisionEncoderConfig( + depth=2, hidden_size=32, intermediate_size=64, num_heads=2, out_hidden_size=16, + num_position_embeddings=64, deepstack_visual_indexes=[0, 1], + ) + encoder = Qwen3OmniMoeVisionEncoder._from_config(config, attn_implementation=attn_implementation).eval() + return VisionEncoderSubmodule(encoder, config=None) + + +def _image(vision): + grid, pixels = torch.tensor([1, 4, 6]), torch.randn(24, 3 * 2 * 16 * 16) + inputs = {"pixel_values": [pixels], "image_grid_thw": [grid]} + return grid, pixels, vision.prepare_inputs("prefill_vision", None, inputs) + + +def test_the_vision_encoder_embeds_an_image_from_its_prepared_inputs_as_it_would_alone(): + vision = _vision() + grid, pixels, prepared = _image(vision) + encoder = vision.vision_encoder + with torch.no_grad(): + ours = vision.forward("prefill_vision", None, **prepared.tensor_inputs) + alone = encoder(pixels, grid_thw=grid[None]) + + assert torch.equal(ours["vision_embeds"][0], alone.pooler_output) and all( + torch.equal(layer, own) for layer, own in zip(ours["deepstack"], alone.deepstack_features, strict=True) + ), "the inputs built before the compiled call embed the image otherwise than the encoder does" + + +def test_a_prompt_with_a_video_keys_only_its_leading_text(model): + out = _prompt( + model, ["text", "video", "text"], {"video_inputs": [torch.zeros(4, 3, 64, 96)]}, + input_metadata={"video_inputs": [{"video_sample_fps": 2.0}]}, + ) + layout = out.metadata["kv_layout"][THINKER_KV]["main"] + + assert [span.walk for span in layout.spans[:layout.keyed_count()]] == ["prefill_text"], ( + "a prompt with a video keyed other than its leading text" + ) + +def _audio(attn_implementation: str = "sdpa") -> AudioEncoderSubmodule: + # The tensors `prepare_inputs` builds depend only on the clip's length and the encoder's + # windows, so a small encoder is enough. + config = Qwen3OmniMoeAudioEncoderConfig( + encoder_layers=1, d_model=32, encoder_attention_heads=2, encoder_ffn_dim=64, output_dim=16, + downsample_hidden_size=8, + ) + encoder = Qwen3OmniMoeAudioEncoder._from_config(config, attn_implementation=attn_implementation).eval() + return AudioEncoderSubmodule(encoder, config=None) + + +def _clip(frames: int, device: str = "cpu") -> dict: + return {"audio_features": [torch.randn(128, frames, device=device)], + "audio_seqlens": [torch.tensor([frames], device=device)]} + + +# shorter than one chunk, shorter than one attention window, and three windows +@pytest.mark.parametrize("frames", [60, 600, 1773]) +def test_the_audio_encoder_embeds_a_clip_from_its_prepared_inputs_as_it_would_alone(frames): + audio, clip = _audio(), _clip(frames) + prepared = audio.prepare_inputs("prefill_audio", None, clip) + # sdpa ignores `max_seqlen`, so comparing the outputs can't test it + assert audio._max_seqlen >= prepared.tensor_inputs["cu_seqlens"].diff().max(), ( + "flash-attn would be told a bound shorter than the clip's longest attention window" + ) + with torch.no_grad(): + ours = audio.forward("prefill_audio", None, **prepared.tensor_inputs) + alone = audio.audio_encoder(clip["audio_features"][0], feature_lens=clip["audio_seqlens"][0]) + + assert torch.equal(ours["audio_embeds"][0], alone.last_hidden_state), ( + "the inputs built before the compiled call embed the clip otherwise than the encoder does" + ) + + +def test_the_audio_encoder_compiles_with_the_default_while_its_output_is_cached(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch, "compile", lambda fn, dynamic, **kwargs: dynamic) + engine = Engine.__new__(Engine) + engine._device = torch.device("cuda") + engine._item_encoders = {"audio_encoder"} + engine._submodules = {"audio_encoder": SimpleNamespace(submodule=_audio(), forward="eager")} + + engine._compile_submodules() + + assert engine._submodules["audio_encoder"].forward is None, ( + "the audio encoder compiled with dynamic=True or stayed eager while its output is cached" + ) + +@requires_cuda +@pytest.mark.parametrize(("node", "seen", "unseen"), [ + ("vision", [(1, 4, 6)], [(1, 8, 10)]), + # The clips of 2, 60 and 50100 frames have a dimension of size 1 inside the encoder, so they compile + # graphs of their own. + ("audio", [1773, 60, 50100, 2], [600, 35]), +]) +def test_an_encoder_compiles_no_graph_for_a_size_inside_the_ranges_it_has_compiled(node, seen, unseen): + pytest.importorskip("flash_attn") + torch._dynamo.reset() + if node == "vision": + submodule, walk = _vision("flash_attention_2"), "prefill_vision" + + def inputs(grid): + return {"pixel_values": [torch.randn(grid[0] * grid[1] * grid[2], 3 * 2 * 16 * 16, device="cuda")], + "image_grid_thw": [torch.tensor(grid, device="cuda")]} + else: + submodule, walk = _audio("flash_attention_2"), "prefill_audio" + + def inputs(frames): + return _clip(frames, "cuda") + + submodule = submodule.to(device="cuda", dtype=torch.bfloat16) + engine = Engine.__new__(Engine) + engine._device = torch.device("cuda") + engine._item_encoders = {node} + engine._submodules = {node: SimpleNamespace(submodule=submodule)} + engine._compile_submodules() + forward = engine._submodules[node].forward + + def encode(size): + prepared = submodule.prepare_inputs(walk, None, inputs(size)) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + forward(walk, None, **prepared.tensor_inputs) + + for size in seen: + encode(size) + frames = counters["frames"]["ok"] + with torch._dynamo.config.patch(error_on_recompile=True): + for size in unseen: + encode(size) + + assert counters["frames"]["ok"] == frames, ( + f"the {node} encoder compiled a graph for a size inside a range it had compiled" + ) diff --git a/test/modular/test_qwen_tts_install.py b/test/modular/test_qwen_tts_install.py index 4490d6a8eb..f6ca8eabdf 100644 --- a/test/modular/test_qwen_tts_install.py +++ b/test/modular/test_qwen_tts_install.py @@ -89,4 +89,5 @@ def floor(extra): for extra in ("all", "qwen3_tts"): assert not floor(extra).contains("4.57.2"), extra - assert floor(extra).contains("4.57.3") and floor(extra).contains("5.19.0"), extra + assert floor(extra).contains("5.19.0"), extra + assert floor("qwen3_tts").contains("4.57.3")