[dghosef]

Python/Rust graph-runtime drift at loop exit: the Rust runtime sends a finishing loop's loop-back edges to other worker graphs. A loop split across worker graphs (BAGEL's CFG-parallel denoise: combine_cfg on rank 0, LLM_cfg_text / LLM_cfg_img on their own ranks, and LLM in a sibling worker graph on rank 0) loops back through those graphs. On the last iteration, the Python runtime drops every loop-back (name, next_node) pair, wherever the dest is. The Rust runtime dropped only the pairs whose dest was in its own worker graph. So after the final step it still sent latents / time_index to LLM, LLM_cfg_text and LLM_cfg_img. Each of them ran one more forward with its loop counter back at 0. Their velocities then parked at combine_cfg, which never ran again.

On main this costs one wasted denoise step per request on every CFG GPU, about 1/49 of the denoise. The image is not affected. A node that reads its loop counter breaks instead: on PR #860's BAGEL walk split, the extra pass hangs image_gen until the request times out.

How it works

1. compile.rs: LoopSpec.loop_back keeps every Loop._loop_back_inputs pair as (signal, dest node name), including pairs whose dest is in another worker graph. Before, it held only the local (signal, NodeId) pairs. Speculation reads the same list, and it is still local-only, as in Python.
2. state.rs, finish_loop: edges are dropped by (name, dest_sym) against that list. This is the check Python makes, so Dest::Local and Dest::External edges are handled alike. It runs inside complete(), before the sibling-graph local resolve and before the TP fanout, so every rank drops the same edges.
3. Deleted Completed.filtered and the return value of register_loop_finish. Nothing read either one, and both returned only the local pairs.

How it was tested

