Qwen3-Omni multimodal encoders migrated into M* and relevant paths (s2s, s2t, i2t, i2s) optimized. by t-avil · Pull Request #150 · mstar-project/mstar

t-avil · 2026-06-30T08:15:04Z

What does this PR do?

Closes #131. Replaces the HF-wrapper AudioEncoderSubmodule / VisionEncoderSubmodule
with from-scratch native M* encoder submodules, matching the already-native
Thinker/Talker/Code2Wav pattern.

Vision (components/vision_encoder.py): native ViT + spatial merge, producing the
final embeds and the DeepStack intermediate features. Patch-embed runs as an
F.linear (the Conv3d kernel==stride shape is pathologically slow in low precision).
Audio (components/audio_encoder.py): native Whisper-style AuT with a varlen
attention primitive (FlashInfer / flash-attn / SDPA backends).
Weight loading: submodule attribute names mirror HF exactly, so the existing
load_weights_from_hf_shards loads the checkpoint unchanged (no remap); load is
completeness-checked (raises on missing keys).
Output contract is unchanged — (merged_embeds, [deepstack...]) feeds
prefill_vision / prefill_audio exactly as before, so the Thinker side is untouched.
HF wrapper kept as an env-gated fallback (MSTAR_QWEN3_NATIVE_{AUDIO,VISION}_ENCODER=0).

Per the issue's "A/B each optimization, keep only what wins":

Cross-request batching: audio encoder varlen-packs concurrent requests (vision
batching is behind MSTAR_BATCH_VISION_PREFILL, opt-in).
Plus GPU log-mel, GPU image preprocess support, and a vision prefill sync-elimination path (MRoPE computed on CPU), each env-gated.

These 2 features are not enabled by default because performance actually degraded.

Encoder block-loop CUDA graph: captured per grid/seqlen layout on the FlashInfer
varlen path, with automatic eager fallback if capture isn't legal. (by default flash-attn varlen is used)
torch.compile: applied to the encoder forward, but the A/B shows it does not
beat eager for these matmul-bound encoders, so it is not relied on as a throughput win.

How was it tested?

Parity + unit tests under test/modular/ (native == HF within tolerance):

test_qwen3_omni_native_encoders_ci.py — CPU/fp32, small random model: state-dict
round-trip, per-layer cosine > 0.9999, all DeepStack levels.
test_qwen3_omni_native_encoders.py — GPU/bf16, real 30B checkpoint: vision (4
resolutions, pooler + DeepStack) and audio (incl. batched packing), cosine > 0.999.
test_qwen3_omni_encoder_graph_parity.py — graph == eager == HF; confirms capture fires.
test_qwen3_omni_gpu_mel_parity.py, test_qwen3_omni_gpu_image_parity.py,
test_qwen3_omni_varlen_backend_parity.py — the optimized GPU/varlen paths vs HF.

Run: pytest test/modular/test_qwen3_omni_*.py -q (the real-checkpoint tests require a
GPU + the 30B weights + flash-attn/flashinfer; the _ci.py test runs CPU-only).

Before/after serving benchmark (4 paths × 6 batch sizes, native vs the HF-wrapper
baseline, env settings, commands and context) lives on the encoders-implemeneted-benchmarked branch .

Benchmark results

Across all four paths, M* sustains ~2× the request throughput (req/s) of vLLM-Omni on
text-output paths (S2T, I2T) and ~3× on speech-output paths (I2S, S2S), measured over a
1→32 concurrency sweep. Systems: M*-new (native encoders, this PR) vs M*-old (the
HF-wrapper baseline) vs vLLM-Omni. Each chart plots four metrics against concurrency.

vLLM reaches the first token faster (TTFT ~1.4× lower at low load, widening to ~2.8×
under load), but its per-token latency degrades as it folds prompt work into decode
steps. M* makes the opposite trade: prefill and decode stay in separate, fixed-shape
graph walks (no mixed continuous batching), so a newcomer waits a beat for an admissible slot
— costing TTFT — but every step stays on the captured fast path. Net: ~1.7–2.4× lower
ITL and ~2× higher request throughput than vLLM. A deliberate TTFT-for-throughput
trade, not a regression.

The encoder is a small slice of the speech pipeline, so the native-encoder change adds
only a ~1–7% throughput bump over the HF-wrapper baseline (M*-old) here. The large
lead over vLLM (~3× req/s, plus lower TTFT and ~3–4× lower ITL) comes from M*'s
in-process, device-direct Talker→Code2Wav pipeline vs vLLM's per-token cross-process
IPC — i.e. the speech-generation path that was already optimized pre-release, not the encoder.

The jump over M*-old comes from batching the audio encoder across requests (old
runs them one at a time). M*-old's apparently-low ITL at higher batch is an artifact of
that serialization — with ~one request in flight, per-token latency is tiny but throughput
collapses; not a real latency win. Versus vLLM, M* leads on throughput and ITL, and on
TTFT at low load (GPU log-mel + no cross-process hop), with vLLM pulling ahead on TTFT
only as concurrency rises.

