Qwen3-Omni: encoders on ragged attention, no flash_attn

README.md

@@ -70,7 +70,7 @@ The PyPI package is `mstar-ai` (`mstar-project` and `mstar-serve` install the sa
 and the CLI are `mstar`. Working from a checkout instead: `uv pip install --torch-backend=auto -e .[all]`.
 Qwen3-TTS also needs `uv pip install --no-deps "qwen-tts==0.1.1"` (see `docs/installation.rst`).
 
-To enable flash-attn support (required for Qwen3-Omni, recommended for BAGEL),
+To enable flash-attn support (recommended for BAGEL),
 ```bash
 # torch 2.10 built for CUDA 13.x (cu13); on a cu12 torch, build from source (docs/installation.rst)
 uv pip install \

benchmark/eval/e2e/cases.yaml

@@ -221,7 +221,7 @@ models:
     config: configs/qwen3omni_colocated.yaml
     family: Omni
     tags: [omni, llm, text, image-in, audio-in, video-in, speech-out, sse, big]
-    requires: [flash_attn, qwen_omni_utils, audioread]
+    requires: [qwen_omni_utils, audioread]
     est_min: 16
     cases:
       - id: text
@@ -313,7 +313,7 @@ models:
     env: {MSTAR_ALLOW_REMOTE: "1"}
     family: Omni
     tags: [omni, rust-frontend, image-in, sse, big]
-    requires: [flash_attn, qwen_omni_utils]
+    requires: [qwen_omni_utils]
     est_min: 14
     cases:
       - id: text
@@ -458,7 +458,7 @@ models:
     optional: true
     family: Omni
     tags: [mooncake, transport, omni, image-in, speech-out, big]
-    requires: [flash_attn, qwen_omni_utils, mooncake.engine]
+    requires: [qwen_omni_utils, mooncake.engine]
     est_min: 10
     cases:
       - id: text
@@ -540,7 +540,7 @@ models:
     case_timeout_s: 900
     family: Omni
     tags: [multi-gpu, tp, omni, image-in, speech-out, big]
-    requires: [flash_attn, qwen_omni_utils, audioread]
+    requires: [qwen_omni_utils, audioread]
     est_min: 16
     conc: [1, 16]
     cases:
@@ -634,7 +634,7 @@ models:
     case_timeout_s: 900
     family: Omni
     tags: [multi-gpu, disaggregated, omni, speech-out, big]
-    requires: [flash_attn, qwen_omni_utils, audioread]
+    requires: [qwen_omni_utils, audioread]
     est_min: 16
     cases:
       - id: text
@@ -879,7 +879,7 @@ models:
     health: true
     family: Omni
     tags: [w2w, transport, shm, disaggregated, omni, speech-out, big]
-    requires: [flash_attn, qwen_omni_utils, audioread]
+    requires: [qwen_omni_utils, audioread]
     est_min: 18
     cases: &qwen3omni_pd_w2w
       - id: text
@@ -917,7 +917,7 @@ models:
     health: true
     family: Omni
     tags: [w2w, transport, disaggregated, omni, speech-out, big]
-    requires: [flash_attn, qwen_omni_utils, audioread, mooncake.engine]
+    requires: [qwen_omni_utils, audioread, mooncake.engine]
     est_min: 18
     cases: *qwen3omni_pd_w2w
 
@@ -931,7 +931,7 @@ models:
     health: true
     family: Omni
     tags: [rdma, disaggregated, omni, speech-out, big]
-    requires: [flash_attn, qwen_omni_utils, audioread, mooncake.engine, rdma_nic]
+    requires: [qwen_omni_utils, audioread, mooncake.engine, rdma_nic]
     est_min: 18
     cases: *qwen3omni_pd_w2w
 

docs/installation.rst

@@ -135,9 +135,8 @@ Model families and some output formats need extra packages, exposed as pip *extr
    * - ``.[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). 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)`_.
+       ``transformers>=5.18``. Its encoders attend through FlashInfer, so it needs no
+       ``flash-attn``.
    * - ``.[orpheus]``
      - Orpheus TTS runtime: ``transformers``, ``flashinfer-python``, ``safetensors``,
        ``einops``, ``huggingface-hub``, ``mooncake-transfer-engine``.
@@ -191,7 +190,7 @@ Model families and some output formats need extra packages, exposed as pip *extr
        Convenient for a machine that serves multiple models; heavier and slower to
        install than a single family's extra. It excludes the separately installed
        TAEHV, ``flash-attn`` and ``qwen-tts`` packages; see below,
-       `flash-attn (Qwen3-Omni, LingBot-World)`_ and `qwen-tts (Qwen3-TTS)`_.
+       `flash-attn (LingBot-World)`_ and `qwen-tts (Qwen3-TTS)`_.
 
 Combine extras as needed (keep ``--torch-backend=auto`` on every install):
 
@@ -225,8 +224,7 @@ Or, in an existing Python 3.12 environment:
 
    If you're just getting started or have the disk/time to spare, ``.[all]`` installs all
    index-hosted model dependencies in one shot. Waypoint still needs the pinned TAEHV
-   command above, Qwen3-Omni still needs ``flash-attn``, and Qwen3-TTS still needs
-   ``qwen-tts`` (see `qwen-tts (Qwen3-TTS)`_):
+   command above, and Qwen3-TTS still needs ``qwen-tts`` (see `qwen-tts (Qwen3-TTS)`_):
 
    .. code-block:: bash
 
@@ -246,8 +244,8 @@ The GPU model families depend on:
 
 - **FlashInfer** (``flashinfer-python``) — paged attention and continuous batching for the
   autoregressive backbones (every model with a ``KV_CACHE`` node runs attention through it).
-- **flash-attn** — used by Qwen3-Omni and by LingBot-World's ``reference_numerics`` mode. **Not installed by any extra**; install it separately
-  (see `flash-attn (Qwen3-Omni, LingBot-World)`_).
+- **flash-attn** — used by LingBot-World's ``reference_numerics`` mode. **Not installed by any extra**; install it separately
+  (see `flash-attn (LingBot-World)`_).
 - **flash-attn-4** — Waypoint's flex-attention ``FLASH`` backend. **Not installed by any
   extra**; install it separately (see `flash-attn-4 (Waypoint, FA3 sm90 kernel)`_).
 - **mooncake-transfer-engine** (0.3.13.post1 or newer) — RDMA tensor transport for
@@ -304,11 +302,11 @@ mstar reads only the speech-tokenizer source files and runs them under transform
 lists them (``transformers==4.57.3``, ``accelerate==1.12.0``, ``gradio``, ``sox`` and
 others) as unmet for ``qwen-tts``; that is expected.
 
-flash-attn (Qwen3-Omni, LingBot-World)
---------------------------------------
+flash-attn (LingBot-World)
+--------------------------
 
-``flash-attn`` is only needed for **Qwen3-Omni** and **LingBot-World** with ``reference_numerics: true``, and it is **not** pulled in by
-``.[qwen3_omni]``, ``.[lingbot_world]`` or ``.[all]`` — you install it as a separate step. The reason: flash-attn
+``flash-attn`` is only needed for **LingBot-World** with ``reference_numerics: true``, and it is **not** pulled in by
+``.[lingbot_world]`` or ``.[all]`` — you install it as a separate step. The reason: flash-attn
 publishes no wheels on PyPI, so ``pip``/``uv`` fall back to compiling it from source, which is
 slow. When a prebuilt GitHub wheel matches your stack, use it to skip the build; otherwise
 build from source (see `No matching wheel — build from source`_ below).

docs/models.rst

@@ -493,6 +493,13 @@ Kokoro notes
 Qwen3-Omni notes
 ----------------
 
+- The audio and vision encoders attend through ragged attention resources
+  (``audio_attn`` within each 8 s window, ``vision_attn`` within each image or
+  video frame group), FlashInfer by default, so no ``flash-attn`` install is
+  needed. A deployment can pick the kernel under ``resources:`` in its YAML,
+  e.g. ``vision_attn: {flashinfer_backend: cutlass}``. Each encoder's block stack
+  is captured as a CUDA graph per token-count bucket; a request past the
+  largest bucket (32768 patches, about 5 minutes of audio) runs it eagerly.
 - Streamed audio starts after 12 codec frames (0.96 s of audio), then 13, then
   25-frame chunks behind 25 frames of left context. Playback stays gapless
   while the Talker produces frames at about 1.5x real time over the first 50

mstar/model/components/aut_submodule.py

@@ -0,0 +1,220 @@
+"""The shared AuT encoder node (Qwen3-ASR, Qwen3-Omni): requests packed into
+one forward, every 8 s attention window of every request one segment of the
+node's ragged-attention step, and the transformer stack captured by packed
+token count.
+
+A model subclasses it for ``prepare_inputs`` only, which turns its request's
+audio into ``(num_mel_bins, frames)`` features (``feature_inputs``).
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from mstar.communication.tensors import NameToTensorList
+from mstar.engine.accelerator_graph_config import (
+    PiecewiseAcceleratorGraphConfig,
+    PiecewiseCallInputs,
+    PiecewiseCaptureShape,
+    PiecewisePackedConfig,
+)
+from mstar.engine.engine import ExecutingBatch
+from mstar.engine.resources import AttentionStep, Segment, SlotLease, SubmoduleStep
+from mstar.model.components.aut_encoder import AuTEncoder, AuTEncoderConfig
+from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs, NodeSubmodule
+
+WINDOW_LABEL = "main"
+# The AuT encoder's transformer stack and projector, as a piecewise capture region
+AUT_ENCODE = "aut_encode"
+
+
+class AuTEncoderSubmodule(NodeSubmodule):
+    """Packed AuT encoder, emitting one ``(num_audio_tokens, output_dim)``
+    ``audio_embeds`` tensor per request. The runner plans the
+    windows' kernel layout before the forward, so nothing is recomputed
+    inside it."""
+
+    MAX_BATCH_SIZE = 8
+    # The engine's blanket torch.compile of ``forward_batched`` traces the
+    # per-request Python around the tensor work and guards on request ids
+    # and batch composition, so batches kept recompiling (1-2 s stalls each,
+    # measured on Whisper's decoder). The encoder's kernels are the ragged
+    # attention and cuBLAS GEMMs, there is nothing for inductor to fuse.
+    disable_torch_compile = True
+
+    # Packed token counts the stack is captured at, and batch sizes (up to
+    # MAX_BATCH_SIZE). 4096 tokens is 8 requests of ~39 s.
+    ENCODE_TOKEN_BUCKETS = [64, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096]
+    ENCODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8]
+
+    @classmethod
+    def max_windows_per_replay(cls, config: AuTEncoderConfig) -> int:
+        """Attention windows per request a captured replay can be handed, to
+        size the node's ragged resource: the largest bucket's tokens make at
+        most this many full windows, and each request adds one partial."""
+        return cls.ENCODE_TOKEN_BUCKETS[-1] // config.window_tokens + 1
+
+    def __init__(self, encoder: AuTEncoder, attn_key: str):
+        super().__init__()
+        self.encoder = encoder
+        self.attn_key = attn_key
+
+    def _param_dtype(self) -> torch.dtype:
+        return self.encoder.conv2d1.weight.dtype
+
+    def feature_inputs(self, features: torch.Tensor) -> NodeInputs:
+        """This node's inputs for one request's ``(num_mel_bins, frames)``
+        features, already on the encoder's device and in its dtype."""
+        num_frames = int(features.shape[-1])
+        return NodeInputs(
+            tensor_inputs={"audio_features": features},
+            kwargs={"num_frames": num_frames},
+            input_seq_len=self.encoder.config.tokens_for_frames(num_frames),
+        )
+
+    def declare_step(
+        self,
+        graph_walk: str,
+        request_ids: list[str],
+        inputs: list[NodeInputs],
+        slot_lease: SlotLease | None = None,
+        piecewise_leases: Mapping[str, SlotLease] | None = None,
+        **kwargs,
+    ) -> SubmoduleStep | None:
+        # leased, the captured stack plans its own attention per replay
+        if (piecewise_leases or {}).get(AUT_ENCODE):
+            return None
+        # one segment per attention window, in the packed order the forward
+        # lays the requests out in
+        segments = [
+            Segment(request_id=rid, label=WINDOW_LABEL, span=window)
+            for rid, inp in zip(request_ids, inputs, strict=True)
+            for window in self.encoder.config.window_lengths(inp.input_seq_len)
+        ]
+        return SubmoduleStep(segments=segments, steps={self.attn_key: AttentionStep(causal=False)})
+
+    def preprocess(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        inputs: list[NodeInputs],
+    ) -> dict[str, torch.Tensor | Any]:
+        lens = [inp.kwargs["num_frames"] for inp in inputs]
+        longest = max(lens)
+        padded = torch.stack([
+            F.pad(inp.tensor_inputs["audio_features"], (0, longest - n))
+            for inp, n in zip(inputs, lens, strict=True)
+        ])
+        return {"audio_features": padded, "feature_lens": lens}
+
+    def get_piecewise_accelerator_graph_configs(
+        self, device: torch.device, autocast_dtype: torch.dtype,
+        tp_world_size: int = 1, **kwargs,
+    ) -> dict[str, PiecewiseAcceleratorGraphConfig]:
+        """The encoder's stack and projector, by batch size and packed token
+        count. At ASR clip lengths (13 tokens a second) its GEMMs are a few
+        dozen rows, so eager it is launch-bound: a 7 s clip's 24-layer stack
+        is ~6.3 ms eager and ~1.9 ms replayed. The conv frontend before it
+        runs eagerly. Batches past the largest bucket run eagerly too."""
+        d_model = self.encoder.config.d_model
+        attn_key = self.attn_key
+
+        def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]:
+            return {"hidden": torch.zeros(shape.total_tokens, d_model, dtype=autocast_dtype, device=device)}
+
+        def declare_step(request_ids: list[str], seq_lens: list[int]) -> SubmoduleStep:
+            # Every window of every request is a segment; the kernel needs
+            # only the boundaries, so they all name the first row's id.
+            # Zero-length rows pad the batch to the bucket.
+            return SubmoduleStep(
+                segments=[Segment(request_id=request_ids[0], label=WINDOW_LABEL, span=n) for n in seq_lens],
+                steps={attn_key: AttentionStep(causal=False)},
+            )
+
+        return {
+            AUT_ENCODE: PiecewisePackedConfig(
+                capture_fn=self._capture_encode,
+                make_static_inputs=make_static_inputs,
+                declare_step=declare_step,
+                lease_before_step=True,
+                total_tokens=self.ENCODE_TOKEN_BUCKETS,
+                capture_batch_sizes=self.ENCODE_CAPTURE_BATCH_SIZES,
+            )
+        }
+
+    def _capture_encode(self, inp: PiecewiseCallInputs) -> dict[str, torch.Tensor]:
+        """The captured region: reads the runner's buffer, never rebinds."""
+        return {"embeds": self.encoder.encode(inp.static_inputs["hidden"], None)}
+
+    @torch.compiler.disable
+    def _replay_encode(
+        self, engine_inputs: ModelInputsFromEngine, hidden: torch.Tensor, windows: list[int],
+    ) -> torch.Tensor | None:
+        """The stack replayed for the whole packed batch, or None if not leased.
+
+        A leased step declares nothing (`declare_step`), so ``engine_inputs.step``
+        is None; a bucket that fits is not enough, as ``can_lease`` may have
+        refused the batch."""
+        runner = engine_inputs.piecewise_runners.get(AUT_ENCODE)
+        bs = len(engine_inputs.request_ids)
+        if runner is None or engine_inputs.step is not None or not runner.can_run(bs, int(hidden.shape[0])):
+            return None
+        # an owned copy: the encoder states leave this worker after the step,
+        # and the buffer is the next replay's
+        return runner.run(
+            static_inputs={"hidden": hidden},
+            request_ids=list(engine_inputs.request_ids),
+            seq_lens=windows,
+            real_bs=bs,
+        ).get("embeds")
+
+    def _encode(
+        self, audio_features: torch.Tensor, feature_lens: list[int],
+        engine_inputs: ModelInputsFromEngine | None = None,
+    ) -> list[torch.Tensor]:
+        enc = self.encoder
+        layout = enc.layout(feature_lens)
+        hidden = enc.frontend(audio_features, feature_lens)
+        embeds = None
+        if engine_inputs is not None:
+            embeds = self._replay_encode(engine_inputs, hidden, layout.window_lengths)
+        if embeds is None:
+            cu = torch.tensor(layout.cu_seqlens, dtype=torch.int32, device=hidden.device)
+            embeds = enc.encode(hidden, cu)
+        return list(embeds.split(layout.tokens_per_request, dim=0))
+
+    def forward(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        audio_features: torch.Tensor,
+        feature_lens: list[int],
+        **kwargs,
+    ) -> NameToTensorList:
+        if audio_features.dim() == 2:
+            audio_features = audio_features.unsqueeze(0)
+        return {"audio_embeds": [self._encode(audio_features, feature_lens, engine_inputs)[0]]}
+
+    def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool:
+        return True
+
+    def max_batch_size(self, graph_walk: str) -> int | None:
+        del graph_walk
+        return self.MAX_BATCH_SIZE
+
+    def forward_batched(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        audio_features: torch.Tensor,
+        feature_lens: list[int],
+        **kwargs,
+    ) -> dict[str, NameToTensorList]:
+        per_request = self._encode(audio_features, feature_lens, engine_inputs)
+        return {
+            rid: {"audio_embeds": [emb]}
+            for rid, emb in zip(engine_inputs.request_ids, per_request, strict=True)
+        }