qwen3_omni: prefix caching for the Thinker by Gaurav-Shah05 · Pull Request #836 · mstar-project/mstar

Gaurav-Shah05 · 2026-10-22T18:35:15Z

This lets Qwen3-Omni reuse the KV of a chat's images, audio and earlier turns.

Follow-up on the prefix caching thread (#210, #401). Qwen3-Omni's Thinker now keys its prompt's text, images and audio, and the replies it generates. A resent image or clip has 1.15x to 2.01x lower TTFT. Requests that ask for audio output stay uncached. The Talker reads one Thinker hidden row per prompt token, and the cache does not keep those rows.

On main, a resent 480x640 image with a new question runs the Thinker's whole prefill again, at 123 ms TTFT. With this PR it takes 61 ms.

Changes

Ten changes:

A submodule can override split_inputs to slice its walk's tensor_inputs, so on a hit the engine runs only that walk's uncached tail.
The load check lets a keyed node stream into another partition, and the conductor turns the cache off for a request that runs that partition.
Like Qwen3.5's, the Thinker's text rows take positions from the position resource's plan, so a cut text walk lands at its keyed position.
The Thinker's split_inputs cuts the Talker masks and deepstack rows with the embeds.
When the cached prefix ends inside an audio clip, the position counter moves past the whole clip.
An item encoder that marks its own dynamic dimensions compiles with only those dimensions dynamic, and every other item encoder keeps dynamic=True.
Qwen3-Omni's audio encoder builds its chunking before the compiled call, so two graphs, split at 1 s, serve clips from 90 ms to 500 s.
Qwen3-Omni's vision encoder builds its grid tensors before the compiled call and marks their patch dimension, so one graph serves every image size.
Qwen3-Omni keys the Thinker's label over its prefill walks and its replies, and declares the checkpoint directory that the cache hashes into its keys.
For a prompt with a video, Qwen3-Omni keys only the leading text, since keying past it has never been run or tested.

Cost

Installing the qwen3_omni extra now needs transformers 5.18 or later, for transformers.vision_utils, and so does the all extra.
The vision encoder's grid tensors take 1.3 to 1.6 ms of host time per image, and its step is 1.2 ms faster than main's at 480x640.
The first image after a boot waits about 1 s longer for the vision encoder's compile (6.1 against 5.1 s).
A video passes its whole patch count as max_seqlen, so flash-attn launches more blocks per frame than it needs. That slows a video's encode by an amount not yet measured.
With the Thinker's cache on, the preprocess worker also keys requests with audio output, though the cache never matches or caches them.
Keying takes 0.2 ms for a 224x224 PNG and 4.6 ms for a 5.8 MB 1200x1600 PNG, mostly hashing the file.
A hit allocates only the prompt's tail, so the pool admits more requests than it can finish.
A request the pool cannot finish waits until the timeout.

Next (out of scope here)

These are open, and none is in this PR:

Caching the Thinker's hidden rows that the Talker reads would let a voice chat reuse its prompt too.
Caching encoder outputs by image or clip digest would let a hit skip the vision and audio encoders, which still run on every send.
Keys for a video's spans would let a resent video reuse its prefix past the leading text.
A per-frame max_seqlen for video would remove the extra flash-attn blocks listed under Cost.

Qwen3.5's prefix caching is #877.

Testing

With main merged in, CI's CPU selection fails the same tests as main. The GPU tests in test_qwen3_omni_prefix_cache.py pass on one H100 with torch 2.12.1+cu129.

On one H100, against main of Oct 22, qwen3omni_colocated (Thinker pool 256 pages), cache on, Thinker greedy, 16 output tokens, streaming, p50, medians over three boots per arm. Another job shared the node at a load average of 3 to 13 on 208 CPUs.

	`main`	This PR	Δ
Same image, new question, 224x224/480x640: TTFT	104.5/123.1 ms	57.1/61.3 ms	1.83x ↓/2.01x ↓
Same image, new question, 1200x1600: TTFT	263.0 ms	176.7 ms	1.49x ↓
Image chat, turn 2, 480x640/1200x1600: TTFT	122.6/271.9 ms	61.5/150.6 ms	1.99x ↓/1.81x ↓
Audio in, new question: TTFT	125.5 ms	109.4 ms	1.15x ↓
Text chat, turn 2, c1/c8: TTFT	41.8/41.5 ms	34.2/34.0 ms	1.22x ↓/1.22x ↓
Shared 2,084-token system prompt, resent alone: TTFT	120 ms	41 ms	2.93x ↓
First send, image 224x224/480x640: TTFT	104.9/120.2 ms	107.1/118.9 ms	1.02x ↑/1.01x ↓
First send, image 1200x1600: TTFT	270.2 ms	263.9 ms	1.02x ↓
First image request after a warm boot: TTFT	5.10 s	6.12 s	1.20x ↑
ITL, c1/c16	6.55/14.24 ms	6.53/14.33 ms	1.00x/1.01x ↑
Boot, warm cache	219.1 s	265.8 s	1.21x ↑

The boot gap follows the Thinker's weight load, which took 37 to 48 s on one main boot and two boots with this PR. The three boots that loaded it in 17 s took 218 to 219 s on both arms. This PR does not change how weights load.

On a later series, the first audio request after a warm boot waited for the encoder's compile on both arms. It took 4.6 to 5.5 s on main and 4.9 to 5.1 s with this PR. For audio first sends, the audio encoder's step with real weights takes 6.9 ms with this PR and 7.2 to 7.4 ms on main.

On main, the first image of a second size waits 6.7 to 7.1 s while the vision encoder compiles again. With this PR no image after the first waits for a compile.

Outputs with the cache on, against cache off:

Same-image resends, and resends whose cached prefix ends inside an image or an audio clip, give the same output, 10 of 10 on each boot.
Turn 2 of image and text chats, eight text chats at c8 included, hits the pages of turn 1's reply.
An audio-in resend hits past the clip and gives the same output.
A request with audio output is never matched or cached.
All of these still hold after merging main of Oct 24.

Under pool pressure, with 32 requests sharing the system prompt above at 1,024 output tokens each, every request times out at 600 s on both arms. Before that, each decodes 837-965 tokens with this PR, against at most one on main.

mstar-review · 2026-10-22T18:39:23Z

Automated review

🟡 No findings on 815b345, light review

Description

The description doesn't say which tests were run or what they returned (ruff, the CPU Core suite, the new test_qwen3_omni_prefix_cache.py and the other modular tests this PR changes), and AGENTS.md requires that for an AI-assisted PR.
The claim that image and audio resends give the same outputs as with the cache off gives no command, prompt set, sample count or torch/CUDA build, so a maintainer can't rerun it.
The benchmark table names the hardware and config but not the harness or the commands, flags and input files, so the numbers can't be reproduced.
It doesn't say that on main no Qwen3-Omni Thinker walk was probed at all, because every walk carried masks_for_talker in tensor_inputs, so text-only chats are newly cached too. The text-chat row of the table implies this but never says it.
It doesn't say that marks_dynamic_dims and the mark_dynamic calls in VisionEncoderSubmodule.prepare_inputs apply whether or not the cache is open, so cache-off deployments and video inputs also change how the vision encoder compiles. It also doesn't say whether video was checked, where cu_seqlens stays unmarked and max_seqlen is now the whole clip's patch count.

What the review checked

The conductor's cache-off check: partition_states[...].is_done is set from get_initial_forward_pass_args before the check reads it, and Qwen3-Omni's Talker reports request_done exactly when audio is not among the output modalities.
streamed_partitions and _refuse_unservable_walks agree. Both ignore only StreamingGraphEdges whose target partition differs from the walk's own.
_skip_cached_prefix: only ThinkerSubmodule overrides split_inputs, so other models' walks with tensor inputs are still never probed.
Text-row positions: RopeManager.plan builds them from the counters in KV plan row order, and padding rows come last. The counter moves in apply_cached_prefix before plan, and pre-plan waits on the previous step's commit, so _position_ids_3d's per-row offsets line up with inputs. Every capture input still carries custom_pos_ids.
The positions of a cut walk: a partly held image or clip keeps its counter at the item's start, and its custom_pos_ids are sliced by the base _split_pos_ids. mrope_pos_advance then moves the counter past the whole item, and a text tail moves it by its own length.
For a single image (t=1), the new image end_pos_base (start_pos + max(h,w) + 1) equals the old vision_pos_ids.max() + 1. The video branch is unchanged.
The audio span in the layout (_get_feat_extract_output_lengths + 2) matches the slots prefill_audio writes, and the position advance matches get_rope_index_audio plus the two sentinels.
declare_step's cg_key_info is still DEEPSTACK exactly when a prefill_vision row is present, matching cg_key_info().
There is no circular import: qwen3_omni_model imports submodules lazily, and the model registry imports model modules lazily, so the top-level transformers.vision_utils import only affects Qwen3-Omni. CI installs transformers without a version pin, and qwen-tts with --no-deps.
The split_inputs override cuts masks_for_talker and deepstack over the same [start, end) as the embeds.

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

dghosef · 2026-10-23T13:59:37Z

Removing the torch.compile increased my audio first send's TTFT from 74ms to 87ms on H200. Do we need bit exactness?

This could be something tunable if/when this pr makes it through
#817

or through environment variables

dghosef · 2026-10-26T21:50:54Z

Thanks. Rereviewing now

It has one entry per frame plus one, so a video's frame count changes its shape; unmarked, the first video recompiles the encoder. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Each video frame is its own attention segment, so the whole clip's patch count overstated it ~t times and flash-attn launched ~t times the blocks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…eligible Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dghosef

Waiting on some results, will approve after. Some other suggestions

dghosef · 2026-10-27T04:29:28Z

            pixel_values,
            grid_thw=grid_thw,
+            # an image is one segment of all its patches, and no video frame's segment is longer
+            max_seqlen=pixel_values.shape[0],


For a video each frame is its own segment, so this overstates max_seqlen by ~t and flash-attn launches ~t times more blocks than it needs. pixel_values.shape[0] // (cu_seqlens.shape[0] - 1) is exact (one item per call), stays symbolic, and is unchanged for images.

pushed a fix

dghosef · 2026-10-27T04:52:14Z

+        if inputs.kwargs or inputs.resource_step_info or (
+            inputs.tensor_inputs and type(submodule).split_inputs is ARNodeSubmodule.split_inputs
+        ):


deciding whether to cut by checking if split_inputs was overridden is fragile. Can you add a flag instead

dghosef · 2026-10-27T04:54:01Z

 from mstar.utils.logging_config import quiet_noisy_loggers
 from mstar.utils.orphan import exit_when_orphaned
 from mstar.utils.profiler import nvtx_enabled, range_pop, range_push
+from mstar.worker.engine_manager import streamed_partitions


streamed_partitions isn't needed for this - Qwen3-Omni's process_prompt can return {"prefix_cache": False} when "audio" in output_modalities and then we can avoid the engine import

I kept the rule in the conductor, since it comes from the graph. If each model's process_prompt had to opt out, a model that forgot would serve hits that leave the partition it streams into short of rows. The worker import is gone though, since streamed_partitions now sits next to the load check in mstar/validation/resources.py.

dghosef · 2026-10-27T05:25:57Z

            **extra_inputs
        }

+    def _position_ids_3d(self, inputs: list[ARNodeInputs]) -> torch.Tensor:


_position_ids_3d is almost identical to qwen3.5's. Factor them out?

process_prompt already has the switch, so the conductor no longer works out which partitions a keyed node streams into.

…positions

… streams into The rule comes from the graph, so a model need not opt out in process_prompt, and it lives beside the load check that lets such a stream load.

dghosef · 2026-10-27T19:29:59Z

0 errors, no TP divergence

TTFT p50 (ms), c1 streaming, qwen3omni_colocated, 3 restarts per arm:

Request	main	PR	speedup
img224 resend	89	48	1.85x
img480 resend	139	94	1.48x
1600x1200 image resend	505	440	1.15x
audio resend	80	46	1.75x
text turn 2	33	27	1.23x
img480 first send	136-145	136-145	same
first image after boot	2.50-2.71 s	5.47-5.64 s	+3 s

ITL p50: c1 4.40-4.44 ms on both. c16 9.78-9.84 ms on main vs 9.53-9.63 ms on the PR.

GPU-bound on both arms (event_sync ~4.2-4.8 ms of a ~4.4 ms step at c1). prepare_inputs goes from 0.18 to 0.28 ms/step since the encoders build their grid inputs on the host.

Correctness:

Partial hits, including ones that end inside the image/clip, match the cold reference 7/7 on every layout.
Speech-out requests stay uncached (0 query / 0 hit tokens), and text-only resends on the same server still hit.
TP2 with prefix_cache_tp: 12,160 of 21,980 tokens hit. PD: 13,245 of 21,990.
Over-admission check (2k shared preamble + image, 24 requests at c16): 0/24 failed on the PR. Main and PR with the cache off time out 16/24 - that's main's existing KV backpressure deadlock, and the cache avoids it here.
After the first video, the vision encoder recompiles (cu_seqlens size 2->4), and an image sent cold then gets a different reply than before the video (2/2 rounds). Marking cu_seqlens dynamic fixes it.

Issues:

The first image after boot is +3 s, not the +1 s in the description, on every round. Rerunning with your latest push to see if this fixed it.
A full cached resend sometimes rewords the reply vs the same prompt sent cold: 1/7 colocated, 2-3/7 under TP2. Main never differs, and cold vs cold is 7/7. My guess is the tail being prefilled separately from its cached prefix. Probably just rounding, not blocking.

Vision encoder max_seqlen sized to one frame: bit-identical outpster eager / 3.5% compiled (flash-attn time -13 to -19%), images unchanged.

Gaurav-Shah05 · 2026-10-27T20:50:12Z

Addressed the rest, and merged main in.

First image: the +1 s was on an H100 (6.1 against 5.1 s).

Rewording: a wrong page or position would break the partial hits too, and those match 7/7, so it's rounding.

dghosef

lgtm

