Stop at max_output_tokens, not one past it by merceod · Pull Request #604 · mstar-project/mstar

merceod · 2026-10-18T13:22:52Z

Qwen3-Omni and Orpheus returned one token more than max_output_tokens (65 for 64 on Qwen3-Omni, where HF generate gives 64), and the Qwen3-Omni Talker one codec frame more than talker_max_output_tokens. The last prefill emits the first token, and for the Talker talker_last_prefill streams the first frame, so a decode loop capped at the whole budget puts out one more than asked for. Main's other TokenStop models already cap decode at max(1, N - 1) (Qwen3.5 and Whisper since #794, Higgs Audio since #859).

The Thinker and Orpheus caps are now max(1, N - 1). The Talker's cap is max(1, T - 1), and its stop check only looks for codec_eos, so the cap alone sets the frame budget. EOS stops are unchanged.

test_model_kv_layouts.py asserts 1 + cap == max_output_tokens for Higgs Audio, Orpheus, Qwen3.5, the Qwen3-Omni Thinker and Whisper at 2, 3 and 64, and the same for the Talker, whose talker_last_prefill row must stop no loop. On main 9 of those 19 fail, here all pass. The Talker test no longer skips without flashinfer, which CPU CI doesn't install. Driven through both graph runtimes with each model's real cap and stop check, main gives N + 1 for Orpheus, the Thinker and the Talker and this branch gives N. ruff is clean. The CPU core paths pass 10340 tests here and 10324 on main, with the same 48 environment failures on both.

On one H200 the Talker streams 11 and 65 frames on main and 10 and 64 here for talker_max_output_tokens 10 and 64, and 2, 3 and 64 frames for 2, 3 and 64 at 4a78dae. The Thinker returns 3, 11 and 65 text tokens on main and 2, 10 and 64 here for max_output_tokens 2, 10 and 64, and Qwen3.5-0.8B returns 1, 2 and 64 on both.

max_output_tokens = 1 on Qwen3-Omni still returns 2 text tokens, since the Thinker has no prefill END_PARTITION.

NSagan271

LGTM, with a note that CI is failing due to the new test requiring a flashinfer dependency. flashinfer either needs to be added to the CI environment, or the test skipped (or the flashinfer import wherever it is used in qwen3-omni can be move within the corresponding function; I think one of the other open PRs does this but I can't recall which one)

NSagan271 · 2026-10-19T00:10:31Z

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Numbers come from one allocation, so treat small deltas as provisional.

Live test: Stop at max_output_tokens, not one past it

Fixed for max_output_tokens ≥ 2 on Qwen3.5, Qwen3-Omni and Whisper, but max_output_tokens = 1 still returns 2 tokens on all three. Qwen3.5 0.8B throughput is within noise.

1x H200 (Qwen3-Omni: 2x, qwen3omni_2gpu.yaml), SHM tensor transport · PR head 8e92be3 vs main 84fd4a6

Claims

Claim	Result
Text requests return exactly max_output_tokens (was one more)	✅ For n = 2, 3, 8 and 64, main returned n+1 and the PR n, on Qwen3.5 0.8B (`ignore_eos`, `/generate` and chat) and on the Qwen3-Omni thinker (text out, no `ignore_eos`, counting prompt). The same for Whisper large-v3 at n = 2, 3, 8 and 16
n = 1	❌ 2 tokens on both arms for all three models (Qwen3.5 `/generate` chunks, chat retokenized; Whisper chunks). The PR's new test covers 2, 3 and 64 only. Hypothesis, from the PR's own accounting: decode step 0 has already emitted the second token when `check_stop` first runs
EOS stops unchanged	✅ Whisper without `ignore_eos` at n=16 stops on EOS at 11 chunks on both arms
Streaming and non-streaming return the same tokens	✅ `/generate` stream chunk count == chat token count (retokenized), for every n on both arms
Orpheus, Higgs Audio	➖ Not run (audio out; one token isn't observable from the client)

Performance

Qwen3.5 0.8B, closed loop, main, PR, PR, main, text tok/s:

arm	c=8	c=32	c=128
main	2915, 2868	7323, 7205	10117, 9547
PR	2619, 2851	6963, 7524	9302, 9794

No regression: at c=32 and c=128 the arms' restart ranges overlap; at c=8 one PR run is 9% low and the other within 0.6% of main. The change is one constant in each check_stop. Host load 24–74.

Client command

python -m benchmark.runner --model qwen3.5 --model-id Qwen/Qwen3.5-0.8B --request-type text_to_text --ignore-eos --output-len-min 128 --output-len-max 384 --profiling-type closed_loop --num-requests <10*c> --num-warmup 1 --max-concurrency <c>. Counts: max_tokens 1, 2, 3, 8, 64, temperature 0, through /generate (chunks) and chat (output retokenized).

Other

CPU tests: test_max_output_tokens.py 24 passed.

Prefill already emitted the first token, so decode step k has emitted k + 2 tokens. The stop checks compared k + 1, which let one extra token through.

merceod · 2026-10-20T11:21:38Z

Rebased onto main. The test/cpu-core.txt line is gone since #581 removed that manifest; test/modular is selected as a whole now, so the new test file is picked up as is. The earlier CPU Core failure was the talker's module-level flashinfer import, which main no longer has. test_max_output_tokens.py: 24 passed locally on the rebased branch.

merceod · 2026-10-25T20:30:38Z

/review

mstar-review · 2026-10-25T20:32:27Z

Automated review

🟡 No findings on 306e1b9, light review

Description

The test drives check_stop only through PythonGraphRuntime, but under AUTO production can run the Rust graph runtime (mstar/graph/runtime/rust.py:406 update_dynamic_loop_iters). The description should say whether the Rust runtime's loop counters were checked to give the same iteration k at decode step k, or that BAGEL's existing + 2 already shows they do.
The motivating claim is parity with HF generate (64 tokens for 64 on Qwen3.5 and Qwen3-Omni). The fix was not re-run end to end on GPU, so the PR has no post-fix server-level token count. One request each against Qwen3.5 and Qwen3-Omni, plus a streaming one to back the streaming/non-streaming claim, would close that gap.

What the review checked

check_stop reads ExecutingBatchRequestInfo.dynamic_loop_iter_counts (mstar/engine/engine.py:232). That comes from InputMetadata, which is copied per batch in _build_input_metadata_and_fwd_info (mstar/worker/worker.py:1476) and incremented for speculated rows, so decode step k sees k on both the normal and the speculative path.
Prefill emits the first token to the client or the downstream stream in every changed model: the Higgs Audio and Whisper prefill nodes have an EMIT_TO_CLIENT new_token edge, and Orpheus prefill streams new_token to snac_decoder (orpheus_model.py:123). So k + 2 tokens are out at decode step k.
BAGEL (bagel/submodules.py:1645) and Cosmos3 (cosmos3/submodules.py:3251-3253) already compare k + 2 against max_tokens, as the description says.
No other LLM check_stop under mstar/model/ still uses the k + 1 comparison against client max_tokens. The Qwen3-Omni talker (line 1617) has its own talker_max_tokens codec budget, which this PR doesn't change.
Existing tests in test/modular/test_qwen3_5_check_stop_batched.py (done=99/max=100, iters=4/max=5) still expect a stop under k + 2, so they don't conflict.
max_output_tokens=1 is not made worse: the decode-step check was already one past, and the conductor's own num_output_tokens >= max_output_tokens check (conductor.py:1660) runs after the prefill walk.
The new test imports only submodule modules that other test/modular tests already import unguarded, so it should import on the CPU Core job.
The 18/24 failures on main fit the test matrix: 6 models x 3 max_tokens values fail, and the 6 EOS cases pass.

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

Gaurav-Shah05

The five counts are right, but the Qwen3-Omni talker still stops one codec frame late, and the new test can be one assertion per model in an existing test (see the two inline comments).

With this PR applied to main, Qwen3.5-0.8B on the Rust graph runtime streamed exactly 2 and 64 tokens for max_output_tokens 2 and 64.

Where the count could live: (Not for this PR)

main's get_loop_caps already caps each decode loop per request on both graph runtimes, and speculation stops at the cap.
- In my runs of your test with each cap at max_output_tokens - 1 and the max_tokens check in check_stop switched off, all six cases stopped at exactly max_output_tokens for 2, 3 and 64 on both runtimes.
So each model could write the count once, in get_loop_caps.
check_stop would then only look for EOS, and a new model couldn't get the budget right in one place and off by one in the other.

Gaurav-Shah05 · 2026-10-27T03:10:48Z

+        # prefill emitted the first token, so decode step k emits token k + 2
        if (not ignore_eos and eos_token_id == token) or \
-                (request_info.dynamic_loop_iter_counts.get("thinker_decode_loop", 0) + 1 >= request_info.max_tokens):
+                (request_info.dynamic_loop_iter_counts.get("thinker_decode_loop", 0) + 2 >= request_info.max_tokens):


The talker has the same bug: talker_last_prefill samples the first codec frame and streams it to Code2Wav, and TalkerSubmodule.check_stop still stops talker_decode_loop at + 1, so the audio runs one frame past talker_max_output_tokens. It should count + 2 like the thinker.

mstar/mstar/model/qwen3_omni/submodules.py

Line 1633 in 11c910b

(request_info.dynamic_loop_iter_counts.get("talker_decode_loop", 0) + 1 >= max_tokens):

Good catch, fixed in ead7b1f. The Talker counts k + 2 now, and 5a6ad55 takes its cap to max(1, T - 1) so the cap test still lines up. Code2Wav gets T frames for talker_max_output_tokens = T (T >= 2).

Gaurav-Shah05 · 2026-10-27T03:11:51Z

+
+@pytest.mark.parametrize("max_tokens", [2, 3, 64])
+@pytest.mark.parametrize("model", list(CASES))
+def test_stops_at_max_output_tokens(model, max_tokens, tmp_path):


Main's test_model_kv_layouts.py already runs each of these five models' check_stop over the loop's iteration counts and checks the stop against its get_loop_caps.

One assertion there, that the prefill token plus the cap equals max_output_tokens, guards the same count in a few lines instead of 205.

Done in c5e53b9, and 6453877 drops this file. By the time I merged main, #858 and #859 had already removed the per-model cap vs check_stop cases for these models, since their check_stop only looks at EOS now and the cap alone sets the budget. So the new test asserts 1 + cap == max_output_tokens directly, one assertion per model at 2, 3 and 64, plus the same check in the Talker's cap test.

#794, #858 and #859 replaced the check_stop code this branch changed with TokenStop, which only checks EOS and leaves the token budget to get_loop_caps. The conflicts take main's side, and the commits after this move what is left of the fix into the caps.

The Thinker's last prefill emits the reply's first text token and talker_last_prefill streams the first codec frame to Code2Wav, so a cap of max_output_tokens (talker_max_output_tokens) let each loop put out one more than asked for. The caps leave room for the prefill's token, as the other TokenStop models' caps do.

talker_last_prefill streams the first frame, so decode step k streams frame k + 2. Counting k + 1 ran the audio one frame past talker_max_output_tokens. The stop check now ends the loop on the same iteration as its cap.

The prefill samples the first token and streams it to the SNAC decoder, so a decode cap of max_output_tokens gave the decoder one token more than asked for.

One assertion per model whose prefill samples the first token (Higgs Audio, Orpheus, Qwen3.5, the Qwen3-Omni Thinker, Whisper), and one more in the Talker's cap test. Orpheus and both Qwen3-Omni loops fail it on main. The Talker test drops its flashinfer skip. The Qwen3-Omni modules import without flashinfer, and CPU CI doesn't install it, so the test never ran there.

test_model_kv_layouts.py now checks the same count with one assertion per model, off each model's own loop cap.

merceod · 2026-11-03T03:58:47Z

Merged main (2394b4d) and reworked for the review.

Main moved under this PR. #794, #858 and #859 put Qwen3.5, Whisper and Higgs Audio on TokenStop, where check_stop only checks EOS and the cap is max(1, N - 1), so those hunks went away in the merge. What main still had off by one is in the caps. The Omni Thinker and Orpheus capped decode at max_output_tokens, and the Talker as you found.

5a6ad55 Thinker cap max(1, N - 1), Talker cap max(1, T - 1)
ead7b1f Talker stop check at k + 2
c3ac825 Orpheus cap max(1, N - 1)
c5e53b9 one assertion per model in test_model_kv_layouts.py, 6453877 drops test_max_output_tokens.py

The new test fails 8 of 25 on main (Orpheus and both Omni loops) and passes here. The Talker test also drops its flashinfer importorskip. The Omni modules import fine without flashinfer and CPU CI doesn't install it, so that test never ran in CI. A script (not in the PR) that runs prefill plus the capped loop through both graph runtimes, with each model's real cap and stop check, gives N + 1 on main for Orpheus, the Thinker and the Talker, and N for all six models here, on Python and Rust. The CPU core paths give 10340 passed here and 10324 on main, with the same 48 environment failures on both (no ftfy or libcuda, HF parity tests that want other weights or transformers).

On the aside, agreed, and main is most of the way there. TokenStop models already count only in get_loop_caps. The Talker, BAGEL's decode and the Cosmos3 reasoner still count in check_stop. Moving them is dropping that arm, keeping the cap at max(1, N - 1) and turning their cap vs stop tests into the 1 + cap check. max_output_tokens = 1 is the one case a cap can't do, since the prefill has to end the partition the way TokenStop.prefill_walks does. I'd do it as a follow-up.

Left alone here, max_output_tokens = 1 on Omni still gives 2 text tokens. The Thinker has no prefill END_PARTITION, and ending at the prefill would skip talker_last_prefill on audio requests. Same follow-up.

Served counts on an H200 to follow.

merceod · 2026-11-03T05:44:51Z

Served counts on one H200, main (2394b4d) against this branch (6453877). Colocated configs, /generate with return_token_ids, greedy Thinker, 3 requests per cell, and all 3 gave the same count.

what the client gets	budget	main	this PR
Omni Talker frames	talker_max_output_tokens 10	11	10
	talker_max_output_tokens 64	65	64
Omni Thinker text tokens (ignore_eos)	max_output_tokens 1	2	2
	max_output_tokens 2	3	2
	max_output_tokens 10	11	10
	max_output_tokens 64	65	64
Omni Thinker tokens in the audio requests	max_output_tokens 400	401	400
Qwen3.5-0.8B text tokens	max_output_tokens 1	1	1
	max_output_tokens 2	2	2
	max_output_tokens 64	64	64

Frames are counted off the streamed PCM. W frames decode to W * 1920 - 555 samples, and every reply came out at a whole number of frames. The token ids match between the arms, this branch just stops one sooner (max_output_tokens 2 gives [33, 878] here and [33, 878, 374] on main). Omni at max_output_tokens 1 still gives 2, the gap I mentioned above.

Gaurav-Shah05

LGTM, with one small point below.

Gaurav-Shah05 · 2026-11-03T06:08:02Z

+        # talker_last_prefill streamed the first frame, so decode step k streams frame k + 2
        if is_eos or \
-                (request_info.dynamic_loop_iter_counts.get("talker_decode_loop", 0) + 1 >= max_tokens):
+                (request_info.dynamic_loop_iter_counts.get("talker_decode_loop", 0) + 2 >= max_tokens):


At talker_max_output_tokens=2 the + 2 arm returns a stop for talker_decode_loop on the talker_last_prefill row, whose walk has no such loop. Both graph runtimes drop it, and the Python runtime logs it as a model bug. Without the arm, get_loop_caps alone gives 2, 3 and 64 frames for 2, 3 and 64 on both graph runtimes in my runs, so TalkerSubmodule.check_stop should check only codec EOS.

Agreed, thanks. 9fcfe15 drops the arm, so TalkerSubmodule.check_stop only looks for codec_eos and the cap from get_loop_caps alone sets the frame budget. 4a78dae moves the test to 1 + cap == T and adds a case that the talker_last_prefill row stops no loop at T = 2. Head is now 4a78dae.

Through both graph runtimes with the real check_stop, T = 2, 3 and 64 give 2, 3 and 64 frames with no dropped stop. On 6453877 the T = 2 run dropped one stop on both runtimes, and the Python one logged it. Served on an H200 at 4a78dae (Rust graph runtime), talker_max_output_tokens 2, 3 and 64 streamed 2, 3 and 64 frames, 3 requests each, and the server log has no "Tried to stop loop" line.

At talker_max_output_tokens 2 the count arm returned a talker_decode_loop stop on the talker_last_prefill row, whose walk has no such loop. Both graph runtimes dropped it, and the Python runtime logged it as a model bug. The loop's cap from get_loop_caps already ends decode at the frame budget on both runtimes.

The Talker test checks 1 + cap == talker_max_output_tokens and that the stop check never ends decode on a frame that is not codec_eos. A new case checks that the talker_last_prefill row stops no loop at a budget of 2.

Main carries #1130 as a squash, so the branch's history already had most of it. The conductor's check_knobs keeps main's extension check and loses the transport knob; test_conductor_wait runs on the Rust transport only.

This comment has been minimized.

Sign in to view

NSagan271 approved these changes Oct 18, 2026

View reviewed changes

Stop decode loops at max_output_tokens, not one past it

306e1b9

Prefill already emitted the first token, so decode step k has emitted k + 2 tokens. The stop checks compared k + 1, which let one extra token through.

merceod force-pushed the users/atindra/max-output-tokens branch from 8e92be3 to 306e1b9 Compare October 20, 2026 11:21

github-actions Bot added model/many size/XS labels Oct 20, 2026

dghosef added the type/bug label Oct 20, 2026

Gaurav-Shah05 requested changes Oct 27, 2026

View reviewed changes

This was referenced Oct 27, 2026

Stop at max_output_tokens, not one past it #906

Closed

Whisper: batch and capture the encoder, capture decode to 64 rows, row outputs for decode #905

Closed

dghosef mentioned this pull request Oct 29, 2026

Muse Glimmer 30B on a shared VLM scaffold (2/3) #951

Open

dghosef assigned dghosef and unassigned dghosef Oct 29, 2026

This was referenced Oct 30, 2026

A stop token sampled at prefill ends the reply, and max_output_tokens counts the prefill's token #794

Merged

Performance 1/7: shared engine, conductor and component changes #858

Merged

merceod added 6 commits November 2, 2026 19:57

Orpheus: the decode cap counts the prefill's token

c3ac825

The prefill samples the first token and streams it to the SNAC decoder, so a decode cap of max_output_tokens gave the decoder one token more than asked for.

test: drop test_max_output_tokens.py

6453877

test_model_kv_layouts.py now checks the same count with one assertion per model, off each model's own loop cap.

mstar-review Bot added model/orpheus model/qwen3_omni and removed model/many labels Nov 3, 2026

merceod requested a review from Gaurav-Shah05 November 3, 2026 04:00

Gaurav-Shah05 approved these changes Nov 3, 2026

View reviewed changes

merceod added 2 commits November 2, 2026 22:19

test: the Talker's cap alone gives the frame budget

4a78dae

The Talker test checks 1 + cap == talker_max_output_tokens and that the stop check never ends decode on a frame that is not codec_eos. A new case checks that the talker_last_prefill row stops no loop at a budget of 2.

mstar-review Bot added size/S and removed size/XS labels Nov 3, 2026

merceod merged commit f3bce74 into main Nov 3, 2026
6 checks passed

merceod deleted the users/atindra/max-output-tokens branch November 3, 2026 06:47

merceod mentioned this pull request Nov 3, 2026

Remove the pyzmq control-mesh transport and MSTAR_RUST_ZMQ #1131

Merged

nick-rui mentioned this pull request Nov 4, 2026

[RFC] Speculative decoding in the resource layer #1186

Open

Conversation

merceod commented Oct 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment has been minimized.

NSagan271 left a comment

Choose a reason for hiding this comment

Uh oh!

NSagan271 commented Oct 19, 2026

Live test: Stop at max_output_tokens, not one past it

Claims

Performance

Other

Uh oh!

merceod commented Oct 20, 2026

Uh oh!

merceod commented Oct 25, 2026

Uh oh!

mstar-review Bot commented Oct 25, 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 Nov 3, 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 Nov 3, 2026

Choose a reason for hiding this comment

Uh oh!

merceod commented Nov 3, 2026

Uh oh!

merceod commented Nov 3, 2026

Uh oh!

Gaurav-Shah05 left a comment

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 Nov 3, 2026

Choose a reason for hiding this comment

Uh oh!

merceod Nov 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

merceod commented Oct 18, 2026 •

edited

Loading