Engine: cut spanned outputs to a per-row valid span after the step by merceod · Pull Request #901 · mstar-project/mstar

merceod · 2026-10-26T23:30:11Z

Batched audio and image nodes pad rows to a bucket and then need each row cut back to its real extent. Doing that inside the forward means a host read of the per-row bounds on the GPU thread, which is a sync in model code and holds up the next step's launch (Kokoro's edge trim paid about 1 ms per 14 ms step that way).

This adds the engine mechanism instead. A submodule declares spanned_outputs(walk) and returns BatchedModelOutput.output_spans, a [bs, 2] int64 device tensor per named output with each row's [start, end) on the last dim. The engine checks the declaration and sizes, the GPU thread queues the spans into the pinned host-row ring behind the step (the same path as the stop rows, so no extra wait), and the main thread cuts each row to a view right after the completion event, before the stop check, the store and any streaming edge see the outputs. The cut is one shared Python step ahead of both graph runtimes; rust/ is untouched. A spanned output that feeds a node this worker can speculate into is refused at startup, since speculation hands outputs on before the cut.

Measured on one B300 with Kokoro (82M) as the first user (its model-side trim moved onto this in a follow-up PR): worker_phases at c=1, the GPU thread's per-step time drops from 14.1 to 8.7 ms and the main thread's wall per step from 14.8 to 14.1 ms; the cut itself costs 0.05-0.06 ms and the staging 0.06 ms. End to end the change is inside the session spread (pooled c=1 audio-s/s 80.9 vs 83.6, c=8 221 vs 218, WER flat at 0.85%).

Tests: new test/modular/test_output_spans.py drives a real Engine on CPU through admit/plan/forward/commit and the real Worker._postprocess_batch up to the store, and fails with the cut disabled; it also covers undeclared and short spans, the per-request merge and the startup refusal. The CUDA staging test passes on GPU. ruff check . clean; test/modular on CPU has the same failing set as main (eval_compare_multigpu, lingbot_world_shell, serving_metrics, wan_vae_decode, z_image_native_equivalence). test/rust: 370 passed, 93 skipped.

A submodule declares spanned_outputs(walk) and returns a [bs, 2] span per row in BatchedModelOutput.output_spans. The gpu thread copies the spans to pinned host memory behind the step; once the main thread has waited on the step's completion event it cuts each row to its span (a view), before the stop check, the store and streaming edges read it. Startup refuses a spanned output that feeds a node the worker speculates into.

mstar-review · 2026-10-26T23:32:42Z

+            else:
+                self.output_spans = {
+                    name: torch.cat([span, other.output_spans[name]])
+                    for name, span in self.output_spans.items()
+                }
+                self.span_request_ids = self.span_request_ids + other.span_request_ids


[P2] The span merge assumes every request returns spans for the same output names, so it can crash or drop spans

On the per-request path (_exec_per_request → merged.update(...)), suppose request A returns spans for {"audio"} and request B returns spans for {"pitch"}. Then other.output_spans[name] raises KeyError, which fails the whole step. If B returns {"audio", "pitch"} instead, B's pitch span is silently dropped and that row reaches the store uncut. That contradicts the comment just above, which says spans can't be dropped on a merge. Merging over the union of names would fix this, but each name would then need its own request-id tuple, because span_request_ids can no longer be shared across names. Alternatively, _real_spans could require the spans to cover exactly spanned_outputs(walk) on every row.

mstar-review · 2026-10-26T23:32:56Z

Automated review

🟡 2 P2 on 5e465bd, light review

Description

The 'Measured on one B300' numbers give no hardware-independent reproduction details: no torch/CUDA build, no exact worker_phases or benchmark command, and no number of runs behind the 'pooled' figures.
The measurements come from a Kokoro trim that lives in a follow-up PR, so nothing in this PR exercises the mechanism on a real model; the description should say which tests and numbers apply to this diff alone.
The CUDA staging test is reported as passing 'on GPU' without naming the GPU or the torch/CUDA build, which AGENTS.md asks for when a change runs on the GPU.

What the review checked