Main brought (since 6cfe39c): #892 Parakeet and Nemotron ASR with the transducer as a graph loop (and the shared mrope_position_ids helper), #917 pytest 9.2 flag proxies, #897 one captured all-to-all for Ulysses, #915 NCCL graph registration off by default, #836 Qwen3-Omni Thinker prefix caching (ARNodeSubmodule.cuts_tensor_inputs, marks_dynamic_dims, the conductor's _streamed_partitions), #909 bench cells, #900 Command A+ batched image tiles. Conflict: - mstar/model/submodule_base.py: main's ARNodeSubmodule.cuts_tensor_inputs beside the stack's token_stop and its stop checks. Independent class attributes; both kept. Ported onto the stack's mechanisms: - mstar/model/parakeet/submodules.py: check_stop_batched drops the step_walk parameter (the stack's three-argument signature) and reads its "active" flags through HostRows.first_values instead of its own row-map loop. Its stop is a flag, not a token, so it has no TokenStop. - test/modular/test_transducer_loop.py: the call without step_walk. The conductor's _streamed_partitions (#836) sits beside _pick_replicas / _register_request; it reads only partition_states and _resource_specs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Brings #858's merge of main: #892 Parakeet / Nemotron ASR, #917, #897 Ulysses all-to-all, #915, #836 Qwen3-Omni Thinker prefix caching, #909, #900. Conflicts: - docs/adding_models.rst: main's marks_dynamic_dims sentence and #859's compile_regions paragraph are independent facts; both kept, in one paragraph. - mstar/model/qwen3_omni/submodules.py (3 hunks, plus auto-merged hunks in the same class): AudioEncoderSubmodule. Two mechanisms for one cause (the HF encoder recompiling per clip length): #859 runs it eagerly (disable_torch_compile) and packs a step's clips into calls of at most MAX_BATCH_FRAMES; #836 builds the encoder's chunk tensors in prepare_inputs with mark_dynamic so the per-clip forward compiles once (marks_dynamic_dims, _max_seqlen). One mechanism: #859's class verbatim. Under disable_torch_compile the engine never compiles the encoder, so #836's concern (a recompiled graph rounding differently from the one whose output the prefix cache holds) cannot arise. The auto-merge had taken #836's kwargs-only forward beside #859's forward_batched, which calls forward with the clips positionally; that keep-both is undone. Main's VisionEncoderSubmodule (marks_dynamic_dims, prepared indices) and Thinker changes (cuts_tensor_inputs, split_inputs, mrope_position_ids, image_span) are kept; #859 did not change those. Imports: both cp_embed_accumulate and mrope_position_ids; the audio helpers only #836's prepare_inputs used are dropped. Tests following from it (test/modular/test_qwen3_omni_prefix_cache.py): - the test that #836's prepared chunk tensors embed a clip as the encoder does alone is deleted with that code; - the compile test now asserts the audio encoder stays eager while its output is cached; - the CUDA no-recompile test keeps its vision case only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…_end_token_id #859's ThinkerSubmodule declares its TokenStop in __init__ from config.im_end_token_id; main's #836 fixture config (merged in 5dada82) had no such field, so the three tests built on it errored at setup. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…igin/main 73c6c17) into merge/860-main Brings main's mstar-project#892 Parakeet / Nemotron ASR, mstar-project#917, mstar-project#897 Ulysses as one captured all-to-all (prefer_all_gather and the all-gather exchange gone, Cosmos3 captures every tier under TP x SP), mstar-project#915, mstar-project#836 Qwen3-Omni Thinker prefix caching, mstar-project#909 and mstar-project#900, with mstar-project#858's and mstar-project#859's resolutions. No conflicts. Files both sides touched (docs/clients.rst, docs/installation.rst, docs/models.rst, sequence_parallel.py, cosmos3/submodules.py, test_accelerator_graph_capture.py, test_flux2_klein_model.py): each merged change equals mstar-project#859's own delta since 3166401, so nothing kept both sides. mstar-project#860's my_rows / gather_rows (Wan2.2's CFG split) sit beside main's new ulysses_attention and use CommGroup.all_gather, which mstar-project#897 left unchanged; nothing in mstar-project#860 passed prefer_all_gather. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Qwen3-Omni: each Thinker decode row emits a Talker mask of its own (R1.1, R2.1) postprocess gave every decode row one cached zeros(1) as its thinker_mask. The store mints a uuid per output and the RDMA/TCP transports register each uuid's address for remote reads, so every row after the first registered the same address again, and the first row's cleanup unregistered an address the others still used. The batched decode now emits the mask as a row output (a row view of the step's clone), a mixed step's unpack gives each decode row a row view of one step tensor, and the eager forward a fresh one per step. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: docstrings say what the audio encoder, Thinker and Talker forwards emit now (R1.13) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Images: read_image refuses an undecodable file, so Qwen3.5's load_image is not a copy of the base one (R1.13, R1.18) The corrupt-file ValueError moves from both load_image bodies into read_image, which every image loader calls. Qwen3.5's override is the decode without the float conversion, and _as_hwc_uint8 drops its float and RGBA branches: load_image always sends uint8 RGB. The video test's image is uint8, as load_image sends it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: the vision grids come from the step metadata alone (R1.19) process_prompt always records each image's and video's host grid, and the schedule always carries them, so the device-grid fallback in _host_grids and the schedule's all-or-nothing guard never ran in serving. Both are deleted; the video tests pass the host grids the way the conductor does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: the API server's GIL switch interval is a model default, not six config lines (R1.26, R2.8) Model.API_SWITCH_INTERVAL_S is the default for api_switch_interval_s, as PREPROCESS_TORCH_THREADS is for preprocess_torch_threads; Qwen3.5 sets 0.001 and its configs drop the key (qwen3_8_27b, the same model class, now gets it too). A config still overrides it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: capture_cudnn_benchmark declares cuDNN's algorithm search for a node's capture warm-up (R1.8) Kokoro set cudnn.benchmark inside its frames capture function, which also runs inside the capture itself and undid the capture's benchmark=False guard (a search there can empty the allocator mid-capture). A node now declares capture_cudnn_benchmark and the engine sets the flag around that node's warm-up and capture only; the capture still turns it off, and eager steps keep the process setting. Kokoro's own scope is deleted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kokoro: one LSTM path and no A/B switches (R1.10) fuse_lstm_directions and decoder_cudnn_benchmark were left-in A/B switches; both are deleted (the latter is the capture_cudnn_benchmark declaration now), and so is the unfused LSTM path: MaskedBiLSTM runs the fused block-diagonal LSTM, built once its weights load, in eval mode with no gradients (it was built after eval() and stayed in training mode). Tests that build KokoroTTS fuse it the way the loader does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice: one backbone path; backbone_graphs gates capture only (R1.7, R1.23, R1.24, R2.7) With backbone_graphs off the backbone declared no attention resource and ran the reference's packed forward through its module-global _CTX (under a lock); on, it ran mstar's layer loop over the same fused weights. The resource is now always declared and the flag only drops the captures, so both settings run encode + head. OmniVoiceBackbone.forward, its _ctx_lock and three of the four required private names go; the reference packed forward lives on in the parity tests as their yardstick (_packed_logits), with the encode-vs- reference GPU test moved beside it. The boot check now also covers the model members encode and head read (_prepare_embed_inputs, audio_heads), and both boot checks stop the worker with SystemExit. build_packed_canvas drops its unused audio_mask_id. The loader drops the reference's packed runner after patching, since nothing in mstar reads it: its 128 MB workspace was dead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Higgs Audio: chunk lengths ride in the step metadata; encoder_graphs gates capture only (R1.11, R1.28) The encoder's prepare_inputs read each chunk's mel length back from the audio_feature_lens tensor with .tolist(). process_prompt already has them on the host, so it records each chunk's encoder output length as an internal kwarg (as Qwen3.5 does its grids) and the prefill step's metadata carries it; the audio_feature_lens tensor and edge are gone. With encoder_graphs off, short clips were also encoded unpadded, a different batching rule and numerics from the captured path. They are now padded and masked either way; the flag only drops the captures. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: the audio encoder sizes its steps from host lengths (R1.11, R1.12, R1.20, R1.21, R2.2) forward_batched read each clip's embedding count back from the device with .tolist() after queuing the encoder, which waited for the whole encoder. A clip's mel frames are its tensor's width, so prepare_inputs derives the frames (input_seq_len) and the embedding count on the host. The batched encoder had no cap either, so a burst of audio requests became one call: a step now takes at most 16 clips (max_batch_size), and forward_batched packs them greedily into calls of at most 30,000 mel frames, since a clip may run 30 min and a call's activations grow with its frames. Refusing an over-budget step in can_batch instead would send it through the engine's per-request path, one eager call per clip. The budget sits in the forward until the scheduler can budget a step's tokens. docs/models.rst notes the batching and that a clip's embeddings depend on what it is batched beside. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: the talker's stop buffer shares its row output's clone (R1.22) The row output was named new_token and the stop buffer layer0_codes, so the engine, which reuses a row clone for a stop buffer of the same name, cloned the step's codes twice. The talker now emits layer0_codes under that name from the start, and postprocess drops its rename. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish S2-Pro, Omnilingual: full fp32 through the engine's precision declarations (R1.25) Fish's full_fp32() and Omnilingual's _full_fp32_matmuls flipped the process-wide TF32 flags inside model code. The codec (with codec_exact_fp32) and the reference encoder now declare float32_matmul_precision='highest' and cudnn_allow_tf32=False, and an fp32 Omnilingual deployment declares the matmul precision; the engine applies them around the forward. Both helpers are deleted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: TokenStop.loopback_signal declares a decode walk's device loop-back; the hooks follow from it (R1.5, R2.3) Higgs copied Qwen3.5's six loop-back hooks verbatim (reads_device_loopback, uniform_row_inputs, _device_loopback, device_loopback_signals, step_is_input_free, inline_client_signals). The walk, the token buffer and the sampler are already the decode loop's TokenStop, so the declaration is one more field on it: loopback_signal, the input signal the walk's token feeds back to. ARNodeSubmodule derives those hooks from token_stop, and the engine keeps the slot masters of token_stop's sampler for a node that sets it. NodeSubmodule.reads_device_loopback is deleted, and both models' prepare_inputs take the loop-back row from the base. Documented in docs/adding_models.rst. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: a vision batch that overflows one replay runs request by request (R1.17, R1.20, R1.21, R2.4) can_batch always packed up to 8 requests, so a step past the top patch bucket or the attention plan's segment ceiling ran the block loop eagerly for all of them, though each would have replayed alone. can_batch now refuses such a batch. docs/models.rst notes the 8-request steps and that an image's embeddings depend on what it is batched beside. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: compile_regions declares the parts of a node the engine compiles on their own (R1.6, R2.5) Qwen3.5's vision encoder compiled its embed and merger lazily on the first request, and did so with no accelerator too; Omnilingual compiled its feature extractor at construction and warmed it from the model loader; Kokoro compiled its frame half in __init__. A submodule now returns those regions from compile_regions(), with optional example arguments, and the engine compiles them (dynamic, under the platform's inductor settings) and warms them at the start of warmup, before capture, on an accelerator only. The three models' own compiles, warm_compile, and the feature-extractor stand-in argument of the wav2vec2 frontend are deleted. Documented in docs/adding_models.rst. Regions compile in torch.compile's default mode, as a whole-forward compile does (the compile level tunes captured forwards), and a node that declares regions with its forward compiled whole fails start-up with a ValueError. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Codecs: one shared Snake kernel, on libdevice's sin (R1.14) Fish S2-Pro's DAC carried a second copy of Qwen3-TTS's Triton Snake kernel. Both now call mstar.model.components.snake, which computes in fp32 for any input dtype and takes sin from libdevice (CUDA's sinf, as torch.sin does) rather than tl.sin, which may lower to the approximate hardware sine and would make Fish's codec_exact_fp32 decode inexact. The Qwen3-TTS GPU test adds large activations, where an approximate sine drifts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictors: Qwen3-Omni and Qwen3-TTS share one set of depth-loop kernels (R1.4, R2.9) Qwen3-Omni's fused code-predictor kernels lived in its own package while Qwen3-TTS's depth loop ran the unfused ops (module q/k norms, a RoPE kernel, two strided cache copies, decode_attn_nhd), and fused_qk_norm_rope in mstar/utils/attention.py had no production caller. The kernels move to mstar/model/components/code_predictor_kernels.py and both depth loops call cp_qk_norm_rope_kv_write and cp_attention. decode_attn_nhd and fused_qk_norm_rope are deleted with test_fused_rmsnorm.py; the kernel tests (now test_code_predictor_kernels.py) check the prologue against fp32 math, and the Qwen3-TTS CPU tests stand in for both kernels. A GPU case runs Qwen3-TTS's depth loop at its real head counts against its layers unfused in fp32 (not yet run on a GPU; its reference math checked on CPU with torch stand-ins for the kernels, where a 100x RoPE theta fails it). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Residual streams start in forward_residual, from a tensor the loop owns (R1.16, R2.6) Every fused-residual layer loop in this PR (the Qwen3-TTS Talker and code predictor, Qwen3-Omni's code predictor) hand-unrolled the first layer, and the two Qwen3-TTS loops started the residual stream at their input, which FlashInfer's fused add + norm then updated in place: the Talker's input_embeds is a packed replay's static input, whose padding rows nothing refreshes, and the 0.6B code predictor's input is the caller's (its projection is the identity). RMSNorm.forward_residual now takes residual None to start a stream, so the loops have no first-layer branch, and each starts from a copy of its input. Starting a stream under TP, from a row-parallel partial sum the start would not reduce, raises ValueError. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: a forward reads its step's piecewise leases off engine_inputs (R1.15) The packed encoders decided whether to replay their lease_before_step region two ways: Qwen3.5's vision encoder inferred the lease from an empty step declaration plus a fitting bucket, Omnilingual and OmniVoice from the bucket alone (correct only because their can_lease mirrors their own eager/packed choice). ModelInputsFromEngine now carries piecewise_leases, the slots this step leased, and all three replay exactly when theirs is there. The ragged config's comment says what max_segments_per_request bounds under one_row_per_step; docs/adding_models.rst shows the new check. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kokoro: the fused LSTM's docstring claims fewer steps and rounding-level agreement, not bit-identity (R1.9) The fusion halves the number of sequential steps; the speedup that implies is measured elsewhere. Its test checks allclose(atol=1e-6), so the docstring says the result matches to float rounding. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStop: check_stop reads a row's first element, as check_stop_batched does merceod, #858 mstar/model/submodule_base.py:1036: check_stop took .item() of the row where check_stop_batched reads [:, 0], so the two disagreed (the per-request one raising) on a stop buffer wider than one element per row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Waypoint harness: serve_rollout's parsers strip the internal id's #<n> suffix too merceod, #858 benchmark/waypoint/benchmark_streaming.py:565: the benchmark's _dit_step_timestamps stripped the suffix, but serve_rollout._dit_schedule, which runs before it, still matched ids exactly, so rows_per_step_histogram came out empty. #860's parser hunk (schedule, overshoot veto, cleanup) moves here beside the benchmark half. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Warmup: repeat and concurrency share one positive-integer check merceod, #858 mstar/api_server/entrypoint.py:1377: the two checks spelled the same predicate twice; one _positive_int keeps them in step. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Warmup: one id form, warmup-<spec>.<level>.<copy>, and the first failure re-raised as is merceod, #858 mstar/api_server/entrypoint.py:1468: three separators (warmup-0, warmup-0-1, warmup-0.1, warmup-0.1-1) depending on level and copy, and an error list built only to raise its first entry. The gather still waits for every copy (return_exceptions) so none is cancelled mid-request. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * API server: api_switch_interval_s must be a positive number, checked before the model is built merceod, #858 mstar/api_server/entrypoint.py:1993: float(interval) took yaml's true as 1 s (logged "True s"), and 0 or "1ms" raised only after the API-server model was built. Validated first now; bool, non-numbers and values <= 0 are a ValueError. Whether the knob should exist (#849) stays the user's call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs/serving.rst: the warmup copy count is the largest node group's replica count merceod, #858 mstar/api_server/entrypoint.py:1425: the docs didn't say that warmup multiplies every entry by the most replicas any node group serves, or that a group with fewer replicas then sees the copies together and may leave its one-request shape cold. A per-replica warmup from the conductor stays the recorded follow-up (P3-4). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ragged attention: the mstar::ragged_attention op body calls the manager's _run merceod, #858 mstar/engine/resources/attn/ragged/flashinfer.py:46: the op repeated _run's wrapper lookup instead of calling it, so wrapper selection lived in two places. (compiler.disable on _run is inert inside an op body.) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Resources: every builder shards heads by EngineResourceInfo.head_shards, set by the engine merceod, #858 mstar/engine/resources/attn/ragged/base.py:31: shard_heads reached only the ragged family; the other six builders (paged, ring and bounded KV, attention, cross-attention, recurrent state, linear attention) divided heads by joint_comm_group.world_size, which can't tell CFG from SP. A model that puts guidance branches on the sp slot and also binds one of those would get its heads halved and fail at append/plan. The engine now sets head_shards once per resource: the lockstep instance's size, or 1 for a node that declares splits_heads = False (a NodeSubmodule attribute, next to the other per-node declarations). RaggedAttentionConfig .shard_heads is deleted. KVManager shards in build, like the others, rather than in __init__; bounded KV refuses head_shards > 1 rather than any multi-rank group. Callers that build EngineResourceInfo by hand over a real TP group (Command A+ reference test, its checkpoint validator, the GPU ragged test) pass head_shards. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * wav2vec2: the frame lengths go up through to_device_nonblocking merceod, #858 mstar/utils/h2d.py:21: wav2vec2's padded-batch mask kept its own pin-then-copy instead of the shared helper. (The sampler's copy stays: it pins straight from a Python list, XPU-aware, in one host allocation.) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Higgs Audio: a short clip is masked and sliced by the one frame count its prompt uses merceod, #859 mstar/model/higgs_audio/submodules.py:154: valid_frames came from the mel width (floor(n/160)) while the slice lengths came from the feature extractor's attention-mask sum (ceil), so whenever the mask sum is one past the width the tower masked a frame the slice then read: at a width of 2 mod 8 that is a whole extra embedding (the stride-2 projector over the last valid row and zeros) the unpadded tower never produces. The step metadata now carries each chunk's mask sum (AUDIO_FEATURE_LENS_KEY), and both the mask and the slice count from it. Which count the HF reference uses was not checked. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: a Thinker decode step emits its layer-0 row, so the row-output store keeps its fast path merceod, #859 mstar/model/qwen3_omni/submodules.py:1128: decode emitted the 2 x hidden concat as a row output and postprocess then rebound thinker_states to a fresh layer-0 slice, which fails store_row_outputs_batch's identity check, so every speech decode step took the per-tensor store. Decode now emits layer_0_embed itself (batched, unbatched and unpacked from a mixed step) and postprocess leaves decode rows alone; the per-step clone halves. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: drop the Talker's unread _tts_eos_embed_cached merceod, #859 mstar/model/qwen3_omni/submodules.py:1254: decode projects _tts_eos_raw inside the captured forward, so the pre-projected EOS embedding had no reader; the field and its projection go. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Higgs Audio: one prefill step built directly, no three-walk scheduler merceod, #859 mstar/model/higgs_audio/higgs_audio_model.py:274: the scheduler could only produce one prefill_audio step, but its routing by prefix, the prefill_schedule / prefill_step kwargs, the empty-schedule raise, the two is_last_prefill derivations and a never-true next-step branch were still there. get_initial_forward_pass_args builds the three edges itself with is_last_prefill True, and the partition step goes straight to decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: the patch grids reach the nodes only as host ints; the unread image_grid_thw edge goes merceod, #859 mstar/model/qwen3_5/qwen3_5_model.py:768: since the grids ride in step metadata, no submodule read the device image_grid_thw edge, yet it was still shipped to both prefill_vision nodes, picked per step and checked for length; and IMAGE_GRIDS_KEY and STEP_GRIDS_KEY were one string under two names. process_prompt keeps the grid tensors local (the layout and the host grids come from them), the walk inputs and picker drop the edge, and the step metadata key gets its own string (step_grids: the picked image and video grids in plan order, not the request's image grids). The TP read-failure test fails the vision walk's vision_embeds read instead. Not run: the interleave tests (need a Qwen3.5 checkpoint). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: drop Qwen3TTSTalkerLayer.forward, which nothing calls merceod, #859 mstar/model/qwen3_tts/components/talker.py:138: the language model runs the fused residual chain over its layers' parts itself, so the layer's own unfused forward had no caller left. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * MoE: the blocks take a fused router through router=, not a flag of their own merceod, #858 mstar/model/components/moe.py:238: fused_softmax_topk was a parameter of four block constructors only to reach the default TopKRouter, while every block already takes router=. A model opts in by passing router=TopKRouter(..., fused_softmax_topk=True). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Components: one length_mask, in model/components/masking.py merceod, #859 mstar/model/higgs_audio/components/audio_tower.py:96: the same [B, size] prefix mask was written seven times under four names (Kokoro's length_mask, Higgs's _prefix_mask, Chatterbox's _non_pad_mask, FastConformer's _time_mask, and inline in wav2vec2, the AuT encoder and EarTTS). Kokoro's moves to model/components and the others import it; Kokoro's masking.py keeps its own channel/feature helpers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Omnilingual ASR: one padded sub-batch generator for both paths; no id-keyed attention cache merceod, #859 mstar/model/omnilingual_asr/submodules.py:277: the packed and the padded transcribe paths each built the same zero-padded sub-batches; _padded_groups yields them for both. _attn_fn's functools.partial cache keyed by id(resource) pinned every resource it saw for nothing (the encoder it feeds isn't compiled, so the callable's identity doesn't matter); the partial is built inline. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice: a row's documents are derived once, on the CanvasItem prepare_inputs builds merceod, #859 mstar/model/omnivoice/submodules.py:365: prepare_inputs, the canvas build and _logits each derived a row's (cond, uncond) document lengths, with the CFG predicate spelled two ways. prepare_inputs now builds the row's CanvasItem and carries it in kwargs; CanvasItem.doc_pair is the one derivation (on does_cfg), read by declare_step, the input_seq_len, the lease's seq_lens in _logits and, through the item, the canvas. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * MoE: a model opts into the fused router by passing it, not through every block's signature merceod, #858 mstar/model/components/moe.py:238 (the #859 side): fused_softmax_topk was threaded through four block constructors and instantiate_moe_block only to reach the default TopKRouter, while every block already takes router=. The parameter is gone from the blocks; Qwen3-Omni's Thinker and Talker pass router=TopKRouter(..., fused_softmax_topk=True). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chatterbox: S3Gen's lengths_to_mask is the shared length_mask too merceod, #859 mstar/model/higgs_audio/components/audio_tower.py:96 ("chatterbox's two"): s3gen_flow kept its own copy of the prefix mask (callers pass integer lengths, so its .long() cast changed nothing). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the get_device compile check joins test_submodule_device_cache.py RA.M1: test_submodule_alloc_dtype.py's get_device test repeated main's test_submodule_device_cache.py except for its no-recompile check, which moves there as a case; the copy goes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Resources: EngineResourceInfo without head_shards shards over its joint group RA.F2: head_shards defaulted to 1 whatever joint_comm_group said, so a resource built by hand over a TP group (Command A+'s reference test and checkpoint validator, the GPU ragged test) kept every head unless the caller remembered to pass it. None now resolves to the group's world size (1 without a group); the engine still passes the node's own count. The three call-site edits 41d22c233 needed are reverted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: TP ranks always split a resource's heads; nodes sharing it must agree RA.F1: with tp > 1, a node declaring splits_heads = False got head_shards = 1 against its TP-split layers, and nodes sharing one spec that disagreed were sized by whichever came first in set order, which can differ across ranks (invariant 11). head_shards is now tp x (sp if splits_heads else 1), the same for every node the resource serves, or load raises ValueError naming them. splits_heads is about SP ranks only, as its docstring and adding_models.rst now say. Test: test_ragged_attention_cpu.py, the tp=2 cases of test_a_node_whose_ranks_split_other_work_keeps_every_head and test_nodes_sharing_a_resource_must_agree_on_splitting_heads fail without it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: test_startup_warmup.py says it covers the API server's process settings RA.F3: the module docstring now names api_switch_interval_s beside warmup_requests, and pytest is imported at the top. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * HostRows.first_values: one row map, no in-order fast path beside it #858 triage row 9 (merceod, submodule_base.py:1055): the in-order prefix path and the row-map path had to agree. One dict of rid -> value off the rows, then one lookup per requested rid; padded rows fall off the zip. Cost (python -P microbench, M-series laptop, in-order batch): 4.0 -> 5.9 us at 64 rows, 1.1 -> 1.2 us at 8; an out-of-order batch gets cheaper. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStop: one per-row verdict holds the walk test; no check_walk flag, no boxing #858 triage row 10 (merceod, submodule_base.py:1069): the walk was tested in check_stop, in check_stop_batched and, behind a check_walk flag, in _token_stops, and check_stop boxed its one row into a 1-tuple and a 1-dict. _token_stops(rid, token, info) -> bool is now the one verdict both checks call per row, and the only walk test (last, so a row that goes on never reads its walk). Whisper overrides it per row (record the token, stop only decode rows, mark stopped); its check_stop_batched, which only pre-filtered a prompt step to word-timestamp requests, goes. Cost (python -P microbench, 64 decode rows on the engine's info wrapper, with row 9): 19.6 -> ~25 us a step, from the per-row call; the pairs loop with check_stop boxing its row kept ~20 us. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: the decode stop is a TokenStop declaration, not a copy of its check Coordinator decision on #858 triage row 9: Qwen3.5's check_stop and check_stop_batched restated ARNodeSubmodule's TokenStop path (decode-only stop on its stop ids unless ignore_eos, or at max_tokens). The LLM node now declares token_stop, as #859 already does, and both methods go. Its stop tests build the real submodule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs/serving.rst: api_switch_interval_s is re-measured once #849 lands B6 (coordinator decision): the knob and its validation stay; its row says it goes if neutral after #849. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the Qwen3.5 loop-cap test checks the real LLM submodule's stop Follow-up to a6a1de8bf (#858 triage row 9): test_model_kv_layouts' Qwen3.5 cap test called LLMSubmodule.check_stop on a stand-in holding only the stop ids; the stop is now the TokenStop the submodule declares, so the test builds the submodule, as #859's version of it does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chatterbox profile_s3gen: the solve inputs use the shared length_mask RA.1: 10368cfb3 deleted s3gen_flow.lengths_to_mask, which the S3Gen profiler still imported. It now calls components.masking.length_mask with the same (lengths, size) arguments. No other user of a name this round deleted is left under benchmark/, scripts/, test/ or docs/. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Components: audio_features and the duplex conformer take the shared length_mask RA.4: the length_mask consolidation left two copies of the prefix mask (the mel frontend's pre-emphasis mask and the Nemotron conformer's pad mask); both call components.masking.length_mask. Kokoro's masking.py now holds only the tail-zeroing helpers, and its docstring says so. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kernels: rms_norm_rope is a shared helper, so it drops its leading underscore RA.5: code_predictor_kernels.py imported the private _rms_norm_rope from mstar/utils/attention.py. It is renamed for every user, and its docstring names both homes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the Higgs loop-back prep test says what it checks on the meta device RA.3: on meta, PinnedStager.copy_ takes its non-accelerator branch, so the test no longer sees the pinned row index; it checks that the prep routes the fresh row through overwrite_rows with no pageable copy of its own. The docstring now says that rather than claiming the pinned path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Omnilingual ASR: a step leases the captured encoder only if its rows fit the segment ceiling #858 triage row 14, #859 side (merceod, engine.py:1552): one replay puts every request's frames under one row, and the encoder attention's max_segments_per_request (max_batch_size by default) is a YAML override; can_lease checked only the frames, so an override below max_batch_size let a full step lease and then fail at plan time. can_lease now also checks the row count against the resource's ceiling, as Qwen3.5's vision encoder does; such a step encodes packed eagerly instead. Test: test_omnilingual_asr_model.py::test_a_step_over_the_audio_budget_is_not_encoded_packed (its 3-row step over a ceiling of 2) fails without it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS, V-JEPA2: replay a leased region exactly when the step leased it #859 triage row 10 remainder (merceod, omnivoice/submodules.py:361): the Qwen3-TTS depth loop and the V-JEPA2 AC block loop chose replay with runner.can_run(batch size), while their declare_step keys off the lease (the region declares CODE_PRED_SAMPLER, or the KV and attention, only when it holds a slot). A batch can_lease refused, or one leased while can_run disagreed, would run one path with the other's declaration. Both now replay iff the region is in engine_inputs.piecewise_leases, as OmniVoice, Omnilingual and Qwen3.5 do. Test: test_qwen3_tts_model.py::test_qwen3_tts_depth_loop_replays_only_on_a_step_that_leased_it fails without it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: Higgs's prefill-step stop case runs prefill rows Follow-up to the #858 merge 75fb1ff64 (#858 triage row 10): the batched stop check no longer skips a step by step_walk; each row's own walk decides, as check_stop always did. The case passed step_walk= "prefill_audio" over rows whose infos said decode, a step the worker never builds; its rows are now prefill rows, one at the token budget. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Delete check_stop_batched's unread step_walk parameter (858 RB.1) No check_stop_batched override reads step_walk any more: each row's walk is on its info. Drop it from the base signatures, the engine's call, and every test caller and fake; the two engine-path tests now pin only that the batched path is taken, and the fakes stop carrying the step_context the engine read only for it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: two blank lines before the new ragged-attention tests (858 RB.2) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: drop the last step_context the engine's stop check no longer reads (858 RB.1) The partway-raise Whisper test's batch fake still carried the step_context that check_stop_for_batch read only for step_walk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: declare_step's docstring states the current replay rule (859 RB.2) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Audio features: the normalize mask uses length_mask too (859 RB.3) The third copy of the length mask in NemoLogMel.batch; same [B, 1, T] bool mask on the same device, with n kept for the divisions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Nits: Kokoro masking's docstring covers reverse_padded; a blank line in the Qwen3-TTS tests (859 RB.4) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictors: the kernel launchers are opaque to Dynamo, so a depth position no longer recompiles The shared code-predictor launchers (cp_qk_norm_rope_kv_write, cp_attention, cp_embed_accumulate) took a Python int depth position and were traced by Dynamo as frames of their own. A captured forward compiles with static ints, so each frame recompiled per depth position and batch bucket: Qwen3-TTS started with 461 recompiles against main's 103, and both frames hit the recompile limit and ran eager anyway. They are now torch.compiler.disable, as main's decode_attn_nhd and apply_rope_pos_ids were; under capture the graph records the same launches. A Qwen3-TTS CPU test compiles the depth forward over four depth positions with Triton launches stubbed opaque and checks neither launcher frame was compiled (four cache entries each before this change). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish S2-Pro: seven decode capture buckets, not nine The compiled AR decode capture is one compile per batch bucket, ~2.7 min each on a cold Inductor cache (the stack's cold start was 1,699 s against main's 153 s). The 64-row cap was measured (c=64 per-request RTF 0.95 -> 0.72, TTFA p95 2.0 -> 0.36 s); the 24 and 48 buckets in between were not, and c=32 was unchanged, as a step costs about the same at 64 rows as at 32. They go, leaving [1, 2, 4, 8, 16, 32, 64]. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * PiecewisePackedConfig: attn_key declares one segment per row (merceod #859 row 10) A packed region whose rows each attend within themselves names its attention resource as `attn_key`, and the config declares one non-causal "main" segment per row; under `one_row_per_step` every length is a segment of the one row. Qwen3.5's vision block loop, Omnilingual's encoder and BAGEL's ViT tower declared exactly that in their own closures, which are deleted. Qwen3-ASR keeps its own: its segments carry the window label. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Warmup: count data-parallel replicas from the worker graphs (merceod #858 row 3) The warmup's replica count re-parsed node_groups from the raw YAML, a mirror of WorkerGraph's ranks // (tp*sp) that floored where WorkerGraph asserts and counted node groups no walk places. It now takes the model's worker graphs, counted as the conductor's _pick_replicas counts them, and the YAML path is gone. get_worker_graphs is safe to call in the API server: only BAGEL's override sets state (three placement flags derived from the same config, idempotent), and nothing the API server process runs after the warmup reads them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the Qwen3-Omni prefix-cache Thinker fixture's config has an im_end_token_id #859's ThinkerSubmodule declares its TokenStop in __init__ from config.im_end_token_id; main's #836 fixture config (merged in 5dada826a) had no such field, so the three tests built on it errored at setup. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice silence trim: the trailing trim reads the lead trim's prefix sums (merceod #859 row 23) remove_silence built three _Segment prefix sums per call, the last over a reversed copy of the lead-trimmed samples. A _Segment now addresses a window of its samples, read backwards when flipped, so `reversed_tail` hands the trailing trim the lead trim's prefix sums with reversed bounds: one build fewer and no reversed copy. Output is unchanged to the sample. The pydub comparison skips without the omnivoice package (as on CI), so a new case checks the trim against the reversed-copy path on random clips at four rates, including rates whose fractional samples per millisecond pad or drop the edge. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 I1.1: one WorkerGraph.replica_ranks for a worker graph's DP replicas The rule "the lockstep instances, or one rank apiece without TP/SP" was written out in the warmup copy count, the conductor's replica pick and assignment, GlobalParallelConfig's instance groups and the placement check's single-instance test. All now read WorkerGraph.replica_ranks. No behaviour change. The tests that faked a WorkerGraph with a SimpleNamespace build a real one, so they follow the property. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 I1.2: serving docs count warmup copies from the worker graphs Since the warmup took its replica count from the model's worker graphs, a node group that no walk places adds no copies; the docs still described the node-group count. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Decoder layers: one run_layers for the deferred-add residual stream; the fused add+norm is an opaque op (merceod #859 row 19) `DecoderLayer.forward_residual` runs a layer with its residual adds left to the norms (`RMSNorm.forward_residual`), and `run_layers` runs a stack of them and the final norm. By default a layer attends through its engine attention after its layer cursor moves; `attend(i, self_attn, h)` stands in for a code predictor's kernels. Qwen3DenseLM, the Qwen3-TTS talker and code predictor and the Qwen3-Omni code predictor use it, and their layer classes are DecoderLayers, so the parameter paths are unchanged. The ops and their order are the same: fp32 CPU outputs, random weights, are torch.equal to the inlined loops they replace. The Omni code predictor's unfused layer forward, its inner model's forward and its own forward had no caller (the inner one could not have run: it called the layer without the keywords it required); they are deleted. OmniVoice keeps its loop: its layers are the reference's own modules, whose norms have no forward_residual and whose MLP is a fused-weight FlashInfer call, so the loop is all it would share. FlashInfer's fused_add_rmsnorm registers its custom op with a no-op (flashinfer/utils.py `register_custom_op` returns the function unchanged in 0.7.2), so Dynamo traced into the wrapper and broke the graph there. Each norm call site then compiled as frames of its own, per norm and shape, which is where Qwen3-TTS 1.7B's extra recompiles (293 against main's 103) sat. `mstar::fused_add_rms_norm` now holds the kernel behind a graph boundary, as `mstar::rms_norm` already did. A CUDA test compiles a three-norm stream with fullgraph=True. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni audio encoder: audio_seqlens is a required input (M-2) The merge brought back the lenient `inputs.get("audio_seqlens", [None])[0]` and an Optional default on forward. The graph always sends audio_seqlens with audio_features (qwen3_omni_model.py's audio encoder node), and preprocess already reshapes it unconditionally, so the read is strict again, as on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni audio encoder: count a clip's embeddings with the encoder's n_window (M-1) prepare_inputs counted a clip's embeddings with _get_feat_extract_output_lengths' default n_window=50, while the prompt layout (Qwen3OmniModel._audio_frames) and the encoder itself use the checkpoint's config.audio_encoder.n_window. forward_batched splits a packed call's output by that count, so an encoder with another window would split at the wrong rows or raise. It now passes the encoder's n_window. A CPU test packs two clips through a 52-frame-window encoder and checks each clip's share against the encoder run on it alone (the default count claims 14 embeddings for the 104-frame clip, the encoder gives 13). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni prefix-cache tests: drop the test that pinned a flag; inline the vision test's constants (M-3) test_the_audio_encoder_stays_eager_while_its_output_is_cached only read back AudioEncoderSubmodule.disable_torch_compile through _compile_submodules, so it pinned a constant rather than behaviour. The vision recompile test's one-line tuple of node, walk and sizes is inlined at its uses. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: a cached Qwen3-Omni audio prefix keeps its first batch's embeddings (M-4) The audio encoder note said a clip's embeddings depend at the last bit on its batch-mates under packing; with prefix caching, a reused audio prefix keeps the embeddings it was first computed with, so the note says so too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * mstar::fused_add_rms_norm: a portable default and a CUDA kernel, so add_rms_norm picks no backend (6080-B) The op had only FlashInfer's kernel as its body, so it could not run off CUDA, and add_rms_norm chose between it and the unfused pair by device and dtype. It now mirrors mstar::rms_norm: the op's default is the unfused add and norm, in place (`residual.add_(x)`, then `x.copy_(norm(residual))`, through mstar::rms_norm or Gemma's fp32 norm), and register_kernel("cuda") runs FlashInfer's fused_add_rmsnorm / gemma_fused_add_rmsnorm. FlashInfer's norm kernels dispatch on fp16 and bf16 only, so the CUDA kernel falls back to the default for other dtypes; that is the one eligibility check left, and it lives with the kernel. add_rms_norm is a reshape around the op. The fake impl is unchanged. No vendor branch: the op picks its kernel by device type exactly as mstar::rms_norm does, so ROCm behaves as before. Gemma's fp32 norm moves from components/norm.py to rms_norm.py (gemma_rms_norm) so the op's default and RMSNorm share it. Behaviour change off CUDA-half: add_rms_norm now updates both inputs in place on every device, as RMSNorm.forward_residual's docstring already states and the CUDA path already did. Values are unchanged; on XPU and CUDA fp32 the in-place form adds one copy of the normed rows per norm. A CPU test checks the default against the unfused add and norm, Gemma included (before, the op imported flashinfer on any device). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the fused add+norm graph-break test compiles with aot_eager (6080-F) backend="eager" stops at Dynamo, so the op's mutates_args declaration was never functionalized. aot_eager runs AOTAutograd's functionalization too, so a wrong mutation declaration fails the test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: the layer loop is run_layers, with the TP group passed to the norms (6080-C) Qwen3_5LanguageModel._forward_fused_residual was a fifth copy of the deferred-add loop. It differed from run_layers only in two ways: the norms take the TP comm_group, and each layer moves its mixer's cursor (attend or mix) to that layer's resource index. run_layers and DecoderLayer.forward_residual now take comm_group and pass it to both norms and the final norm. Qwen3.5 passes its cursor move through the existing attend hook (_attend). The hand-written loop is deleted. The TP ops are the same: RMSNorm.forward_residual already took comm_group, and it reduces only when residual is given. The first layer's input norm started the stream without a group (the embeddings are a full tensor), and that is still what happens. The guard that raised ValueError when a stream started under TP is removed. It refused exactly this case, a TP stack whose first input is the full embedding output, and the norm cannot tell a full tensor from a partial sum. The docstring states the contract. The cursor move now runs after the layer's input norm rather than before it. That is host-side state the norm does not read. CPU fp32 check (throwaway harness, not committed): tiny Qwen3.5 config, random weights, stub mixers. The new forward is torch.equal to the deleted loop, and the cursor calls come in the same sequence. TP2 needs a GPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictors: one cp_attend for Qwen3-TTS and Qwen3-Omni (6080-D) The two forward_depth_unrolled attend closures were the same code. The only difference was that Omni called F.linear(h, attn.qkv_proj.weight) where TTS called attn.qkv_proj(h). Both are ParallelAttention modules on a trivial comm group, and Omni's has no qkv bias, so qkv_proj(h) is that same F.linear. cp_attend in code_predictor_kernels.py is the shared function, and each predictor binds its cache arguments with functools.partial. cp_attend is plain Python, so its projections still compile; the launchers it calls keep their torch.compiler.disable. Omni's leftover locals go, and so does its comment that FlashInfer's rope needs bf16 (the Triton prologue replaced that rope). The Qwen3-TTS comment that said Omni used a different attention primitive and deferred sharing until #199 now points at the shared path. The TTS CPU fakes patch the kernels module, where cp_attend looks the launchers up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni code predictor: docstrings describe the dense-KV Triton path; dead helpers go (6080-E) forward_depth_unrolled's docstring and section header called it SDPA and a replacement for a paged-FlashInfer path. It attends through cp_attend's Triton kernels, and the eager fallback is gone. consolidate_stacked_weights said the per-head lm_head modules serve that eager path; they stay only as the weight loader's targets. get_embedding, get_lm_head, the codec_embedding alias, the FlashInfer _apply_rope and _num_residual had no callers anywhere in the repo (grep over mstar/, test/, benchmark/, scripts/) and are deleted. A plan reference ("Option P") goes from the lm_head_weight comment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice silence test: say the shared-prefix-sum case guards the flipped path on CI (7C-1) The pydub comparison skips without the omnivoice package, as on CI, so this case is what checks _Segment.reversed_tail there; the docstring says so. 7c270c703's saving, measured (inv 10): remove_silence on 24 kHz speech-like clips (mid_sil=500, lead_sil=100, trail_sil=100, OmniVoice's settings), parent 30fb70b4f and 7c270c703 interleaved in one process on the Eras CPU box (r7i.4xlarge, OMP_NUM_THREADS=2), 550 timed calls each, two runs: 10 s clip 6.73 -> 5.46 ms median (-1.27 ms, -19%), 30 s clip 13.7-17.6 -> 11.4-15.1 ms (-2.3 to -2.5 ms). A 10 s OmniVoice request takes about 130 ms end to end at c=1, so the saving is about 1% of it and over 1 ms; the change stays. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: wrap the packed-clip test's long line (M-1) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the host-sync test's stand-in audio encoder has an n_window (M-1) AudioEncoderSubmodule.prepare_inputs now reads the encoder's n_window, as the real Qwen3OmniMoeAudioEncoder carries; the stand-in gets the checkpoint's 50. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 859 R2.1, R2.2: fix stale code-predictor docstrings; note the fused add-norm fallback R2.1: the code predictor runs in the autocast dtype, not float32, and its scratch KV cache is allocated once and never zeroed; it needs no clearing because every position a depth step reads is written earlier in that step. R2.2: say at mstar::fused_add_rms_norm's default kernel that it is the portable fallback (XPU, CUDA fp32), opaque to Inductor, and that an XPU model adopting run_layers can register a fused kernel there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Warmup: a replicated entry runs one request alone first Warmup sent each entry `copies` times at once, `copies` being the most replicas any worker graph has. Where only a cheap node is replicated (a VAE and text encoder x2 beside one DiT SP2 instance), every copy reached the one DiT together, it batched them, and its bs=1 shape stayed cold for the first client. Past one copy, each level now runs one request (`warmup-{i}.{l}.solo`) to its end before the `copies * concurrency` batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * merceod #859 row 11: the patch grids become host tuples once, in process_prompt process_prompt turns each image's and video's grid into a (t, h, w) tuple of Python ints where it builds the lists, sizes the spans from those tuples, and puts the same tuples in the metadata. The wire codec carries tuples of ints through model_kwargs and step_metadata exactly, so the schedule's tuple(...) and the submodules' _host_grids re-conversions were no-ops; both go, along with the host() helper and media's 2-D grid guard (every grid is (3,)). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * merceod #858 row 15: one pinned H2D helper, capability-gated and list-aware to_device_nonblocking decides pinning with DeviceCapability.of(device).pin_memory instead of a cuda string check, so XPU takes the non-blocking pinned path too, and takes a Python sequence plus a dtype, built straight into pinned memory (one host allocation, one async copy). The hand-rolled copies switch to it: the sampler's keep_mask, nemotron_duplex's device_counts (deleted) and chunk copy, and the callers that built a pageable tensor only to hand it to the helper (parakeet, wav2vec2, fastconformer, audio_features). audio_features pads into pinned memory on the same capability check. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 859 R3.1: Higgs valid_frames go up from a list Pass the host list of chunk frame counts straight to to_device_nonblocking with dtype=torch.long, so it is built once in pinned memory instead of as a pageable tensor that the helper then pins again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 R3.1: warmup runs its solo once per entry, only when replicas are uneven The one-request-alone warmup ran once per concurrency level and also in fully replicated deployments, where every worker graph already sees one copy. It now runs once per entry, before the levels, and only when some placed worker graph has fewer replicas than the copies sent. _data_parallel_replicas returns the most and fewest replicas from the same worker graphs, and _run_warmup_requests takes that pair. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 R3.2: the warmup solo logs its own outcome The solo used to share the batch's try, so a failed solo skipped the batch and the warning blamed the concurrency level, and the success line's time included the solo. The solo now has its own try (from R3.1) and its own done line, and the batch's time covers the batch alone. The startup line says when each warmup runs once alone before its copies. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 R3.3: to_device_nonblocking says an already-pinned source is read in place The docstring said the caller need not keep the source, but an already-pinned CPU tensor is copied from in place: writing to it before the copy lands changes what arrives. Say so, and point a reused buffer at PinnedStager. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * code predictor: no per-position recompile of cp_attend (GPU recompile count) A captured forward compiles with dynamic=False, so every compiled frame that holds a depth step's cache_pos (or passes it to a call at a graph break, partials included) guards on it and recompiles at each of the 16 depth positions and each batch bucket. Before run_layers this could not happen: the launcher calls sat in forward_depth_unrolled's own layer loop, and Dynamo runs a frame with a graph break inside a loop eagerly, so only the norm and MLP frames it called were compiled. With run_layers the breaks landed in loop-free frames that hold the ints: cp_attend and its resume frame, DecoderLayer.forward_residual (through the partial's args) and forward_depth_unrolled (rebuilding the partial at the break). Every frame holding them now runs eager: - cp_attend is torch.compiler.disable'd; its qkv/o projections are cuBLAS GEMMs at -O1 compiled or not. - run_layers calls attend in its own loop rather than through DecoderLayer.forward_residual (deleted; it had no other caller), so an opaque attend sends the loop to eager as the old depth loop did. The ops and their order are unchanged for every caller. - forward_depth_unrolled (Qwen3-TTS, Qwen3-Omni) skips its own frame with torch.compiler.disable(recursive=False). Its callers already broke at it. Only the norm and MLP frames compile, none keyed on a position. The test compiles the talkers' depth-loop shape over 2 and 4 positions for both code predictors and requires the same frame count; at 3abebd178 it is 11 vs 19 (Qwen3-TTS) and 17 vs 31 (Qwen3-Omni). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * run_layers: one compiled frame per layer besides attention run_layers runs eager (its loop breaks at the layer cursor), so every compiled frame it calls is dispatched on each uncaptured forward. Calling the input norm, the post norm and the MLP separately made that three frames per layer, where DecoderLayer.forward_residual and its resume frame had made two. DecoderLayer.forward_mlp_residual runs the post-attention add+norm, the MLP and the next norm's add+norm (the next layer's input norm, or the final norm) as one frame holding no layer index. The ops, their order, the fused add+norm placement and comm_group are unchanged; the first layer's input norm runs once before the loop. The test counts the compiled graphs one forward of a Qwen3 dense stack runs: each extra layer may add attention's graphs plus one. At d55935f03 it adds 5 (attention's 2 plus input norm, post norm, MLP). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: _attend runs eager, so no frame guards on the layer index run_layers calls Qwen3_5LanguageModel._attend from its eager loop, so _attend was a top-level compiled frame. It reads _cursor_attr[i] and _resource_idx[i] before the graph break at set_layer_idx, which guards it on i: one empty-graph compile per layer per compile wrapper, counted as recompiles and toward MSTAR_RECOMPILE_LIMIT. It is torch.compiler.disable(recursive=False) now, so the mixer it calls still compiles as its own frame. No other attend binds i in a compiled frame: the default path sets the cursor in run_layers' eager loop, and cp_attend is disabled. The depth-loop test becomes test_layer_loop_compiles_no_frame_per_step, parametrized over Qwen3-TTS, Qwen3-Omni (now including its opening two-token pass at cache_pos 0) and Qwen3.5's layer loop with stub mixers. It counts the frames Dynamo compiles rather than graphs: _attend's per-layer compiles have empty graphs, so a graph count is equal at d55935f03 (5 vs 5) while compiles are 11 vs 9 there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * cp_attend docstring: the cuBLAS claim holds up to -O1 At -O2 max-autotune may pick a Triton GEMM template for the compiled projections, so eager is not equivalent there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the per-layer frame bound compares the two-layer difference whole (D5) (calls4 - calls2) // 2 rounded an odd difference down, so one stray frame per two layers could pass. Compare the difference with twice the bound. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: run_layers' compile tests live with the Qwen3 LM, on one counter (D4) The Qwen3.5 layer-loop case moves from the Qwen3-TTS module into test_qwen3_lm.py beside the per-layer frame test, and both modules count through one dynamo_counts fixture in conftest (frames compiled, and graphs one compiled call runs) instead of a helper each. The code-predictor depth-loop cases stay with their builders in the TTS module. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * run_layers: an empty layer stack is the final norm again (D3) Starting the stream at layers[0]'s input norm raised IndexError on an empty stack, which had returned the final norm before 76ba991e2. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * code predictor: cp_attend's projections compile again (D1) cp_attend was torch.compiler.disable'd whole, so its one-row qkv and o projections ran on eager cuBLAS. At the default -O1 (coordinate descent) Inductor rewrites an M==1 mm into a fused Triton reduction, and Qwen3-TTS measured 25-30% slower at -O0 at c=1 (docs/serving.rst). It is disable(recursive=False) now: its own frame, which binds cache_pos and layer_idx, stays eager, and the projections compile as int-free frames of their own. The launchers stay disabled. The docstrings said the projections were cuBLAS GEMMs either way; they now say what compiles. The new test requires both projections' weights in a compiled graph of the depth loop, for both code predictors; it fails at 9f26410bd. The per-position compile test still passes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M1: copy an already-pinned host tensor to the device in place The merge of main took #904's pageable path for every CPU tensor, so a source that is already pinned (NemoLogMel.batch pads into pinned memory) paid a second pinned allocation and a full-batch memcpy per call. Read a pinned source in place again, and restore the docstring caveat that the caller must not write to it until the copy lands. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M2: size video decode threads from usable_cpus, apart from the torch pool video_decode_threads re-read CPU affinity beside main's usable_cpus(), which also reads the cgroup quota, and was capped by torch's intra-op pool, which main #903 shrank to a quarter of the CPUs (at most 4). That cut V-JEPA2 and Cosmos3 decodes from 8 FFmpeg threads to 4, or 1 on a small host. FFmpeg's frame threads are not torch intra-op threads, so decode is now min(8, usable_cpus()). The docs paragraph that said both that the quota is read and that it is not now says one thing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M3: test GLM-5.2's stop id at prefill The merge ported GLM-5.2's prefill stop rules onto TokenStop without a test: a stop id sampled on the prefill walk ends the reply at the first decode step, whose token postprocess drops, and both stop checks agree; with ignore_eos the prefill sets nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the projection-compile test counts through dynamo_counts (E2) It had its own backend, reset and torch.compile beside the fixture. The fixture's DynamoCounts gains traced, the data_ptrs of the tensors its compiled graphs take, and the test reads that. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M4: declare a prefill-sampled first token on TokenStop GLM-5.2 and Command A+ each hand-built the same rule beside TokenStop: a stop id sampled by a prefill walk has no loop to stop, so the reply ends at the loop's first step, whose token postprocess drops, and max_tokens counts the prefill's token. TokenStop.prefill_walks now declares those walks; the base keeps the flag in request_state and applies both rules in check_stop and check_stop_batched, and ended_at_prefill() tells postprocess to drop the closing token. GLM-5.2's _token_stops / check_stop overrides and Command A+'s _ended_at_prefill set, check_stop and cleanup_request are deleted. Command A+ now sets the flag for max_tokens 1 too (no decode step follows, and removal frees it) and keeps it until removal instead of clearing it on the stopping step, as GLM-5.2 did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M1: model the mocked host tensor in the pin test as pageable test_to_device_nonblocking_pins_where_the_device_capability_says builds its host from Mock(spec=torch.Tensor), whose is_pinned() is truthy, so it now took the already-pinned path. It is about the pageable source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N1: copy a pinned source in place only when it is contiguous .to() of a pinned non-contiguous tensor packs it into a pageable temp first, so the copy blocks the caller. Such a source goes through the pinned buffer like a pageable one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N2: read usable_cpus once per process for video decode threads video_decode_threads re-read the affinity and the cgroup files on every decode. Cache it, as image_io.PNG_WORKERS does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N3: say a generated_key counter must count the prefill token itself prefill_walks adds the prefill's token to the loop count only when TokenStop counts loop iterations; a submodule's own generated_key counter is read as is. No model declares both today. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N4: the engine drops the closing token of a reply that ended at prefill Declaring TokenStop.prefill_walks still made each model's postprocess delete token_stop.buffer when ended_at_prefill(). Engine._postprocess_batch now does it after postprocess, the one place every forward path's rows are postprocessed; a submodule without prefill_walks pays one getattr per batch. GLM-5.2's and Command A+'s own deletes go. Tests that drove postprocess by hand (GLM-5.2's engine-cycle driver and both models' prefill-stop cases) now call Engine._postprocess_batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N4: say prefill_walks needs the loop to emit the token as buffer The engine drops token_stop.buffer on the closing step, which removes the token only where the loop's client edge carries that output (GLM-5.2, Command A+). Qwen3.5's decode emits text_inputs, which this would miss. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 / run_layers: capture works again (fail_on_recompile) Every capture runs its forward under torch.compiler.set_stance( "fail_on_recompile") (capture_into_graph). torch.compiler.disable( recursive=False) skips only its wrapper's code; the decorated function's own frame never gets a cache entry, so Dynamo's callback runs on every call, and convert_frame lets it run eager only because its caller is the wrapper. The capture stance's fail_callback has no such check: it raises on the first warm call, whatever the shapes. So every bucket of a forward that reaches one fails to capture: Qwen3.5 through _attend (1438ea127), Qwen3-TTS and Qwen3-Omni through forward_depth_unrolled (d55935f03) and cp_attend (a9bf02141). The undecorated _attend of 0fe8c812d captured because it was compiled at warm-up, one entry per layer index. No frame is recursive=False now. Each per-layer or per-step int is read only in fully torch.compiler.disable'd cursor methods, called from loops that already run eager, and the attention runs as its own int-free module frame: - run_layers' hook is set_cursor(i, self_attn), which replaces the default self_attn.attend.set_layer_idx(i); the loop then calls self_attn(h) on both paths. Qwen3.5's is _set_cursor. - cp_attend is gone. A code predictor's layers are DepthAttention (a ParallelAttention: qkv and o projections, compiled) over the predictor's DepthKV, a cursor holding the frame-local KV cache and the depth step (bind_step, set_layer_idx, and the two kernel launches). run_layers moves it per layer through the default path. - forward_depth_unrolled takes only the input embeddings; the depth loops in the Talkers bind the step on cp.depth first, so cache_pos never enters a compiled frame. dynamo_counts now makes its second, warm call under the capture stance, as the runner warms a bucket and then captures it. With only that change on da5514d26, the Qwen3.5 per-layer test and both code-predictor tests (per-position, projections compile; TTS and Omni) fail with the same "Detected recompile ... '_attend' / 'forward_depth_unrolled'" errors the GPU runs logged; they pass here, with the per-layer, per-position and projection checks unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: a DepthKV binding serves one pass over the layers DepthKV kept the previous step's cache, positions and cache_pos, so a depth step whose loop skipped bind_step wrote and attended over the previous step's slots without a sign. The last layer's call now drops the binding after its launches, inside the eager __call__ (no new graph break; a capture has already recorded the cache's pointers, and the caller owns its storage), and a call with none bound raises ValueError. The leading dim of the cache is the layer count, as both Talkers allocate it. The bind…