Combines both effects: cross-request audio-encoder batching drives the gain over
M*-old (up to ~3.4× throughput), on top of the already-optimized in-process
speech-generation walk. Against vLLM, M* leads on every metric — ~3× req/s, ~2× lower
TTFT, ~3–4× lower ITL — again from avoiding the per-token IPC vLLM pays across its
three-process pipeline.

Checklist

ruff check . passes
Added or updated tests / docs where relevant

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013jbAqykVTX655bFQYs5Vya

Shorten verbose comments/docstrings added by the encoder work and drop the obsolete encoder-async note (and its now-unused import); no code changes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013jbAqykVTX655bFQYs5Vya

zip(..., strict=False), drop unused math import, split a multi-statement line, wrap >120-char lines, and sort test imports so ruff check . passes. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013jbAqykVTX655bFQYs5Vya

NSagan271 · 2026-07-01T07:00:42Z

+        self._cg_max_keys = int(os.environ.get("MSTAR_ENCODER_CG_MAX_KEYS", "16"))
+        self._cg_warmed = False
+
+    def _maybe_cg_warmup(self, input_features, feature_lens):


I think it would make sense to refine / clean up the existing PiecewiseCudaGraphRunner implementation to work in this case instead of re-implementing a version of it here. We implemented the piecewise cuda graph runner to be able to handle a similar use case of capturing a function that is internal to the forward pass, but it has some potential issues:

The piecewise cuda graph config is not in the submodule base class, and I'm not sure if it is documented

The current config is a raw dictionary instead of a dataclass, so it's hard for a model author to parse what the config actually needs.

The piecewise cuda graph runner is set as a property on the submodule, but it would be potentially cleaner to have it stored on the engine level and passed in through ModelInputsFromEngine.

The current piecewise cuda graph runner only supports a batch of equal-size inputs. It can be extended to also have a FlashInferPackedCudaGraphConfig-style option, to make it to work for the vision encoder as well.

What the actual config defines is pretty vjepa2-specific and can also be simplified. Right now, the model author has to specify a a fn_factory that takes in a static cache manager + static_pos_bufs, and returns a tensor -> tensor function.

A more generalizable way could be to have any function that should be captured by a piecewise cuda graph take in: static_inputs: dict[str, torch.Tensor], static_cm: BatchedCacheManager | None=None, and **kwargs, and return a dict[str, torch.Tensor]. The config can potentially include: (1) a handle to the function that is captured, (2) a handle to a function that takes in the sequence length and batch size, and returns static inputs, (3) kwargs that will get passed into the captured function, (4) a boolean for whether to plan attention (and/or an optional function for actually planning the attention).

(minor/future): right now, only one function per submodule can be captured. We can extend this by returning a dictionary of label (just some arbitrary string label for the function) to piecewise cuda graph runner config.

@t-avil @merceod your call on whether to apply these changes in this PR / or if my idea about the piecewise cuda graph runner makes sense; I can also write up an Issue for this.

@t-avil Actually, I might take a look at this separately and make a branch off of main

@NSagan271 Sounds great, I agree on extending the generalization on PiecewiseCudaGraphRunner. Does make sense to make it a separate change.

@t-avil Implemented in #154, which I'm about to merge. Lmk if this works for your purposes! If not, I can always modify in another PR.

Will check asap - been working on the second version of this pr (vllm-omni released a huge update pretty recently)

We started to do much better on text generation with your version of PiecewiseCudaGraphRunner adopted.

Thank you so much for generalizing it; looks good to adopt it on my side for now

…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).

…nd-rolled capture Both encoders carried their own torch.cuda.CUDAGraph capture — a per-layout graph cache, side-stream warmup, static-buffer copy, eager fallback — duplicating PiecewiseCudaGraphRunner. That duplication is what review of mstar-project#150 flagged. They now declare PiecewisePackedConfig from get_piecewise_cuda_graph_configs() and the shared runner owns capture/replay; the old implementation is deleted rather than left behind a flag. The captured regions are unchanged (_layer_loop_tail, _block_loop_tail); only who captures them changes. Keying moves from one graph per EXACT cu_seqlens layout to (segments, total_tokens) buckets, so cu_seqlens — and vision's cos/sin — become static input buffers copied per replay instead of constants baked into the graph. Enabling change to the runner (~47 lines, additive): it only planned attention when uses_kv_cache=True, so a region running FlashInfer *ragged* varlen with no KV cache could not use it at all. Two optional hooks, make_attn_state(shape) and plan_attn_fn(state, shape, seq_lens), let such a region own a per-bucket wrapper built with use_cuda_graph=True + fixed indptr buffers and re-plan it OUTSIDE the graph before each replay; _flashinfer_varlen skips its own plan() when external_plan is set, since a host-side plan inside a replayed region is illegal. Gated on both hooks, so KV-cache configs and the V-JEPA2 migration are byte-identical. Capture buckets are MEASURED (MSTAR_ENCODER_CG_PROBE=1), not estimated: vision i2t is segments=1 / 576..1024 tokens but i2s reaches segments=4 / ~4096, and audio s2t is segments 1..7 / 36..426 while s2s reaches segments 1..16 / ~1057. Sizing from the text paths alone left multi-image and long-audio requests with no bucket, silently falling back to eager — correct, but a hidden perf cliff. Instrumentation makes that visible: per-path counters, a one-shot log per distinct layout, and a WARNING the first time a layout fits no bucket. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

