[NSagan271]

Most check_stop policies follow one of two forms: AR models check the iter count against a maximum, plus an EOS check if ignore_eos is not set, and diffusion/flow models just do the iter count check. This PR ports both of these common checks into the Rust runtime, with the option for a submodule-level override.

What and why

Declared stop policies:

• The hook. NodeSubmodule.get_stop_policies(graph_walk) -> dict[str, StopPolicy, ...] | None (dict from loop name -> policy). The default is None, which keeps calling the submodule's check_stop.
• StopKind.AR and StopKind.ITER encode the AR-style and diffusion-style stop checks described above. StopKind.EOS checks only the token, for a loop whose cap (Model.get_loop_caps) is its budget, as main's TokenStop models' are.
• Where they run. The worker declares each local walk's policies to the graph runtime once, at startup. Both runtimes implement GraphRuntime.check_stops; the per-row decision is the policy's own (StopPolicy.stops), and each runtime keeps a request's limits (StopLimits) per (node, walk) on its request info, registered whenever the worker receives the request's forward-pass info, so they go when the request is removed. The runtime reads the policies' token column from the step's stop buffers (BatchedModelOutput.check_stop_buffers), so the fast path needs both declared policies and batched output. A step the runtime can't check (no declared policies, a step that mixes walks, or no stop buffer for the token column) still calls check_stop.
• Models. Orpheus declares AR. A TokenStop (Performance 1/7: shared engine, conductor and component changes #858) with runtime_check=True derives its policy (ARNodeSubmodule.get_stop_policies): EOS on its loop's walk, while its prefill_walks keep check_stop, which can end the partition there (A stop token sampled at prefill ends the reply, and max_output_tokens counts the prefill's token #794). Qwen3.5 sets it; the other TokenStop models are unchanged here. Qwen3.5 and Orpheus return their token as a stop buffer. Chatterbox and Qwen3-TTS keep check_stop: they count their own steps in request_state, not by the loop's iteration count. Whisper keeps check_stop too: since Whisper/Qwen3-ASR PR 5: serving performance #568 its stop check also keeps each token for the word-timestamp align walk.
• The non-custom stop checks now live in the graph runtime.

Postprocess as begin / resume. Inside the runtime, a step's postprocess is three phases: check stops (and stop loops), route, send. The worker drives it as begin_postprocess(request) then resume_postprocess(...), and the runtime stops for the worker only where a phase needs Python, named by PostprocessNeed:

• STOPS: the runtime's own check failed or missed a row's limits, so the model's check_stop decides; resume with the stops and the rows it failed. A step the runtime doesn't check never stops here: the worker runs check_stop before beginning and passes its stops in (PostprocessRequest.stops with runtime_stops=False), and the runtime takes them as final.
• STORE: the outputs have to be stored as tensors before routing; resume with them. A step that returns log-probs (Sampler: return the sampled token's log-prob and ids on /generate #914) always takes this break, since the store reads each row's log-probs for the route.
• SEND: the send needs what only Python has (forward-pass info when a walk completes, the final publish of stopped rows, tensors to register, local streams, new-token counts, profiling); resume with the SendInput.

PhasedPostprocess is the Python implementation over each runtime's phase methods; the Rust runtime runs the same flow natively (rust/src/graph/runtime/postprocess.rs). It holds the GIL, except around the route and send of an inline step, which touch nothing in Python; that time is recorded as rust.postprocess.nogil. This replaces route_and_send_inline and the worker's fast-path fallbacks: _postprocess_batch is now one begin / resume loop with a handler per break. Before it begins, the worker hands the engine the step's host outputs (Engine.step_finished; today that keys the step's tokens for prefix reuse), so a request that stops on this step is up to date first.

Instrumentation. Under MSTAR_PHASE_TIMING, worker.msg.<type> times each worker message type, with stage laps inside _add_new_request and _remove_request.

• Renamed postprocess phases, for anyone comparing worker_phases runs across this PR. worker.postprocess.runtime is all the time inside the runtime's begin / resume calls. Each break the worker handles is timed under its own name: check_stop (the model's check_stop, before the runtime or at a STOPS break), store_tensors and send_prep. prematerialize and tail are the stages before and after. The old check_stop_tail, stop_loops, route_outputs, register_outputs, route and send stages are gone. rust.route_and_send_inline.nogil becomes rust.postprocess.nogil.