The spans survive capture and replay: AcceleratorGraphRunner returns the captured BatchedModelOutput whole as slot.static_outputs, so output_spans comes back with each replay.
_real_spans returns a view of the spans without copying, which is safe: _stage_host_rows copies it to pinned memory on the execution stream before the completion event, ahead of any later replay.
The cut swaps entries in per-rid lists that _merge_per_rid builds fresh for each step (cloned values or row-view lists), so it never changes captured static outputs.
_crop_to_spans runs after _wait_for_step and before _prematerialize_for_check_stop, the stop check, inline emit and the store; the store's row-view fast path is turned off whenever a row is cut.
_refuse_speculated_spans applies the same source-side and target-side enable_async_scheduling checks and the same direct node.outputs edges that PythonGraphRuntime._speculate_walk and ingest_for_speculation use to pick speculation targets.
exec_and_postprocess has one caller, on the GPU thread, so _stage_host_rows stages the spans on every accelerator path; on CPU the cut reads the CPU output_spans directly.
The cut runs on the Python side of the worker before the store, so the Rust graph runtime and transport need no matching change (invariant 5).

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

Gaurav-Shah05

The span merge on the per-request path drops a request's span or fails the whole step (see the comment on BatchedModelOutput.update).

On one H100 with configs/qwen3tts_1p7b.yaml at c32, a model that returns no spans pays nothing. Three alternating launches per build:

	main	this PR	Δ
Worker step time	10.51 ms	10.44 ms	1.0x
Main-thread CPU per step	7.00 ms	7.04 ms	1.0x
Audio-s/s	135.7	134.9	1.0x

One mechanism with #886:

Its withheld_outputs decides the same thing per row, which part of a step's output is data. Qwen3-TTS's EOS frame is a row whose span is empty, and the Talker can compute that span on the device from its sampled token.
In my run an empty span reached the store as a zero-width tensor, which a streaming edge still counts as an item and StreamBuffer._collate can't torch.stack with the real frames.
If an empty span left the row's output out of the store instead, Qwen3-TTS could be the first user of spans, with one mechanism for both. What do you think?

Gaurav-Shah05 · 2026-10-27T03:40:49Z

