Qwen3-Omni multimodal encoders migrated into M* and relevant paths (s2s, s2t, i2t, i2s) optimized.

benchmark/dataset.py

@@ -485,12 +485,14 @@ def __init__(
             )
         self.prompts = prompts
 
-        raw = load_dataset(
-            "ethz/food101",
-            split=split,
-            cache_dir=cache_dir,
-            trust_remote_code=True,
-        )
+        # ethz/food101 is now distributed as Parquet; newer `datasets` rejects
+        # `trust_remote_code`, so don't pass it (fall back only on older
+        # `datasets` that still require a loading script).
+        try:
+            raw = load_dataset("ethz/food101", split=split, cache_dir=cache_dir)
+        except (TypeError, ValueError):
+            raw = load_dataset("ethz/food101", split=split, cache_dir=cache_dir,
+                               trust_remote_code=True)
         # Shuffle with a fixed seed so we get class-diverse images while
         # remaining deterministic across runs (so the same N images appear in
         # the same order on every benchmark invocation).

benchmark/runner.py

@@ -257,8 +257,17 @@ def _pct(values, p):
                 "output_bytes": dict(m.output_bytes),
             })
 
+        import dataclasses
+
+        def _ls(x):
+            return dataclasses.asdict(x) if x is not None else None
+
+        def _lsmap(d):
+            return {k: _ls(v) for k, v in (d or {}).items()}
+
         payload = {
             "system": "ours",
+            "inference_system": self.config.inference_system.value,
             "model": getattr(self.config.model, "__class__", type(self.config.model)).__name__,
             "request_type": self.config.request_type.value,
             "profiling_type": self.config.profiling_type.value,
@@ -273,6 +282,15 @@ def _pct(values, p):
             "jct_p95_ms": _pct(jcts_ms, 95),
             "jct_p99_ms": _pct(jcts_ms, 99),
             "request_throughput": (agg.request_throughput or 0.0),
+            "text_token_throughput": agg.text_token_throughput,
+            "audio_seconds_throughput": agg.audio_seconds_throughput,
+            "audio_duration_mean_s": agg.audio_duration_mean_s,
+            "batch_size": agg.batch_size,
+            "max_concurrency": agg.max_concurrency,
+            "ttft": _lsmap(agg.ttft),
+            "itl": _lsmap(agg.itl),
+            "e2e_latency": _ls(agg.e2e_latency),
+            "rtf": _ls(agg.rtf),
             "per_request": per_request,
         }
 

mstar/engine/cuda_graph_config.py

@@ -215,6 +215,13 @@ class PiecewiseCudaGraphConfig(ABC):
     #     seq_lens for BATCHED, packed indptr for PACKED).
     uses_kv_cache: bool = False
     plan_fn: Callable[["BatchedCacheManager", PiecewiseCaptureShape], None] | None = None
+    # (5) attention planning without a KV cache, for regions running FlashInfer
+    #     ragged varlen: the plan is host-side, so it cannot live in the graph.
+    #     make_attn_state builds one wrapper per bucket (needs use_cuda_graph=True
+    #     so replanning reuses the captured index buffers); plan_attn_fn replans it
+    #     before each replay. Ignored when uses_kv_cache is True.
+    make_attn_state: Callable[[PiecewiseCaptureShape], Any] | None = None
+    plan_attn_fn: Callable[[Any, PiecewiseCaptureShape, list[int]], None] | None = None
     # Whether the runner advances cache seq_lens (Python-only, post-replay) after
     # each ``run``. True suits the common case where the captured region consumes
     # its planned tokens once per step. Set False when the caller advances the

mstar/engine/cuda_graph_runner.py

@@ -93,6 +93,9 @@ class PiecewiseGraphData:
     static_cache_manager: BatchedCacheManager | None
     dummy_rids: list[str]
     shape: PiecewiseCaptureShape
+    # Per-bucket no-KV attention wrapper, replanned outside the graph per replay.
+    # None for KV-cache configs, which plan through static_cache_manager.
+    attn_state: Any = None
 
 @dataclass
 class CudaGraphData:
@@ -2506,6 +2509,13 @@ def warmup_and_capture(self) -> None:
     def _capture_one(self, shape: PiecewiseCaptureShape) -> None:
         static_inputs = self.config.make_static_inputs(shape)
         static_cm, dummy_rids = self._setup_cache_manager(shape)
+        # Built once so the graph captures its index buffers; replanned, never
+        # rebuilt, per replay.
+        attn_state = (
+            self.config.make_attn_state(shape)
+            if static_cm is None and self.config.make_attn_state is not None
+            else None
+        )
 
         fn = self.config.capture_fn
         if self.config.compile:
@@ -2517,15 +2527,21 @@ def _capture_one(self, shape: PiecewiseCaptureShape) -> None:
             )
 
         def run_fn():
