feat(qwen3-omni): multimodal encoders migrated into M* and optimized for text gen by t-avil · Pull Request #185 · mstar-project/mstar

t-avil · 2026-07-23T23:44:37Z

POC Single-config Qwen3-Omni serving that beats vLLM-Omni 0.24 on text (#131)

This PR implements #131. vLLM-Omni is one of the most heavily contributed serving paths for this model, and the bar moved mid-effort: the goal shifted from matching the 0.21 paper version to matching-or-beating the newest 0.24, which substantially improved text generation. This follows up on the work in #150 and carries the text-generation result across that gap — one served configuration that edges ahead of 0.24 on text at serving concurrency while keeping M*'s speech lead.

In one line: ~6–25% faster than vLLM-Omni 0.24 on text throughput at serving batch sizes, and ~2–3× on speech (request throughput).

Results vs vLLM-Omni 0.24 (B32, 2× H200, continuous batching, closed-loop, natural-EOS)

path	this work	vLLM-Omni 0.24	ratio
i2t (tok/s)	1842	1745	1.06×
s2t (tok/s)	855	665	1.29×
i2s (req/s)	2.34	1.14	2.05×
s2s (req/s)	14.8	4.9	3.03×
t2s (req/s)	1.59	0.83	1.91×

Honest scope. i2t wins every batch (+6% to +25%; strongest at B4–B16). s2t wins at B4/B8/B32 (+24–29%) and is within noise of 0.24 at low concurrency (B1/B2 ≈ 0.98×) and a wash at B16 (≈1.00×). Speech wins every batch: i2s ~2.0–2.6×, s2s ~2.5–3.2×, t2s ~1.6–2.05× (request throughput — on audio-seconds/s the speech margin is smaller). B32 text cells are host-load sensitive (M* decode is host/GIL-bound). The
vllm024 and upmain reference series should be re-run through your own pipeline before defending those cells.

Full raw data, the 3-series charts, env, and reproduce steps live on the companion branch (linked, not a separate PR):

https://github.com/t-avil/mstar/tree/encoders-implemented-v2-benchmarked/benchmarks/qwen3-omni-131

What's implemented

Every optimization is an MSTAR_* env flag, default-OFF and byte-identical when off; fp8/custom-ops are bounded-to-rounding. CPU parity tests pin the gates.

Replicated-encoder placement + output-modality routing (conductor) — encoders as BF16 DP-replicas on both ranks; each request's encode routed to the idle rank by output modality (text→rank 0, speech→rank 1). Recovers both per-modality topology optima under one served config, no topology swap.
Multiprocess preprocess pool — off-process multimodal preprocess; removes the i2t first-token (TTFT) blow-up.
Host-floor decode stack — sidecar/integer/deferred check-stop, ordered + batched emit, lower-overhead routing/send, on-device batched position prep, cached sampler config. Cuts the per-step Python/GIL decode floor (M* decode is host-bound).
fp8 grouped-GEMM MoE + torch.library custom ops (compiled Thinker graph stays intact under fp8).
Chunked + captured-mixed prefill with wider vision/prefill capture grids.
Occupancy-auto-gated s2t audio-prefill merge.

Note:This is a fairly large change. I'm happy to adopt this PR either partially (keeping only the changes that align with the roadmap) or iteratively(splitting it into multiple PRs and merging the changes one by one).

The archived implementation combined all of these features, so in this draft pr I wanted to demonstrate what the final result could look like if these optimizations were adopted together

…text (mstar-project#131) Adds one served Qwen3-Omni configuration (configs/qwen3omni_2gpu_dpenc.yaml) that beats vLLM-Omni 0.24 on text throughput while keeping M*'s 2-3x speech request throughput lead. Follows up on mstar-project#150; measured on 2x H200, continuous batching, closed-loop, natural-EOS. Mechanism and features (every feature is an MSTAR_* env flag, default-OFF and byte-identical when off; fp8/custom-ops are bounded-to-rounding): * Replicated-encoder placement + output-modality routing (conductor): the encoder node group runs as BF16 DP-replicas on both ranks; each request's encode is routed to the idle rank by output modality (text->rank 0, speech->rank 1), recovering both per-modality topology optima under one served config. * Multiprocess preprocess pool: off-process multimodal preprocess removes the image-to-text first-token blow-up. * Host-floor decode stack: sidecar/integer/deferred check-stop, ordered + batched token emit, lower-overhead routing/send, on-device batched position prep, cached per-request sampler config -- cut the per-step Python/GIL decode floor. * fp8 grouped-GEMM MoE + torch.library custom ops (compiled Thinker graph stays intact under fp8). * Chunked + captured-mixed prefill with wider vision/prefill capture grids. * Occupancy-auto-gated speech-to-text audio-prefill merge. Adds CPU parity tests for the feature gates. ruff-clean; no design/markdown docs (the full method write-up, benchmark data, and charts live on the companion encoders-implemented-v2-benchmarked branch).

…d-v2 # Conflicts: # mstar/worker/worker.py

… PLR0913 ruff 0.16 enables PLR0917 by default; it fires across the existing codebase (kernels, samplers, tests). The repo already ignores every other arg-count rule (PLR0911/0912/0913/0915), so ignore its positional-args sibling too.

NSagan271

I went through the PR and flagged anything that stood out to me, but it was too large for me to look through everything thoroughly. I do think that this PR is too long to review / validate by itself, and would prefer for it to be broken down into several smaller PRs. In addition, the length of the comments and large number of test files makes it hard to read the code.

There is also a lot of qwen3omni-specific code added to the system (data preprocess worker, worker, micro scheduler, etc.) that needs to be properly abstracted into features that any model could, in theory, take advantage of. I flagged such code where I saw it, but I did not have enough time to look through every single change in this pass, so just keep an eye out.

Also, there are a lot of different changes here, and they conceptually make sense, but I'm unclear specifically which changes result in big performance wins, and which just add complexity to the system. I think there are too many changes to do a full A/B on everything, but it would be nice to know what are the critical pieces that should get merged more quickly (because I think it could take a while to properly design/review/validate/refine/test everything).

cc @merceod

P.S. (we can keep this PR around for paper rebuttals if needed, but in its current state it needs a lot of work to be able to be merged into main)

NSagan271 · 2026-07-24T13:32:11Z

There are a large amount of test files; can you commit a subset?

NSagan271 · 2026-07-24T14:10:41Z

+        # Requires a real model with a tokenizer (postprocess); skipped for the
+        # dummy/None model so tests and dummy configs are untouched.
+        self.detok_client: DetokClient | None = None
+        detok_on = os.environ.get("MSTAR_DETOK_PROC", "0") == "1"


Environment variables are now documented in docs/environment_variables.rst, can you add the environment variables for this PR there? Also, if some features help situationally (e.g., PREPROC_PROC helps for multimodal-processing-heavy models expected run under high batch sizes, but wouldn't be useful for, e.g., pi0.5), that should be documented.

NSagan271 · 2026-07-24T14:17:18Z

-            if result.request_id not in self.per_request_reading_tensors:
-                logger.warning(
-                    "Late result chunk for cleaned-up request %s, ignoring",
+            # Tolerant decrement: a LATE chunk can land after cleanup_request


Nit: some of the comments can be made more concise, and especially the docstrings at the beginning of new files.

NSagan271 · 2026-07-24T14:20:29Z

+            # popped this rid's counter (in-flight SHM read completing in the
+            # main-thread-pop -> worker-thread-cleanup window; ordered emit
+            # widens it by holding+late-flushing chunks). A bare subscript
+            # here raised KeyError and aborted the whole drain loop, dropping


To make debugging these kinds of issues easier in the future, can you wrap the self.preprocess_worker.get_result_chunks() call in entrypoint.py in a try... except...?

NSagan271 · 2026-07-24T14:56:33Z

+        # present): a CUDA-hidden child would produce a CPU mel that is not
+        # bit-identical, so such requests stay inline (byte-identical). i2t
+        # (image+text) is unaffected — its preprocessing is CPU-deterministic.
+        self._serve_uses_gpu_mel = (


This feels like a very model-specific flag; the data worker should be as model-agnostic as possible to keep the abstractions clean (and avoid accumulating model-specific code, which would get messy to deal with).

What self._serve_uses_gpu_mel is actually doing here is to determine whether preprocessing outputs should go to the child pool or stay in the current process. I think that could be made better-abstracted by moving def _should_offload(self, input: PreprocessInput) -> bool: to the model base class (with a default to return True --- the preproc_enabled gate can remain in the data worker) and probably a different name, like should_offload_preprocess. And a model like Qwen3-omni can override this with the specific "if GPU mel then don't offload, otherwise ok to offload" logic.

NSagan271 · 2026-07-24T16:29:37Z

+
+
+class FlashInferSplitMixedWrapper:
+    """Split attention for a captured thinker_mixed step (MSTAR_MIXED_SPLIT_ATTN).


Nit: can you make this docstring model-agnostic?

NSagan271 · 2026-07-24T16:36:10Z

                    must_yield_away = (
                        consecutive_spec_steps >= max_consecutive_spec
                        or must_yield_for_fairness
                    )
-                    if not must_yield_away:
+
+                    # MSTAR_ADMIT_FASTPATH (fix #4): arrival-triggered admission.


The long comments make these changes pretty hard to parse. Can you heavily reduce the comments here? And also, it might make sense to break some logic into helper functions (but I've not looked too closely).

NSagan271 · 2026-07-24T16:43:54Z

+        #   admitted via the normal ready path) so the guess is uniform decode,
+        #   not heterogeneous. This is what keeps the folded mixed step INSIDE
+        #   the chain instead of breaking it.
+        if batch.graph_walk == "thinker_mixed":


I do think mixed batches are a good idea, but it needs a proper model-agnostic design to be integrated into our system. Right now, it adds a lot of long, model-specific code to worker.py and micro_scheduler.py, which makes everything harder to read/debug, and also harder for another model to implement the same optimizations.

I haven't read through the current mixed batch code here thoroughly enough to create a full design, but I think the model should have a function that declares a list of MixedBatchConfig dataclasses, where the dataclass includes walks that can be merged together, the name of the merged walk, any other fields that can fit well into a config (default on the base class: empty list; no mixed batch).
It will probably need function handles to any specific logic that can't just be config-driven, but I've not scoped that out.

If you pull the mixed batch logic out into another PR, I can help on the design end.

NSagan271 · 2026-07-24T16:45:43Z

+        etype = engine.engine_type()
+        if etype == EngineType.STATELESS:
+            return True
+        if etype == EngineType.KV_CACHE and batch.graph_walk.startswith("prefill"):


The string matching logic is brittle; it works and can be a good default, but having the option for a model-specific override would be more robust.

NSagan271 · 2026-07-24T16:57:48Z

        self.nodes = set(self.worker_graph.section.get_nodes().keys())
        self.loops = set(self.worker_graph.section.get_loops().keys())

+        # _native_bridge stays None (no C++ state-machine bridge on this path);


I don't see _native_bridge being passed as non-None anywhere in this PR. Also, in terms of adding a C++ state machine, the repo already has an optional Rust dependency, and a draft PR to port some of the graph logic to Rust exists in #171 (though I have not closely reviewed the PR). So instead of this native bridge, we can either get #171 merged sooner, or you can coordinate with the author of that PR to integrate whatever features are needed?

t-avil · 2026-07-25T23:57:06Z

Thanks for the detailed review, @NSagan271!

I was mainly looking for feedback on which features are worth adopting (which you did -thank you!), so I really appreciate you taking the time to go through everything. The goal was to demonstrate the combined result and get feedback on which optimizations best fit the roadmap (as mentioned in the PR description).

I'll start breaking it into smaller, feature-focused PRs and leave this one open as a reference for now.

That said, do you think #150 would be worth merging first to close #131, while keeping this PR open until we've safely extracted the useful pieces?

cc @merceod

NSagan271 · 2026-07-26T03:46:23Z

@t-avil Sounds good! I do agree on getting #150 merged. Can you look into my comments on the piecewise cuda graph runner for that? I'm pushing that abstraction for this because it's a good opportunity to refine the piecewise runner to be more generally useful, so open to feedback on changing the abstraction.

t-avil requested a review from merceod July 23, 2026 23:54

t-avil changed the title ~~feat(qwen3-omni): single-config serving that beats vLLM-Omni 0.24 on …~~ feat(qwen3-omni): multimodal encoders migrated into M* and optimized for text gen Jul 24, 2026

Merge remote-tracking branch 'upstream/main' into encoders-implemente…

5025c10

…d-v2 # Conflicts: # mstar/worker/worker.py

merceod requested review from Gaurav-Shah05, NSagan271 and stephen-dwq July 24, 2026 05:24

NSagan271 requested changes Jul 24, 2026

View reviewed changes



		class FlashInferSplitMixedWrapper:
		"""Split attention for a captured thinker_mixed step (MSTAR_MIXED_SPLIT_ATTN).

Uh oh!

Conversation

t-avil commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

POC Single-config Qwen3-Omni serving that beats vLLM-Omni 0.24 on text (#131)

Results vs vLLM-Omni 0.24 (B32, 2× H200, continuous batching, closed-loop, natural-EOS)

What's implemented

Uh oh!

NSagan271 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

t-avil commented Jul 25, 2026

Uh oh!

NSagan271 commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

t-avil commented Jul 23, 2026 •

edited

Loading

NSagan271 left a comment •

edited

Loading

NSagan271 commented Jul 26, 2026 •

edited

Loading