• New: test/rust/test_runtime_parity.py::test_the_last_iteration_sends_no_loop_back_to_another_worker_graph, parametrized over both runtimes and over where step runs: on a peer worker (the wire path) or in a sibling worker graph on this worker. It builds the loop with the real _divide_into_worker_graphs, runs it for all 3 iterations, and checks that the last iteration routes no loop-back to step (not staged for the peer's read, not ingested into the sibling graph) and holds no reference for it.
  • On origin/main: [rust-peer] and [rust-sibling] fail ([True, True, True] != [True, True, False]), and both [python-*] cases pass.
  • With the fix, all four pass.
• A one-off harness, not committed: main's real BAGEL image_gen_cfg worker graphs on main's 3-GPU configs/bagel_cfg_parallel.yaml, with three runtimes delivering every peer send over the real ZMQ transport for Rust. Per-node loop counters:
  • On origin/main, Python: LLM, LLM_cfg_text, LLM_cfg_img and combine_cfg each run k = 0..48, and vae_decoder runs once.
  • On origin/main, Rust: LLM, LLM_cfg_text and LLM_cfg_img each run k = 0..48 and then 0.
  • With the fix, both runtimes run k = 0..48.
• CPU box, linux, with mstar_rust built from the branch:
  • pytest test/rust/ on 532c6ed: 539 passed, 5 skipped.
  • cd rust && cargo test --release --no-default-features --features python on 532c6ed: 77 passed.
  • The CPU Core set on 532c6ed (this branch with main 6690cff merged): 8613 passed, 950 skipped and 2 failed. Both failures, the two cases of test_eval_compare_multigpu.py::test_dump_stacks_retries_with_sudo_only_when_asked, also fail on main (fadeb27). No failure is new. The head, 608b272, only adds the merge of main fadeb27, and CI's CPU Core job passes on it.
  • ruff check .: clean.
• GPU, H200s, this PR (608b272) against main (fadeb27), each on both graph runtimes (Python with MSTAR_RUST_GRAPH=0), over the SHM transport:
  • Configs: configs/bagel_cfg_parallel.yaml on 3 GPUs; configs/qwen3omni_2gpu.yaml and configs/orpheus.yaml on 2. Qwen3-Omni ran in its own venv with torch 2.10.0+cu130, torchaudio 2.10.0 and the flash-attn 2.8.3 wheel from docs/installation.rst, because the default venv's torch 2.12.1 and torchaudio 2.11.1 fail to load libtorchaudio.
  • BAGEL CFG-parallel, 8 seeded t2i requests per arm (3 at c=1, 4 at c=4, 1 resend), with forwards counted from a temporary log line per executed batch. On main, the Rust runtime runs LLM, LLM_cfg_text and LLM_cfg_img 50 times per request (k = 0..48 and then 0). Main's Python runtime, and this PR on both runtimes, run each of them 49 times. combine_cfg runs 49 times and vae_decoder once in every arm. The images are bit-identical across all 4 arms. No request failed, and the server was idle after every drain.
  • The new parity test, rerun on the same two commits: main's build fails [rust-peer] and [rust-sibling] and passes both [python-*] cases. This PR passes all four (with test/modular/test_graph_runtime.py: 213 passed, 1 skipped).
  • Streaming and dynamic loops: Qwen3-Omni (thinker, talker, Code2Wav) and Orpheus (LLM → SNAC), 22 requests per arm (16 at c=4, 4 at c=1, 2 resends). Both are 22/22 in all 4 arms, with no empty audio, and the server was idle after every drain.
  • Not tested: TP (none of these configs use it), and the TCP transport.
• GPU, H200, cold, with deterministic Inductor, on an earlier revision of this PR: PR Performance 3/7: image, video and world models #860's BAGEL walk split (d7e6e68), whose nodes read the loop counters. There the extra pass hangs image_gen until the 600 s timeout on the Rust runtime, on both the 2- and 3-GPU CFG-parallel layouts. With the fix cherry-picked, t2i, edit, emptywin and think all pass on both layouts. Forward counts and KV publishes then equal the Python runtime's.

Follow-up, not in this PR: the finished-iteration filter covers only the outputs of the node that completes a worker graph's last iteration. If one worker graph owns two or more nodes of a split loop, an earlier node can still send a loop-back to a peer on the final iteration. Both runtimes have this gap (rust/src/graph/state.rs:563, mstar/model/base.py:773), so parity holds. BAGEL is unaffected, because combine_cfg is its graph's only loop node.

🤖 Generated with Claude Code

[mstar-review]

Automated review

🟢 No findings on f5191ce, light review

What the review checked

• finish_loop (rust/src/graph/state.rs:622-623) now drops edges by (name, dest_sym), the same check as Python's filtered_signals in mstar/graph/base.py:782 and graph_io.py:102. It covers Dest::Local and Dest::External edges alike.
• dest_sym is set to the interned dest name for every edge kind (compile.rs:59), and to the node name for re-emitted external inputs (state.rs:688), so it lives in the same namespace as the interned loop_back dest names.
• _divide_into_worker_graphs copies the original loop's full _loop_back_inputs into each worker-graph Loop (mstar/model/base.py:384), and rust.py:143 passes it through unchanged, so the cross-graph pairs reach compile_one.
• Speculation's next-iter check (state.rs:808) compares against g.node(d).name for a local d, so it stays local-only as before.
• Filtering for nested loops still works: each finish_loop in the cascade filters the whole out against its own loop's pairs, which matches Python unioning outer.filtered_signals.
• The filter runs inside complete(), before the sibling-graph resolve and TP fanout in runtime.rs, so every rank drops the same edges (invariant 11).
• Completed.filtered and the return value of register_loop_finish had no other readers in rust/src. The only caller (runtime.rs:926) already ignored the return value.
• The new parity test only uses helpers that exist (_divide_into_worker_graphs, _released, _ready, _ingest_block, _spec), and the module is guarded by importorskip on mstar_rust.

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

[NSagan271]

🤖 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: #931 at 608b272 vs main fadeb27, H100 80GB (2-GPU allocations)

Verdict: the fix works and nothing regressed. On BAGEL CFG-parallel, base's Rust runtime runs one extra pass on the CFG branches. With the PR, every loop node runs exactly as many times as combine_cfg, matching the Python runtime. Each model and layout below ran on both runtimes with 0 tracebacks and 0 failed requests. Main has moved 2 commits since fadeb27; neither touches these files.

Each arm's own mstar_rust was built from its tree. Forwards were counted per (worker, node, walk, request) by a probe wrapping Worker._execute_on_gpu_thread. Everything ran at -O0, over SHM. Each BAGEL cell is 7 seeded t2i requests (3 at c=1, then 4 at c=4).

The claim: BAGEL CFG-parallel

This allocation had only 2 GPUs, so configs/bagel_cfg_parallel.yaml ran with LLM_cfg_text and LLM_cfg_img both on rank 1. That keeps both paths the new test covers: the loop-back to a peer worker, and to LLM in a sibling worker graph on rank 0.

image_gen_cfg forwards per request	LLM	LLM_cfg_text	LLM_cfg_img	combine_cfg	vae_decoder
main, Rust	50 ×7	50 ×7	50 ×5, 49 ×2	49	1
main, Python	49	49	49	49	1
PR, Rust	49	49	49	49	1
PR, Python	49	49	49	49	1

• The c=1 images are bit-identical across all 4 arms (3/3).
• At c=4, 2 of 4 images differ between any two arms, including main-Python vs PR-Rust, neither of which has the extra pass. The differences are small, PSNR 36–39 dB, consistent with batch composition. I read them as run-to-run noise, not the PR.
• The PR's new parity test, run against each arm's build: main's fails [rust-peer] and [rust-sibling], and passes both [python-*]. The PR's passes all 4.

Other loop layouts (both runtimes; main vs PR on Rust)

Model, layout	What was checked	Result
Qwen3.5-0.8B, 1 GPU (AR decode)	LLM forwards per request vs tokens returned; greedy text	✅ forwards = tokens on all 96 requests in all 4 arms; c=1 text 32/32 identical; c=16 41–52/64, same band as a main-Rust vs main-Python control (41/64)
Qwen3.5-9B, TP2 (qwen3_5_9b_tp2.yaml)	same, on both ranks	✅ forwards = tokens on both ranks, all 96 requests; c=1 text 32/32; no hang
BAGEL PD-disaggregated (bagel_pd_disaggregated.yaml)	t2i forward counts and images	✅ image_gen 49 + vae_decoder 1 per request in every arm; c=1 images bit-identical
Orpheus PD-disaggregated (orpheus_pd_disaggregated.yaml, LLM → SNAC streaming)	benchmark.runner TTS, 24 requests at c=4	✅ 24/24 completed, 0 failed, in every arm

Not covered

• The 3-GPU layout of bagel_cfg_parallel.yaml as shipped (this allocation had 2 GPUs). PR Performance 3/7: image, video and world models #860's walk split, where the extra pass hangs image_gen.
• The TCP and RDMA transports. Throughput: the change only touches loop finish, so I didn't benchmark it.
• The follow-up named in the description (several loop nodes in one worker graph): no shipped config has that layout.

[NSagan271]

LGTM