…mstar-project#154) Brings in the generalized PiecewiseCudaGraphRunner so the Qwen3-Omni encoders can adopt it instead of carrying their own capture engine, which is what review of this PR asked for. Merged at the exact main commit that introduced it rather than cherry-picked, so the runner shares history with main and does not merge back as a duplicate patch. Auto-merged with no conflicts. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

The encoders carried their own torch.cuda.CUDAGraph capture — a per-layout graph cache, side-stream warmup, static-buffer copy, eager fallback — duplicating the runner generalized in mstar-project#154. That duplication is what review of this PR flagged. They now declare PiecewisePackedConfig from get_piecewise_cuda_graph_configs() and the shared runner owns capture and replay; the hand-rolled path is deleted, not left behind a flag. The captured regions are unchanged (_layer_loop_tail, _block_loop_tail); only who captures them changes. Keying moves from one graph per EXACT cu_seqlens layout to (segments, total_tokens) buckets, so cu_seqlens — and vision's cos/sin — become static input buffers copied per replay instead of constants baked into the graph. One addition to the runner (~47 lines, additive): it only planned attention when uses_kv_cache=True, so a region running FlashInfer *ragged* varlen with no KV cache could not use it at all. make_attn_state(shape) and plan_attn_fn(state, shape, seq_lens) let such a region own a per-bucket wrapper built with use_cuda_graph=True + fixed indptr buffers and re-plan it OUTSIDE the graph before each replay; _flashinfer_varlen skips its own plan() when external_plan is set, since a host-side plan inside a replayed region is illegal. Gated on both hooks, so KV-cache configs and V-JEPA2 are byte-identical. Capture buckets are measured, not estimated (MSTAR_ENCODER_CG_PROBE=1): vision i2t is segments=1 / 576..1024 tokens but i2s reaches segments=4 / ~4096, and audio s2t is segments 1..7 / 36..426 while s2s reaches segments 1..16 / ~1057. Sizing from the text paths alone left multi-image and long-audio requests with no bucket, silently falling back to eager — correct, but a hidden perf cliff. Instrumentation makes that visible: per-path counters, a one-shot log per distinct layout, and a WARNING the first time a layout fits no bucket. Benchmarked on GPUs 6+7 over i2t/s2t/i2s/s2s x B=1..32 (24/24 cells, 0 failed); charts and numbers on encoders-implemeneted-benchmarked under benchmarks/qwen3-omni-native-encoders-piecewise/. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

Resolves the one conflicted file, mstar/model/qwen3_omni/qwen3_omni_model.py (7 regions), where mstar-project#196's prompt/schedule rework overlaps the native-encoder work. Everything else auto-merged. Resolved in upstream's favour: mstar-project#196's correctness fix survives intact — the _MM_SPLIT_SENTINEL, the partition into text-before / modality / text-after with each half tokenized separately, and num_thinker_prefill_steps derived from the built schedule rather than len(input_modalities). Our vLLM prompt-layout TOKEN-SLICE path is kept only for AUDIO-ONLY requests, where it exists to preserve BPE merges across the split by slicing already-tokenized ids; anything involving vision, including mixed audio+image, takes mstar-project#196's path. That audio-only narrowing is enforced in BOTH process_prompt and the schedule builder's early-return, deliberately: mstar-project#196 also yields two text spans, so without a vision guard on the early-return a mixed audio+image request would match it and silently lose its vision walk. One adaptation was required rather than copying mstar-project#196 verbatim. It computes num_mm_inputs from pil_images/np_audios, but this branch only populates those lists when GPU preprocess is OFF, and MSTAR_GPU_IMAGE_PREPROCESS / MSTAR_GPU_MEL are ON by default — so the counts would be zero for every image request, the sentinel would never be inserted, and every vision request would silently skip the split. Modality presence is therefore measured from the RAW inputs, which preserves mstar-project#196's actual behaviour under this branch's defaults. Validated on GPUs 6+7 before merging: i2t and s2t both 8/8 completed, 0 failed, and generated image descriptions confirm prefill_vision runs. A resolution of this conflict on a base lacking mstar-project#183 previously raised KeyError: NodeAndGraphWalk(node='thinker_decode_loop', graph_walk='prefill_vision') on every vision request; mstar-project#183 is in main, and that error does not occur here. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

NSagan271

Some comments. Sorry I'm being so hard on design here---I'm using these few initial model integration/optimization PRs as an opportunity to improve M*'s extensibility.

NSagan271 · 2026-07-31T06:29:34Z


        if outputs.new_token_outputs:
-            name_to_count: dict[str, int] = {}
+            name_to_new_token: dict = {}


Is reverting name_to_count to name_to_new_token (and overall changing counts sent to the conductor back to materializing full tokens) an intentional change or an unintended side-effect of a merge?