+                self.span_request_ids = other.span_request_ids
+            else:
+                self.output_spans = {
+                    name: torch.cat([span, other.output_spans[name]])


This is a bug on the per-request path: the merge keeps only the first spanned request's names. In my run a later row's extra spanned output reached the store uncut, and disjoint names raised KeyError here and failed the whole step. Requiring a step that returns any spans to return one for every declared output, checked in Engine._real_spans, keeps the names aligned.

Fixed in 219cc2c: Engine._real_spans now refuses a step that spans some declared outputs but not all ("returned spans for ['audio'] ... but none for ['level']"), so every forward on a walk carries the same names, and BatchedModelOutput.update refuses a merge of spans with different names instead of dropping one or raising a KeyError. Two tests cover both; the docs sentence says a step that returns spans returns one for every declared output.

Gaurav-Shah05 · 2026-10-27T03:41:06Z

+        (``NodeSubmodule.spanned_outputs``). Called once the step's completion
+        event has fired, so the spans staged behind the step are on the host.
+        The cut is a view, and it replaces the tensor in its list in place, so
+        an output the submodule also put under a second name is cut there too.


Nit: This holds only for a name added in postprocess or unpack_packed_outputs. A list the forward returns under two names gets one list per name from Engine._merge_per_rid, so the second name keeps its full width, as it did in my run.

Right, the in-place cut only reaches a second name that postprocess or unpack_packed_outputs gave the same list. The docstring now says that and that names the forward returned separately keep their own widths (219cc2c).

Co-authored-by: mstar-review[bot] <345662112+mstar-review[bot]@users.noreply.github.com>

Gaurav-Shah05

lgtm, applied the bot's suggestion

…spans

… users/atindra/output-spans

…ath reads its failed rids

Main since 73c6c17: #883 mstar_rust in the Linux wheels, #901 output spans (BatchedModelOutput.output_spans, NodeSubmodule.spanned_outputs, the worker's crop), #904 DiT seed noise drawn at ingest (on_request_added, NoiseStager deleted, to_device_nonblocking as one memcpy into pinned memory), #903 API-server torch pool from the CPU quota and pooled image encodes, #924 robotics harness, #819 Inductor cache under ~/.cache, #881 CommGroup.agree / JointGroups.agree (agree_across_ranks and flagged_on_any_rank deleted), #724 CUDA IPC KV in one 2D memcpy a page, #387 GLM-5.2 (MLA, fp8 MoE with NoAuxTcRouter, the DSA resource), #921 vjepa2 AC crop, #922 Cosmos3 contiguous positions, #923 dense attention falls back when fa3-fwd has no kernel for the GPU. Conflicts: - mstar/utils/h2d.py, to_device_nonblocking: the stack's signature (data: Tensor | Sequence, device, dtype=None) and DeviceCapability pin_memory gate, with main's tensor path (#904): a CPU tensor is written into a fresh pinned buffer with one memcpy rather than pin_memory()'s copy_ over the intra-op pool. Every CPU tensor now goes through a fresh buffer, so the stack's "an already-pinned source is read in place" caveat (R3.3) no longer applies and is dropped from the docstring. test_accelerator_graph_backend's pinning test checks the pinned torch.empty and its non-blocking copy instead of pin_memory(). - mstar/model/components/diffusion/denoise_loop.py: imports, both (main's copy, the stack's functools). Main code on a superseded path, ported here: - GLM-5.2 (glm52/submodules.py): its own check_stop / check_stop_batched (four arguments, step_walk, its own row-map loop) / _stops -> token_stop = TokenStop(decode_loop, new_token, eos_token_ids, sampler) with a _token_stops override for its two differences (a stop token at prefill ends the reply at the first decode step; max_tokens counts the prefill's token), and a check_stop override for that closing step, whose token postprocess drops. test_glm52_model's helper and test_model_kv_layouts' GLM cap test build the real submodule (over a stub language model) instead of object.__new__, which would leave token_stop unset. - PinnedStager.to_device (#387), a second "host values -> new device tensor" helper beside to_device_nonblocking: deleted with its test (test_pinned_stager_to_device.py); its callers, GLM-5.2's prefill seq_lens and the DSA resource's page tables, call to_device_nonblocking(list, device, dtype=torch.long). test_dsa_resource spies on that instead. The other files both sides touched (engine.py, accelerator_graph_runner.py, kv/manager.py, attn/base.py, ragged/base.py, communication.py, model/base.py, submodule_base.py, entrypoint.py, conductor.py, moe.py, image_io.py, vjepa2_model.py, docs, tests) auto-merged into disjoint hunks: per file the merged delta equals main's own. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The codec's prepare_inputs dropped the EOS frame with a boolean index, which is a device sync on the GPU thread every codec step (2-3 ms on a B300) and keeps a codec step from overlapping the decode step before it. The codec can't drop it by count on the host because it doesn't know which window ends the stream; only the Talker's check_stop does. The stop check now returns the outputs it withholds: `check_stop` may return `StopCheck(stops, withheld)` (a NamedTuple, beside the plain `set[str]` every other model keeps returning) and `check_stop_batched` a `BatchedStopCheck` keyed by rid, so a row's token is read once, the batched form has a counterpart (the per-row dicts `_LazyPerRid` skips are never built), and `ignore_eos` can withhold the EOS frame without stopping. `Engine.check_stop_for_batch` writes `batch.withheld_outputs` and returns the stops as before; the worker pops the withheld names from the per-rid outputs before `_store_and_route`, so the route carries no tensor for that (row, signal) and both graph runtimes see the same `RouteInput` (`rust/` untouched; both count stream items from the descriptors they route, so a withheld row is not in the done count either). All three codecs are on it: Qwen3-TTS's Talker withholds `codec_tokens` on EOS; Qwen3-Omni's Talker its codec_eos frame, and Code2Wav drops the `eos_mask.any()` read (a final window with one new frame now stays on the captured graph instead of going eager); Fish S2 Pro's AR the frame whose head index is `eos_code`, and the codec drops `codes[codes[:, 0] < eos_code]`. Each codec has the empty-final-chunk guard: a chunk whose only frame was EOS becomes an empty window that emits nothing. An output cut to an empty span (#901) is left out of the store the same way, so a streaming edge gets no item for it instead of a zero-width one. Setup for every number below: one B300 (sm_103, driver 595.91.07 / CUDA 13.2), torch 2.12.1+cu132, flashinfer 0.7.2.post1, server `MSTAR_SHM_ARENA=1 python -m mstar.cli.main serve qwen3_tts_1p7b --gpus 0 -O2` from the arm's tree, client `benchmark/tts_speech_bench.py --engine mstar` (voice vivian, English, warmup 5c) on the first 200 Seed-TTS EN sentences, WER with `benchmark/tts_speech_wer.py` (whisper-large-v3). - worker_phases at c=8 (`python -m benchmark.worker_phases.server` around the same serve command, untimed): gpu_thread.prepare_inputs mean 1.24 -> 0.43 ms, iter_total mean 7.18 -> 6.44 ms - audio-s/s, A/B from one checkout with `git apply -R`/`git apply`, A B A B, two server sessions per arm, unseeded c=1/8/32: c=1 14.2 -> 14.4, c=8 54.8 -> 58.1, c=32 106.9 -> 115.3 (the two-session ranges touch; the phase table is the firmer evidence) - outputs: seeded c=1 WAVs byte-compared against main, 200/200 identical in three branch server processes (an A/A control between two main processes is 199/200); the same check passed on the colocated and `qwen3tts_1p7b_split.yaml` configs with both the Rust and the Python graph runtime (20/20 each), 0 errors Also in here: the docs claimed a bf16 codec halves codec time; on B300 cuDNN picks `implicit_convolve_sgemm<bf16>` for the convs and bs>=8 gets slower with the same WER, so the default stays fp32 and the comment says why. And a YAML note on why the Talker stays on FA2 there: trtllm-gen replays a stale block table under capture and faults, cudnn is a wash. The trtllm-gen part is an engine-level issue, filed as #898. Tests: the stop-check return form and the engine path (`test_withheld_outputs.py`), one CPU test per codec through the real `check_stop`, StreamBuffer and policy (EOS withheld, the final chunk off a window boundary carries the last frames' audio, an EOS-only stream ends with an empty chunk that decodes nothing, the audio still stitches), the zero-tensor stream parity test on both runtimes, and #629's hand-built step carrying the new field. Touched modules plus `test/rust/test_runtime_parity.py`: 881 passed, 3 skipped; `test/rust`: 428 passed, 99 skipped; `ruff check .` clean; the CPU Core suite's failing set is main's. ## Since the review Codec host time on the GPU thread (`prepare_inputs` of the codec walk, mean ms, main to branch; one untimed `MSTAR_PHASE_TIMING=100` session per arm, B300): | model | c=1 | c=8 | c=32 | |---|---|---|---| | Qwen3-TTS codec_chunk | 2.42 to 0.19 | 3.37 to 0.25 | 1.66 to 0.34 | | Qwen3-Omni code2wav_chunk | 2.89 to 0.15 | 2.33 to 0.17 | | | Fish S2 Pro codec_chunk | 5.44 to 0.21 | 6.78 to 0.22 | 5.49 to 0.27 | Wall equalled CPU on main: the old sync spin-waited. Served, main against branch, A B A B, two sessions per arm (server pinned to 8 cores, `-O2`, Rust graph runtime; Qwen3-TTS and Fish on tts1's 200 Seed-TTS sentences, Omni on 48 two-sentence questions with the Thinker at temperature 0; WER by whisper-large-v3), audio-s/s mean over the sessions: | cell | main | branch | TTFA p50 ms main / branch | WER % main / branch | |---|---|---|---|---| | Qwen3-TTS c=1 / 8 / 32 | 14.5 / 58.7 / 125.9 | 14.3 / 60.8 / 126.1 | 26 / 48 / 77 to 25 / 47 / 80 | 1.19 / 1.4 / 1.1 to 1.19 / 1.3 / 1.0-1.7 | | Fish S2 Pro c=1 / 8 / 32 | 4.9 / 17.1 / 31.3 | 5.1 / 18.8 / 33.6 | 108 / 202 / 400 to 100 / 190 / 346 | 1.85 / 2.2 / 1.1 to 1.85 / 1.6 / 0.9 | | Qwen3-Omni speech c=1 / 8 | 11.3 / 39.2 | 11.6 / 39.2 | 186 / 320 to 178 / 329 | 0.9-1.4 / 1.0-3.1 to 1.4-3.4 / 1.4-1.6 | No errors in any cell. Seeded c=1 outputs are byte-identical to main on Qwen3-TTS and Fish (200/200 over six session pairs, an A/A pair included); Omni is not reproducible across its own processes on main either (0/48), so WER is its guard, and its high cells are one garbled sampled reply per arm with reply endings matching on both. Qwen3-TTS c=32 splits into a 120 and a 138 audio-s/s mode per process on both arms (five sessions each). The Python graph runtime (`MSTAR_RUST_GRAPH=0`) runs all three with no errors; its seeded Qwen3-TTS and Fish outputs are identical to the Rust runtime's. Not verified: a split config with the codec on another worker (that the done count leaves out withheld rows is from both runtimes' code, and the parity test covers the local stream only); a withheld output that is also an inline client signal would still be sent (no model withholds one). The Fish c=8/32 and Qwen3-TTS c=8 gains come from two sessions each with non-overlapping ranges; the GPU thread no longer blocking on the codec's device read is the likely cause (the wait moves to `event_sync`).

mstar-review Bot added area/engine area/model-shared area/worker size/M labels Oct 26, 2026

mstar-review Bot reviewed Oct 26, 2026

View reviewed changes

Comment thread docs/adding_models.rst Resolved

mstar-review Bot added type/feature type/perf labels Oct 26, 2026

Gaurav-Shah05 requested changes Oct 27, 2026

View reviewed changes

Gaurav-Shah05 mentioned this pull request Oct 27, 2026

Qwen3-TTS: withhold the EOS frame from the codec stream #886

Merged

Engine: a step that returns spans returns one for every declared output

219cc2c

merceod requested a review from Gaurav-Shah05 October 28, 2026 05:27

Apply suggestion from @mstar-review[bot]

eb391a4

Co-authored-by: mstar-review[bot] <345662112+mstar-review[bot]@users.noreply.github.com>

Gaurav-Shah05 approved these changes Oct 28, 2026

View reviewed changes

merceod added 3 commits October 27, 2026 22:53

Merge remote-tracking branch 'origin/main' into users/atindra/output-…

6ec814d

…spans

Merge remote-tracking branch 'origin/users/atindra/output-spans' into…

1a584e5

… users/atindra/output-spans

test_output_spans: the fake worker carries a scheduler, as the step p…

b1ade7a

…ath reads its failed rids

merceod merged commit fd3b7e1 into main Oct 28, 2026
6 checks passed

merceod deleted the users/atindra/output-spans branch October 28, 2026 07:11

dghosef mentioned this pull request Nov 4, 2026

Let a submodule emit outputs from the host rows at the stop check #1116

Open

Navigation Menu

Repository navigation

Engine: cut spanned outputs to a per-row valid span after the step - #901

Engine: cut spanned outputs to a per-row valid span after the step#901
merceod merged 6 commits into
mainfrom
users/atindra/output-spans

merceod commented Oct 26, 2026

Uh oh!

mstar-review Bot Oct 26, 2026

Uh oh!

Uh oh!

mstar-review Bot commented Oct 26, 2026

Uh oh!

Gaurav-Shah05 left a comment

Uh oh!

Gaurav-Shah05 Oct 27, 2026

Uh oh!

merceod Oct 28, 2026

Uh oh!

Gaurav-Shah05 Oct 27, 2026

Uh oh!

merceod Oct 28, 2026

Uh oh!

Gaurav-Shah05 left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

merceod commented Oct 26, 2026

Uh oh!

mstar-review Bot Oct 26, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mstar-review Bot commented Oct 26, 2026

Automated review

Description

Uh oh!

Gaurav-Shah05 left a comment

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

merceod Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

merceod Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants