[merceod]

A Qwen3-Omni request with no audio output now decodes on its own walk, thinker_text_decode: its token stays on the device, nothing goes to the Talker, and the stop check reads one host copy. Both decode walks also get a KV plan cache that matches, a captured decode step builds its embeds and rotary tables in the graph, and a decode batch of 8 or 16 rows routes and aligns its experts in one kernel. Against main b06ae391f (#970 and #1140 in), on one B300, two or three sessions each alternated: T2T +5.3 / +1.1 / +4.2% tok/s at c=1/16/64 (+2.0 / -0.3% at c=1/16 without one slow main session), A2T +3.4 / +10.2% RTFx at c=1/16, T2S +3.8% and A2S +4.0% audio-s/s at c=1, 0 failed.

Merged with main. f7fceb005 merges main b06ae391f into the branch (no rebase). #970 and #1140 changed the two files this touches, and where both did the same thing, main's version stays: the Thinker's fused add + RMSNorm (keep_states now drops main's layer-0 view and layer-N add), a decode row as its token alone with its position from the plan, and a mixed step's decode rows embedded in one call in preprocess, which now also takes loop-back rows (_decode_tokens). #1140's per-pair GEMVs keep decode batches of up to 7 rows: route_align hands fused_experts an alignment, which turned that path off, so the router now aligns only where fused_experts reads one (1efd5666b). This PR's walk, plan-cache fix, in-graph tables, fused route from 8 rows and B300 tiles stay. The text-decode test moved to the Rust runtime with #1130, and main's new tile-table trace test resets this branch's table cache (d117982ee).

Merged with #986. ad1b4e8cf merges main f609cac54. fused_moe/kernels.py and test_fused_moe_config.py are main's (the registry lookup replaces this branch's table read and nearest-size search). runner.py keeps main's one get_config read plus runs_per_pair. The B300 tiles moved into the registry as mstar/kernels/tables/moe_gemm/NVIDIA_B300_SXM6_AC.json, with the same tiles at every count. Tests: ruff clean; GPU modules 1038 passed, 6 skipped, 2 failed (main's two); CPU Core fails main's ids. T2T on one B300, one session per arm: 325 vs 312 tok/s at c=1 and 1840 vs 1809 at c=16, 0 failed, nothing recompiled after startup.

How this sits on #859

#859 declared what used to be each model's own hooks. ARNodeSubmodule.device_loopback = DeviceLoopback(walk, signal, sampler, buffer) gives uniform_row_inputs, device_loopback_signals, inline_client_signals and loopback_inputs; ARNodeSubmodule.token_stop = TokenStop(...) gives check_stop and check_stop_batched. For the Thinker it also dropped the per-row host syncs in decode, made decode's thinker_states the layer-0 row, and set fused_softmax_topk on the routers. The Thinker's layer loop itself didn't change, it isn't on run_layers.

What this adds is the walk #859 couldn't declare for thinker_decode: the rows-only path needs every output of a walk to stay on the device, and thinker_decode streams states to the Talker on every step. So the text walk declares:

self.device_loopback = DeviceLoopback(
    walk=THINKER_TEXT_DECODE, signal="text_inputs", sampler=THINKER_SAMPLER, buffer="new_token",
)
self.token_stop = TokenStop(
    loop="thinker_decode_loop", buffer="new_token", stop_ids=frozenset({config.im_end_token_id}),
    sampler=THINKER_SAMPLER, walk=THINKER_DECODE_WALKS,
)

On top of the declarations the walk has its graph (_get_thinker_forward picks it from audio_output), keep_states=False in _decode so the two Talker states go, and _preprocess_decode, which reads through loopback_inputs. Prefill now returns row outputs, so a mixed step's text-decode rows find their stop check and inline emit. A captured text-decode step always builds its rotary tables in the graph, as Qwen3.5's decode has since #859, so ingraph_decode_rope plays no part here.

Compared with this PR's first version, these went away in favour of the declarations: _stops, check_stop, check_stop_batched, uniform_row_inputs, device_loopback_signals, inline_client_signals, _device_loopback, reads_device_loopback, _ingraph_rope and the hand-written slot read in _preprocess_text_decode. submodules.py went from +198/-62 to +145/-36 against main then; against today's main it's +178/-84, since its decode paths replace the ones #970 rewrote. step_is_input_free stays, since the engine still reads it (engine.py:2969).

TokenStop.walk also takes a set now. Both decode walks end one loop, and the field took one walk or None. None would stop correctly, because no prefill walk carries this loop and the graph runtime drops the stop (check_dyn_loop). But that's the runtime's model-bug path, and prefill rows sample im_end all the time, since a prompt section tends to end a user turn: two per A2T request in our runs, one Tried to stop loop ... potential logical bug error each (528 over one guard server's 264 A2T requests on the Python runtime before #1130, 0 with the set; the Rust runtime drops them without a log). A set keeps what main does for a prefill row, which is name no stop. TokenStop.stops_on is the one place that reads the field, and no other declarer changes.

qwen3omni_pd_2gpu.yaml is #859's config and placed only thinker_decode on the decode rank, so it gains the new walk next to qwen3omni_pd_disaggregated.yaml. A test checks every config that splits the Thinker by walk.

The rest of the change

step_is_input_free on both decode walks. The KV plan cache checks step.segments is cache.step_segments at commit (kv/manager.py:3036). A step redeclared before the forward is a different object, so the epoch bumped every step and _fast_decode_plan almost never matched. Fast plans over a T2T phase pass: 74 of 36,036 on main, 34,101 of 36,084 here (measured on 53c47cb8c, before the merge). On requests with audio output it matches about half of the Thinker's steps at c=1 on one GPU: 878 fast plans over 1,740 thinker_decode steps, against 1 over 1,796 on main. Why the rest miss isn't traced.

One kernel for routing and alignment. route_align moves from Command A+ to mstar/utils/fused_moe/route.py and gains a softmax and any top-k. fused_softmax_topk, set on both Omni routers since #859, is the one flag, and TopKRouter.route is the one place that reads the split: on CUDA, a batch of up to 128 renormalized slots whose block can consume an alignment (one GPU, no masked rows) takes route_align, unless fused_experts would run it as #1140's per-pair GEMVs, which read no alignment (runs_per_pair, the check fused_experts itself makes: up to 7 tokens at the Omni expert sizes). Anything else takes router_topk, and with the flag off it's softmax, torch.topk and the renormalize. For the Thinker's decode buckets that means 1, 2 and 4 rows run main's per-pair path and 8 and 16 take route_align. The Talker has norm_topk_prob=False, so it stays on router_topk. GLM's and Command A+'s routers don't set the flag and don't reach any of this.

So the baseline is router_topk, the two alignment kernels and the id cast. Per decode step at T2T from a node-traced nsys window, one server each on 1 x B300, before the merge (main is 53c47cb8c, which had no per-pair path; 640 and 491 steady steps at c=1, 383 and 414 at c=8), as count x µs each = µs a step:

kernel	c=1 main	c=1 this	c=8 main	c=8 this
_router_topk_kernel	48 x 2.44 = 117		48 x 2.61 = 125
moe_align_block_size_kernel	48 x 3.10 = 149		48 x 3.21 = 154
count_and_sort_expert_tokens_kernel	48 x 1.33 = 64		48 x 1.33 = 64
id cast, int64 to int32	48 x 0.99 = 47		48 x 1.00 = 48
_route_align_kernel		48 x 3.69 = 177		48 x 5.66 = 272
routing and alignment	377	177	392	272
router GEMM	48 x 1.87 = 90	48 x 1.99 = 96	48 x 4.97 = 239	48 x 4.99 = 239
fused_moe_kernel (B300 tiles)	96 x 12.89 = 1237	96 x 9.69 = 931	96 x 45.09 = 4329	96 x 39.01 = 3745
act_and_mul_kernel	48 x 1.78 = 86	48 x 1.76 = 84	48 x 2.21 = 106	48 x 2.41 = 116
moe_sum_reduce_kernel	48 x 1.91 = 92	48 x 1.87 = 90	48 x 2.28 = 109	48 x 2.13 = 102
MoE block	1882	1378	5174	4474

The route saves 200 µs a step at c=1 and 120 at c=8; the tiles save 307 and 584. The decode graph goes from 1167 kernels to 1026.

Since #1140, main runs up to 7 tokens as per-pair GEMVs, and the merge keeps that path where both apply. On B300, the Thinker's MoE block alone (router GEMM, route, experts) at E=128, H=2048, I=768, 48 calls over 8 weight sets captured in one graph and replayed, median of 5 x 100 replays, three runs within 0.1 µs (scripts/moe_micro.py in the run directory):

tokens	main's route, µs a step	route_align + grouped GEMM, B300 tiles	change
1	1574	1479	-6.0%
2	1980	1985	+0.3%
4	3297	2908	-11.8%
8	4637	4524	-2.5%
16	6626	6772	+2.2%

At 1 to 7 tokens main's route is the per-pair GEMVs and at 8 and 16 it's router_topk plus the align kernels into the same grouped GEMM. End to end it's a trade, two sessions each alternated on GPU 0, the merge before that commit (f7fceb005, fused route at every size up to 16 rows) against this head:

cell	fused route from 1 row	this (per-pair up to 7)	fused vs this
T2T c=1 tok/s	310.8, 310.6	315.1, 318.1	-1.9%
T2T c=4 tok/s	762.1, 766.8	712.3, 711.0	+7.4%
A2T c=1 RTFx	58.6, 58.9	59.8, 60.8	-2.6%

The micro and the cells agree at 4 rows and not at 1, which I haven't traced. The fused route also moves greedy tokens at 1 to 4 rows: its two servers agree 32/32 with each other and 5 or 6 of 32 with main or this head (A2T 192 of 200), where main and this give each other 20 to 32 (A2T 200). Dropping the runs_per_pair check in TopKRouter.route is the one-line way back to it, and a cut between 2 and 4 rows looks like the B300 optimum (a guess from these two cells). I kept main's where the two overlap: its per-pair tiles were swept on H100, and the B300 tiles that make the grouped GEMM fast here don't load on Hopper, so which one wins there isn't measured. At 16 tokens route_align, one program over 128 slots, is a little slower than the two align kernels in this setup.

In-graph tables on both decode walks. A captured thinker_decode step takes ids and builds its embeds and rotary tables in the graph, through the same _decode as the text walk: only keep_states differs, and the layer-0 row and mask go to the Talker as before. A decode row brings its id and nothing else, which #970 also made main's thinker_decode row do; a mixed step embeds its decode rows in preprocess, once a step rather than once a row, and reads their positions off the plan, with main's code, which now also reads a loop-back row's token off the sampler's slot master (_decode_tokens). What this adds over main is that a captured decode step builds its embeds and rotary tables in the graph, where main's builds them in preprocess. Measured before the merge, T2T pushed through thinker_decode in a scratch build, before and after: eager kernels a step 12.1 to 0.1 at c=1 and 29.4 to 1.9 at c=8, copies 20 to 16, tok/s 278 / 1119 / 2913 to 283 / 1121 / 2943 at c=1/8/32. The Talker and Code2Wav set the audio-output cells, where it doesn't show:

before → after the port, two sessions each	audio-s/s	change	TTFA p50 ms
T2S c=1	17.1, 17.2 → 17.8, 17.4	+2.4%	126, 126 → 125, 130
T2S c=16	94.4, 97.2 → 97.2, 93.9	-0.3%	273, 268 → 266, 276
A2S c=4	27.8, 26.7 → 27.1, 26.6	-1.4%	285, 293 → 280, 293
A2S c=16	59.7, 59.5 → 59.7, 57.9	-1.4%	469, 447 → 447, 467

One graph for both decode walks. Sending text requests down thinker_decode costs too much and can't be declared away: the rows-only path needs every output signal of the walk on the device (Worker._all_on_device), and thinker_decode streams thinker_states and thinker_mask to the Talker. Keeping both walks and letting them replay one capture does work, and I left it out (measured before the merge): it saves six captures but no measurable GPU memory or warm boot (95.3–95.9 GB, 165–174 s across these arms), and it makes the text walk carry two clones.

T2T tok/s, one session	text walk	through thinker_decode, states shipped	change	one capture, two walks	change
c=1	299	283	-5.6%	297	-0.6%
c=8	1131	1121	-0.9%	1138	+0.6%
c=32	3011	2943	-2.2%	3024	+0.4%

The kernels aren't the cost of the first: the graph gets 2 more, +12 µs a step at c=1 and +26 at c=8 (0.4%). The rest is the host's per-request outputs, postprocess and the stream send. The one-capture build matches the head on every greedy set (T2T 32/32, A2T 32/32, T2S and A2S text 8/8, 1,778 decode steps bit-identical), and text and audio requests at once ran with 0 failed.

Tuned tiles for B300. moe_gemm registry entries for the Thinker's and Talker's expert sizes (mstar/kernels/tables/moe_gemm/NVIDIA_B300_SXM6_AC.json), swept with the pre-#986 tuner and converted to the registry's format. They serve the grouped GEMM: prefill, and decode from 8 rows.

Gate against main

Main is b06ae391f, with #970, #1140 and #1127 (resident planning, which Omni doesn't declare, so it stays off for both arms). 1 x B300 SXM6 (sm_103), torch 2.12.1+cu132, FlashInfer 0.7.2.post1, -O2, qwen3omni_colocated.yaml, knobs kv_plan_cache: 1, ingraph_decode_rope: true, gc_gen0_threshold: 50000, prefix caching off, server on 8 cores, client on 4, an idle-priority spinner on each of those cores so they don't park. Two server sessions per arm in the order main, this, this, main, on one Inductor cache, and a third main session for T2T c=1/16 and A2T c=16, run after the route check above. T2T is benchmark.runner, closed loop, --ignore-eos, 257 tokens a request, 64/160/640 requests at c=1/16/64. A2T is 200 LibriSpeech test-clean clips over /generate, greedy, capped at 256. T2S and A2S are 32 of those prompts and clips answered in speech, greedy Thinker, seeded Talker. Each cell warms 5 waves first. ITL is (e2e - TTFT) / (tokens - 1), p50 per request.

cell	main	this	change	TTFT p50 ms main	this	ITL p50 ms main	this
T2T c=1 tok/s	282, 311, 312	318, 317	+5.3%	26, 24, 25	22, 24	3.44, 3.11, 3.10	3.06, 3.07
T2T c=16 tok/s	1719, 1799, 1793	1789, 1790	+1.1%	56, 55, 59	54, 47	9.04, 8.68, 8.69	8.73, 8.75
T2T c=64 tok/s	2955, 2955	3079, 3077	+4.2%	2732, 2724	2621, 2597	10.92, 10.94	10.48, 10.54
A2T c=1 RTFx	58.3, 58.2	60.3, 60.2	+3.4%	38, 39	38, 38	3.41, 3.42	3.30, 3.29
A2T c=16 RTFx	249.1, 247.1, 248.6	277.4, 269.8	+10.2%	98, 98, 94	78, 77	14.36, 13.48, 13.80	13.08, 12.98

cell	main	this	change	TTFA p50 ms main	this
T2S c=1 audio-s/s	17.13, 16.46	17.46, 17.40	+3.8%	128, 131	125, 125
A2S c=1 audio-s/s	14.77, 14.44	15.30, 15.09	+4.0%	142, 142	137, 138

Failed requests: 0. This is steady across its sessions (T2T c=1 315 and 318 in two more, in the route check above). Main's first session is 9% under its other two at T2T c=1, which looks like a different compile on the same cache (a guess, not traced). Leaving it out, T2T is +2.0% at c=1 and level at c=16 (-0.3%), and A2T c=16 is +10.4%. TTFT at c=64 is 2.6 to 2.7 s on both arms; I haven't looked at where it queues. For scale, one session of the old base 53c47cb8c on this harness read 241 / 1628 tok/s at T2T c=1/16 and 205 RTFx at A2T c=16, so main itself moved +29 / +10 / +21% since then.

The c=1 gain is smaller than against 53c47cb8c (+18%) for one reason I can point to: a 1-row decode step now runs main's per-pair MoE on both arms, where before the merge the route and the tiles saved this PR 507 µs a step at c=1 (the nsys table above). That #970's own host work also overlaps with this PR's is a guess I haven't measured. A2T c=16 keeps most of its gain, +10.2% against +12.8% before.

The rest of this section ran before the merge, against 53c47cb8c, and isn't rerun.

worker_phases at T2T, untimed, on the second GPU, mean ms a step. event_sync is the GPU wait, so the CPU side is what shrinks, and at c=32 the wait takes up the difference:

phase	c=1 main	this	c=8 main	this	c=32 main	this
iter_total	4.02	3.35	7.59	6.95	11.11	10.35
iter_total.cpu	1.79	1.46	5.01	4.70	7.60	7.86
worker.postprocess.event_sync	0.83	0.85	3.74	3.85	5.30	6.61
worker.gpu_thread.total.cpu	1.32	0.84	1.67	1.07	2.67	1.36
worker.gpu_thread.prepare_inputs	0.19	0.04	0.44	0.15	1.15	0.32
engine.preprocess	0.28	0.10	0.30	0.13	0.47	0.15
engine.collect_outputs	0.29	0.01	0.30	0.02	0.45	0.05

With the route and tiles off, the walk and the plan cache alone gave T2T 249 / 959 tok/s at c=1/8 in an untimed guard, where main gave 220 / 936 and this PR 287 / 1043 in the same setup. That's one run of 32 requests each, so read it for direction only.

vLLM-Omni 0.30.0 the same day as that gate, stock deploy (qwen3_omni_moe.yaml: Thinker on GPU 0, Talker and Code2Wav loaded on GPU 1 and idle for text, Triton MoE, FlashInfer attention, prefix caching off), same cores and the same benchmark.runner cells at temperature 0 and 256 tokens. It goes through /v1/chat/completions with the harness's system message and M* through /generate, so the prompts aren't byte-equal. It holds 247 GB on the Thinker's GPU against M*'s 96 GB.

T2T	vLLM-Omni (s1, s2)	main (mean)	this (mean)	this vs vLLM-Omni	TTFT p50 ms vLLM-Omni	this
c=1 tok/s	267, 271	250	297	+10.2%	22, 22	26, 25
c=8 tok/s	1140, 1133	1028	1138	+0.1%	64, 53	45, 46
c=32 tok/s	3057, 3090	2782	3026	-1.5%	150, 144	79, 66

Before the merge, main sat 7 to 10% behind vLLM-Omni on T2T, and this closed that at c=8 and c=32 and passed it at c=1. Main has moved since (#970), so that column is old.

Correctness

Greedy, exact text at c=1 on the merged head, one Inductor cache: the 32 T2T prompts (natural EOS, 256 tokens max) and the 200 A2T clips, on the gate's servers unless the row says otherwise:

compare	T2T, of 32	A2T, of 200
main vs main, three servers	20, 21, 23	200
main vs main, one cold compile (two guard servers)	5	7 of 8
this vs main, six pairs	20, 20, 23, 23, 32, 32	200 each (four)
this vs itself, the gate's two servers and the two in the route check above	32, 21	200, 199
this, same prompts through both decode walks (T2T vs T2S, A2T vs A2S), 32 each, each server	32	32

Counted in Thinker decode steps up to each prompt's first differing token: this vs main matches 7,344 of 7,344 against one main server, 6,537 of 7,316 against the second and 6,196 of 7,340 against the third, and main's servers give each other 6,196 to 6,537. The A2T steps match on every pair, 5,605 of 5,605. So the merged head moves no token that main doesn't move between its own boots. That isn't the "route and tiles move bf16 rounding" picture from before the merge: a 1-row step now takes main's per-pair route on both arms, and the fused route at 1 to 4 rows was what moved the tokens (the merge without 1efd5666b gives 5 or 6 of 32 against main, above). The tiles still run this PR's prefill, and on these sets that moved nothing.

Before the merge, at f4cbbeb0b, not rerun: the route and tiles off vs main 32/32 and 32/32, loop-back on vs MSTAR_DEVICE_LOOPBACK=0 32/32 and 32/32, before vs after the port 32/32 and 32/32, and the 200-clip A2T WER at c=8 (5.30–5.37% on main, 5.30–5.54% here).

Audio output. T2S and A2S text match as T2T and A2T do above (A2S 32/32 on every pair). PCM bytes are 0/32 on every pair, main vs main and this vs itself included: the seeded Talker's audio still differs across server processes, as before the merge.

Mixed steps with audio. The 28 short prompts and 4 of 2,600–6,500 tokens at c=16 staggered 150 ms, answered in speech, on two servers per arm (the first of each a cold compile, and the first one here before 1efd5666b): 32/32 with audio, 0 failed, about 2,020 s of audio a pass on both.

Stalls. None in 13 servers at this merge, 0 of 7 on main and 0 of 6 here. The old head's one stall was a runtime dynamo compile of the Talker's long prefill row in fused_experts. Here that compile happens on both arms, once per server, only on the long prompts answered in speech: 20 recompiled frames on main (4 of them torch_dynamo_resume_in_fused_experts_*) and 18 here (2), over 4–10 s each, and every one finished. The gate servers recompiled nothing after startup. So it isn't specific to this PR, and the one stall stays unexplained.

Tests

• ruff check .: passes.
• CPU Core at 1efd5666b (CUDA hidden, xdist, 2 x B300 host): 11003 passed, 997 skipped, 18 failed; main b06ae391f on the same host 10946, 974, 17. The same 17 by id; the 18th (test_eval_compare_multigpu::test_a_head_that_hangs_is_cut_at_the_cap_with_stacks_and_fails) fails alone on main too. d117982ee changes a CUDA-only test. The TP lockstep sim and the piecewise signature module, run serially: 27 passed.
• GPU modules on CUDA at the pushed head (test_fused_moe_route, test_fused_moe_config, test_fused_moe_decode, test_fused_moe_align, test_moe_router, the Omni modules, test_worker_device_loopback, test_inline_tokens, test_stop_policies, mstar/model/command_a_plus/tests): 989 passed, 6 skipped, 2 failed. The two failures are test_fused_experts_matches_naive_and_precomputed_alignment[1] and [3], which fail the same way on main: since Per-model perf and serving fixes for Qwen3.5, Qwen3-Omni, Fish, OmniVoice and others #1140, fused_experts runs 1 and 3 tokens per pair without an alignment and grouped with one, so the two aren't bit-equal.
• test_fused_moe_route.py gains a test that a batch fused_experts runs per pair takes router_topk and the next size up takes route_align. The rest as before: the split at the cap, no renormalize and no alignment target taking router_topk, a block offering its weights only where it can use the alignment, and a default TopKRouter, GLM's NoAuxTcRouter and CommandAPlusRouter refusing any fused route. test_qwen3_omni_text_decode.py covers both decode walks: ids only on a routed row, both captures taking ids, in-graph tables, the states and mask on thinker_decode, no host read, a mixed step's planned positions, and the stop on either walk through the Rust runtime.
• Command A+: not served, it's 437.5 GB with tp4/tp8 configs on a 2-GPU box. language_model.py has no diff; its tests are in the GPU run above. route_align moved vs pre-move on B300 at E=128, K=8 (before the merge): outputs equal on 1020/1020 calls per size, 42.7 / 45.3 / 49.7 µs a call pre-move vs 43.7 / 46.6 / 50.9 moved at 1 / 8 / 16 tokens, and the extra µs is the Python launch, which a captured graph doesn't pay.

Not in this PR

• Not rerun at this head: I2T, V2T, the HF fp32 reference set, the KV plan-cache counts, worker_phases, nsys, and the 2-GPU configs. Those ran before the merge.
• thinker_decode still computes layer_n_hidden (main's residual add now), which decode doesn't read. Dropping it touches the audio path, so it's a separate change.
• Where the per-pair path should hand over to the fused route on B300 (+7.4% at T2T c=4 for the fused route, -1.9% at c=1, above), and whether route_align should stop at 64 slots. Both want a check on Hopper too.
• A Command A+ end-to-end run: it needs a tp4 node.
• One capture replayed by both decode walks: measured free above and left out.

On H200

NSagan271's live test of the first version (101f9fe against its merge base 9b90e81, H200, Rust runtime, torch 2.10, median of 3 servers per arm) is the Hopper picture. The B300 tile tables don't load there, so what carries over is the text walk, the plan cache and route_align on 1 GPU (it is off under TP).

cell	1 GPU	thinker TP2
I2T c=1	+13%	+1%
I2T c=4	+5%
I2T c=16	+5%	+1%
A2T c=16	−9.5% in one round, not in a rerun of six fresh servers	+7%
T2S c=1	+8%
T2S c=16	+7%	−1%
A2S c=4	+5%	+6%

The A2T c=16 dip came with p99 inter-token gaps of 185–239 ms in a round where A2T ran after the I2T cells on each server, and did not show when A2T ran first; both arms show the same occasional gaps. An I2T-then-A2T run on one server with this head is queued on our side. Greedy replies there were identical across repeats on every server, and more reproducible across restarts than main on single images (21/21 vs 11/21).

[merceod]

Merged main (through #861) at 101f9fe. The conflict was the return_token line in postprocess; beyond it, thinker_text_decode counts as a decode walk in the row-length hook chunked prefill (#916) added, so it reports one token like thinker_decode.

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human.

Live test: #1111 at 101f9fe vs its merge base 9b90e81, H200 (sm_90)

Verdict: works, and most of the speedup carries over to Hopper. The B300 tile tables don't load here, so the gain comes from the text-only decode walk, the plan-cache fix and route_align (route_align only on 1 GPU, since it is off under TP). There were no failures in any phase. Outputs differ from main only where main also differs between its own restarts.

Setup:

• Qwen3-Omni on qwen3omni_colocated.yaml (1 GPU, route_align active) and on qwen3omni_thinker_tp2.yaml (3 GPUs).
• Three fresh servers per arm, in the order main, PR, PR, main, PR, main.
• Rust runtime, torch 2.10, FlashInfer 0.7.2. Benchmarks use benchmark.runner closed loop, without ignore-EOS.

Throughput, median of 3 servers (text tok/s, or audio s/s for speech)

cell	1 GPU main	1 GPU PR	Δ	thinker TP2 main	TP2 PR	Δ
I2T c=1	200	226	+13%	168	170	+1%
I2T c=4	456	480	+5%
I2T c=16	1060	1114	+5%	1105	1117	+1%
A2T c=16	392	355	−9.5%*	383	411	+7%
T2S c=1	11.7	12.6	+8%
T2S c=16	57.0	60.7	+7%	69.8	69.3	−1%
A2S c=4	14.9	15.7	+5%	18.8	20.0	+6%

* The A2T c=16 dip didn't reproduce. In the first round, all three PR servers read below all three main servers (347–359 vs 382–395). The slow tail came from p99 inter-token gaps of 185–239 ms vs 53–64 ms, and the p50 was unchanged. In that round A2T ran after the I2T cells on each server, while four other jobs shared the node.

I reran it as the first workload on six fresh servers, three runs back to back each:

run	main	PR
1st	336, 403, 356	361, 375, 441
2nd	440, 422, 403	380, 405, 473
3rd	435, 404, 427	458, 486, 468

Both arms show the same occasional 150–180 ms p99 gaps. I don't have an explanation for the first round. It could be an order effect that only shows after the I2T cells, or node contention. One way to check is to run I2T then A2T on the same server.

TP2 gains are smaller than the PR's table. I2T c=1 is +1% here, against +24% on B300. Under TP the eager router stays, which the PR states, and the H200 decode step may be bound elsewhere. I haven't profiled it.

Correctness

• Greedy replies: 24 prompts covering audio, image, two images, video, text, and text around attachments, each sent twice per server. All 24 repeats were identical on every server.
  • Across arms, every case that differs also differs between two main servers. Text and text-around-attachment prompts matched on all 18 cross-arm pairs.
  • The PR is more reproducible across restarts than main on single images: 21/21 vs 11/21. That is consistent with route_align breaking ties toward the lowest expert every time (hypothesis).
• Text and speech replies in flight together: these exercise the new thinker_text_decode walk next to the audio walk. No errors; speech came back for 6/6 on every server. Text matched solo on 7–12 of 12 for both arms, the same spread.
• Speech abort storm: 3 rounds of 32 requests with 16 dropped. All kept and fresh requests finished on every server and both configs, and /health answered 200 afterwards.
• GPU tests on 1× H200: the PR tree passed 43 (test_fused_moe_route, test_fused_moe_config, test_fused_moe_fp8, Command A+ test_fused_kernels); the main tree passed 21 (no route test there). The Command A+ kernel tests pass after route_align moved.
• No recompiles on PR servers: the PR servers logged no recompile lines (main: 3 at 1 GPU, 6 at TP2).

Not covered

• B300 tile tables, since this is H200.
• Command A+, GLM and other MoE models end to end. Only their kernel tests ran. The get_config lookup changes how the nearest batch size is found but picks the same entry, smaller on a tie.
• CPU suites, which CI covers.

[merceod]

Thanks for the Hopper run. The branch is being brought onto main after #859 (its device loop-back and token-stop declarations replace the hooks this branch carried on its own), so the table will be refreshed on that base, and I'll put the H200 columns next to the B300 ones.

On TP2: under TP on H200 only the text-only decode walk and the plan-cache fix are active (no tile tables, route_align off), so a smaller gain there is in line with what the PR says; I don't have a per-piece split at TP2 on B300 either. The A2T c=16 rerun reads like the same p99 gaps on both arms. I'll run I2T then A2T on one server with the rebased branch to check the order effect.

[merceod]

Ported onto main after #859 at cea7d03: the walk now declares DeviceLoopback and TokenStop (the only base change is TokenStop.walk taking a set, since both decode walks end one loop). Description refreshed with the B300 gate on that base (T2T +18/+11/+9%, A2T +15/+9/+9%) and your H200 table; the I2T-then-A2T order check is queued.

[NSagan271]

Mainly looks good, just a few comments.

[NSagan271]

Can you also port embed/RoPE-in-graph to thinker_decode (or do they not carry over?) This might also allow one graph to be captured for both decode walks (though then the text decode would pay the cost of computing and returning the thinker states; not sure if that would degrade throughput)

[merceod]

They carry over, done in 9fe581fe0. A captured thinker_decode step takes ids and builds its embeds and rotary tables in the graph, through the same _decode as the text walk. Only keep_states differs, and the layer-0 row and mask go to the Talker as before.

Parity. One B300, c=1, one Inductor cache: 1,778 Thinker decode steps over 8 T2S and 8 A2S are bit-identical before vs after (token, thinker_states bytes, mask), T2T 32/32, A2T 32/32, T2S and A2S text 8/8. WAV bytes and greedy codec tokens are 0/16, but so is before vs before: the Talker's own step inputs differ across boots from step 2 to 18, with the same text conditioning going in. I haven't traced why.

What it buys. On the Thinker alone, T2T pushed through thinker_decode in a scratch build: eager kernels a step 12.1 to 0.1 at c=1 and 29.4 to 1.9 at c=8, tok/s 278 / 1119 / 2913 to 283 / 1121 / 2943 at c=1/8/32. On audio output it doesn't show, the Talker and Code2Wav set those cells:

two sessions each	before audio-s/s	after	change	TTFA p50 ms before	after
T2S c=1	17.1, 17.2	17.8, 17.4	+2.4%	126, 126	125, 130
T2S c=16	94.4, 97.2	97.2, 93.9	-0.3%	273, 268	266, 276
A2S c=4	27.8, 26.7	27.1, 26.6	-1.4%	285, 293	280, 293
A2S c=16	59.7, 59.5	59.7, 57.9	-1.4%	469, 447	447, 467

One graph. Two ways to read it, and I measured both, T2T tok/s, one session each:

	text walk	text requests on thinker_decode, states shipped	change	one capture, both walks replay it	change
c=1	299	283	-5.6%	297	-0.6%
c=8	1131	1121	-0.9%	1138	+0.6%
c=32	3011	2943	-2.2%	3024	+0.4%

Merging the walks is out: it's over 2% at c=1 and c=32, and the declaration can't be met, because the rows-only path needs every output signal of the walk on the device (Worker._all_on_device) and thinker_decode streams thinker_states and thinker_mask to the Talker. Computing the states isn't the cost, that's 2 more kernels, +12 µs a step at c=1 and +26 at c=8. It's the per-request outputs, postprocess and the stream send on the host.

One capture with two walks is your suggestion taken literally, and it's free: replay_graph_walks on thinker_decode's config, the text walk stays rows-only and never reads the states. Same tokens as the head on every set above, and text and audio at once ran with 0 failed. I left it out because it saves six captures and no measurable GPU memory or warm boot (95.3–95.9 GB, 165–174 s across these arms), and the text walk would carry two clones for it. It's a 6-line patch, so say if you'd rather have the one config and it goes in.

[NSagan271]

The description says no Omni router sets fused_softmax_topk, but the Thinker and Talker both do since #859. So the baseline this replaces is the fused softmax + top-k kernel, not softmax → torch.topk → renormalize → cast. Could the description's kernel accounting be updated to match?

That also leaves TopKRouter with two fused routing paths: fused_softmax_topk inside forward, and route_aligned picked separately in _route. Can route_aligned be the fused_softmax_topk path's branch for batches up to ROUTE_ALIGN_MAX_SLOTS, with the existing fused kernel as the fallback above that? Then there would be one flag and one place to read the routing logic. If there's a reason to keep them separate, a sentence in the docstring would help.

[merceod]

You're right about the baseline, and yes. Folded in f469d63ea: TopKRouter.route is the one place that reads it. With fused_softmax_topk on CUDA, up to ROUTE_ALIGN_MAX_SLOTS renormalized slots whose block can consume an alignment (one GPU, no masked rows) take route_align, anything else takes router_topk, and with the flag off it's the eager chain. forward calls it with nothing to align against, and _route passes gate_up_proj only where the block can use the result. The Talker has norm_topk_prob=False, so it stays on router_topk.

This also fixes something: route_aligned didn't read the flag, so a default TopKRouter in a ParallelSparseMoeBlock would have taken it. Now it doesn't. GLM's and Command A+'s routers don't set the flag, and test_a_router_without_the_flag_routes_as_it_always_did refuses any fused route for them and for a default TopKRouter.

No token moved, and the head boots off the same Inductor cache in the usual 170 s: the aligned path and the fallback trace as they did before the fold.

Per decode step at T2T, nsys, one server each on one B300, main 53c47cb8c, 491+ and 383+ steady steps at c=1 and c=8, count x µs each = µs a step:

kernel	c=1 main	c=1 this	c=8 main	c=8 this
_router_topk_kernel	48 x 2.44 = 117		48 x 2.61 = 125
moe_align_block_size_kernel	48 x 3.10 = 149		48 x 3.21 = 154
count_and_sort_expert_tokens_kernel	48 x 1.33 = 64		48 x 1.33 = 64
id cast, int64 to int32	48 x 0.99 = 47		48 x 1.00 = 48
_route_align_kernel		48 x 3.69 = 177		48 x 5.66 = 272
routing and alignment	377	177	392	272
router GEMM	48 x 1.87 = 90	48 x 1.99 = 96	48 x 4.97 = 239	48 x 4.99 = 239
fused_moe_kernel (B300 tiles)	96 x 12.89 = 1237	96 x 9.69 = 931	96 x 45.09 = 4329	96 x 39.01 = 3745
act_and_mul_kernel	48 x 1.78 = 86	48 x 1.76 = 84	48 x 2.21 = 106	48 x 2.41 = 116
moe_sum_reduce_kernel	48 x 1.91 = 92	48 x 1.87 = 90	48 x 2.28 = 109	48 x 2.13 = 102
MoE block	1882	1378	5174	4474

So against the real baseline the one kernel saves 200 µs a step at c=1 and 120 at c=8. The tiles save 307 and 584. The description has this table now.

[NSagan271]

Command A+ was affected by these changes; can you re-benchmark it if you haven't already?

[merceod]

Not end to end: it's 437.5 GB with tp4/tp8 configs and this is a 2-GPU box. Here's what I ran instead.

The diff. With softmax off, which is the default and what Command A+ passes, _route_align_kernel computes the same sigmoid weights with the same rounding. The new BLOCK_K only differs from K when top-k isn't a power of two, and Command A+'s is 8, so the row mask, slot numbering and vector widths come out as they were. top_vals starts at -inf instead of 0, which only shows past K. language_model.py has no diff: same route_align arguments, same router_topk past 128 slots, and decode.py has no diff either. Two neighbours it could feel: _nearest_tuned picks the same entry as the old min for every M up to 20,000 on all five tables, GB200's tp4 table included, and fused_experts reads the same get_config values later.

The kernel. Moved vs main's copy on a B300, one process, alternated on the same logits, at the HF config's sizes (128 experts, top-k 8, block 16 at tp4 and tp8, and in GB200's table). µs per call, tp4 (tp8 is within 1 µs):

tokens	slots	pre-move	moved	past the launch, pre-move	moved	outputs equal
1	8	42.7	43.7	8.3	8.3	1020/1020
8	64	45.3	46.6	10.3	10.6	1020/1020
16	128	49.7	50.9	15.1	15.0	1020/1020
32	256	74.7	76.8	39.8	40.7	1020/1020
64	512	339.2	347.0	304.3	311.0	1020/1020

The extra µs inside the cap is the Python launch, two more kernel arguments, which a captured graph doesn't pay. 32 and 64 tokens are past the cap with the assert lifted, for the comparison only: the call site takes router_topk there, on both sides.

The tests. mstar/model/command_a_plus/tests on CUDA: 46 passed, 4 skipped (HF and metadata gated). I added one that the re-export is the moved kernel and stays a sigmoid route by default.

If you want the end-to-end number, command_a_plus_tp4.yaml on your node, this branch vs main, is the run.

[merceod]

Thanks, taking all three.

1. The in-graph embed and tables port to thinker_decode as well; I'll measure the audio-output cells with it. On one graph for both walks: DeviceLoopback needs every output of its walk to stay on the device, and thinker_decode ships the layer-0 row to the Talker each step, so one graph would mean every text request also emits and ships those states. I'll put a number on that cost and answer with it.
2. Right, both Omni routers set fused_softmax_topk since Performance 2/7: speech and language models #859, so the baseline is the fused softmax + top-k kernel plus the two alignment kernels. I'll fold route_aligned under that flag for batches up to ROUTE_ALIGN_MAX_SLOTS with the fused kernel above it, and redo the accounting from a per-kernel trace against that baseline.
3. Command A+ is 437 GB of bf16 weights with tp4/tp8 configs, which the 2-GPU box here can't serve. What I can give you from here is the moved kernel measured at its router sizes against the pre-move copy, plus its kernel tests; if you can run the end-to-end on your node, that would settle it.

[NSagan271]

@merceod I set off a run for command A+, though it's on H200 so it won't exercise the new Blackwell block sizes and such. Also a follow-up (I can take) that would be complementary with this PR: if the Engine can have some outputs for a graph walk loop back via on-device buffers and others be streamed / relayed through SHM, then we would theoretically not need a separate graph walk for text decode.

[merceod]

Thanks for setting off the Command A+ run; the B300 tiles only load on a matching GPU name, so H200 is the right check for the move itself.

On the follow-up: agreed, that's the cleaner end state. The split exists only because DeviceLoopback wants the walk's whole output set on the device while thinker_decode relays the layer-0 row to the Talker through SHM. With per-output routing in the engine (some outputs loop back on the device, others relay), one walk would serve both, and for a decode row the states are just the layer-0 row, so text requests would pay close to nothing for computing them. The run on pod A now measures what the states cost when they are shipped, which bounds that from above. Glad for you to take the engine side; this PR can drop the extra walk once it's there.

[merceod]

All three are addressed at f4cbbeb0b, with main merged in (#1110, #1129). thinker_decode takes ids and builds its tables in the graph now, route_aligned is folded under fused_softmax_topk as TopKRouter.route, and Command A+ has a diff statement, a kernel microbenchmark and its tests. I couldn't run it end to end on this box, so that one's yours if you want it.

Nothing moved tokens: 1,778 Thinker decode steps are bit-identical before vs after, T2T 32/32, A2T 32/32. Against 53c47cb8c on one B300, one session each: T2T +18.3 / +10.4 / +7.7% at c=1/8/32, A2T +15.8 / +12.8% at c=1/16, T2S +6.9 / +3.7% at c=1/16, A2S +10.6 / +4.6% at c=4/16, 0 failed. Description's updated.

One open item: one server of 8 at this head stalled for good on the mixed audio set at c=16 staggered, 0% CPU, in a runtime compile of the Talker's 1,017-token prefill row (torch_dynamo_resume_in_fused_experts_at_128). The 7 others and 6 on main, some on an empty Inductor cache, ran clean, and I had no stack on the one that stalled. So I can't say whether it's this PR: fused_experts' get_config read moved here, and main's 0 of 6 doesn't rule it out. If you see one, py-spy dump is what I'm missing.

[merceod]

On the A2T-after-I2T question from the live test (#1111 (comment)):

Ran the I2T-then-A2T check on one B300 with the 1-GPU colocated config. The branch had moved, so this is the head f4cbbeb0b (cea7d03 plus the thinker_decode rope commit, the one-flag fused route and a main merge) against main 53c47cb8c. Four fresh servers in the order main, PR, PR, main. The first two run I2T c=16 three times and then A2T c=16 three times, and the last two run A2T c=16 three times as their first workload. Each pass is 5 warmup waves and then 200 measured requests (LibriSpeech test-clean clips, or 200 image questions), greedy, closed loop. ITL is every gap between streamed tokens in the pass, pooled.

server, order	cell	pass 1 / 2 / 3	TTFT p50 ms	ITL p50 ms	ITL p99 ms	gaps >150 ms
main, I2T then A2T	I2T tok/s	1487 / 1494 / 1500	84 / 85 / 82	8.24 / 8.21 / 8.22	27.2 / 27.6 / 26.8	0 / 3 / 3
	A2T RTFx	206.6 / 219.3 / 222.3	102 / 98 / 101	14.19 / 13.70 / 13.06	113.4 / 45.0 / 43.2	1 / 7 / 3
PR, I2T then A2T	I2T tok/s	1604 / 1599 / 1604	73 / 73 / 73	7.67 / 7.67 / 7.68	23.7 / 23.7 / 23.7	0 / 0 / 0
	A2T RTFx	236.5 / 234.2 / 257.2	86 / 84 / 85	11.72 / 11.70 / 11.49	110.5 / 39.9 / 33.0	6 / 2 / 4
PR, A2T first	A2T RTFx	224.0 / 241.1 / 241.5	89 / 90 / 89	11.39 / 10.96 / 10.75	116.1 / 117.3 / 38.7	17 / 3 / 11
main, A2T first	A2T RTFx	207.7 / 224.3 / 234.4	103 / 94 / 91	12.24 / 12.71 / 13.01	113.9 / 40.6 / 38.2	2 / 8 / 11

No order effect on either arm. On the PR, A2T after I2T averages 242.6 RTFx against 235.5 when it runs first. Main reads 216.1 against 222.1, and both differences are well inside the 10-13% spread between passes on a single server. The PR leads main in both orders, +12% after I2T and +6% with A2T first. The long gaps show up on both arms: across the three A2T passes, a server logged 11 or 21 gaps over 150 ms on main and 12 or 31 on the PR. Every server's first A2T pass has a p99 around 110-116 ms, whichever arm or order, so that tail goes with the first pass on a server and not with what ran before it. There were no recompiles after startup and no cgroup throttling during any pass.

So I can't reproduce the -9.5% here, and I'd put round 1 down to the shared node. Caveats: this is one server per arm and order, not three, and it's B300, where the tuned tiles load, so it says nothing about the H200 path without them.

[merceod]

Merged main b06ae391f into the branch as f7fceb005 (a merge commit, no rebase) and pushed d117982ee. #970 and #1140 had rewritten the Thinker's decode in the same two files, components/thinker.py and submodules.py. Where both did the same thing I kept main's: the fused add + norm (keep_states now drops main's layer-0 view and layer-N add), the decode row as its token alone, and the one-call embed of a mixed step's decode rows, which now also takes loop-back rows. The text walk, the plan-cache fix, the in-graph tables and the B300 tiles stay as they were. Two test fixes ride along: the text-decode test's runtime check moved to the Rust runtime (#1130), and main's new tile-table trace test resets this branch's table cache.

One thing merged clean but needed a call. #1140's per-pair GEMVs only run when fused_experts isn't handed an alignment, and route_align hands it one, so the merge alone switched main's path off for 1, 2 and 4-row decode steps. I kept main's there (1efd5666b: the router asks the same check fused_experts makes), so route_align now serves 8 and 16 rows. On B300 that's a trade. Two sessions each, the merge without that commit vs with it: T2T c=4 762/767 vs 712/711 tok/s (+7.4% for the fused route), T2T c=1 311/311 vs 315/318 (-1.9%), A2T c=1 -2.6%. The fused route at 1 to 4 rows is also what moved greedy tokens against main before: without the commit it's 5-6/32 vs main, with it 20-32/32, same as main vs main. It's a one-line flip if you'd rather take the c=4 number on B300. I didn't, since the per-pair tiles were tuned on H100, where this PR's tiles don't load.

What each part gives now, on one B300 with the same harness and pins, from the PR's old base 53c47cb8c:

cell	main 53c47cb8c	main b06ae391f (3 sessions)	this (2 sessions)	main moved	this on top
T2T c=1 tok/s	241	282, 311, 312	318, 317	+29%	+2.0%
T2T c=16 tok/s	1628	1719, 1799, 1793	1789, 1790	+10%	-0.3%
A2T c=16 RTFx	205	249, 247, 249	277, 270	+21%	+10.4%

"Main moved" is against main's two matching sessions, and so is "this on top": main's first session was 9% under the other two at T2T c=1 and decoded different greedy tokens from them, so it looks like a different compile on the same cache (a guess). The rest of the gate, main vs this, two sessions each alternated: T2T c=64 2955/2955 vs 3079/3077 tok/s (+4.2%), A2T c=1 58.3/58.2 vs 60.3/60.2 RTFx (+3.4%), T2S c=1 17.1/16.5 vs 17.5/17.4 audio-s/s (+3.8%), A2S c=1 14.8/14.4 vs 15.3/15.1 (+4.0%). 0 failed anywhere. So #970 and #1140 took most of what this did for c=1 text, and what's left of this PR is mostly the plan cache and the 8 and 16-row route at batch, which A2T c=16 still shows. That split is my reading of the code paths, not something I profiled.

Greedy at c=1: this matches main on 7,344 of 7,344 T2T decode steps against one main server, and on 6,196 to 6,537 against the other two, the same range main's own servers give each other. A2T is 200/200 clips on every pair. The two decode walks give the same text on every server. WAV bytes don't repeat across server processes on either arm.

Stalls: 0 of 14 servers at the merge and around it (5 main, 8 this branch including the pre-fix merge, 1 old main). The runtime compile of the Talker's long prefill row in fused_experts, where the old head stalled once, happens on main too, once per server, only on long prompts answered in speech, takes 4–10 s, and finished every time.

CPU Core fails the same ids as main (one more, test_eval_compare_multigpu::test_a_head_that_hangs_is_cut_at_the_cap_with_stacks_and_fails, fails alone on main too). The GPU modules fail only test_fused_experts_matches_naive_and_precomputed_alignment[1] and [3], which fail on main since #1140 (per pair without an alignment vs grouped with one). The description's tables are replaced.

[merceod]

Merged main f609cac54 (#986's kernel registry) into the branch as ad1b4e8cf (a merge commit, no rebase) and pushed it. Three files conflicted:

• fused_moe/kernels.py and test_fused_moe_config.py are now main's. This branch's only edits there were its own table cache (_tuned_tables) and nearest-size search (_nearest_tuned), and the registry's constant _tile_rows and traced row search replace both. The tests reset through registry.reset_lookups() as Add a kernel registry with per-GPU tuned MoE tiles #986's do.
• fused_moe/runner.py keeps main's one get_config(..., dtype, device, tuned) read before the alignment, so block_m = config["BLOCK_SIZE_M"] again. It also keeps runs_per_pair, which means no route alignment where the experts run per pair. I dropped this branch's second read after the alignment, since the registry lookup doesn't break the graph anymore.
• Not a conflict, but needed: the B300 tiles were in the old fused_moe/configs/ directory, which nothing reads now. They're in mstar/kernels/tables/moe_gemm/NVIDIA_B300_SXM6_AC.json, converted with Add a kernel registry with per-GPU tuned MoE tiles #986's rule for the H200 tables, and they give the same tiles as the old lookup at every count from 0 to 4096. The two route tests on the old table API are gone; main's test_get_config_takes_the_row_up_to_m_and_down_overrides and test_compiled_fused_experts_takes_the_tables_tiles cover the same ground.

Tests:

• ruff check . is clean.
• GPU modules on a B300 at ad1b4e8cf: 1038 passed, 6 skipped, 2 failed. That covers fused MoE config/registry/decode/route/align, the router, the Omni modules (771 passed), loopback, inline tokens, stop policies and Command A+. The 2 failures are test_fused_experts_matches_naive_and_precomputed_alignment[1] and [3], the same ids main fails. test_compiled_fused_experts_is_one_graph passes, and so does the symbolic-batch trace on every shipped table, the B300 one included.
• CPU Core matches main's 18 failures by id. A 19th, test_eval_compare_multigpu::test_a_base_that_hangs_is_reported_broken_not_flagged, passes alone on both trees.

Served T2T on one B300, one session per arm (main f609cac54 first):

	c=1 tok/s	c=16 tok/s
main	311.7	1808.7
this	325.1	1840.3

That's +4.3% at c=1 and +1.7% at c=16, 0 failed. Neither server recompiled anything after startup, and no fused_experts resume frames showed up, since #986's constant lookup removed that graph break.

Main has moved again since (#978, #1159, #932). It merges clean, but #932 reworks the KV plan cache this PR's step_is_input_free fix feeds, and I haven't run the two together.