Oops, great catch; yep - thats a post-merge side effect. Fixed

NSagan271 · 2026-07-31T06:31:20Z

+                        audio_seqlens, req_token_counts=None, **kwargs):
+        # Thread the (possibly absent) piecewise runner onto the encoder for this
+        # forward; the encoder's forward reads it via ``_piecewise_runner``.
+        self.audio_encoder._piecewise_runner = engine_inputs.piecewise_runners.get("layer_loop")


Nit: in my opinion, it is cleaner to pass the piecewise runner in as an input rather than setting a field on the audio encoder directly.

NSagan271 · 2026-07-31T08:06:40Z

+    return {"ws": ws, "wrapper": wrapper, "cu_obj": None}
+
+
+def make_fi_graph_state(device, num_segments):


I think this has revealed another gap in the system: the stateless engine should have some functionality to provide a flashinfer wrapper and do attention planning. I think the make_attn_state is good to have for flexibility, but the use case for this audio and vision encoder is one that the engine side should be able to handle. I haven't fully scoped this out in my mind, but @t-avil if you agree, can I branch off of your branch to implement it tomorrow?

Yes, please !

NSagan271 · 2026-07-31T08:34:13Z

+            # [prefill_text(prefix), prefill_audio, prefill_text(suffix)], so the
+            # audio walk's BOS/AUDIO/EOS embeddings land between them.  Slicing
+            # avoids retokenizing across the boundary (which can shift BPE
+            # merges), unlike the #196 separate-tokenization path.


A few questions:

What does vLLM-Omni do for mixed-modality inputs? Why does the slice lose vision information in the case of audio + vision inputs?

How big of a difference does the separate-tokenization path make? I think exact parity is good, but just making sure the advantage is worth the complexity.

TLDR we followed what bagel already did... and probably it is worth to open up a new issue.

vLLM writes the prompt once and leaves a blank where each attachment goes, then fills each blank in the position it sits - so N attachments in any order just means N blanks, and audio->image->audio needs no special handling. We do the opposite: we chop the text and hand the pieces to the model with the attachments in between

Doing it properly means moving to per-item placeholders plus an order-preserving adapter, which touches the intake, the schedule builder, and both models. It seems to be a real refactor, though one I'd expect to remove more code than it adds, and it would fix ordering, multiple files, and single-step multimodal batching in one go. I've kept the audio-only restriction here as a correctness guard and would rather open a separate ticket for the refactor.

For the separate-tokenization path, I'm okay with removing it. I only kept it because I noticed occasional one-token drifts from tokenizing the prompt (around 4% of the time) in separate pieces, but they're pretty minor in practice. If we're doing the follow-up adapter refactor anyway, I think we can remove this path and simplify everything around a single implementation.

That sounds good to me, can you open up a new issue for this?

NSagan271 · 2026-07-31T08:43:37Z