END_PARTITION (#794). The worker takes it out of the model's stops in _check_stops, which both paths that call check_stop go through: a row inside a loop fails, as on main, and any other ends its partition, through the request's partition_done or, at a STOPS break, on the runtime directly. A step whose stop check withheld an output (StopCheck.withheld) routes through STORE, never inline, as main's fast path refused it.

Docs. docs/adding_models.rst documents get_stop_policies and TokenStop.runtime_check.

Evidence

1× H200, main 9c05c56 (this PR's base) against this PR; Rust graph + ZMQ, SHM arena, prefix cache off, closed loop, benchmark/worker_phases at period 100. ABBA, two servers per arm, 0 failed requests. Qwen3.5 0.8B as configured in configs/qwen3_5_0.8b.yaml (--ignore-eos, 256–512 output tokens); Orpheus with configs/orpheus_colocated.yaml (LLM and SNAC on one worker).

Postprocess work (ms per worker iteration, both servers pooled): worker.postprocess_batch less the waits in it, the wait for the step (completion_event_sync) and the stop check's host copy (prematerialize).

config	main	this PR	change
Qwen3.5, mixing off, c=64	0.416	0.384	−8%
Qwen3.5, mixing off, c=128	0.634	0.612	−3%
Qwen3.5, mixing on, c=64	0.457	0.422	−8%
Qwen3.5, mixing on, c=128	0.770	0.725	−6%
Orpheus, c=32	0.847	0.812	−4%
Orpheus, c=64	1.435	1.206	−16%

• The model's check_stop stage falls from 0.11–0.15 ms (Qwen3.5) and 0.15–0.31 ms (Orpheus) per iteration to 0.01–0.03 ms; what remains are the steps the runtime doesn't check (prefill and mixed steps). The runtime's own begin / resume span is 0.20–0.37 ms, and holds the route and send the worker ran before.

Throughput:

config	main	this PR
Qwen3.5, mixing off, c=64 (tok/s)	20463 / 20457	19921 / 20754
Qwen3.5, mixing off, c=128	26218 / 27110	27239 / 25680
Qwen3.5, mixing on, c=64	19956 / 20254	20434 / 20441
Qwen3.5, mixing on, c=128	26758 / 26764	26847 / 26906
Orpheus, c=32 (audio-sec/s)	61.4 / 61.9	61.1 / 62.8
Orpheus, c=64	85.6 / 84.2	87.7 / 88.7

• Qwen3.5 is within noise; Orpheus at c=64 is 4% faster.

Per-request message cost (the new worker.msg.* spans, measured before the rebase onto main; Qwen3.5 c=128, 2,306 requests): about 0.68 ms of main-thread CPU per request.

stage	ms per request
new request	0.37 (reading the prompt tensor 0.18, engine 0.11, runtime 0.05)
removal	0.24 (engine 0.13, the "Request cleanup complete" INFO log line 0.07)
input signals	0.06

These are follow-up targets, not changed here.

Tests

• New tests:
  • test/modular/test_stop_policies.py: each model's declared policies agree with its own check_stop, Qwen3.5's as derived from its TokenStop;
  • test/rust/test_runtime_check_stops.py: both runtimes, over the token column the worker's PostprocessRequest hands them, through a speculated decode loop, stop what the policy does; a recycled handle has no limits until its own are registered.
  • test/modular/test_postprocess_request.py: the worker's PostprocessRequest routes a combined batch under its walks' output signals, and carries the model's stops when the runtime doesn't check them.
  • test/rust/test_postprocess.py: the Rust runtime's native begin / resume matches PhasedPostprocess over the same runtime: every break, the frames sent and the loop counters, across runtime and model stops, inline and stored outputs, int32/int64 tokens, and host rows in or out of order.
• Mutation checks. An off-by-one in the count check fails all 24 policy tests; ignoring ignore_eos fails 14 of 24.

Not in this PR

• Steps that mix walks still call check_stop in Python.
• The per-request costs above: the prompt read and the engine's add and remove. The INFO log line is a marker in test/waypoint/serve_rollout.py, so lowering it needs a replacement signal there.

AI assistance was used to write this change and its description.

🤖 Generated with Claude Code

[mstar-review]

[P2] loop_count was added between check_dyn_loop and its doc comment, so check_dyn_loop's doc now sits on loop_count.

The rustdoc for loop_count now begins "Whether this request's current walk actually contains the loop. A stop for one it does not is a model bug", which isn't what it does, and check_dyn_loop is left with no doc.

Suggested change

	/// Whether this request's current walk actually contains the loop. A stop
	/// for one it does not is a model bug: logged by the caller and dropped.
	/// `rid`'s iteration of loop `lp` in partition `p`, as
	/// `get_dynamic_loop_iters` reports it (the last worker graph with the
	/// loop wins); 0 when no live worker graph has it.
	fn loop_count(&self, rid: u32, p: Sym, lp: Sym) -> u32 {
	let Some(part) = self.info(rid).and_then(|i| i.partitions.get(&p)) else { return 0 };
	let mut count = 0;
	for &wg in &part.walk_worker_graphs {
	let Some(state) = self.state(wg, rid) else { continue };
	if let Some(&lid) = self.graphs[wg as usize].loop_by_name.get(&lp) {
	count = state.loop_iter(lid);
	}
	}
	count
	}
	fn check_dyn_loop(&self, rid: u32, partition: &str, loop_name: &str) -> bool {
	/// `rid`'s iteration of loop `lp` in partition `p`, as
	/// `get_dynamic_loop_iters` reports it (the last worker graph with the
	/// loop wins); 0 when no live worker graph has it.
	fn loop_count(&self, rid: u32, p: Sym, lp: Sym) -> u32 {
	let Some(part) = self.info(rid).and_then(|i| i.partitions.get(&p)) else { return 0 };
	let mut count = 0;
	for &wg in &part.walk_worker_graphs {
	let Some(state) = self.state(wg, rid) else { continue };
	if let Some(&lid) = self.graphs[wg as usize].loop_by_name.get(&lp) {
	count = state.loop_iter(lid);
	}
	}
	count
	}
	/// Whether this request's current walk actually contains the loop. A stop
	/// for one it does not is a model bug: logged by the caller and dropped.
	fn check_dyn_loop(&self, rid: u32, partition: &str, loop_name: &str) -> bool {

[mstar-review]

Automated review

🟡 3 P2 on 5c6e575, light review

Four of the five earlier findings are fixed (combined-walk output signals, Orpheus/Whisper stop buffers, stream_partition_done after a dropped row, nogil accounting), and the GIL-release P1 is addressed: the Rust postprocess now holds the GIL except for an inline step's route and send. One new P2 is posted, on a misplaced Rust doc comment.

Description

• The description still doesn't say that the worker.postprocess.* phase names changed: check_stop_tail, stop_loops, route_outputs, register_outputs, route and send are gone, and runtime, tail and send_prep (as a break stage) replace them; rust.route_and_send_inline.nogil is also gone. This matters to anyone comparing worker_phases runs across this commit.
• Test commands and their results are missing: ruff check ., the CPU Core pytest command, cargo test plus pytest test/rust/ against a rebuilt extension, and pass/skip counts. The Tests section lists new modules but not that they were run, or on which commit.
• There is no statement of which GPU or integration tests were run on the head commit. The benchmark numbers are stated to predate the last commit, and the Whisper WER check isn't tied to a commit.
• Per the repo's PR-splitting guidance (a note, not a blocker): the postprocess begin/resume restructure and the per-message instrumentation could have landed separately from the stop-policy feature.

What the review checked

• The combined-walk batch down to one walk now routes under _walk_signals(batch, walks[0]); test_postprocess_request.py covers it.
• A batch still labelled with a combined walk can't take the runtime check: combined walk names are rejected as real walks (Model combined-walk validation), so they are never in my_walks or _stop_token_signal, and the step falls back to check_stop.
• After relabel_if_uniform or with no request_walks, walks[0] equals the declared batch_N.graph_walk, so check_stops is never asked about an undeclared walk.
• An AR policy only runs in the runtime with a token column: _runtime_stop_column returns runtime_stops=False when _host_column is None, and the Rust stops_over errors (falling back to STOPS) without a covering int32/int64 column.
• Python _keep_rows and Rust keep_rows both filter stream_done and stops against dropped rows; the new modular test and the Rust unit test exercise this.
• The Rust pp_run holds the GIL for the stop loops and for a stored step's route and send (sent try-first via send_frames), and releases it only around an inline route and send. Python _postprocess_step sums nogil_ns over crossings and records it once at DONE.
• Orpheus and Whisper now return check_stop_buffers={"new_token": ...}. Whisper's row_outputs row views are the same tensor[i:i+1] slices as its old per-rid dict (Engine._collect_outputs).
• The runtime check skips check_stop_for_batch, which has no side effects beyond registering a failure on a raise; that is not needed when the runtime decides.
• Stop limits live on per-request runtime info in both runtimes and are dropped in remove_request (invariant 14); max_tokens is typed int.
• The loop count is read at postprocess time with step N+1 already speculated; test_runtime_check_stops.py drives both runtime kinds through a speculated decode loop and compares against the policy on the worker's own counts.
• Engine.get_stop_policies follows the mgmt is None guard pattern of its neighbours, and _declare_stop_policies runs once at startup.

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

[stephen-dwq]

lgtm; i can see no races and optimizations are correct