* Tests: the Qwen3-Omni prefix-cache Thinker fixture's config has an im_end_token_id #859's ThinkerSubmodule declares its TokenStop in __init__ from config.im_end_token_id; main's #836 fixture config (merged in 5dada826a) had no such field, so the three tests built on it errored at setup. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL parallel CFG: a branch node sitting out a step builds only what its forward reads merceod #860 row 12. Under image_gen_cfg the step's branch count was known in prepare_inputs, but a node sitting the step out (cfg_img all of t2i, cfg_text outside the CFG interval) still seeded step-0 latents, built the sentinel frame, three per-label position-id tensors and the shape inputs, only for the forward to return a zero-row velocity. prepare_inputs now counts the branches first and returns the slim NodeInputs (the count, requires_cfg) when the node sits out. forward checks the skip ahead of the single-branch dispatch and sends sat_out_velocity(config, device), which no longer needs latents: same (0, latent width) shape and the latents' dtype rule as before, so the v_<label> edge each runtime routes is unchanged. declare_step still returns None for the step, and no image_gen_cfg capture exists (the LLM captures only decode and prefill_text), so no captured input shape moves. _sits_out / _cfg_label replace three NODE_TO_CFG_LABEL lookups, the per-node output-name dict and the _NODE_TO_CFG_LABEL alias. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice silence trim: the trailing trim reads the lead trim's prefix sums (merceod #859 row 23) remove_silence built three _Segment prefix sums per call, the last over a reversed copy of the lead-trimmed samples. A _Segment now addresses a window of its samples, read backwards when flipped, so `reversed_tail` hands the trailing trim the lead trim's prefix sums with reversed bounds: one build fewer and no reversed copy. Output is unchanged to the sample. The pydub comparison skips without the omnivoice package (as on CI), so a new case checks the trim against the reversed-copy path on random clips at four rates, including rates whose fractional samples per millisecond pad or drop the edge. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 I1.1: one WorkerGraph.replica_ranks for a worker graph's DP replicas The rule "the lockstep instances, or one rank apiece without TP/SP" was written out in the warmup copy count, the conductor's replica pick and assignment, GlobalParallelConfig's instance groups and the placement check's single-instance test. All now read WorkerGraph.replica_ranks. No behaviour change. The tests that faked a WorkerGraph with a SimpleNamespace build a real one, so they follow the property. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan VAE decoders: one FastDecoder setup in wan_vae, used by Wan2.2, Cosmos3 and LingBot merceod #860 row 10. The shared Wan VAE decoder was set up four times (Wan2.2, Cosmos3 whole-clip, Cosmos3 streaming, LingBot): a to_channels_last before a DecoderStep, and a decode_causal call under autocast-off, with only Wan2.2 switching layout per request under inference_mode(False). wan_vae.FastDecoder now owns it: the layout switch at construction (before the step compiles), the DecoderStep, decode (a whole clip from an empty cache) and decode_stream (a caller-owned cache). Both calls put the VAE back in the decode layout first, outside inference mode, and use_stock_layout switches it to the loaded layout for a stock decode, so Wan2.2's per-request switch and everyone else's construction-time switch are one mechanism. A VAE already in the right layout skips the switch. wan_vae.decode_dtype(fp32_env) is the cuDNN-gated decode dtype (bf16 from 9.16, the model's *_VAE_DECODE_FP32 override) that Wan2.2 and Cosmos3 each spelled out. Numerics per model are unchanged: each caller keeps its cast where it was (Wan2.2 and Cosmos3 whole-clip at construction, the Cosmos3 streaming node lazily, LingBot bf16 unless vae_decode_fp32), its deterministic flag (False only on the Cosmos3 streaming step), its compile switch, latents_per_call and cudnn_benchmark scope. The Cosmos3 streaming node's step compiles only when the warm cache is on, as before; its stateless re-decode never runs it. Deleted: the three to_channels_last + DecoderStep copies, Wan2.2's per-request inference_mode(False) block and both inline decode_causal calls, the two cuDNN-version dtype bodies (and the TODO asking for this), Cosmos3's _fast_step / _step and LingBot's _decode_step. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 I1.2: serving docs count warmup copies from the worker graphs Since the warmup took its replica count from the model's worker graphs, a node group that no walk places adds no copies; the docs still described the node-group count. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Decoder layers: one run_layers for the deferred-add residual stream; the fused add+norm is an opaque op (merceod #859 row 19) `DecoderLayer.forward_residual` runs a layer with its residual adds left to the norms (`RMSNorm.forward_residual`), and `run_layers` runs a stack of them and the final norm. By default a layer attends through its engine attention after its layer cursor moves; `attend(i, self_attn, h)` stands in for a code predictor's kernels. Qwen3DenseLM, the Qwen3-TTS talker and code predictor and the Qwen3-Omni code predictor use it, and their layer classes are DecoderLayers, so the parameter paths are unchanged. The ops and their order are the same: fp32 CPU outputs, random weights, are torch.equal to the inlined loops they replace. The Omni code predictor's unfused layer forward, its inner model's forward and its own forward had no caller (the inner one could not have run: it called the layer without the keywords it required); they are deleted. OmniVoice keeps its loop: its layers are the reference's own modules, whose norms have no forward_residual and whose MLP is a fused-weight FlashInfer call, so the loop is all it would share. FlashInfer's fused_add_rmsnorm registers its custom op with a no-op (flashinfer/utils.py `register_custom_op` returns the function unchanged in 0.7.2), so Dynamo traced into the wrapper and broke the graph there. Each norm call site then compiled as frames of its own, per norm and shape, which is where Qwen3-TTS 1.7B's extra recompiles (293 against main's 103) sat. `mstar::fused_add_rms_norm` now holds the kernel behind a graph boundary, as `mstar::rms_norm` already did. A CUDA test compiles a three-norm stream with fullgraph=True. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * LingBot fast numerics: compile module-level pointwise chains once, not closures per block merceod #860 row 18. Off reference numerics, _fast_numerics setattr'd a torch.compile closure over six bound methods (_pre, _post_self, _norm3, _ffn, _qkv, _q) onto every block and attention module: 180 compiled wrappers on the 30-layer DiT, each tracing its own graph and counting against Dynamo's per-code cache limit. The chains are now module-level functions of the modules they run (_qkv(attn, x), _pre(norm1, ...), _ffn(norm2, ffn, ...), ...), each compiled once (functools.cache) with the options 98da67e10 set (dynamic=False, emulate_precision_casts). Blocks and attention modules take one constructor flag, fast (was fa4: flash-attn-4 and the compiled chains are both the off-reference path), and call _pointwise(fast, fn, *args); with fast off it calls fn eagerly, the same ops the methods ran, so the reference path is unchanged. Deleted: _fast_numerics and the six swappable methods. Dynamo guards on the modules passed in by their structure, not their identity: with a CompileCounter backend, three blocks' worth of calls trace six frames, one per chain. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 CFG-parallel: gather through the group, drop gather_rows (860 I1.1) gather_rows was CommGroup.all_gather behind a None-group convention the rest of sequence_parallel does not use. The dit now holds a trivial group when none is passed (one GPU), so my_rows takes a CommGroup like the other helpers and the forward calls sp_group.all_gather directly; a trivial group's all_gather returns its input, as gather_rows(None, x) did. The collective and the guidance_scale > 1 condition it runs under are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL parallel CFG: the sat-out velocity drops its dtype logic (860 I1.2) sat_out_velocity read the autocast dtype inside the forward and called that the latents' dtype rule, but step-0 latents were fp32 (prepare_inputs runs outside autocast), so the claim did not hold. Nothing reads the dtype: the edge carries zero elements and combine_cfg checks only shape[0]. Build the tensor in the default dtype, say so in the docstring, and stop the 6055b60f4 test asserting fp32. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan VAE FastDecoder: one layout switch path; Cosmos3 reads the VAE's layout (860 I1.4) FastDecoder.__init__ called to_channels_last directly, outside _use_layout's inference_mode(False), though the class says every switch runs there. It now goes through _use_layout(True), which converts the same weights in the same way (and skips a VAE a sibling node already converted, as the idempotent call did). Cosmos3's decoder nodes kept a _channels_last flag beside their FastDecoder's. _vae_input now asks wan_vae.is_channels_last(self.vae): the flag was set exactly when a FastDecoder converted this node's VAE, and Cosmos3 never switches it back, so z's layout is unchanged in every config. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni audio encoder: audio_seqlens is a required input (M-2) The merge brought back the lenient `inputs.get("audio_seqlens", [None])[0]` and an Optional default on forward. The graph always sends audio_seqlens with audio_features (qwen3_omni_model.py's audio encoder node), and preprocess already reshapes it unconditionally, so the read is strict again, as on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni audio encoder: count a clip's embeddings with the encoder's n_window (M-1) prepare_inputs counted a clip's embeddings with _get_feat_extract_output_lengths' default n_window=50, while the prompt layout (Qwen3OmniModel._audio_frames) and the encoder itself use the checkpoint's config.audio_encoder.n_window. forward_batched splits a packed call's output by that count, so an encoder with another window would split at the wrong rows or raise. It now passes the encoder's n_window. A CPU test packs two clips through a 52-frame-window encoder and checks each clip's share against the encoder run on it alone (the default count claims 14 embeddings for the 104-frame clip, the encoder gives 13). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni prefix-cache tests: drop the test that pinned a flag; inline the vision test's constants (M-3) test_the_audio_encoder_stays_eager_while_its_output_is_cached only read back AudioEncoderSubmodule.disable_torch_compile through _compile_submodules, so it pinned a constant rather than behaviour. The vision recompile test's one-line tuple of node, walk and sizes is inlined at its uses. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: a cached Qwen3-Omni audio prefix keeps its first batch's embeddings (M-4) The audio encoder note said a clip's embeddings depend at the last bit on its batch-mates under packing; with prefix caching, a reused audio prefix keeps the embeddings it was first computed with, so the note says so too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * mstar::fused_add_rms_norm: a portable default and a CUDA kernel, so add_rms_norm picks no backend (6080-B) The op had only FlashInfer's kernel as its body, so it could not run off CUDA, and add_rms_norm chose between it and the unfused pair by device and dtype. It now mirrors mstar::rms_norm: the op's default is the unfused add and norm, in place (`residual.add_(x)`, then `x.copy_(norm(residual))`, through mstar::rms_norm or Gemma's fp32 norm), and register_kernel("cuda") runs FlashInfer's fused_add_rmsnorm / gemma_fused_add_rmsnorm. FlashInfer's norm kernels dispatch on fp16 and bf16 only, so the CUDA kernel falls back to the default for other dtypes; that is the one eligibility check left, and it lives with the kernel. add_rms_norm is a reshape around the op. The fake impl is unchanged. No vendor branch: the op picks its kernel by device type exactly as mstar::rms_norm does, so ROCm behaves as before. Gemma's fp32 norm moves from components/norm.py to rms_norm.py (gemma_rms_norm) so the op's default and RMSNorm share it. Behaviour change off CUDA-half: add_rms_norm now updates both inputs in place on every device, as RMSNorm.forward_residual's docstring already states and the CUDA path already did. Values are unchanged; on XPU and CUDA fp32 the in-place form adds one copy of the normed rows per norm. A CPU test checks the default against the unfused add and norm, Gemma included (before, the op imported flashinfer on any device). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the fused add+norm graph-break test compiles with aot_eager (6080-F) backend="eager" stops at Dynamo, so the op's mutates_args declaration was never functionalized. aot_eager runs AOTAutograd's functionalization too, so a wrong mutation declaration fails the test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: the layer loop is run_layers, with the TP group passed to the norms (6080-C) Qwen3_5LanguageModel._forward_fused_residual was a fifth copy of the deferred-add loop. It differed from run_layers only in two ways: the norms take the TP comm_group, and each layer moves its mixer's cursor (attend or mix) to that layer's resource index. run_layers and DecoderLayer.forward_residual now take comm_group and pass it to both norms and the final norm. Qwen3.5 passes its cursor move through the existing attend hook (_attend). The hand-written loop is deleted. The TP ops are the same: RMSNorm.forward_residual already took comm_group, and it reduces only when residual is given. The first layer's input norm started the stream without a group (the embeddings are a full tensor), and that is still what happens. The guard that raised ValueError when a stream started under TP is removed. It refused exactly this case, a TP stack whose first input is the full embedding output, and the norm cannot tell a full tensor from a partial sum. The docstring states the contract. The cursor move now runs after the layer's input norm rather than before it. That is host-side state the norm does not read. CPU fp32 check (throwaway harness, not committed): tiny Qwen3.5 config, random weights, stub mixers. The new forward is torch.equal to the deleted loop, and the cursor calls come in the same sequence. TP2 needs a GPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictors: one cp_attend for Qwen3-TTS and Qwen3-Omni (6080-D) The two forward_depth_unrolled attend closures were the same code. The only difference was that Omni called F.linear(h, attn.qkv_proj.weight) where TTS called attn.qkv_proj(h). Both are ParallelAttention modules on a trivial comm group, and Omni's has no qkv bias, so qkv_proj(h) is that same F.linear. cp_attend in code_predictor_kernels.py is the shared function, and each predictor binds its cache arguments with functools.partial. cp_attend is plain Python, so its projections still compile; the launchers it calls keep their torch.compiler.disable. Omni's leftover locals go, and so does its comment that FlashInfer's rope needs bf16 (the Triton prologue replaced that rope). The Qwen3-TTS comment that said Omni used a different attention primitive and deferred sharing until #199 now points at the shared path. The TTS CPU fakes patch the kernels module, where cp_attend looks the launchers up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni code predictor: docstrings describe the dense-KV Triton path; dead helpers go (6080-E) forward_depth_unrolled's docstring and section header called it SDPA and a replacement for a paged-FlashInfer path. It attends through cp_attend's Triton kernels, and the eager fallback is gone. consolidate_stacked_weights said the per-head lm_head modules serve that eager path; they stay only as the weight loader's targets. get_embedding, get_lm_head, the codec_embedding alias, the FlashInfer _apply_rope and _num_residual had no callers anywhere in the repo (grep over mstar/, test/, benchmark/, scripts/) and are deleted. A plan reference ("Option P") goes from the lm_head_weight comment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice silence test: say the shared-prefix-sum case guards the flipped path on CI (7C-1) The pydub comparison skips without the omnivoice package, as on CI, so this case is what checks _Segment.reversed_tail there; the docstring says so. 7c270c703's saving, measured (inv 10): remove_silence on 24 kHz speech-like clips (mid_sil=500, lead_sil=100, trail_sil=100, OmniVoice's settings), parent 30fb70b4f and 7c270c703 interleaved in one process on the Eras CPU box (r7i.4xlarge, OMP_NUM_THREADS=2), 550 timed calls each, two runs: 10 s clip 6.73 -> 5.46 ms median (-1.27 ms, -19%), 30 s clip 13.7-17.6 -> 11.4-15.1 ms (-2.3 to -2.5 ms). A 10 s OmniVoice request takes about 130 ms end to end at c=1, so the saving is about 1% of it and over 1 ms; the change stays. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: wrap the packed-clip test's long line (M-1) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the host-sync test's stand-in audio encoder has an n_window (M-1) AudioEncoderSubmodule.prepare_inputs now reads the encoder's n_window, as the real Qwen3OmniMoeAudioEncoder carries; the stand-in gets the checkpoint's 50. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 859 R2.1, R2.2: fix stale code-predictor docstrings; note the fused add-norm fallback R2.1: the code predictor runs in the autocast dtype, not float32, and its scratch KV cache is allocated once and never zeroed; it needs no clearing because every position a depth step reads is written earlier in that step. R2.2: say at mstar::fused_add_rms_norm's default kernel that it is the portable fallback (XPU, CUDA fp32), opaque to Inductor, and that an XPU model adopting run_layers can register a fused kernel there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * LingBot-World: drop the fast config from this stack; #926 owns it #926 brings LingBot-World's fast path through the ragged attention resource, a deterministic block compile and a bf16 channels-last decode, under the same lingbot_world_1_3b_fast.yaml name. This removes everything here that existed only for that config, as introduced by 29bf80124 (its LingBot fast parts), 7b38d578d, fac2fb744, bc110b94d, e8519d791, 35526b3b9, dbc9bb679 and f312336d2: - configs/lingbot_world_1_3b_fast.yaml and its rows in docs/models.rst and docs/installation.rst; - reference_numerics, the flash-attn-4 attention path and its import check, the compiled block pointwise chains and emulate_precision_casts, the one-row timestep embedding, the cuDNN-autotuned condition encoder; - vae_decode_compile / vae_decode_cudnn_benchmark, the bf16 weight cast, and the named vae_decode_* model kwargs (back to main's **kwargs); - their tests (the bool-settings test, the kwargs test, the compiled-chain test). config.py, lingbot_world_model.py and test_lingbot_world_shell.py are main's again. Kept, as they serve the default config: 45e67ab32's shared FastDecoder (eager, contiguous by default, the same decode as main's DecoderStep path), the per-chunk camera modulation (99a7b64b6), the single-row e/e0 broadcast and the no-op RoPE concat skip, all bit-identical on the default config. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * LingBot-World: back to main exactly The LingBot-World code on this stack now matches main: the per-chunk camera modulation, the single-row e/e0 broadcast, the RoPE-concat skip and the decode through wan_vae.FastDecoder.decode_stream go, and the dit decodes through main's DecoderStep / decode_causal again (run_controls.py too). wan_vae's FastDecoder and decode_dtype stay for Wan2.2 and Cosmos3. git diff origin/main over mstar/model/lingbot_world, configs/lingbot* and the LingBot tests is empty. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Warmup: a replicated entry runs one request alone first Warmup sent each entry `copies` times at once, `copies` being the most replicas any worker graph has. Where only a cheap node is replicated (a VAE and text encoder x2 beside one DiT SP2 instance), every copy reached the one DiT together, it batched them, and its bs=1 shape stayed cold for the first client. Past one copy, each level now runs one request (`warmup-{i}.{l}.solo`) to its end before the `copies * concurrency` batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * merceod #859 row 11: the patch grids become host tuples once, in process_prompt process_prompt turns each image's and video's grid into a (t, h, w) tuple of Python ints where it builds the lists, sizes the spans from those tuples, and puts the same tuples in the metadata. The wire codec carries tuples of ints through model_kwargs and step_metadata exactly, so the schedule's tuple(...) and the submodules' _host_grids re-conversions were no-ops; both go, along with the host() helper and media's 2-D grid guard (every grid is (3,)). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * merceod #858 row 15: one pinned H2D helper, capability-gated and list-aware to_device_nonblocking decides pinning with DeviceCapability.of(device).pin_memory instead of a cuda string check, so XPU takes the non-blocking pinned path too, and takes a Python sequence plus a dtype, built straight into pinned memory (one host allocation, one async copy). The hand-rolled copies switch to it: the sampler's keep_mask, nemotron_duplex's device_counts (deleted) and chunk copy, and the callers that built a pageable tensor only to hand it to the helper (parakeet, wav2vec2, fastconformer, audio_features). audio_features pads into pinned memory on the same capability check. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * merceod #860 row 15: share ragged_for between DenoiseLoopSubmodule and Wan2.2 Wan22DitSubmodule's dit attention callable (_attention_fn) was a copy of DenoiseLoopSubmodule.ragged_for. Move ragged_for, its per-binding memo and the rebind reset into RaggedAttentionMixin beside RaggedAttentionFn in components/diffusion/attention.py; both classes now inherit it. Wan sets attn_resource_key = SELF_ATTN and calls ragged_for(VIDEO_SPAN), so it gets the same partial(resource.run, label="video") from the same resource key, memoized per binding as before. Wan's copy, its bind_node_resources override, its _self_attn_fn field and its one-time "self-attention through the ragged resource" log line are deleted, as are two test lines that stubbed that field. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 859 R3.1: Higgs valid_frames go up from a list Pass the host list of chunk frame counts straight to to_device_nonblocking with dtype=torch.long, so it is built once in pinned memory instead of as a pageable tensor that the helper then pins again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 R3.1: warmup runs its solo once per entry, only when replicas are uneven The one-request-alone warmup ran once per concurrency level and also in fully replicated deployments, where every worker graph already sees one copy. It now runs once per entry, before the levels, and only when some placed worker graph has fewer replicas than the copies sent. _data_parallel_replicas returns the most and fewest replicas from the same worker graphs, and _run_warmup_requests takes that pair. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 R3.2: the warmup solo logs its own outcome The solo used to share the batch's try, so a failed solo skipped the batch and the warning blamed the concurrency level, and the success line's time included the solo. The solo now has its own try (from R3.1) and its own done line, and the batch's time covers the batch alone. The startup line says when each warmup runs once alone before its copies. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 858 R3.3: to_device_nonblocking says an already-pinned source is read in place The docstring said the caller need not keep the source, but an already-pinned CPU tensor is copied from in place: writing to it before the copy lands changes what arrives. Say so, and point a reused buffer at PinnedStager. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL parallel CFG: the guided window is its own walk, decided at intake (merceod #860 row 3) On a CFG-parallel deployment every image step ran image_gen_cfg's three branch nodes and combine_cfg. Each node counted the step's branches from its own worker's loop counter, and a branch outside the window sat the step out: a zero-row velocity, a None declare_step, and a host check in combine that the counts agreed. The guard on those counter reads failed every CFG-parallel request ("loop counter reads 0 on a later iteration"): under the Rust graph runtime the branch nodes are dispatched once more after the loop's last iteration with the counter at 0 (MSTAR_RUST_GRAPH=0 passes; that runtime fix is separate). Main never read a counter on these nodes. t(k) decreases with the step, so a request's guided steps are one window [a, b) with one branch count, known at intake. process_prompt resolves it (cfg_window, an internal kwarg) and the conductor runs the request's phases (image_gen_phases): image_gen_head for k < a on LLM alone, image_gen_cfg2/3 for the window on the branch nodes, then image_gen from b (which decodes), or image_decode when the window reaches the last step. A request without guidance runs image_gen and never reaches the branch ranks. One builder makes every denoise walk from (branches, decode), each with its own loop, which get_loop_caps caps at the phase's steps; latents and time_index cross phases as persisted edges, and get_partition_forward_pass_args moves the request on, so every rank follows the conductor's walk. Each phase carries its branch count in step_metadata; only single-GPU image_gen with the window inside it still reads its loop counter (and keeps the seeded-latents guard). Deleted: sat_out_velocity, cfg_branch_skipped, CFG_BRANCH_LABELS, _sits_out, the sat-out forward and prepare_inputs branches (row 12's 6055b60f4/28dc0d77d slim inputs), declare_step's None return, combine's zero-row check and its per-step branch count, and image_gen_cfg. The CFG-parallel configs name the new walks. Tests: the sit-out cases become a phase-schedule test against a copy of the old per-step count (interval, scales, timesteps/shift, both placements) and one test of the conductor's walk sequence. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: stage host index lists straight into pinned memory; tests name the image_gen_cfg2 walk (860 I2.1, I2.3) The ViT cu_seqlens and the LLM text_indexes pass their lists to to_device_nonblocking with dtype=, one pinned allocation each instead of a pageable tensor that is then pinned. No other caller builds a torch.tensor from a host list for it. The image_gen_cfg walk no longer exists; the tests that named it now name image_gen_cfg2. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL test: the conductor walk accumulates persist signals and drops what is unpersisted (860 I2.2) The conductor appends each walk's persist signals and removes the uuids a walk unpersists; the test replaced them per walk, so it passed even if _get_unpersist_tensors stopped returning the fed latents and time_index. It now accumulates like the conductor and asserts every fed uuid is unpersisted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: name a window's CFG walk from CFG_WALKS; IMAGE_GEN_LOOP is test-local (860 I2.4) image_gen_phases looked the CFG walk up by formatting its branch count into a name; it now reads CFG_WALK_BY_BRANCHES, the inverse of CFG_WALKS, so the walk names live in DENOISE_WALKS alone. IMAGE_GEN_LOOP had no production reader and moves into the one test that uses it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL CFG-parallel: publish only the KV labels the request's CFG walk reads (860 I2.5) The window is known at intake, so the main LLM's publish labels follow the request's phases on a CFG-parallel deployment: cfg_text for a 2-branch window, cfg_text and cfg_img for 3, and nothing for a guided request whose window has no step. A t2i default request no longer exports cfg_img at prefill and after a think decode, which no node of its walks reads. The labels come from the conductor's request config, so every rank of the LLM instance publishes the same set; the KV layout, the consumers' needed labels and the single-node and split-LLM deployments are unchanged (a remote LLM's single-node walks read every branch label). The conductor-walk test's intake moves into a helper that the new test shares. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * graph runtime (Rust): drop 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 peers) loops back through them. On its last iteration Python drops every loop-back (name, next_node) pair whatever the dest. Rust compiled only the pairs whose dest was in its own worker graph and dropped only Dest::Local edges, so it still sent latents/time_index to the peers (and to LLM, a sibling worker graph on the same worker). Each branch then ran one more pass with its loop counter back at 0. Compile every loop-back pair as (signal, dest node name), one list for the exit filter and speculation, and drop edges by (name, dest_sym) when the loop finishes, as Python does. Delete Completed.filtered and register_loop_finish's return value: neither was read, and both returned only the local pairs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * code predictor: no per-position recompile of cp_attend (GPU recompile count) A captured forward compiles with dynamic=False, so every compiled frame that holds a depth step's cache_pos (or passes it to a call at a graph break, partials included) guards on it and recompiles at each of the 16 depth positions and each batch bucket. Before run_layers this could not happen: the launcher calls sat in forward_depth_unrolled's own layer loop, and Dynamo runs a frame with a graph break inside a loop eagerly, so only the norm and MLP frames it called were compiled. With run_layers the breaks landed in loop-free frames that hold the ints: cp_attend and its resume frame, DecoderLayer.forward_residual (through the partial's args) and forward_depth_unrolled (rebuilding the partial at the break). Every frame holding them now runs eager: - cp_attend is torch.compiler.disable'd; its qkv/o projections are cuBLAS GEMMs at -O1 compiled or not. - run_layers calls attend in its own loop rather than through DecoderLayer.forward_residual (deleted; it had no other caller), so an opaque attend sends the loop to eager as the old depth loop did. The ops and their order are unchanged for every caller. - forward_depth_unrolled (Qwen3-TTS, Qwen3-Omni) skips its own frame with torch.compiler.disable(recursive=False). Its callers already broke at it. Only the norm and MLP frames compile, none keyed on a position. The test compiles the talkers' depth-loop shape over 2 and 4 positions for both code predictors and requires the same frame count; at 3abebd178 it is 11 vs 19 (Qwen3-TTS) and 17 vs 31 (Qwen3-Omni). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * run_layers: one compiled frame per layer besides attention run_layers runs eager (its loop breaks at the layer cursor), so every compiled frame it calls is dispatched on each uncaptured forward. Calling the input norm, the post norm and the MLP separately made that three frames per layer, where DecoderLayer.forward_residual and its resume frame had made two. DecoderLayer.forward_mlp_residual runs the post-attention add+norm, the MLP and the next norm's add+norm (the next layer's input norm, or the final norm) as one frame holding no layer index. The ops, their order, the fused add+norm placement and comm_group are unchanged; the first layer's input norm runs once before the loop. The test counts the compiled graphs one forward of a Qwen3 dense stack runs: each extra layer may add attention's graphs plus one. At d55935f03 it adds 5 (attention's 2 plus input norm, post norm, MLP). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: _attend runs eager, so no frame guards on the layer index run_layers calls Qwen3_5LanguageModel._attend from its eager loop, so _attend was a top-level compiled frame. It reads _cursor_attr[i] and _resource_idx[i] before the graph break at set_layer_idx, which guards it on i: one empty-graph compile per layer per compile wrapper, counted as recompiles and toward MSTAR_RECOMPILE_LIMIT. It is torch.compiler.disable(recursive=False) now, so the mixer it calls still compiles as its own frame. No other attend binds i in a compiled frame: the default path sets the cursor in run_layers' eager loop, and cp_attend is disabled. The depth-loop test becomes test_layer_loop_compiles_no_frame_per_step, parametrized over Qwen3-TTS, Qwen3-Omni (now including its opening two-token pass at cache_pos 0) and Qwen3.5's layer loop with stub mixers. It counts the frames Dynamo compiles rather than graphs: _attend's per-layer compiles have empty graphs, so a graph count is equal at d55935f03 (5 vs 5) while compiles are 11 vs 9 there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * cp_attend docstring: the cuBLAS claim holds up to -O1 At -O2 max-autotune may pick a Triton GEMM template for the compiled projections, so eager is not equivalent there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the per-layer frame bound compares the two-layer difference whole (D5) (calls4 - calls2) // 2 rounded an odd difference down, so one stray frame per two layers could pass. Compare the difference with twice the bound. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: run_layers' compile tests live with the Qwen3 LM, on one counter (D4) The Qwen3.5 layer-loop case moves from the Qwen3-TTS module into test_qwen3_lm.py beside the per-layer frame test, and both modules count through one dynamo_counts fixture in conftest (frames compiled, and graphs one compiled call runs) instead of a helper each. The code-predictor depth-loop cases stay with their builders in the TTS module. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * run_layers: an empty layer stack is the final norm again (D3) Starting the stream at layers[0]'s input norm raised IndexError on an empty stack, which had returned the final norm before 76ba991e2. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * code predictor: cp_attend's projections compile again (D1) cp_attend was torch.compiler.disable'd whole, so its one-row qkv and o projections ran on eager cuBLAS. At the default -O1 (coordinate descent) Inductor rewrites an M==1 mm into a fused Triton reduction, and Qwen3-TTS measured 25-30% slower at -O0 at c=1 (docs/serving.rst). It is disable(recursive=False) now: its own frame, which binds cache_pos and layer_idx, stays eager, and the projections compile as int-free frames of their own. The launchers stay disabled. The docstrings said the projections were cuBLAS GEMMs either way; they now say what compiles. The new test requires both projections' weights in a compiled graph of the depth loop, for both code predictors; it fails at 9f26410bd. The per-position compile test still passes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M1: copy an already-pinned host tensor to the device in place The merge of main took #904's pageable path for every CPU tensor, so a source that is already pinned (NemoLogMel.batch pads into pinned memory) paid a second pinned allocation and a full-batch memcpy per call. Read a pinned source in place again, and restore the docstring caveat that the caller must not write to it until the copy lands. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M2: size video decode threads from usable_cpus, apart from the torch pool video_decode_threads re-read CPU affinity beside main's usable_cpus(), which also reads the cgroup quota, and was capped by torch's intra-op pool, which main #903 shrank to a quarter of the CPUs (at most 4). That cut V-JEPA2 and Cosmos3 decodes from 8 FFmpeg threads to 4, or 1 on a small host. FFmpeg's frame threads are not torch intra-op threads, so decode is now min(8, usable_cpus()). The docs paragraph that said both that the quota is read and that it is not now says one thing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M3: test GLM-5.2's stop id at prefill The merge ported GLM-5.2's prefill stop rules onto TokenStop without a test: a stop id sampled on the prefill walk ends the reply at the first decode step, whose token postprocess drops, and both stop checks agree; with ignore_eos the prefill sets nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the projection-compile test counts through dynamo_counts (E2) It had its own backend, reset and torch.compile beside the fixture. The fixture's DynamoCounts gains traced, the data_ptrs of the tensors its compiled graphs take, and the test reads that. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M4: declare a prefill-sampled first token on TokenStop GLM-5.2 and Command A+ each hand-built the same rule beside TokenStop: a stop id sampled by a prefill walk has no loop to stop, so the reply ends at the loop's first step, whose token postprocess drops, and max_tokens counts the prefill's token. TokenStop.prefill_walks now declares those walks; the base keeps the flag in request_state and applies both rules in check_stop and check_stop_batched, and ended_at_prefill() tells postprocess to drop the closing token. GLM-5.2's _token_stops / check_stop overrides and Command A+'s _ended_at_prefill set, check_stop and cleanup_request are deleted. Command A+ now sets the flag for max_tokens 1 too (no decode step follows, and removal frees it) and keeps it until removal instead of clearing it on the stopping step, as GLM-5.2 did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * M1: model the mocked host tensor in the pin test as pageable test_to_device_nonblocking_pins_where_the_device_capability_says builds its host from Mock(spec=torch.Tensor), whose is_pinned() is truthy, so it now took the already-pinned path. It is about the pageable source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N1: copy a pinned source in place only when it is contiguous .to() of a pinned non-contiguous tensor packs it into a pageable temp first, so the copy blocks the caller. Such a source goes through the pinned buffer like a pageable one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N2: read usable_cpus once per process for video decode threads video_decode_threads re-read the affinity and the cgroup files on every decode. Cache it, as image_io.PNG_WORKERS does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N3: say a generated_key counter must count the prefill token itself prefill_walks adds the prefill's token to the loop count only when TokenStop counts loop iterations; a submodule's own generated_key counter is read as is. No model declares both today. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N4: the engine drops the closing token of a reply that ended at prefill Declaring TokenStop.prefill_walks still made each model's postprocess delete token_stop.buffer when ended_at_prefill(). Engine._postprocess_batch now does it after postprocess, the one place every forward path's rows are postprocessed; a submodule without prefill_walks pays one getattr per batch. GLM-5.2's and Command A+'s own deletes go. Tests that drove postprocess by hand (GLM-5.2's engine-cycle driver and both models' prefill-stop cases) now call Engine._postprocess_batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * N4: say prefill_walks needs the loop to emit the token as buffer The engine drops token_stop.buffer on the closing step, which removes the token only where the loop's client edge carries that output (GLM-5.2, Command A+). Qwen3.5's decode emits text_inputs, which this would miss. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 / run_layers: capture works again (fail_on_recompile) Every capture runs its forward under torch.compiler.set_stance( "fail_on_recompile") (capture_into_graph). torch.compiler.disable( recursive=False) skips only its wrapper's code; the decorated function's own frame never gets a cache entry, so Dynamo's callback runs on every call, and convert_frame lets it run eager only because its caller is the wrapper. The capture stance's fail_callback has no such check: it raises on the first warm call, whatever the shapes. So every bucket of a forward that reaches one fails to capture: Qwen3.5 through _attend (1438ea127), Qwen3-TTS and Qwen3-Omni through forward_depth_unrolled (d55935f03) and cp_attend (a9bf02141). The undecorated _attend of 0fe8c812d captured because it was compiled at warm-up, one entry per layer index. No frame is recursive=False now. Each per-layer or per-step int is read only in fully torch.compiler.disable'd cursor methods, called from loops that already run eager, and the attention runs as its own int-free module frame: - run_layers' hook is set_cursor(i, self_attn), which replaces the default self_attn.attend.set_layer_idx(i); the loop then calls self_attn(h) on both paths. Qwen3.5's is _set_cursor. - cp_attend is gone. A code predictor's layers are DepthAttention (a ParallelAttention: qkv and o projections, compiled) over the predictor's DepthKV, a cursor holding the frame-local KV cache and the depth step (bind_step, set_layer_idx, and the two kernel launches). run_layers moves it per layer through the default path. - forward_depth_unrolled takes only the input embeddings; the depth loops in the Talkers bind the step on cp.depth first, so cache_pos never enters a compiled frame. dynamo_counts now makes its second, warm call under the capture stance, as the runner warms a bucket and then captures it. With only that change on da5514d26, the Qwen3.5 per-layer test and both code-predictor tests (per-position, projections compile; TTS and Omni) fail with the same "Detected recompile ... '_attend' / 'forward_depth_unrolled'" errors the GPU runs logged; they pass here, with the per-layer, per-position and projection checks unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: a DepthKV binding serves one pass over the layers DepthKV kept the previous step's cache, positions and cache_pos, so a depth step whose loop skipped bind_step wrote and attended over the previous step's slots without a sign. The last layer's call now drops the binding after its launches, inside the eager __call__ (no new graph break; a capture has already recorded the cache's pointers, and the caller owns its storage), and a call with none bound raises ValueError. The leading dim of the cache is the layer count, as both Talkers allocate it. The bind_step argument shapes move from Qwen3-Omni's forward_depth_unrolled docstring into DepthKV.bind_step's; the stale "Position IDs come in directly" bullet goes, the Omni class docstring calls depth the cursor it is, and the recursive=False explanation lives only in run_layers. Tests: a second forward_depth_unrolled without rebinding raises (fails at 16d62450f: DID NOT RAISE). The node's bind_node_resources leaves each DepthAttention on the predictor's DepthKV (fails with the empty bind_resources override removed). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: the first layer consumes a DepthKV binding The guard no longer depends on the cache layout. bind_step marks the binding fresh and set_layer_idx(0), the start of a pass, raises ValueError unless it is, then consumes it; the last-layer clear and the None check in __call__ go. bind_step's docstring says the cache is not cleared between frames (a depth step reads only what its frame wrote) and that the submodule allocates it. The code-predictor tests share _cpu_predictor (config, bs, calls) for their identity-norm, fake-kernel setup. The bound-once test still fails at 16d62450f (DID NOT RAISE). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Rust loop exit: test the sibling-graph path and the routed edge directly The split-loop test now also runs with step in a sibling worker graph on this worker, and asserts the last iteration routes no x to step: not staged for a peer's read, and not ingested into the sibling graph. Correct the LoopSpec / LoopArg comments, which said both lists are kept verbatim, and inline the single-use resolve closure. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * LingBot-World: call the ragged resource through resource.run, as the scaffold does Main's #926 added LingBot with RaggedAttentionCallable, which this PR deletes in favour of functools.partial(resource.run, label=...); the module failed to import (test_lingbot_world_shell collection, lingbot startup validation). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test_subgroup_creation_order: the stub worker graph carries replica_ranks Main's #705 test stubs a WorkerGraph; this PR's GlobalParallelConfig reads WorkerGraph.replica_ranks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test_withheld_outputs: the fake check_stop_batched drops step_walk, as the base did Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test_row_output_store: a withheld output stays out of the row store Both Talkers now withhold their EOS frame through the row-store fast path; pin that it reads the popped per-row outputs, not the row views. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: the Cosmos3 two-GPU numbers name their hardware and predate main's Edge changes Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStop: leave the token budget to the loop cap TokenStop's stop check re-checked max_tokens, which every adopter's get_loop_caps already enforces: the runtime caps the decode loop at the same iteration (Python and Rust both stop the loop and refuse to speculate past max_iters), and the conductor derives the cap and request_info.max_tokens from one max_output_tokens, also when it cuts that to fit the context. So the budget is checked once, by the cap. generated_key existed only to count that budget and goes with it. Tests that pinned the stop check ending a row at its budget now expect the row to go on; the cap-vs-stop-check comparisons for the TokenStop adopters are deleted (there is no second count to compare), and the two hand-driven loops apply the request's cap as the worker does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStop: withhold a prefill-ended reply's closing token A reply whose prefill sampled a stop id closes its decode loop on the loop's first step, whose token comes after the reply's last. The engine dropped that token after postprocess by probing the submodule for `token_stop`: a model concept in engine code, and three routes missed it (a rows-only step skips the per-request postprocess, the inline client emit reads the host rows, and a row left empty by the pop was skipped by the stop check, stalling it). The closing step now says so through the existing primitive: check_stop returns StopCheck({loop}, withheld=(buffer,)) and check_stop_batched a BatchedStopCheck. The engine's pop and its `getattr` are gone, and the flag is read on the main thread after prefill's check_stop. The worker's inline emit honours withheld outputs per row: such a row gets None, which both graph runtimes (Python and Rust) route as that row's tensors, of which it has none. That also closes main's gap where a withheld output still rode inline. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Warmup docs: the solo run happens only when replicas are uneven docs/serving.rst said every run of a warmup entry past one copy starts with one request alone; _run_warmup_requests sends that solo request only when some worker graph has fewer replicas than the most, once per entry. Say that, and that otherwise bs=1 warms only when concurrency includes 1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * video_decode_threads: the cap of 8 is per decode The docstring said the bound keeps concurrent decodes from oversubscribing; each decode starts its own threads, so it doesn't. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ragged attention: name each manager uniquely in the op registry The mstar::ragged_attention op found its manager by resource key in a process-wide registry, so two live engines with the same key (tests, harnesses) handed one engine's compiled graph the other's plans. The op key is now the resource key plus a number unique in the process. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Warmup: a replica count that fails leaves warmup at one copy _data_parallel_replicas ran outside the warmup's try, so a config the worker graphs can't be built from (or no worker graphs: max([])) stopped startup, against the promise that a failing warmup is logged and skipped. It now logs and returns DataParallelReplicas(1, 1). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * one_row_per_step docs: what declare_step receives Under the flag the region's declare_step gets the first request alone with every request's seq_lens, so max_segments_per_request is a per-step ceiling there. Say so where the flag is documented. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: MSTAR_PY_SWITCH_INTERVAL_SEC gets its worker-scheduling row The worker reads MSTAR_PY_SWITCH_INTERVAL_SEC at the top of Worker.run, but docs/environment_variables.rst had no row for it. It gets one under Worker scheduling, naming Worker.run as its reader. The variable applies to workers only: the API server's interval stays the serving config's api_switch_interval_s, whose measurements covered the API server alone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Comments: the decode budget is the loop cap, not check_stop Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: a prefill_walks token counts toward the budget, so the loop cap is one fewer Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: say which MSTAR_PY_SWITCH_INTERVAL_SEC values the worker ignores The worker passes the value to sys.setswitchinterval, which refuses only values of 0 or below, so `nan` and `inf` are applied rather than ignored with a warning, as the row implied. The row now names what is ignored: a value float() cannot parse, or one of 0 or below. _env_number's docstring cited this variable as raising a too-small value to a minimum; it ignores one instead, so only MSTAR_RECOMPILE_LIMIT is left as the precedent. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: point api_switch_interval_s to the workers' switch-interval variable The key sets the API server's GIL switch interval only. The worker processes read theirs from MSTAR_PY_SWITCH_INTERVAL_SEC, whose row already points back to this key; the serving config row now says so too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ragged attention: update a test docstring for unique op keys test_a_compiled_region_attends_through_a_rebuilt_resource still said the op names its manager by resource key, with a rebuilt resource taking over that key. Each manager now registers under its own op key, the resource key plus a number unique in the process. The docstring now says what the test checks: a region compiled against one manager, then called with a manager rebuilt under the same resource key, attends through the rebuilt one's plan while the old one is still registered. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * PiecewisePackedConfig: attn_key is resolved when the region declares, not by rewriting declare_step attn_key used to overwrite the config's public declare_step field with a bound method in __post_init__. A dataclasses.replace of such a config then passed both fields and raised, and a shallow copy kept declaring through the original's method. The runner now asks the config for the region's step (region_step), which answers from attn_key or declare_step on the instance it is called on. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: a forward gets a lease_before_step region's runner only on a step that leased it A region that takes its slot before the outer declaration must replay only on a step that leased it: unleased, the step's own declare_step covers the region's resources, and a replay would declare, plan and commit them a second time. Each forward used to enforce this itself, against ModelInputsFromEngine.piecewise_leases (Qwen3.5 vision, Qwen3-TTS, V-JEPA2, Omnilingual, OmniVoice) or against engine_inputs.step being None (Qwen3-ASR, nemotron duplex). The engine now leaves an unleased region out of the runners it hands the forward (SubmoduleManagement.runners_for), so a forward replays whenever its runner is there. The piecewise_leases field on ModelInputsFromEngine and the per-model checks and their tests go. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Device loop-back is its own declaration, DeviceLoopback, not a TokenStop field TokenStop.loopback_signal bundled two concerns into the stop policy: which token ends the loop, and how a decode walk reads its token back on the device. Declaring it also made the walk's step input-free, which silently turned on pre-plan reuse for Whisper, Qwen3-ASR and Higgs decode, whose measured numbers predate that. Loop-back is now ARNodeSubmodule.device_loopback, a DeviceLoopback(walk, signal, sampler, buffer); the loop-back hooks (uniform_row_inputs, device_loopback_signals, inline_client_signals) follow from it, and one base read, loopback_inputs, replaces the four copies of the slot-master read and the routed-row overwrite (Whisper passes its timestamp rule state as a declared loop-back state). step_is_input_free is again each node's own answer: Qwen3.5 decode says so, the others keep the default. A walk is now a required field, so the engine's check for a missing one goes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Stop check: the step closing a reply that ended at prefill withholds the loop-back signal too When a prefill walk samples a stop id, the loop's first step only closes the reply, and its outputs are withheld. It withheld the token's buffer (new_token), but a device loop-back node's client edge carries its loop-back signal (text_inputs), inline off that buffer: the worker nulls an inline value only when the signal itself is withheld, so the token sampled after the stop id reached the client. On the tensor path the routed signal leaked the same way. The closing step now withholds the node's DeviceLoopback signal beside the buffer, so prefill_walks and device loop-back compose. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chatterbox T3, Higgs Audio, Qwen3-ASR: a stop token sampled at prefill ends the reply These prefills sample the reply's first token, but their TokenStop declared no prefill_walks. A stop id sampled there was either ignored (Higgs, Qwen3-ASR: only a decode row could stop) or a stop for a loop not yet entered, which the runtime drops (T3), so decode ran on to the cap. They now list their prefill walks: the loop's first step closes the reply and its outputs are withheld. As TokenStop says, that first token counts toward the budget, so Higgs and Qwen3-ASR now give the loop max_output_tokens - 1 and a capped reply is max_output_tokens long, not one more (T3 already did). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: a captured decode step stages no rotary tables, whatever its rows bring A captured decode builds its rotary tables inside the graph from the planned positions, so preprocess skipped them, but only on the loop-back branch. A step whose rows all brought their token routed (each request's first decode step, or every step with MSTAR_DEVICE_LOOPBACK=0) still built the 3D positions and cos/sin, and the capture held two static buffers for them, all unread. preprocess now returns before them on any captured decode step, and the capture test reads engine_inputs.captured instead of probing the step's context for a slot lease. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * mstar::fused_add_rms_norm: the CUDA kernel takes FlashInfer's fused path only when x, residual and weight share a dtype FlashInfer's fused add+RMSNorm reads the residual and the weight in x's dtype, but the kernel branched on x's dtype alone. mstar::rms_norm casts its weight; this op, which run_layers hands every decoder stack, did not, so a model whose norm weight or residual stream was in another dtype would have its buffers read as the wrong type. Any mix now takes the unfused add and norm, which casts the weight. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * One cuDNN benchmark scope, beside the engine's precision scope The engine's private _cudnn_benchmark, added for capture_cudnn_benchmark, was a line-for-line copy of wan_vae.cudnn_benchmark. The helper now lives once, in mstar.engine.fp32_precision next to fp32_precision_scope, the other process-wide backend flag the engine sets around a node's warm-up and capture. wan_vae imports it from there, so Cosmos3's VAE decode keeps calling it through wan_vae. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: the audio encoder is the AuT node Qwen3-ASR runs, packed and captured Qwen3-Omni's audio encoder ran transformers' Qwen3OmniMoeAudioEncoder eagerly, about 10 ms a call and mostly launches, with its own step splitter around it. That is the same AuT architecture, on the same thinker.audio_tower weights, that Qwen3-ASR already runs natively: every attention window of every clip is a segment of a ragged attention resource, and the 32-layer stack replays a CUDA graph by packed token count. The audio node is now that node (Qwen3ASREncoderSubmodule), built without a spectrogram config so it takes the API server's log-mel features as they arrive, over an AuTEncoder loaded from the checkpoint and a RaggedAttentionSpec sized for the longest clip (MAX_AUDIO_SECONDS). Omni's audio config is AuTEncoderConfig. The HF wrapper, its frame-budget splitter, its return_dict probe and their tests go; the host-sync test now checks the node Omni runs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ragged attention: a one-row-per-step region's ceiling is max_segments_per_step A region that replays a whole step as one row (one_row_per_step) captures at batch size 1, so its plan held max_segments_per_request segments for every request of the step: Qwen3.5's vision tower and Omnilingual set a "per request" ceiling that meant per step, Omnilingual's to its max batch size. RaggedAttentionConfig.max_segments_per_step now names that ceiling (and a deployment can override it like the others); when set, it sizes every captured plan of the resource. Both adopters declare and read it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: the vision node boots on the image processor alone The vision node's pixel normalizer loaded the video processor at worker boot, image-only deployments included, to check that the two processors normalize alike. A checkpoint with no video_preprocessor_config.json gets the video processor class's defaults, so one whose image normalization differs from those failed to boot even when it would never see a video. The check now runs where the video processor is first loaded, on the first video: the normalizer reads the image processor only, and such a checkpoint serves images and refuses videos. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the refused video processor is read into a name (ruff B018) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision: the compile regions warm at start-up; the docs say when model code still compiles The vision encoder's compile_regions declared no example calls, so a deployment without warmup_requests compiled its patch embed and merger on the first image request. They now warm on the smallest image the processor makes, in the dtypes a step brings, so the dynamic graphs exist before serving. The docs said not to call torch.compile in model code, which several models still do for compile options compile_regions does not take (static shapes, a mode, fullgraph); they now say to prefer compile_regions and when that does not fit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: Qwen3-Omni's audio encoder loads no transformers config to check The audio encoder is now built from mstar's own reading of config.json (AuTEncoderConfig), with no AutoConfig call that could run checkpoint code, so the remote-code test checks the vision encoder's loader alone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kokoro: say that cuDNN benchmarking covers both captured halves The engine turns cudnn.benchmark on around every capture warm-up of the node, so cuDNN times the convolutions of the text half's buckets as well as the frame half's. The comment on capture_cudnn_benchmark and the models doc named only the frame half. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kokoro: put the fused BiLSTM in the module tree MaskedBiLSTM.fuse kept the fused LSTM in a plain list, outside the module tree, beside the two unidirectional LSTMs it was built from. A later .to(), .half() or state-dict load then reached only the unused copies, and the fused weights that run stayed behind on the old device or dtype. fuse now registers the fused LSTM as a submodule and deletes fwd and bwd, so one copy of the weights exists and every module-tree operation moves it. The checkpoint still loads into fwd and bwd, before fusing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Split speech configs: run their two workers under MPS The Chatterbox, OmniVoice and Qwen3-TTS 0.6B split configs put two workers on one GPU, and their throughput numbers were measured with the workers' kernels overlapping under CUDA MPS: without it the Chatterbox split measured the same as one worker. The YAMLs left starting the daemon to the user. Set `mps: true`, which has the conductor run a private MPS daemon for the deployment, as the ASR split configs do, and point the models doc at it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * One residual code predictor for Qwen3-Omni and Qwen3-TTS Both Talkers carried their own copy of the code predictor: the layer, inner-model and predictor classes, a depth-loop driver, and a frame-local KV scratch allocated in each submodule. The two drivers had drifted. Qwen3-Omni's opened the depth sequence with one two-token pass, sliced a constant position table and fused the embedding lookup, sum and code write into one kernel; Qwen3-TTS still ran a separate opening pass, filled a position tensor per depth step and summed the embeddings with three kernels a step. ResidualCodePredictor, beside the kernels they already shared, now holds the layers, the stacked heads, the input projection (the identity on Qwen3-Omni and the 0.6B) and one depth loop that takes a sampling callable. Both Talkers build it and call its depth loop. Its DepthKV owns the scratch, sized once from the node's largest batch, and the prologue kernel takes each token's RoPE position from its cache slot, which is what both callers passed, so no position tensors remain. For Qwen3-TTS this removes one pass over the layers and about 45 small kernels per frame. Its embedding sum is now accumulated in fp32 and rounded once, as Qwen3-Omni's was, instead of rounded after every add, so greedy codes can differ at near-ties. Qwen3-Omni's numerics are unchanged: its sum is rounded to the model dtype inside the step instead of when the next step reads it. The wrapper of the embed-accumulate kernel now also asserts contiguous tokens, since the loop takes any sampling callable, and the loop copies the first embedding even when it is already fp32. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS test: name the row-output test…

Gaurav-Shah05 added 9 commits October 22, 2026 03:34

engine: probe a walk whose submodule cuts its own tensor_inputs

4419d97

conductor: cache off when a keyed walk streams to a running partition

199d9ea

qwen3_omni: thinker cuts its inputs and places a cut walk's positions

c8d15e9

engine: an item encoder may opt out of its dynamic=True compile

b8985ab

qwen3_omni: run the audio encoder eager as an item encoder

0970ba5

qwen3_omni: the vision encoder compiles one graph for every image size

7757d3e

qwen3_omni: the thinker keys its prompt, images, audio and replies

0c6cabd

test: prefix caching for the thinker

fd22735

comment pass

815b345

mstar-review Bot added area/conductor area/engine area/worker model/qwen3_omni size/M labels Oct 22, 2026

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

test pass

644b3da

dghosef self-assigned this Oct 23, 2026

Gaurav-Shah05 added 10 commits October 24, 2026 13:57

qwen3_omni: build the audio encoder's chunking before the compiled call

a738c43

engine: drop the item encoder's eager opt-out

d814ac2

test: Qwen3-Omni's encoders reuse their graphs for sizes in range

635bb0e

Merge remote-tracking branch 'origin/main' into prefix-cache-omni

1a21c9f

worker: find streamed partitions from keyed labels

40df852

qwen3_omni: drop the keyed stream that kv_labels replaces

b4014c9

test: the prefix tests declare kv labels and layouts

1621ed5

qwen3_omni: a prompt with a video keys its leading text

a9d75f1

qwen3_omni: the layout takes an image's span from image_span

027cf1e

qwen3_omni: a video prompt keys one span, with the reason

975c7ca

dghosef and others added 4 commits October 26, 2026 21:20

qwen3_omni: mark the vision encoder's cu_seqlens dynamic

1937be9

It has one entry per frame plus one, so a video's frame count changes its shape; unmarked, the first video recompiles the encoder. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

qwen3_omni: size the vision encoder's max_seqlen to one frame

863e0da

Each video frame is its own attention segment, so the whole clip's patch count overstated it ~t times and flash-attn launched ~t times the blocks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

docs: Qwen3-Omni's Thinker counts prefix-cache hits too

00b6fe0

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

prefix-caching skill: tensor inputs a split_inputs override cuts are …

02381ed

…eligible Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dghosef reviewed Oct 27, 2026

View reviewed changes

Gaurav-Shah05 added 5 commits October 27, 2026 14:07

engine: a submodule that cuts its tensor_inputs says so with a flag

fada3f7

qwen3_omni: a request with audio output turns the prefix cache off

fdd6c43

process_prompt already has the switch, so the conductor no longer works out which partitions a keyed node streams into.

model: Qwen3.5 and Qwen3-Omni share their 3D position ids

9f6c6b7

test: the tensor_inputs flag, the audio-output opt-out and shared 3D …

de5fd16

…positions

Merge remote-tracking branch 'origin/main' into prefix-cache-omni

f38c386

mstar-review Bot added area/model-shared model/qwen3_5 and removed area/worker labels Oct 27, 2026

Gaurav-Shah05 added 2 commits October 27, 2026 14:45

conductor: cache off for a request that runs a partition a keyed node…

4722f7c

… streams into The rule comes from the graph, so a model need not opt out in process_prompt, and it lives beside the load check that lets such a stream load.

test: a request whose talker runs is not cached

65a1f31

dghosef approved these changes Oct 27, 2026

View reviewed changes

Gaurav-Shah05 merged commit 07d23c5 into main Oct 27, 2026
6 checks passed

Gaurav-Shah05 deleted the prefix-cache-omni branch October 27, 2026 21:21

vasilevklart mentioned this pull request Oct 27, 2026

test, ci: restore torch's flag proxies by assignment (pytest 9.2), and pin pytest #917

Merged

dghosef mentioned this pull request Oct 28, 2026

Performance 2/7: speech and language models #859

Merged

merceod mentioned this pull request Nov 1, 2026

Qwen3-Omni: one Thinker prefill walk and one Talker prefill step #1093

Draft

Navigation Menu

Repository navigation

qwen3_omni: prefix caching for the Thinker - #836

qwen3_omni: prefix caching for the Thinker#836
Gaurav-Shah05 merged 32 commits into
mainfrom
prefix-cache-omni

Gaurav-Shah05 commented Oct 22, 2026 •

edited

Loading

Uh oh!

mstar-review Bot commented Oct 22, 2026

Uh oh!

dghosef commented Oct 23, 2026

Uh oh!

dghosef commented Oct 26, 2026

Uh oh!

dghosef left a comment •

edited

Loading

Uh oh!

dghosef Oct 27, 2026 •

edited

Loading

Uh oh!

dghosef Oct 27, 2026

Uh oh!

dghosef Oct 27, 2026

Uh oh!

Gaurav-Shah05 Oct 27, 2026

Uh oh!

dghosef Oct 27, 2026

Uh oh!

dghosef commented Oct 27, 2026

Uh oh!

Gaurav-Shah05 commented Oct 27, 2026

Uh oh!

dghosef left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Gaurav-Shah05 commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Cost

Next (out of scope here)

Testing

Uh oh!

mstar-review Bot commented Oct 22, 2026

Automated review

Description

Uh oh!

dghosef commented Oct 23, 2026

Uh oh!

dghosef commented Oct 26, 2026

Uh oh!

dghosef left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dghosef Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

dghosef Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

dghosef Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

dghosef commented Oct 27, 2026

Uh oh!

Gaurav-Shah05 commented Oct 27, 2026

Uh oh!

dghosef left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Gaurav-Shah05 commented Oct 22, 2026 •

edited

Loading

dghosef left a comment •

edited

Loading

dghosef Oct 27, 2026 •

edited

Loading