-                aud_out["input_features"]
-            )
+            for img in raw_image_inputs:
+                pv, grid_thw = _gpu_image_preprocess(


How computationally intensive is _gpu_image_preprocess (it looks fairly lightweight but I'm just checking to see if it could possibly interfere with GPU work on whatever the designated GPU is). Also, I'm pretty sure that img is by default a CPU tensor if it comes in directly from the data worker, so there is still a CPU roundtrip without changes to the data worker.

Good catch. It was intended to be forced onto cuda; I'm not sure where that got lost.

Image preprocessing

The preprocessing function is now device-agnostic. Previously, the data worker always passed it a CPU tensor, so it only ever executed on the CPU. I also renamed it to _image_preprocess, since the previous name implied a device placement that the call path never guaranteed.

I intentionally avoided refactoring the worker logic, so the CPU→GPU copy still exists. Even so, moving preprocessing to the GPU still provides a noticeable improvement—the gains come despite the round-trip, not because it was removed.

Batch i2t req/s (CPU → GPU) i2t TTFT p50 (CPU → GPU)

B1 0.60 → 0.64 (+8%) 228 → 128 ms

B8 2.17 → 2.38 (+10%) 337 → 222 ms

B32 3.90 → 4.46 (+14%) 517 → 339 ms

Audio is the bigger one. Log-mel had the same problem and it was capping speech
throughput at ~4.5 req/s from B4 to B32 — flat across 8× concurrency. On GPU:

batch s2t req/s s2t TTFT p50 s2t tok/s

B8 4.5 → 9.4 1341 → 237 ms 108 → 224

B32 4.5 → 16.0 ~6700 → 413 ms 110 → 394

Separately, there was a dtype round-trip: decode gives uint8, the base class converted to float [0,1], and preprocessing converted straight back to uint8 to resize.

PiecewiseCudaGraphRunner only planned attention when uses_kv_cache=True: both _capture_one's plan() and run()'s step 2 are guarded on static_cache_manager. A captured region that runs FlashInfer *ragged* varlen attention owns no KV cache but still needs a host-side replan with the real per-request seq_lens before every replay, so it had no way to use the runner and had to hand-roll its own capture. That is exactly the duplication flagged in review of mstar-project#150. Adds two optional config hooks, mirroring the existing KV plan path: make_attn_state(shape) -> one persistent capture-safe wrapper per bucket, built at capture time and threaded into capture_fn as attn_state= plan_attn_fn(attn_state, shape, seq_lens) -> replan before each replay Purely additive and gated on both hooks being set, so KV-cache configs and the V-JEPA2 migration are byte-identical. Replay planning reuses _replay_seq_lens, which zero-pads to shape.bs: the planned indptr sums to the real token count, leaving the bucket's zeroed pad tail in no segment. Ragged attention is block-diagonal, so the pad tail cannot leak into real segments. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

Both encoders hand-rolled their own torch.cuda.CUDAGraph capture — a per-layout graph cache, side-stream warmup, static-buffer copy, eager fallback — duplicating PiecewiseCudaGraphRunner. That duplication is what review of mstar-project#150 flagged. Now each submodule returns a PiecewisePackedConfig from get_piecewise_cuda_graph_configs() and the runner owns capture and replay. The captured regions are unchanged (_layer_loop_tail, _block_loop_tail); only who captures them changes. Keying moves from one graph per EXACT cu_seqlens layout to (segments, total_tokens) buckets, so cu_seqlens — and vision's cos/sin — become static input buffers copied per replay instead of constants baked into the graph. The ragged FlashInfer wrapper is built once per bucket with use_cuda_graph=True + fixed indptr buffers and re-planned outside the graph before each replay; _flashinfer_varlen skips its own plan() when external_plan is set, since a host-side plan inside a replayed region is illegal. The legacy path stays reachable behind MSTAR_ENCODER_LEGACY_CG=1 so one build can A/B both capture paths instead of comparing across builds. Instrumentation, because "piecewise did not regress" and "piecewise never ran" are otherwise indistinguishable in results.json: per-path counters, a one-shot log per distinct layout, and a WARNING the first time a layout fits no bucket. MSTAR_ENCODER_CG_PROBE=1 forces eager and logs every layout, so capture buckets can be sized from measured (segments, total_tokens) instead of guesswork — the shipped defaults do NOT cover batch 16/32 and are sized in a follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

…le settings Review follow-ups on mstar-project#150 (NSagan271's worker.py question). Revert the new_token_counts -> new_tokens change across worker.py, node_manager_utils.py, ipc_format.py and conductor.py. It was not a merge side-effect: it first appears in this branch's own 35b9940, and neither upstream/main nor the merge parent 68950c3 has it. Reverting fixes two problems at once. First, correctness. buffer_new_tokens() was called from INSIDE the `for signal in outputs.new_token_outputs` loop while being handed the cumulative name->tokens dict, and the buffer extends rather than replaces. A request emitting N distinct new-token signals therefore counted signal 0 N times, signal 1 N-1 times, and so on. conductor.py uses that running total as a stop condition (num_output_tokens >= max_output_tokens), so affected requests terminated early and returned truncated output -- while also reporting inflated throughput. Single-signal requests (i2t, s2t) run the loop once and were unaffected, which is why this survived testing; the paths that emit both Thinker text and Talker codec (s2s, t2s) were not. Second, cost. Counting went from tensor.numel() -- metadata, free -- to tensor.cpu().numpy().tolist(), a blocking device-to-host copy per new-token tensor, per request, per step, for every model. The batched _d2h_new_tokens helper written to amortize that was never called and prematerialized_new_tokens was never passed by any caller, so the fast path was unreachable. The only consumer of the token values is conductor.py, which takes len() of them, so the counts the old format already carried are sufficient. Separately, stop imposing qwen3-omni's compile requirements on every model. stateless_engine had been changed to compile all submodule forwards with dynamic=False; that suits an encoder driven by fixed piecewise capture buckets but forces per-shape recompiles on any variable-shape submodule sharing the enc_dec config. The engine now defaults to dynamic=None and reads an opt-in `torch_compile_dynamic` attribute, which the two native encoder submodules set to False -- so qwen3-omni keeps the behavior it was benchmarked with and other models keep Inductor's default. Also restores the kv_cache_engine._compile_submodules docstring, which had been rewritten to claim mode="max-autotune-no-cudagraphs" and "must be called BEFORE CUDA graph capture" while the code does neither (it still passes dynamic=None with the default mode, and warmup calls it after capture); and makes _clone_tensor_input copy mutable non-tensor values instead of sharing them with the original. test/modular/test_qwen3_omni_*.py is unchanged at 47 passed / 11 failed / 10 skipped. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

… encoder Addresses NSagan271's review comment on mstar-project#150. The submodules were reaching into the encoder and assigning ``encoder._piecewise_runner`` before each forward, which the encoder then read back with ``getattr(self, "_piecewise_runner", None)``. That makes a per-call input look like module state: the encoder appears to own a runner it does not, the attribute silently persists after the call, and nothing in the signature says the forward depends on it. Both native encoders now take ``piecewise_runner`` as an explicit keyword argument, defaulting to None (the eager path), and the four submodule call sites pass ``engine_inputs.piecewise_runners.get(...)`` directly. No behavioral change; test/modular/test_qwen3_omni_*.py stays at 47 passed / 11 failed / 10 skipped. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

…encoder Addresses NSagan271's review comment on mstar-project#150: "these attention primitives can be shared model components instead of living in audio_encoder.py". The varlen attention stack -- four SDPA variants, the FlashInfer ragged path, the capture-override plumbing, head-dim padding and backend selection -- was defined in the Qwen3-Omni audio encoder, and the vision encoder imported it from there via ``audio_encoder as AE``, reaching through a sibling model file for eleven symbols. Nothing in it is audio-specific. It now lives in mstar/model/components/varlen_attention.py, which is the package explicitly documented for model-agnostic building blocks. The piecewise-capture telemetry (note_encoder_path / encoder_path_counts / note_encoder_layout) moves alongside it to mstar/model/components/encoder_telemetry.py, for the same reason and because it was the other half of what the vision encoder was reaching across for. Both encoders now import from components/ and the AE alias is gone. The duplicated "is capture legal" predicate -- ``_FLASHINFER_AVAILABLE and _VARLEN_BACKEND == "flashinfer"``, spelled out in both encoders against two private globals of another module -- is now ``capture_legal_backend()``. As the reviewer noted, BAGEL's ViT carries its own near-duplicate ``_sdpa_varlen``; that is left alone here but flagged in the new module's docstring as the next thing to fold in. Test-only fallout: the three test modules that reached into audio_encoder for these globals now import the components modules directly. Unexpected but verified side effect: the full-suite result goes from 64 passed / 10 skipped to 70 passed / 4 skipped. The six newly-running tests are the GPU/bf16 encoder-vs-HF parity tests, which were being skipped by a collection-order effect -- their ``import flash_attn`` guard evaluated False during a full-suite run despite flash_attn being importable. Running that file alone yields 6 passed on both the old and the new code, so this is coverage recovered, not a guard weakened. The remaining 4 skips are the always-skipped audio_output_parity scaffold. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012fTosh9fKWQuGLYcHT2h1m

_dump_obj / _tensor_dump_dir wrote intermediate tensors and token ids to MSTAR_DUMP_DIR. Removed with their three call sites.

Places the audio block inside the user turn ahead of the instruction and pins audio M-RoPE h/w to temporal, so the 3D position ids match HF get_rope_index. Drops MSTAR_VLLM_PROMPT_LAYOUT; the legacy bare-block layout is gone.

Drops MSTAR_GPU_MEL and MSTAR_GPU_IMAGE_PREPROCESS along with the HF-CPU image branch and its pil_images round-trip. The log-mel CPU fallback stays, keyed on torch.cuda.is_available() alone.

…ll paths MSTAR_VLLM_AUDIO_SENTINELS selected alternate audio marker ids (151669/151670); MSTAR_BATCH_VISION_PREFILL allowed multi-request vision prefill. Both defaulted off and neither had test coverage.

MSTAR_QWEN3_NATIVE_{AUDIO,VISION}_ENCODER duplicated the native_audio_encoder / native_vision_encoder config fields. The fields remain; removing the override drops the last _envflag caller and the helper.

…constants Buckets and backend selection are measured values, not deployment settings, and a typo in MSTAR_VARLEN_BACKEND silently selected an SDPA fallback and disabled graph capture with no error. Now CAPTURE_TOKENS_{AUDIO,VISION}, CAPTURE_BATCH_SIZES_{AUDIO,VISION} and _VARLEN_BACKEND. Also drops MSTAR_ENCODER_CG_PROBE and MSTAR_VISION_GRAPH_ALIGN.

buffer_new_tokens() was called inside the `for signal in outputs.new_token_outputs` loop while being passed the cumulative name->tokens dict, and the buffer extends rather than replaces, so a request emitting N distinct new-token signals counted signal 0 N times. conductor.py uses that running total as a stop condition (num_output_tokens >= max_output_tokens), so affected requests ended early with truncated output and inflated throughput. Single-signal paths (i2t, s2t) run the loop once; s2s and t2s emit both Thinker text and Talker codec. Counting had also moved from tensor.numel() to tensor.cpu().numpy().tolist(), a blocking device-to-host copy per tensor, per request, per step, for every model. The batched _d2h_new_tokens helper was never called and prematerialized_new_tokens was never passed. conductor.py only takes len() of the values, so counts suffice.

StatelessEngine compiled every submodule forward with dynamic=False, forcing a recompile per distinct shape on variable-shape submodules sharing the enc_dec config. Defaults to None; the two native encoder submodules set torch_compile_dynamic = False.

It claimed mode="max-autotune-no-cudagraphs" and "must be called BEFORE CUDA graph capture". The code passes dynamic=None with the default mode, and warmup calls it after capture.

Values that were neither tensors nor sequences were returned by reference, so a clone shared them with its original. Dicts are recursed into; anything else is copied.

Replaces assignment of encoder._piecewise_runner before each forward with an explicit keyword argument defaulting to None.

They targeted encoder._cg_cache, _cg_warmed and MSTAR_ENCODER_CUDA_GRAPH, none of which exist since the move to PiecewiseCudaGraphRunner, so both raised AttributeError before asserting anything. They now build a runner from the encoder's get_piecewise_cuda_graph_config and require encoder_path_counts() to show *.piecewise == 1 and *.eager == 0 -- a captured graph existing does not mean the forward replayed it, and a layout fitting no bucket falls back to eager silently. Also restores flash_attn_varlen_func and _FLASH_ATTN_AVAILABLE, which the third test stubbed and left set (8 varlen failures later in the same session), and removes its CUDA gate since it runs on CPU.

Cosine was computed in float32 over ~36M values and returned 1.0045, above 1.0, so the gate compared nothing; now float64, with the threshold tightened from 0.999 to 0.9999. max-abs over that many values is a tail statistic set by one pixel at the kernel's worst rounding boundary, so the 99.99th percentile is the magnitude gate and max is a looser ceiling. Adds smooth-gradient inputs alongside uniform noise. Measured at 3000x2000: content max p99.99 mean cos(float64) uniform noise 0.2039 0.1098 0.0021 0.9999293 smooth gradient 0.0078 0.0078 0.00033 0.9999950 The residual is CPU-vs-CUDA bicubic on uint8; the GPU path runs the same torchvision kernel in the same order as HF.

…onents Both were defined in qwen3_omni/components/audio_encoder.py, and vision_encoder imported eleven symbols from it as `audio_encoder as AE`. Neither is audio-specific. Adds capture_legal_backend() to replace the predicate spelled out in both encoders against two private globals of another module. BAGEL's ViT still carries a near-duplicate _sdpa_varlen; noted in the new module's docstring.

The function has never run on a GPU. data_worker constructs with a hardcoded device="cpu" (data_worker.py:86), PreprocessWorker takes no device parameter, and _preprocess_loop's device argument is never passed, so load_image always returns a CPU tensor. Every op runs on the input's own device, which makes the name a claim about placement that the call path does not support. Also corrects the module comment, which cited ~175 ms of CPU round-trip as the dominant I2T TTFT cost. That figure is the HF processor on a ~3000px image; at the sizes actually served it is ~7 ms. Measured against the HF processor: image HF proc (CPU) ours (CPU) ours (H200) 512x512 6.87 ms 1.29 ms 0.26 ms 1024x768 16.68 ms 5.98 ms 0.27 ms 3000x2000 140.56 ms 147.93 ms 1.23 ms Only the CPU column applies today. The docstring no longer claims the input arrives on the GPU. Test file and test renamed to match.

_audio_mel_gpu resolved its device as `waveform.device if waveform.is_cuda else torch.device("cuda")`, so a CPU waveform -- which is the only kind data_worker produces -- was copied to cuda:0 and processed there. That opened a CUDA context on a model GPU from the API-server process (~94 MiB measured) and put per-request kernels on a device the worker owns, for preprocessing that is not on the model's critical path. It now follows the input device, matching _image_preprocess, so the data worker is uniformly CPU. The cost is small and the function still beats the HF CPU feature extractor either way: clip HF FE (CPU) ours (CPU) ours (H200) 5 s 3.50 ms 2.13 ms 0.39 ms 30 s 6.84 ms 3.35 ms 0.48 ms Since it is now faster than the HF path on CPU as well, the torch.cuda.is_available() gate is gone and the HF feature-extractor fallback is only used when the processor itself is missing.

The base class decodes to uint8 and then converts to float in [0, 1]. _image_preprocess resizes on uint8 to match HF's image processor, so it converts straight back -- two full passes over the tensor that cancel, with a 4x larger tensor carried in between (69 vs 17 MiB for a 3000x2000 image). Overriding load_image to skip the conversion is worth ~35-40% of preprocessing time, measured against the HF processor on the same inputs: image HF (CPU) ours float32 ours uint8 512x512 1.62 ms 1.19 ms 0.75 ms 1024x768 10.76 ms 6.41 ms 3.73 ms 3000x2000 140.00 ms 153.99 ms 125.00 ms It also flips the large-image case: with the float round-trip we were slower than the processor we replaced. The base class still returns float, since other models consume that range directly; _image_preprocess already accepted either dtype.

Both are device-agnostic and the data worker hands them CPU tensors, so both were running on CPU. Adds a thin ``_gpu`` wrapper for each and points the two call sites at it, so the device is a visible choice at the call site rather than something the core guesses. For log-mel this restores the previous behaviour: it used to force torch.device("cuda") and 0ac383a changed it to follow the input device, which made it CPU. That was a large regression, hidden by an isolated microbenchmark showing log-mel at ~0.05% of per-request latency. Measured end to end instead, CPU log-mel was capping speech throughput at ~4.5 req/s from B4 to B32 -- flat across an 8x concurrency increase: s2t req/s CPU -> GPU TTFT p50 CPU -> GPU tok/s CPU -> GPU B8 4.54 -> 9.43 1341 -> 237 ms 108 -> 224 B32 4.47 -> 16.00 ~6700 -> 413 ms 110 -> 394 Image preprocessing had never run on a GPU in any configuration: i2t req/s CPU -> GPU TTFT p50 CPU -> GPU B1 0.60 -> 0.64 228 -> 128 ms B8 2.17 -> 2.38 337 -> 222 ms B32 3.90 -> 4.46 517 -> 339 ms The cause is not arithmetic -- per call this is ~0.2 ms of resize and normalize. Preprocessing runs in one thread, so on CPU that thread is held for the whole call while on GPU it enqueues and returns, which is why the gap widens with concurrency. It also holds despite the CPU->GPU copy: decode is CPU-only, so the copy is still paid. Cost is ~13 MiB transient plus a ~94 MiB CUDA context in the API-server process. Measured on configs/qwen3omni_2gpu.yaml with no MSTAR_* flags, GPUs 2,3, n and warmup matched per batch; valid arm-vs-arm only. The device is hardcoded to "cuda" and should come from config before this is correct under PD-disaggregation.

t-avil and others added 4 commits June 30, 2026 07:31

Qwen3-Omni native encoders + perf optimizations

b64d286

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013jbAqykVTX655bFQYs5Vya

Merge branch 'mstar-project:main' into encoders-implemeneted

508a7b4

NSagan271 reviewed Jul 1, 2026

View reviewed changes

t-avil mentioned this pull request Jul 23, 2026

feat(qwen3-omni): multimodal encoders migrated into M* and optimized for text gen #185

Draft

t-avil and others added 3 commits July 30, 2026 23:31

t-avil requested a review from NSagan271 July 31, 2026 05:42

NSagan271 requested changes Jul 31, 2026

View reviewed changes

t-avil added 11 commits August 1, 2026 00:41

qwen3-omni: remove _dump_obj tensor dumping

9b03ec9

_dump_obj / _tensor_dump_dir wrote intermediate tensors and token ids to MSTAR_DUMP_DIR. Removed with their three call sites.

qwen3-omni: always use the vLLM prompt layout

441c756

Places the audio block inside the user turn ahead of the instruction and pins audio M-RoPE h/w to temporal, so the 3D position ids match HF get_rope_index. Drops MSTAR_VLLM_PROMPT_LAYOUT; the legacy bare-block layout is gone.

qwen3-omni: always run GPU log-mel and GPU image preprocessing

e285f47

Drops MSTAR_GPU_MEL and MSTAR_GPU_IMAGE_PREPROCESS along with the HF-CPU image branch and its pil_images round-trip. The log-mel CPU fallback stays, keyed on torch.cuda.is_available() alone.

qwen3-omni: remove the unused audio-sentinel and batched-vision-prefi…

ec004a5

…ll paths MSTAR_VLLM_AUDIO_SENTINELS selected alternate audio marker ids (151669/151670); MSTAR_BATCH_VISION_PREFILL allowed multi-request vision prefill. Both defaulted off and neither had test coverage.

qwen3-omni: remove the native-encoder env override

6488472

MSTAR_QWEN3_NATIVE_{AUDIO,VISION}_ENCODER duplicated the native_audio_encoder / native_vision_encoder config fields. The fields remain; removing the override drops the last _envflag caller and the helper.

engine: fix the _compile_submodules docstring

e107880

It claimed mode="max-autotune-no-cudagraphs" and "must be called BEFORE CUDA graph capture". The code passes dynamic=None with the default mode, and warmup calls it after capture.

model: copy mutable non-tensor values in _clone_tensor_input

73613c2

Values that were neither tensors nor sequences were returned by reference, so a clone shared them with its original. Dicts are recursed into; anything else is copied.

qwen3-omni: pass piecewise_runner as a forward argument

d47b76b

Replaces assignment of encoder._piecewise_runner before each forward with an explicit keyword argument defaulting to None.

t-avil added 7 commits August 1, 2026 00:41

t-avil requested a review from NSagan271 August 1, 2026 08:40

NSagan271 mentioned this pull request Aug 3, 2026

Ragged prefill wrapper #203

Open

2 tasks

t-avil mentioned this pull request Aug 6, 2026

Order-preserving multimodal input adapter #209

Open

		return {"ws": ws, "wrapper": wrapper, "cu_obj": None}


		def make_fi_graph_state(device, num_segments):

Batch	i2t req/s (CPU → GPU)	i2t TTFT p50 (CPU → GPU)
B1	0.60 → 0.64 (+8%)	228 → 128 ms
B8	2.17 → 2.38 (+10%)	337 → 222 ms
B32	3.90 → 4.46 (+14%)	517 → 339 ms

batch	s2t req/s	s2t TTFT p50	s2t tok/s
B8	4.5 → 9.4	1341 → 237 ms	108 → 224
B32	4.5 → 16.0	~6700 → 413 ms	110 → 394

Uh oh!

Conversation

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

Uh oh!

What does this PR do?

How was it tested?

Benchmark results

Checklist

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 Jul 30, 2026 • 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!

NSagan271 left a comment

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!

Uh oh!

Uh oh!

Uh oh!

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

Image preprocessing

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

t-avil commented Jun 30, 2026 •

edited

Loading

t-avil Jul 30, 2026 •

edited

Loading