+            extra = {} if attn_state is None else {"attn_state": attn_state}
             return fn(
                 static_inputs=static_inputs,
                 static_cm=static_cm,
+                **extra,
                 **self.config.forward_kwargs,
             )
 
         def plan():
             if static_cm is not None:
                 self._plan(static_cm, shape)
+            elif attn_state is not None and self.config.plan_attn_fn is not None:
+                # Bucket's own partition sums to total_tokens, so capture sees the
+                # widest indptr any replay can ask for.
+                self.config.plan_attn_fn(attn_state, shape, list(shape.seq_lens))
 
         plan()
 
@@ -2557,6 +2573,7 @@ def plan():
             static_cache_manager=static_cm,
             dummy_rids=dummy_rids,
             shape=shape,
+            attn_state=attn_state,
         )
 
     def _setup_cache_manager(
@@ -2787,6 +2804,15 @@ def run(
                 data.shape,
                 seq_lens=self._replay_seq_lens(data.shape, seq_lens, real_bs),
             )
+        elif data.attn_state is not None and self.config.plan_attn_fn is not None:
+            # _replay_seq_lens zero-pads to shape.bs, so the indptr sums to the
+            # real token count and the pad tail lands in no segment; ragged
+            # attention is block-diagonal, so it cannot leak into real segments.
+            self.config.plan_attn_fn(
+                data.attn_state,
+                data.shape,
+                self._replay_seq_lens(data.shape, seq_lens, real_bs),
+            )
 
         # --- 3: replay ---
         data.graph.replay()

mstar/engine/stateless_engine.py

@@ -543,16 +543,24 @@ def _apply_torch_compile(self, node_name: str, submodule: NodeSubmodule) -> None
                 self.config.name, node_name,
             )
             return
+        # ``dynamic`` defaults to None (Inductor decides). A submodule whose
+        # forward only ever sees a fixed set of shapes -- e.g. one driven by
+        # piecewise capture buckets -- can set ``torch_compile_dynamic = False``
+        # to force static specialization. Leaving the default is important for
+        # variable-shape submodules, which would otherwise recompile per shape.
+        dynamic = getattr(submodule, "torch_compile_dynamic", None)
         try:
             if hasattr(submodule, "forward"):
                 submodule.forward = torch.compile(
                     submodule.forward,
                     fullgraph=False,
+                    dynamic=dynamic,
                 )
                 logger.info(
-                    "StatelessEngine[%s]: torch.compile applied to %s.forward",
+                    "StatelessEngine[%s]: torch.compile applied to %s.forward (dynamic=%s)",
                     self.config.name,
                     node_name,
+                    dynamic,
                 )
             # forward_batched is intentionally left eager — Inductor would pay
             # a ~30s one-shot trace cost for dynamic varlen shapes on the

mstar/model/components/encoder_telemetry.py

@@ -0,0 +1,50 @@
+"""Piecewise-capture telemetry, shared by encoders.
+
+A captured replay and a silent eager fallback produce identical results, so the
+path taken is counted rather than inferred; ``encoder_path_counts()`` lets a
+benchmark assert it measured the path it meant to.
+"""
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+_ENCODER_PATH_COUNTS: dict[str, int] = {}
+_SEEN_LAYOUTS: set[tuple] = set()
+_SEEN_LAYOUTS_CAP = 512          # bounded: a long run must not leak keys
+_WARNED_NO_BUCKET: set[str] = set()
+
+
+
+def note_encoder_path(path: str) -> None:
+    _ENCODER_PATH_COUNTS[path] = _ENCODER_PATH_COUNTS.get(path, 0) + 1
+
+
+def encoder_path_counts() -> dict[str, int]:
+    """Snapshot of {path: count}, e.g. {"vision.piecewise": 96, "vision.eager": 0}."""
+    return dict(_ENCODER_PATH_COUNTS)
+
+
+def note_encoder_layout(kind: str, n_seg: int, total_tokens: int, fitted: bool) -> None:
+    """Log once per distinct layout; WARN the first time one fits no bucket —
+    the only visible signal that buckets are mis-sized, since the fallback is
+    otherwise silent."""
+    if not fitted and kind not in _WARNED_NO_BUCKET:
+        _WARNED_NO_BUCKET.add(kind)
+        logger.warning(
+            "%s encoder: NO capture bucket fits segments=%d total_tokens=%d — "
+            "falling back to eager. Widen CAPTURE_BATCH_SIZES_%s / "
+            "CAPTURE_TOKENS_%s; piecewise numbers from this run are "
+            "NOT measuring the captured path.",
+            kind, n_seg, total_tokens, kind.upper(), kind.upper(),
+        )
+    key = (kind, n_seg, total_tokens)
+    if key in _SEEN_LAYOUTS:
+        return
+    if len(_SEEN_LAYOUTS) < _SEEN_LAYOUTS_CAP:
+        _SEEN_LAYOUTS.add(key)
+    logger.info("%s encoder layout: segments=%d total_tokens=%d bucket=%s",
+                kind, n_seg, total_tokens, "HIT" if fitted else "MISS")
+
+