Repository navigation
Wan2.2 step: one-pass adaLN per frame, fused q/k norm+RoPE, no host syncs; Cosmos3: keep the dense attention plan - #963
Merged
Merged
Conversation
(cherry picked from commit bada56b)
Every token of a latent frame shares its timestep, so the time embedding and the six modulation vectors are computed once per frame (bit-identical rows) instead of per token. The norm+modulation+gated-residual steps go through components/diffusion/modulation.py, which runs a single-pass Triton kernel inside the compiled DiT and the plain eager ops otherwise.
The slice-assign rope made Inductor write eight fp32 intermediates per block. Eager now uses the shared stacked rotation (same bits), and the compiled DiT runs q and k through one kernel.
The step index comes from the loop counter, the UniPC tables go up through pinned memory, and the corrector solve skips its host-side info check, so the worker queues each step behind the previous one.
mstar-review
Bot
added
area/model-shared
area/resources
model/cosmos3
model/wan22
size/M
labels
Automated review🟡 No findings on Description
What the review checked
Advisory, not a merge gate. Rules: |
Gaurav-Shah05
requested changes
Oct 30, 2026
Gaurav-Shah05
left a comment
Gaurav-Shah05
left a comment
Collaborator
There was a problem hiding this comment.
Wan's step is 1.10x faster on H100 too, but on H100 the per-frame time embedding breaks the port's bitwise test, and the sync-free step has no test that guards it.
I ran Wan22 at 832x480x33 (20 steps) and Cosmos3-Edge on its stock config at 832x480x121 on one H100, two servers per arm, alternated against the branch's base.
Apart from the stock row, the Wan rows add compile_dit_sizes: ["832x480x33"] to configs/wan22.yaml, since the stock config compiles only 1280x704x81 and runs this size eager, without either kernel.
| Base | #963 | Δ | |
|---|---|---|---|
| Wan c=1 DiT step | 137.7 ms | 124.8 ms | 1.10x ↓ |
| Wan c=1 DiT step, stock config | 243.4 ms | 220.3 ms | 1.10x ↓ |
| Wan c=1 p50 | 3.95 s | 3.68 s | 1.07x ↓ |
| Wan c=2 | 16.6 videos/min | 17.8 videos/min | 1.08x ↑ |
| Cosmos3-Edge c=1 DiT step | 225.0 ms | 223.9 ms | 1.005x ↓ |
Cosmos3-Edge's videos are identical to the base's. Wan's are not:
- on the stock config they sit at 25 dB PSNR from the base's, against 43 dB between two base servers.
- That is a different trajectory, not a worse one, since the base's own compiled and eager outputs are 21 dB apart.
Direction for a later PR. Your follow-up moves Wan's UniPC port into components/diffusion for Cosmos3's four generation paths.
DenoiseLoopSubmodulealready takes the step from the loop counter and stages each request'ssigmasandtimestepson the device once.- It skips the extra step the async worker can schedule past a request's last one, and batches requests of the same size into one forward.
- With UniPC as a scheduler
DenoiseLoopSubmodulecarries, Wan and Cosmos3 would get all of that, and Wan at c=2 would run one forward per step instead of two. What do you think?
| batch, frames = timestep.shape | ||
| sincos = _sinusoidal_timestep_embedding(timestep.flatten(), self.freq_dim) | ||
| temb = self.time_embedder(sincos.unflatten(0, (batch, seq_len))) | ||
| temb = self.time_embedder(sincos.unflatten(0, (batch, frames))) |
Collaborator
There was a problem hiding this comment.
On H100 test_wan22_native_dit_equivalence.py fails at this embedder (temb differs by up to 9.8e-4, and the base passes), since cuBLAS picks a different kernel for the second Linear at 18 rows than at 7020. The bit-for-bit claims in the module docstring and in test_sequential_cfg_trajectory_is_bitexact_vs_oracle should go, and the native test should hold per-frame rows to per-token ones within that rounding, with the blocks checked bitwise on the same temb and modulation rows.
Collaborator
Author
There was a problem hiding this comment.
Done, in f2b23ab. The bit-for-bit claims are gone from the module docstring and the sequential test, which is now test_sequential_cfg_trajectory_matches_oracle, same for the no-CFG one. test_wan22_native_dit_equivalence.py holds per-frame rows to per-token rows within that rounding (temb 2e-3, modulation rows 3.2e-2) and checks patchify, all 30 blocks and the head bitwise, with the wrapped module handed the native rows expanded per token.
On this box (B300) 18 and 7020 rows are bitwise under the default numerics at the four timesteps I tried, in four fresh processes, and the row counts I swept from 2 to 8192 agree at t = 937; only M=1 differs. With NVIDIA_TF32_OVERRIDE=0 the two row counts do differ: temb by 3.1e-5, modulation by 3.9e-3. Switching TF32 off in the time embedder alone reproduces your number, 9.77e-4 on temb and 1.56e-2 on the modulation rows, so my guess is that H100 is choosing TF32 at one row count and not the other. I can't check that here, so it's a hypothesis. The native test passes in all three setups, with blocks at max_abs 0, which also exercises the row hand-over for real since the rows differ in the last two.
I also recorded an oracle on the B300, since this PR had never run the reference suite, and that turned up two things:
- Under the suite's numerics main is 0.0 on the sequential and no-CFG trajectories, and this branch is 2.06 and 0.48. So the difference shows up here too, just not under the default numerics.
- With the embedder run at the reference's row count the branch goes back to 0.0 on all of them, so the time embedding is the only departure. I kept that as
test_sequential_cfg_trajectory_is_bitexact_on_per_token_rows, so the suite still has a 0.0 trajectory gate on everything else.
SEQ_TRAJECTORY_MAX_ABS now equals PER_STEP_MAX_ABS (4.5): 2.06 here, 3.65 with temb's fp32 rows perturbed by 1.1e-4, the size of a TF32 kernel change. NO_CFG_TRAJECTORY_MAX_ABS is 1.0 (0.48, 0.76 perturbed). The I2V test in that suite still expected a per-token timestep from this PR's first commit, also fixed. I didn't run any of this on H100, so if the bounds are too tight there, tell me what you see.
| # FlashInfer's plan is a device kernel plus a stream sync. | ||
| wrapper = self._wrapper(slot_key.label, slot_key.slot) | ||
| layout = (tuple(cu_q), tuple(cu_k), causal) | ||
| if self._planned.get(slot_key) != layout: |
Collaborator
There was a problem hiding this comment.
The wrapper holds the schedule, so this check could live in _RaggedPrefillBase._plan. The ragged managers behind FLUX.2-klein, Z-Image and LTX-2.5 replan the same wrappers every step, and a comment in their plan already asks for this cache. Moving the check there is for another PR, but what do you think?
Collaborator
Author
There was a problem hiding this comment.
Yes, I think that's the right home, and it should be its own PR. The wrapper holds the schedule either way, so the check moves down cleanly and FLUX.2-klein, Z-Image and LTX-2.5 get it for free. I'd want their numbers measured on that PR since it changes three more models' planning. I left this one as it is here.
| return int(time_index.reshape(-1)[0].item()) | ||
|
|
||
|
|
||
| def _to_device_without_sync(t: torch.Tensor, device: torch.device) -> torch.Tensor: |
Collaborator
There was a problem hiding this comment.
mstar/utils/h2d.py to_device_nonblocking already does this for any host tensor, and DenoiseLoopSubmodule stages its schedule with it. forward should call it, and this helper should go.
Collaborator
Author
There was a problem hiding this comment.
Done, in 93bc40c. forward calls mstar/utils/h2d.py's to_device_nonblocking and the branch's helper is gone.
| num_steps = int(step_metadata["num_inference_steps"]) | ||
| guidance_scale = float(step_metadata["guidance_scale"]) | ||
| k = int(time_index.item()) | ||
| k = _loop_step(fwd_info, time_index) |
Collaborator
There was a problem hiding this comment.
After this line nothing on the served path reads time_index: prepare_inputs and forward both take k from the loop counter, and the edge and its + 1 are carried only for _loop_step's .item() fallback. The edge and the fallback should go, with the tests passing the counter as test_step_prep_host_syncs.py does, so a missing counter fails instead of putting the sync back.
Collaborator
Author
There was a problem hiding this comment.
Done, in 1af7dfe. The edge is gone from the graph, prepare_inputs, forward and the seed, and _loop_step is a plain fwd_info.dynamic_loop_iter_counts[DENOISE_LOOP_NAME], so a missing counter is a KeyError, with test_wan22_step_prep_fails_without_the_loop_counter on it. The strict read depends on the runtime reporting the counter at iteration 0, so I ran it on a live server on both graph runtimes (MSTAR_RUST_GRAPH auto and 0): 20 steps per request, overshoot at 20 vetoed, no errors. The tests pass the counter the way test_step_prep_host_syncs.py does.
| return torch.linalg.solve(torch.stack(rows), torch.stack(b)) | ||
| # solve_ex: the same factorization as solve, without solve's host-side | ||
| # check of the LAPACK info (a device sync every step) | ||
| return torch.linalg.solve_ex(torch.stack(rows), torch.stack(b)).result |
Collaborator
There was a problem hiding this comment.
With this back to linalg.solve, or time_index.item() or the pageable table copies back in forward, every CPU test still passes. forward should have a CUDA test like test_iteration_zero_stages_its_schedule_and_seeds_without_a_host_sync, run with a stub transformer under torch.cuda.set_sync_debug_mode("error"), which fails on each of the three.
Collaborator
Author
There was a problem hiding this comment.
Done, in 4090bb6: test_step_prep_host_syncs.py::test_wan22_forward_takes_no_host_sync, T2V and I2V, stub transformer, four steps so it covers the predictor alone, both corrector orders and the last step, under set_sync_debug_mode("error"), and it checks the result against a host run. I mutated each one in a scratch copy on GPU 1 and it fails with called a synchronizing CUDA operation on all of them: linalg.solve back, .item() on a device step tensor in forward, a pageable sigmas copy, a pageable timesteps copy.
| assert torch.equal(x_out, plain) and torch.equal(out, plain_mod) | ||
|
|
||
|
|
||
| def test_eager_rms_norm_rope_is_the_modules_bitwise(): |
Collaborator
There was a problem hiding this comment.
Nit: This test builds its expected value from the same two calls rms_norm_rope's eager branch makes, so it can't fail. Comparing against the slice-write rotation removed from dit.py would test that rms_norm_rope's eager branch gives the same bits as before on Wan's tables.
Collaborator
Author
There was a problem hiding this comment.
Done, in 3f49241. It compares rms_norm_rope's eager branch against the slice-write rotation that was in dit.py, on Wan22RoPE3D's tables, on CPU and CUDA. It fails on a flipped rotation and on rotating in bf16 instead of fp32 (mutated in a scratch copy, CPU).
Collaborator
Author
|
Taking all six: the bit-for-bit claims go and the native test holds per-frame rows to per-token ones within cuBLAS's rounding with the blocks checked bitwise on the same rows; a CUDA test under |
Nothing on the served path read the edge anymore, it was only there for the .item() fallback. A caller with no counter now gets a KeyError instead of the sync.
forward with a stub transformer under sync-debug mode, so linalg.solve, a host read of a device tensor or a pageable table copy fails it.
cuBLAS can pick another kernel at 18 rows than at 7020. The native test holds those rows to that rounding and checks the blocks bitwise on the same rows. The reference suite prices the trajectory and keeps a 0.0 gate with the embedder at the reference's row count. Also fixes its I2V test, which still expected a per-token timestep.
Collaborator
Author
|
Pushed ce81c45 (merge of main at 7eac23c) and 93bc40c..3f49241 for the six points; description rewritten with the B300 gate, the trajectory analysis and the new bounds. One thing to know: this branch still carries #944's earlier commit, and #944's branch has moved since, so it needs a merge once #944 lands. Agreed on the direction, and I'd do it as the next PR once this lands, not here. The gate after this rework, on 1x B300 against main at 7eac23c + #944's commit, same setup as yours with Re-requesting. |
The two bounds were set from the worst case seen: one sat below a bf16 step at temb's top, and the other let a moved bf16 boundary through.
Gaurav-Shah05
approved these changes
Oct 30, 2026
Gaurav-Shah05
left a comment
Gaurav-Shah05
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. I merged main and pushed two test commits on top:
- Embedder check. It's held with
torch.equalon the reference's per-token timesteps, as on main, in place of the two hand-picked bounds (2e-3 sat below a bf16 step at the top oftemb's range, and 3.2e-2 passed a moved bf16 cast). - No-CFG gate. It runs on per-token rows through the sequential test's wrapper, now one helper, and is back at 0.0.
They are not tested on GPUs, but should be good to merge after they pass.
This was referenced Oct 31, 2026
This was referenced Nov 1, 2026
merceod
pushed a commit
that referenced
this pull request
Nov 3, 2026
* API server: the queue gauge counts inputs not being prepared, pooled or inline (R3.1, R3.2)
The data-worker thread publishes how many of its taken-up inputs are not
being prepared (waiting_inputs), and queue_depth adds that to the queue,
so one request preparing reads the same with or without a pool and the
gauge no longer reads the thread's private deque. The four descriptions
of the gauge now say the same thing, and that client-stream appends count.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server: validate the class PREPROCESS_WORKERS value; fix a stale comment (R3.3, R3.4)
A class value that is not a positive integer now stops startup naming the
model class, instead of a config value building a pool for it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server: one positive-int check for preprocess_workers; fix the startup comment (R4.2, R4.3)
_positive_int moves from entrypoint.py into data_worker.py (entrypoint
already imports from it) and checks both the class PREPROCESS_WORKERS and
the config's preprocess_workers, replacing their two inline copies. The
messages are unchanged.
The comment above the check claimed it runs first, but check_media_limits()
runs before it; it now gives the reason: nothing to close yet.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server: without a pool, the queue gauge's count is the deque's length (R4.4)
_publish_waiting() runs before every hand-off and rescanned the whole
_preparing deque each time, k^2/2 per pass. Without a preprocess pool no
entry there is ever being prepared (the inline one is popped first), so
len() is exact; the scan stays only with a pool.
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>
* Tests: pin the pooled queue gauge once settled, and its count during a hand-off (R4.1)
The pool case of test_queue_depth_counts_requests_waiting_for_the_data_worker
could pass on the moment both requests were still in in_queue; it now also
reads the gauge after the wait, when the pool has prepared #1 and it waits
behind "hold". test_a_prepared_request_waits_while_a_peer_cannot_take_more
records waiting_inputs at each hand-off, which pins the publish before it.
Checked against mutated code: counting only unprepared entries reads 1 at
the new assert (the old _wait still passed); dropping the publish before the
hand-off records [2, 1].
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>
* Docs: say which outputs a stop check may withhold
A withheld output has to leave by a streaming edge that no node this worker
speculates into reads, and cannot be the loop-back signal. The startup check
enforces this for TokenStop.withhold; a StopCheck returned from check_stop
(the Qwen3-Omni Talker, Fish S2 Pro's AR) is not checked, so the paragraph
that introduces StopCheck now states the constraint.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Worker, engine: the stop check's host outputs are Mappings in every signature
_prematerialize_for_check_stop returns _LazyHostOutputs, which copies an
output name the first time it is read, not a dict of host copies of every
tensor; its annotation and docstring now say so. Engine.extend_prefix_chains
is annotated with the Mapping its body already checks for, and the resource
runner's per-request fallback checks Mapping too, so one contract is checked
one way (today only plain dicts reach that fallback).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: one idle poll for the API server's stand-in transports; drop dead ones
The API server's message loop now waits in poll_for_messages instead of
sleeping, so a stand-in that returns at once made a loop running on a thread
spin on the GIL. api_server_helpers.idle_poll waits out the timeout as the
socket's poll does, and every stand-in for the API server's loop uses it;
their mis-indented and blank-padded copies go with them.
The data-worker thread and the Rust-frontend bridge never poll, so the
poll_for_messages entries on their stand-ins are removed. The conductor
stand-in in test_kv_layout_intake.py had none, so each pass of the main loop
raised AttributeError and was logged; it now answers as the other conductor
stand-ins do.
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>
* Qwen3-Omni audio encoder: a step whose rows all failed batches without raising (R1.18)
The engine asks `can_batch` with whatever rows survived `prepare_inputs`, so
a step whose rows all failed reached `layout([])`, whose `max()` of no chunk
widths raised and turned the rows' own failures into a step exception.
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>
* Qwen3-Omni audio encoder on the shared AuT encoder (R1.1, R1.10, R1.11, R1.5)
Omni's encoder node drove transformers' Qwen3OmniMoeAudioEncoder internals
from mstar code, a second implementation of `components/aut_encoder.py`,
which Qwen3-ASR already runs packed, windowed and captured. Omni now builds
its tower as `AuTEncoder`, loaded from `thinker.audio_tower.*` as Qwen3-ASR
loads it, and its config is `AuTEncoderConfig` (the Omni-only
`AudioEncoderConfig` is gone).
The frontend improvements move into the shared component, so Qwen3-ASR
gets them too:
- the trim back to each chunk's valid rows gathers by host-built indices
(a pinned, non-blocking upload) instead of a pageable H2D plus boolean
indexing, both host syncs;
- the chunk count pads with empty chunks to a power of two (cuDNN plans
each new conv shape);
- token counts follow the checkpoint's `n_window` instead of 100 frames;
- the eager path passes the layout's host boundaries, not a device tensor.
The submodule keeps its packing policy (one replay a step, `can_lease` and
the all-or-nothing `can_batch` past the largest bucket). Deleted with the HF
copy: `AudioLayout`, `layout`, `_frontend`, `_layers`, the `_attn_fns` memo,
the hand-written region declaration (now `attn_key=AUDIO_ATTN`), the
double-clone of the leased output, the per-row layout lists rebuilt three
to four times a step, the HF FA2 construction and its stale comments.
GPU-unverified: compare audio_embeds against the base on real weights.
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 audio encoder: the capture check reads the attention resource's ceiling (R1.2)
`_fits_capture` compared a step's windows with the class constant, while the
captured wrapper is sized by the resource's `max_segments_per_request`, which
a deployment may override under `resources: audio_attn:` (docs/serving.rst).
An override below the constant let a step lease a replay whose plan raised
for too many segments, failing every clip in the step; one above it was
ignored. It now reads the bound resource's ceiling, as the Qwen3.5 vision
tower and Omnilingual do.
The declared ceiling is derived from the encoder's geometry (the top
bucket's full windows plus a short last one per clip) instead of a literal
that assumed `n_window_infer` = 800.
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>
* AuT encoder: clamp an fp16 residual stream as HF does; Omni's audio attention follows the engine dtype (R1.8)
`AuTEncoderLayer` dropped HF's fp16 guard (clamp the residual stream to
max - 1000 after each layer), so under `autocast_dtype: float16` a long
clip's stream could overflow to inf and hand the Thinker inf/NaN audio
embeddings; Qwen3-ASR and, now, Qwen3-Omni run this layer. The clamp is a
static dtype check, so the bf16 graph is unchanged.
Qwen3-Omni's AUDIO_ATTN declared `dtype=torch.bfloat16` whatever the
deployment dtype, so an fp32 or fp16 deployment ran the encoder's attention
through a bf16 plan. It now leaves the dtype to the engine, as Qwen3-ASR's
AuT attention and the Qwen3.5 vision tower do.
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-Omni: drop the audio_seqlens tensor nothing reads (R1.5)
The captured encoder takes a clip's frame count from its features' width,
so `audio_seqlens` was shipped to the encoder node, and waited on, with no
reader. `process_prompt` no longer builds it (the prompt layout counts the
encoder rows from the same width), the encoder node and the audio walk's
edges drop it, and the comments that still sent it to the encoder are
corrected.
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>
* Prefill ladder override: one checked override point in ARNodeSubmodule; Omni keeps its default ladder (R1.3)
Qwen3.5 and Qwen3-Omni each overrode `PREFILL_TOKEN_BUCKETS` in their own
submodule constructor, unchecked: `[]` silently meant the default, a repeat
was captured twice, and a scalar raised a TypeError at build. The override
is now one `ARNodeSubmodule` constructor argument that both submodules pass
through, and a ladder that is not a non-empty list of distinct positive
token counts raises ValueError when the worker builds the submodule.
Qwen3-Omni's default ladder goes back to 128-2048. The 8192 ladder was
measured on 2-GPU audio and text traffic, so `configs/qwen3omni_2gpu*.yaml`
set it under `model_kwargs`, and every other Omni deployment keeps its
captures, graph memory and startup time.
The serving-config override the engine applies per (node, walk), with no
model code, remains the north star; not built here.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* configs/qwen3_5_9b_dp2.yaml: the 9B prefill ladder (R1.4)
Each DP2 replica is the single-GPU 9B LLM of qwen3_5_9b.yaml, but the config
left out its 3072-7168 prefill buckets, which docs/models.rst says every
qwen3_5_9b*.yaml adds, so prefill and mixed steps of 2049-8192 tokens padded
to the next power of two on DP2.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni Talker on DecoderLayer and run_layers; Thinker layers start the residual stream through forward_residual (R1.7)
The Talker layer had exactly `DecoderLayer`'s shape but a hand-written
`(mlp_out, residual)` contract and its own deferred-residual loop beside
`run_layers`; it is now a `DecoderLayer` and its stack runs through
`run_layers`, as the code predictor's does. Parameter names are unchanged.
The Thinker keeps its own loop (deepstack, the accept-layer tap and
`token_valid` don't fit `run_layers`) but drops the `residual is None`
branch that duplicated `forward_residual`'s own.
Not done here (follow-up, needs a GPU): passing `comm_group` so TP configs
fuse all-reduce + add + RMSNorm, which needs `o_proj` and the TP MoE to hand
back partial sums and `init_allreduce_fusion` at build.
GPU-unverified: Talker codes / audio output against the base.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test_multimodal_ordering: run the Qwen3-Omni cases in CPU Core (R1.9)
Four Qwen3-Omni cases skipped on `importorskip("flashinfer")`, which CPU
Core never installs, although the Omni modules import fine without it
(test_qwen3_omni_prefix_cache imports them unguarded and runs there). So
the only test of the rewritten, host-sync-free `talker_text_mask` never
ran. All four now run.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Accelerator graph runner: `_stage`'s docstring says what a replay's padded tail holds (R1.13)
It said trailing slots keep capture-time contents, but a forward whose
residual stream starts at a static input (Qwen3-Omni's Thinker and Talker,
`run_layers` users such as Qwen3DenseLM) updates that buffer in place, so
the tail holds the last replay's stream. Real rows are unaffected: attention
plans exclude the tail and the other ops are row-wise.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni Thinker: one mechanism for the Talker masks' host sync; the vision prefill's mask writes don't sync (R1.15)
With `talker_text_mask` built from device ops, skipping the masks for a
text-only request saved only a few small kernels a text row, a second
mechanism for the sync the rewrite already removed, and added an optional-key
path in `split_inputs`. Every prefill row builds its masks again.
`prefill_vision` still marked its sentinels with a list index, a pageable
upload, and filled deepstack through a boolean mask, which reads the mask
on the host: two scalar writes and a slice now, as the audio path does.
Not here (follow-up): the masks' consumers (`Thinker.postprocess`'s
`selected[inclusion_mask]`, `_get_talker_embeds`'s `.any()` and boolean
indexing) still read them on the host, and the vision walk keeps its
`.item()` reads.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni Thinker decode: padding rows take the broadcast positions; decode rows skip unused reads (R1.14)
A decode row takes the plan's positions like a text row, but the
`thinker_decode` capture template still carried `custom_pos_ids`, and a
replay's padding rows are clones of it, so every padded decode step (a batch
not in 1, 2, 4, ..., 32) built its positions row by row (a cast per real row
and a cat) instead of the one broadcast view. The template is now a text
row. Decode rows also no longer read the position counter, which only the
audio and vision walks use.
GPU-unverified: decode capture and padded-batch replay.
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>
* Tests: pin the decode-row interleave, multi-window encoder parity and the replay's windows (R1.17)
- test_qwen3omni_vision_mixed: a mixed step's decode tokens embed in one call
and must land back in row order beside the prompt rows; a misordered
iterator would hand decode rows each other's embeddings unnoticed.
- test_aut_encoder: the HF-parity cases (multi-chunk clips, 26-row windows,
packed requests) ran only under transformers 4, where HF's SDPA path
ignores windows and needs the masking fixture; transformers 5 windows by
itself, so the fixture now applies on 4.x only and the cases run in CPU
Core, against the shared frontend Qwen3-Omni now runs too.
- test_qwen3_omni_prefix_cache: the packed-clip parity covers a ten-chunk,
two-window clip, and the replayed path, whose windows must be the step's.
Each fails under a mutation of what it pins (reversed decode order, an
off-by-one trim, one window for the replay).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* ARNodeSubmodule: check a ladder's entries before hashing them (R1.3)
A ladder of unhashable entries reached `set()` first and raised TypeError
instead of the ValueError the check exists to give.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test_trust_remote_code: the Qwen3-Omni audio encoder reads no HF config (R1.1)
Built as the shared AuT encoder from the model's own parsed config, the
audio encoder no longer calls `AutoConfig.from_pretrained`, so the case
checking that call's `trust_remote_code` goes; the vision encoder's stays.
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 after what it checks
The clone prefill returns row outputs as well, and the test asserts it,
so drop "except for the clone prefill" from its name and docstring.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: describe the Qwen3-Omni 2-GPU PD config and OmniVoice's token buckets
configs/qwen3omni_pd_2gpu.yaml and the backbone_token_buckets model
kwarg were added without a line in the models doc.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice: check the packed runner at boot, drop the unused plan dtype
The loader deletes the reference's packed runner (`_fi_runner`, a 128 MB
workspace mstar never uses), but the boot check did not cover it, so an
omnivoice release that renamed it would fail with a bare AttributeError
rather than the SystemExit the other private names get. The boot check
now requires it.
OmniVoiceBackbone also kept a plan_dtype that only the parity test read;
the test plans the reference runner with the model's own dtype instead.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice test: compare the silence trim with pydub at full scale
The reference converts float audio to int16 with a clip, as the port
does, so a +-1.0 sample stays at the int16 limits rather than wrapping.
Add a full-scale burst between silences to the edge cases checked
against the reference, where omnivoice is installed.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor: check the sampled tokens with is_contiguous
The embed-accumulate kernel reads row b's token at offset b. A one-row
tokens tensor is read correctly whatever its stride, which the
stride(0) == 1 check refused; is_contiguous is the exact condition.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT frontend: pad only the last convolution launch (R1.11)
The step's chunk count was padded to a power of two, so a long step convolved
up to twice its chunks (Omni's 630-chunk top step as 1024; a Qwen3-ASR step
of eight 20-minute clips, 9600 chunks, as 16384). The convs already run in
`conv_chunksize` launches, and only the last one's batch varies, so only it
pads: up to a power of two, at most `conv_chunksize`. Steps of up to 256
chunks pad exactly as before, the conv shapes stay a handful, and a long
step adds at most one partial launch's padding.
GPU-unverified: audio_encoder phase time on 60 s clips at c=16 (Omni) and
on long-form Qwen3-ASR batches, against the previous head.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: serving keys for the speech split configs and the switch interval
The mps row and the MPS paragraph named only the ASR split configs; the
Chatterbox, OmniVoice and 0.6B Qwen3-TTS split configs set mps: true now
too. The api_switch_interval_s row said anything but a positive number
fails at startup, while an explicit null keeps Python's 5 ms (the model's
API_SWITCH_INTERVAL_S applies only when the key is absent). The resource
override example also lists max_segments_per_step among a ragged attention
resource's keys.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: Qwen3-Omni's audio encoder is the AuT node
The Qwen3-Omni notes still described the HF encoder's frame budget (16
requests, a new call past 300 s of audio). The audio encoder now runs
Qwen3-ASR's AuT node on the API server's log-mel features: up to 8
requests a step in one packed forward, the stack replayed by packed token
count up to 4096 tokens, eager past that.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS, Omnilingual, ragged attention: comments name the new owners
Qwen3-TTS's resource docstring still said the Talker submodule owns the
code predictor's frame-local KV scratch, which its DepthKV owns now.
Omnilingual's capture check, the ragged captured wrapper's sizing and the
can_lease docs still spoke of a per-request segment ceiling where the code
reads the per-step one (max_segments_per_step, through max_segments_for).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Comments and docs: a shard_over_sp=False ragged resource splits its heads over TP only
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* docs/models.rst: the Omni audio encoder's eager and sub-second cases as the code has them (R2.1)
A step also packs, and a clip also replays, only within the attention
plan's `audio_attn.max_segments_per_request` windows, so a clip under
8192 rows but over that ceiling runs the layers eagerly too. A clip
under 1 s encoded alone differs from HF's in its last conv row, which
the windowed attention mixes into every embedding of the clip (one
window), not only in its last rows.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Head sharding: one per-resource switch, shard_over_sp; NodeSubmodule.splits_heads goes
Every resource now shards its heads over all TP x SP ranks of the node's
lockstep group (EngineResourceInfo.head_shards, a property), and a ragged
resource that each SP rank runs whole says so with
RaggedAttentionConfig(shard_over_sp=False). The node-level splits_heads flag,
the engine's per-node fold and its agreement check, and their docs and tests
are removed.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni Thinker prefill capture: padding rows take the broadcast positions (R2.2)
The `prefill_text` capture template, which `prefill_text`, `prefill_audio`
and the mixed walk replay, still carried `custom_pos_ids`, and a replay's
padding rows are built from it, so every padded text-only prefill or
mixed step built its MRoPE positions row by row (a cast per real row and
a cat) instead of the one broadcast view. The template is now a text
row, as the decode one is since R1.14; an audio row brings its own
positions.
GPU-unverified: prefill capture and padded text/mixed replays.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* GLM-5.2 prefill ladder through ARNodeSubmodule's override; the base declares the ladder (R2.3)
`model_kwargs.prefill_token_buckets` had two mechanisms: Qwen3.5 and
Qwen3-Omni pass it to `ARNodeSubmodule.__init__`, which refuses an empty
or repeated ladder, while GLM-5.2 read its config's copy at capture with
`or`, so `[]` silently meant the default and a repeat was captured
twice. GLM-5.2 now passes it to the same constructor and its capture
reads `self.PREFILL_TOKEN_BUCKETS`; the `or` fallback is gone.
`ARNodeSubmodule` declares `PREFILL_TOKEN_BUCKETS`, the attribute the
override replaces and the capture configs read.
Test: test_glm52_prefill::test_a_bad_prefill_ladder_is_refused ([] and
[64, 64]) did not raise before and raises ValueError now; the module's
helper builds the submodule through its constructor.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: the frontend takes the step's clips as a list; Omni and Qwen3-ASR stop padding them (R2.4)
Both encoder submodules zero-padded every clip of a step to the longest
and stacked them, and the frontend sliced each clip straight back out.
Omni's features arrive in fp32, so a long clip batched with short ones
allocated B x 128 x T_max x 4 B per step (about 0.5 GB for one 10 min
clip beside fifteen 2 s clips) for nothing. The frontend now takes the
clips as they are, `(num_mel_bins, frames)` each, and reads each one's
length off its shape; the padded stack, the `feature_lens` argument,
Qwen3-ASR's `num_frames` kwarg and its 2-D `forward` special case go.
The convolution input is built from the same frames in the same order,
so the embeddings are unchanged.
GPU-unverified: the transient memory saving (no CPU-observable change).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: one checkpoint loader, `AuTEncoder.from_checkpoint`, for Qwen3-Omni and Qwen3-ASR (R2.5)
The two models loaded the shared tower with the same six steps (build on
meta, cast, `to_empty`, stream `thinker.audio_tower.*`, `load_weights`,
`eval`), copied line for line. The sequence is now a classmethod on the
encoder, which owns the parameter layout, and both models call it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT frontend: pad the last convolution launch to a multiple of 32 chunks, not a power of two (R2.6)
The power-of-two rule could nearly double a launch: 257 chunks ran as
500 and 300 as 500, about 4 GFLOP per empty chunk, and Qwen3-ASR pays
it too since the frontend is shared. A multiple of 32 (8 for a launch
of at most 8 chunks) wastes at most 31 chunks at the cost of about 17
conv shapes instead of 7, each a one-time cuDNN plan.
GPU-unverified: the audio encoder's phase time on Omni (5 x 60 s clips)
and Qwen3-ASR (8 x 37 s), power of two vs multiple of 32.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: pack a step only up to the largest captured token count
Qwen3-Omni's HF audio encoder split a step into calls of at most 30,000
mel frames. On the AuT node nothing bounded a step, so eight long clips
ran as one eager forward over every clip's tokens. can_batch now packs a
step only while it fits the largest capture bucket (4096 tokens); past it
the clips run one by one, each replayed when it fits a bucket, so an
eager step holds one clip's activations. The same node serves Qwen3-ASR,
whose steps up to 4096 tokens are unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: count tokens by the checkpoint's own chunk, for both models
tokens_for_frames counted 13 tokens per 100-frame chunk whatever the
config's n_window, while the frontend cuts chunks of 2 * n_window frames
and Qwen3-Omni's prompt layout counted with transformers' helper at the
checkpoint's n_window. It now counts each full chunk_frames chunk by the
convolutions' arithmetic, and the layout calls it, so the prompt, the
encoder's split and its frontend share one count. Unchanged for the
released checkpoints (n_window 50).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Ragged attention: refuse a per-request ceiling on a resource sized per step
A resource with max_segments_per_step sizes its plan from that alone, so
a deployment override of max_segments_per_request was accepted and then
read by nothing. apply_yaml_overrides now refuses it, naming the per-step
key. Its docstring counts the ceilings it takes correctly.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Test: the engine hands a forward only the regions its step leased
The case called SubmoduleManagement.runners_for directly, so reverting the
engine's call to it went unnoticed. It now drives Engine._drive_step and
reads the runners the forward receives.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Cosmos3: take cudnn_benchmark from fp32_precision, drop the wan_vae re-export
The scope moved to mstar.engine.fp32_precision and wan_vae kept a
re-export for Cosmos3's two decode paths. They import it from its owner
now, the re-export goes, and its test moves beside the other fp32
precision tests.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni: drop audio_seqlens, which the audio encoder no longer reads
The AuT node takes a clip's frame count off its unpadded log-mel
features, so the audio_seqlens tensor rode to the encoder unread. The
prompt layout, its one reader, counts the same frames off the features'
shape, so the tensor, its graph input and its routing go.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder node: move it beside the encoder, one loader and spec for both models
Qwen3-Omni imported its audio encoder node from the qwen3_asr package and
repeated its loader and attention spec line for line. The node is now
AuTEncoderSubmodule in mstar/model/components/aut_encoder_submodule.py,
with AUT_ATTN, aut_attention_spec and load_aut_encoder_submodule, and
both models call them. It takes a LogMelSpectrogram instead of
Qwen3-ASR's config. Behaviour is unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: a Qwen3-Omni audio step over 4096 tokens runs its clips one by one
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT frontend: pad the last convolution launch to the smaller of a power of two and a multiple of 32 (R3.1)
The multiple-of-32 rule (R2.6) ran a last launch of 9-16 chunks as 32,
where the power-of-two rule had run it as 16, so a single 9-16 s
Qwen3-ASR clip convolved 2-3.5x its chunks. Taking the smaller of the
two roundings (still at least 8, at most conv_chunksize) pads no launch
more than either rule did: 9-16 -> 16, 17-64 as before, 65-96 -> 96,
and at most 31 empty chunks, for 18 conv shapes instead of 17.
GPU-unverified: the encoder's time on Qwen3-ASR at c=1 with 10 s and
14 s clips, multiple of 32 vs this rule (CUDA events around
AuTEncoder.frontend, or worker_phases on the encoder node).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* ARNodeSubmodule: narrow the PREFILL_TOKEN_BUCKETS comment to what holds (R3.2)
The comment said every submodule that captures prefill names its ladder
PREFILL_TOKEN_BUCKETS. BAGEL's LLM (PREFILL_TEXT_TOKEN_BUCKETS), the
Qwen3-Omni Talker (TALKER_PREFILL_TOKEN_BUCKETS) and Cosmos3's DiT (an
instance ladder with an env-var override) don't, and an override passed
to one of them would be ignored. The comment now states the obligation
instead: a submodule that takes `prefill_token_buckets` must read its
capture ladder off `self.PREFILL_TOKEN_BUCKETS`, as Qwen3.5, the Omni
Thinker and GLM-5.2 do.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Name the TP x SP head-sharding test for the kinds it covers; reflow a docs paragraph
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: say a step holds up to 4096 tokens; import its loader at the top
The token budget applies to Qwen3-ASR as well as Qwen3-Omni, so the ASR
docs and the node's docstrings now say a larger step runs its clips one
by one. Both models import load_aut_encoder_submodule beside
aut_attention_spec, from the module they already load at the top.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 9B DP2 config: give the GDN pool prefix snapshot slots (RM.1)
The dp2 config was written before #877 and has no `gdn_state.prefix_slots`,
so the recurrent pool can't serve a cached prefix and the engine keeps the
LLM node's KV prefix cache closed with it. Each replica runs at world size 1,
so it takes the same 128 snapshot slots as qwen3_5_9b.yaml.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan 2.2: keep the dit's heads whole over SP with shard_over_sp=False
The node-level splits_heads flag is gone; the dit's one ragged attention
declares the same thing on the resource. Under SP each rank still runs
every head (the ranks split CFG branches), and TP still splits them.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: run a guided window as its own walk on one node too
A single-node (or LLM-disaggregated) deployment ran a guided request as one
image_gen walk and read each step's branch count off the loop counter against
a CFGWindow in step metadata, with a guard for a counter that disagreed with
the seeded latents. The CFG-parallel layouts already ran the window as its own
walk with a fixed count. Every layout now does: a guided request runs
image_gen_head (if the window starts late), then image_gen_guided2 or
image_gen_guided3 on LLM alone, then image_gen (or image_decode when the window
reaches the last step). Each step runs the same branches as before.
DenoiseWalk gains `split` (a node per branch or not), ImageGenPhase loses its
`branches`, and the per-step count, its guard, the CFG_BRANCHES kwarg and the
cfg_window / image_gen_branches step metadata are deleted. The disaggregated
configs and the docs' examples list the new walks for the image LLM.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Say the api_switch_interval_s numbers predate #849's output-path rewrite
The row still said the knob would be re-measured once #849 lands. It has
landed and changed the message-thread and data-worker path the knob
tunes, so the numbers are from before that change and the re-measure is
still pending.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the head-sharding paragraph from the submodule attributes section
It was there to document the removed splits_heads attribute; what is
left repeats the RaggedAttentionConfig row.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the TP x SP head-sharding test that main's sizing already passes
Head sizing now equals main's, so the six cases pinned nothing this
branch changes. The shard_over_sp=False test stays; the helper loses its
recurrent branch, which only the deleted test used.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: hold guidance KV only for the branches a request's window runs
A request is guided when its CFG window has a step, not when a scale is above
1: a window with no step (a cfg_interval that excludes every timestep) runs no
guidance, so it no longer prefills or decodes cfg_img or forks cfg_text.
`_requires_cfg` is gone; process_prompt resolves the window before the layout
and the conductor reads it from the stamped window.
The layout reserves a branch label's denoise scratch only where the window
runs that branch: cfg_img's only with image guidance, and neither branch's
for an unguided request, which no layout denoises on them any more. A guided
text-to-image request no longer holds a latent-sized cfg_img scratch it never
writes.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the shard_over_sp sizing test, which main's sizing already passes
Main already sizes a shard_over_sp=False ragged attention over the TP
group, so the test pinned nothing this branch changes. The file's
coverage goes back to main's plus the manager registry test.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: comments state the contract instead of the code they replaced
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL docs: give the measured output shift of guided text-to-image
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test_kv_admission: name the CFG walk BAGEL defines
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Cosmos3 and LingBot decode through wan_vae.FastDecoder and decode_dtype
wan_vae.FastDecoder and decode_dtype had one user, Wan2.2, while Cosmos3's
two decoder nodes and LingBot-World kept their own copies of the cuDNN
>= 9.16 dtype gate, of the channels-last + DecoderStep setup and of the
autocast-off decode_causal call. All of them now use the shared helpers,
and the copies (_vae_decode_dtype, _fast_step, the streaming node's _step,
LingBot's inline gate and _decode_step) are deleted.
Numerics and switches are unchanged: Cosmos3 keeps COSMOS3_VAE_DECODE_FP32,
its compile switch, its cast before the layout switch, the streaming
step's non-deterministic compile (built only with the warm cache), its
latents per call and cudnn_benchmark scope; LingBot keeps no override
(decode_dtype's variable is now optional), and bf16 and channels-last only
off reference numerics. Every node keeps the layout it set up, so
FastDecoder's per-call layout check is a no-op for them.
FastDecoder's autocast-off scope now follows the latents' device instead of
naming cuda.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan 2.2: drop the FFN GELU epilogue, which has no measurement of its own
MLP(gelu_epilogue=) and Wan22DiT's fuse_ffn_gelu ran the FFN GELU in the
up-projection's cuBLASLt epilogue through the private
torch._addmm_activation. Nothing measured it apart from the ragged
attention change it shipped with, so it goes, as unmeasured performance
paths do: MLP is back to its plain form and the DiT builds without the
flag.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3 text encoder: refuse a row without text_len instead of never replaying
The shared encoder node picks its captured (masked) graph from each row's
text_len step metadata. A model that reused the node without setting it
got every batch run eager, with no error. cg_key_info now raises
ValueError for such a row; FLUX.2-klein and Z-Image set it at intake.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Image output knobs: an explicit null is the default, and only integers pass intake
image_output_format let png_compress_level: null through, and encode_image
then failed on int(None) after the denoise and decode. encode_image now
reads null as the default level, as it already did for output_compression.
Intake also refused nothing that int() accepted, so True passed as 1 and
50.9 was truncated to 50; it now takes only integers (not bools) in range,
as BAGEL's width and height check does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* V-JEPA2: say the B=1 guard routes eager steps only, and drop "as before" comments
A step that leases the AC rollout capture replays forward_batched at any
batch size, B=1 included, whatever can_batch returns, so can_batch's B=1
guard covers eager steps only; its comment now says so. Two comments
described the code by comparison with an earlier version ("as before the
capture", "same as before") and now state what the code does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* wan_vae.decode_dtype: say that an override of 0 forces bf16
The override variable is two-way, as Wan2.2's and Cosmos3's were before it
was shared (and as the Wan fast-decode test uses it): set, it forces fp32,
or bf16 when it is 0/false/no/off. The docstring read as if 0 meant "not
forced".
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Higgs Audio: declare text_output so inline decode tokens skip the per-row postprocess
Higgs's decode token rides inline, but without a declared text_output
Model.postprocess_inline built a tensor and called postprocess for every
row on every step. Declaring TextOutput(self._detokenizer), with no stop
ids as before, gives it the memoized id-to-bytes path Qwen3-ASR and
Whisper use; postprocess hands text to the base class. Bytes are
unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5: say the API_SWITCH_INTERVAL_S gain predates #849
docs/serving.rst says the number was measured before #849 and awaits a
re-measure; the comment pointing at it now says the same.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Conductor: build its instance table from WorkerGraph.replica_ranks
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Image output knobs: serve the whole numbers main served
Intake took only Python ints, so a whole-number float (90.0) or a number in
a JSON string ("90"), which main read through int() and served, became a
400. It now accepts any whole number int() reads in range, and still refuses
bools, fractions such as 50.9, infinities and out-of-range values.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: build only the window walks a deployment can run
A guided window runs on the branch nodes on CFG-parallel deployments and on
LLM alone elsewhere, never both, yet CFG-parallel builds still placed the
unsplit image_gen_guided2/3 walks and listed them in every request's
needed-label map. LLM_GRAPH_WALKS is now keyed by the layout, and the walk
graphs and the label map both read it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: say why a split LLM still publishes both branch labels
The comment dated from when a remote image_gen walk's branch count varied
per step. Each walk now fixes its count; what keeps cfg_img published on a
split deployment is a remote decode reading it and the single-node denoise
walks naming their labels per request rather than per walk.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL and Cosmos3: drop loop-counter leftovers from comments and names
denoise_loop's docstring still named the step's loop counter, which BAGEL
no longer reads; declare_step's walk-branch narrowing kept a requires_cfg
guard that active_labels already implies, and its guided-row set shadowed
the module's guided(). Cosmos3's decode-dtype comment now gives both
directions of its override.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Image output: check only the chosen format's compression knob
Intake refused a request whose other-format knob was out of range (a JPEG
request with png_compress_level 10, a PNG request with output_compression
150), which main served because encode_image never reads it. Check only
the knob the format applies.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: say LLM-disaggregated, not split, in the branch-label comment
"Split" means CFG-parallel in the neighbouring code (DenoiseWalk.split,
LLM_GRAPH_WALKS).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* TP read-failure test: derive the loop-back hooks from device_loopback = None
The stand-in stubbed the deleted _device_loopback and restated the three
hooks DeviceLoopback now derives; clearing the declaration gives the same
answers through the base class.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* TP read-failure test: drop the stale "emit" from the stand-in's comment
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Merge-up fixes: DeviceLoopback in the withhold check, chunk-sink stand-in poll, postprocess_s on inline failure
- _refuse_unroutable_withholds read TokenStop.loopback_signal, which main
moved to ARNodeSubmodule.device_loopback; it now reads the signal from
the submodule's DeviceLoopback, and its test declares one.
- test_chunk_sink: its request profile stub lacked output_postprocess_s,
so every pass raised before relaying a chunk; it uses a real
RequestProfile, its transport's poll ends the pass (the loop no longer
sleeps), and an error in the pass is raised instead of retried.
- An inline value whose batched postprocess_inline call failed is charged
its share of that call as well as its own retry.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Withhold check: the loop-back signal is exempt only on its own walk
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Withhold check test: the loop-back exemption is scoped to its walk
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-ASR/Omni AuT node: pin the replay's windows and packed-clip parity on the shared node
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Attention: select_last_hidden is one eager boundary, so prefill buckets don't recompile its frame
The disabled `qo_indptr_buf` call broke the graph inside `select_last_hidden`,
leaving a resume frame that holds both the bucket's static token count
(`hidden`) and its row count (`qo_indptr`). Captured forwards compile with
dynamic=False, so that frame took one cache entry per (tokens, rows) bucket:
Qwen3.5's 9 token buckets by 7 prefill/mixed row counts made 63, and the
13-step prefill ladder made 91, past the recompile limit of 84.
Disabling the whole method moves the break to its call site, where the
caller already breaks; no compiled frame holds the pair. Same for the MLA
and XPU managers, which share the shape.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL takes its prefill's last rows through select_last_hidden; test docstrings
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* GDN prefill run: state the zeroed-tail precondition the mixed wrapper also meets
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice: key the backbone's step outputs by the engine's row ids
The packed backbone keyed its outputs by CanvasItem.request_id, which
prepare_inputs now fills from fwd_info.request_id, the wire request id. The
engine looks each row's outputs up by its worker-local handle, found none,
and every OmniVoice request failed in postprocess with KeyError 'c_logits'.
forward_batched now keys rows by engine_inputs.request_ids, as the other
batched submodules do, and CanvasItem carries no request id at all, so the
canvas can't be keyed by the wrong one again. The unread request_id kwarg
goes with it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni 2-GPU config: keep the 8192 Thinker ladder off the 80 GB default
qwen3omni_2gpu.yaml, the CLI default and the config sized for 80 GB GPUs,
ran the Thinker's prefill ladder to 8192. On an H100 its 4096-8192 captures
don't fit beside the 1,024-page KV pool (59 of 86 Thinker buckets captured),
so those steps run eagerly and a deployment that requires accelerator graphs
refuses to start. The default goes back to the model's 128-2048 ladder, as on
the base. The 8192 ladder, measured on 2xH200, stays in
qwen3omni_2gpu_large_kv.yaml, which is already the H200 variant; its header,
the default's, the submodule comment and the Qwen3-Omni notes say so.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni 2-GPU config: don't claim the remaining captures fit 80 GB
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Pin video decode threads to the CPUs the plan leaves under pinning
A case with a pinned CPU fails if video_decode_threads goes back to
usable_cpus(); the MSTAR_STEP_THREAD_CORE row names FFmpeg's per-decode
threads among what the pins shrink.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Ragged attention: a wrapper skips replanning the layout it last planned
The dense manager kept the (cu_q, cu_k, causal) each of its wrappers was
last planned with and skipped an unchanged plan; the ragged self- and
cross-attention managers replanned every step, so a denoise loop paid
FlashInfer's plan (a kernel and a stream sync) once per step for the same
segments. The skip now lives in the wrapper both use
(_RaggedPrefillBase._plan): each wrapper is one (label, slot) or one
(bucket, slot, label), so the schedule it keeps is the one the step reading
it planned. The dense manager's own table goes.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan2.2 docs: date the attention and CFG-parallel timings as before #963
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Conductor: leave the idle wait to its own change
The conductor's 1 ms nap becomes a wait on the poll in a change of its
own, which waits until the next timed check instead of a fixed 1 ms. This
branch's capped poll, its failed-pass back-off, the conductor stand-ins'
poll_for_messages and their test go, so those files are main's again.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server tests: the failed-wait test pins its back-off; stand-in waits as main has them
The message loop always passes a computed timeout, so the 1 ms defaults
left by the fixed poll go: the chunk-sink stand-in takes main's again and
the failed-wait stand-in takes it as required. That test now asserts the
10 ms back-off, without which a persistent wait error would spin.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Preprocess pool: at most the CPUs the plan leaves unpinned
Every other API-server pool is sized from the CPU plan; the preprocess
pool took its declared or configured size uncapped, so under
step_thread_core a model declaring 8 workers could run them, each with its
own FFmpeg threads, on fewer free CPUs.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Data worker: drain routes before an input only when it prepares here
The drain keeps a call's held items from waiting behind a preprocess that
may take seconds. An input prepared on the pool only stores and sends when
it is handed on, and each pass already drains routes at its top, so it no
longer takes a drain of its own.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 vision: compile the block loop, so a step past the captured buckets runs compiled
disable_torch_compile left the block loop uncompiled whenever it did not
replay: a long video (over 16384 patches) ran the blocks' fp32 rope, norms
and activations as eager kernels, slower than the whole-forward compile it
replaced. The block loop is now a compile region like the embed and the
merger, dynamic in the patch count; it compiles in its capture's warm-up.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server: drop the api_switch_interval_s knob
Its gain was measured before #849 and #939 changed the message-thread and
data-worker path, and the re-measure it was waiting on never ran. The
serving-config key, its validation, docs rows and test go; workers keep
main's MSTAR_PY_SWITCH_INTERVAL_SEC unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the GIL switch-interval knob's model defaults (the knob is removed in #858)
Reverts #859's additions on top of #858's api_switch_interval_s: the
Model.API_SWITCH_INTERVAL_S class default, Qwen3.5's 0.001 override,
the model argument to entrypoint._set_switch_interval, the doc wording
and the test fakes that carried the attribute. Every file #858 owns is
back to #858's text, so #858's removal of the knob merges in cleanly.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: the TP2 publish test's KV managers hold one rank's head
KVManager shards its config in build, so the test's direct constructor
shards it the same way and each rank moves its own shard.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server: drop warmup_requests' concurrency option
No config sets it, so its parsing, docs and tests covered a dead path. Each
entry still goes out once per data-parallel replica, now as
warmup-<spec>.<copy>.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 vision: a request too large for a replay gets an encoder step of its own
Up to 8 requests share an encoder step, and a step that cannot replay as
one runs them back to back (can_batch). For requests that are each past the
captured buckets (long videos) that only held every request's embeddings,
and so their prefills, until the last one was encoded: under load the
prefills bunched into waves (TTFT +30-50% on 60 s video at c=16/32, decode
almost never interleaved with a prefill). The vision node now splits its
steps by capture key, a request too large for a replay being a group of its
own, and the engine lets a node whose only capture is a piecewise region
split that way.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: Canary's decode cap leaves the prefill's token in the budget
#794's cap-versus-stop-check agreement tests for Canary and Canary-Qwen assumed
a stop check that counts max_tokens; TokenStop leaves the budget to the loop's
cap, so they pin the cap instead.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: last_prefill_flag for a prompt prefilled over several passes; GLM-5.2 test comment names the loop cap
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 vision split: review cleanups
- Engine.capture_group asks Engine.splits_by_capture_key; the module helper
existed only for a test fake, which now builds a real Engine.
- The split docstrings say a piecewise region is not tied to a walk, and that
on a piecewise-only node the key only groups the scheduler's rows.
- Stale 'eager' comments on the vision tower: past the top bucket the block
loop runs compiled, uncaptured.
- docs/models.rst: drop the one-boot measurement of a different overlay.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Chatterbox, Higgs Audio tests: a stop id at prefill ends the partition
#858 moved TokenStop's prefill_walks onto END_PARTITION; the step closing
a reply that ended at prefill no longer exists.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Orpheus test: the stop check's fake info names its walk
TokenStop now reads each stopping row's walk to choose between its loop
and END_PARTITION.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 9B DP2 config: the replica pick follows #920's router
The conductor now weighs a predicted prefix hit beside the live requests.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Review wording: the vision split's capture-only scope, Qwen3.5 configs' start-up compile, Qwen3-Omni's encoder header
- The Qwen3.5 vision split applies while the block loop is captured, and
requests that fit alone but not together still run one after another in
one step; cg_key_info's and capture_group's docstrings say a
piecewise-only node's key is a scheduler group, not a bucket.
- The Qwen3.5 configs' warmup comments: the vision regions comp…
merceod
pushed a commit
that referenced
this pull request
Nov 3, 2026
* docs/models.rst: the Omni audio encoder's eager and sub-second cases as the code has them (R2.1)
A step also packs, and a clip also replays, only within the attention
plan's `audio_attn.max_segments_per_request` windows, so a clip under
8192 rows but over that ceiling runs the layers eagerly too. A clip
under 1 s encoded alone differs from HF's in its last conv row, which
the windowed attention mixes into every embedding of the clip (one
window), not only in its last rows.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Head sharding: one per-resource switch, shard_over_sp; NodeSubmodule.splits_heads goes
Every resource now shards its heads over all TP x SP ranks of the node's
lockstep group (EngineResourceInfo.head_shards, a property), and a ragged
resource that each SP rank runs whole says so with
RaggedAttentionConfig(shard_over_sp=False). The node-level splits_heads flag,
the engine's per-node fold and its agreement check, and their docs and tests
are removed.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni Thinker prefill capture: padding rows take the broadcast positions (R2.2)
The `prefill_text` capture template, which `prefill_text`, `prefill_audio`
and the mixed walk replay, still carried `custom_pos_ids`, and a replay's
padding rows are built from it, so every padded text-only prefill or
mixed step built its MRoPE positions row by row (a cast per real row and
a cat) instead of the one broadcast view. The template is now a text
row, as the decode one is since R1.14; an audio row brings its own
positions.
GPU-unverified: prefill capture and padded text/mixed replays.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* GLM-5.2 prefill ladder through ARNodeSubmodule's override; the base declares the ladder (R2.3)
`model_kwargs.prefill_token_buckets` had two mechanisms: Qwen3.5 and
Qwen3-Omni pass it to `ARNodeSubmodule.__init__`, which refuses an empty
or repeated ladder, while GLM-5.2 read its config's copy at capture with
`or`, so `[]` silently meant the default and a repeat was captured
twice. GLM-5.2 now passes it to the same constructor and its capture
reads `self.PREFILL_TOKEN_BUCKETS`; the `or` fallback is gone.
`ARNodeSubmodule` declares `PREFILL_TOKEN_BUCKETS`, the attribute the
override replaces and the capture configs read.
Test: test_glm52_prefill::test_a_bad_prefill_ladder_is_refused ([] and
[64, 64]) did not raise before and raises ValueError now; the module's
helper builds the submodule through its constructor.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: the frontend takes the step's clips as a list; Omni and Qwen3-ASR stop padding them (R2.4)
Both encoder submodules zero-padded every clip of a step to the longest
and stacked them, and the frontend sliced each clip straight back out.
Omni's features arrive in fp32, so a long clip batched with short ones
allocated B x 128 x T_max x 4 B per step (about 0.5 GB for one 10 min
clip beside fifteen 2 s clips) for nothing. The frontend now takes the
clips as they are, `(num_mel_bins, frames)` each, and reads each one's
length off its shape; the padded stack, the `feature_lens` argument,
Qwen3-ASR's `num_frames` kwarg and its 2-D `forward` special case go.
The convolution input is built from the same frames in the same order,
so the embeddings are unchanged.
GPU-unverified: the transient memory saving (no CPU-observable change).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: one checkpoint loader, `AuTEncoder.from_checkpoint`, for Qwen3-Omni and Qwen3-ASR (R2.5)
The two models loaded the shared tower with the same six steps (build on
meta, cast, `to_empty`, stream `thinker.audio_tower.*`, `load_weights`,
`eval`), copied line for line. The sequence is now a classmethod on the
encoder, which owns the parameter layout, and both models call it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT frontend: pad the last convolution launch to a multiple of 32 chunks, not a power of two (R2.6)
The power-of-two rule could nearly double a launch: 257 chunks ran as
500 and 300 as 500, about 4 GFLOP per empty chunk, and Qwen3-ASR pays
it too since the frontend is shared. A multiple of 32 (8 for a launch
of at most 8 chunks) wastes at most 31 chunks at the cost of about 17
conv shapes instead of 7, each a one-time cuDNN plan.
GPU-unverified: the audio encoder's phase time on Omni (5 x 60 s clips)
and Qwen3-ASR (8 x 37 s), power of two vs multiple of 32.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice: code2wav never looks up the MASK id; unmask knobs must be finite (R1.1)
A canvas that still held MASK cells went into the codec decode, whose
codebooks end one id before the MASK id: on the GPU that lookup is a
device-side assert that kills the worker and every request on it. The
host check that used to refuse such a canvas went with the move of the
output shaping to the API side. The decode now reads code 0 in those
cells, and the NaN poison still makes postprocess refuse the waveform,
so only that request fails.
process_prompt now checks the unmask loop's float knobs are finite
numbers (one parser, also used where the conductor reads them): an
infinite reveal score ties every cell, which is how a client knob could
leave MASK cells. The refusal message names both causes of a non-finite
waveform instead of labelling a codec NaN as masked cells.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice: bound the output-shaping durations, take 0/1 for postprocess_output (R1.12)
output_shaping checked the shaping kwargs' types, not their ranges:
pad_duration 1e5 passed process_prompt and then had the API server
allocate tens of GB of padding after the whole unmask loop, and NaN or
inf failed only after generation. Both durations must now be finite and
between 0 and 10 seconds. postprocess_output takes 0 and 1 again, as JSON
clients send booleans (the base read it through bool(); the move to
postprocess had made them a 400).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: a decode row reads its text token by its loop iteration (R1.4)
Decode prepare_inputs read the streamed text condition at a per-request
counter and advanced it. prepare_inputs runs before admit, and a refused
admit prepares the step again, so each refusal skipped one text token
(a dropped sub-word, possibly an early EOS, no error). Mixed steps made
this reachable for decode rows: a step carrying a new prompt's pages is
refused under KV pressure, and every decode row in it was re-prepared.
The index is now the step's own iteration of talker_decode_loop
(InputMetadata.dynamic_loop_iter_counts, incremented for a speculated
step), which a refused admit leaves alone; the counter is gone.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: pack a step only up to the largest captured token count
Qwen3-Omni's HF audio encoder split a step into calls of at most 30,000
mel frames. On the AuT node nothing bounded a step, so eight long clips
ran as one eager forward over every clip's tokens. can_batch now packs a
step only while it fits the largest capture bucket (4096 tokens); past it
the clips run one by one, each replayed when it fits a bucket, so an
eager step holds one clip's activations. The same node serves Qwen3-ASR,
whose steps up to 4096 tokens are unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: count tokens by the checkpoint's own chunk, for both models
tokens_for_frames counted 13 tokens per 100-frame chunk whatever the
config's n_window, while the frontend cuts chunks of 2 * n_window frames
and Qwen3-Omni's prompt layout counted with transformers' helper at the
checkpoint's n_window. It now counts each full chunk_frames chunk by the
convolutions' arithmetic, and the layout calls it, so the prompt, the
encoder's split and its frontend share one count. Unchanged for the
released checkpoints (n_window 50).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Ragged attention: refuse a per-request ceiling on a resource sized per step
A resource with max_segments_per_step sizes its plan from that alone, so
a deployment override of max_segments_per_request was accepted and then
read by nothing. apply_yaml_overrides now refuses it, naming the per-step
key. Its docstring counts the ceilings it takes correctly.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Test: the engine hands a forward only the regions its step leased
The case called SubmoduleManagement.runners_for directly, so reverting the
engine's call to it went unnoticed. It now drives Engine._drive_step and
reads the runners the forward receives.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Cosmos3: take cudnn_benchmark from fp32_precision, drop the wan_vae re-export
The scope moved to mstar.engine.fp32_precision and wan_vae kept a
re-export for Cosmos3's two decode paths. They import it from its owner
now, the re-export goes, and its test moves beside the other fp32
precision tests.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni: drop audio_seqlens, which the audio encoder no longer reads
The AuT node takes a clip's frame count off its unpadded log-mel
features, so the audio_seqlens tensor rode to the encoder unread. The
prompt layout, its one reader, counts the same frames off the features'
shape, so the tensor, its graph input and its routing go.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder node: move it beside the encoder, one loader and spec for both models
Qwen3-Omni imported its audio encoder node from the qwen3_asr package and
repeated its loader and attention spec line for line. The node is now
AuTEncoderSubmodule in mstar/model/components/aut_encoder_submodule.py,
with AUT_ATTN, aut_attention_spec and load_aut_encoder_submodule, and
both models call them. It takes a LogMelSpectrogram instead of
Qwen3-ASR's config. Behaviour is unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: a Qwen3-Omni audio step over 4096 tokens runs its clips one by one
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor kernels: fit cp_linear's tiles to the device's shared memory (R1.2)
cp_linear's tiles were tuned on an H200 (227 KB of shared memory a
block). At 3 to 64 rows the SwiGLU tile buffers x and both weights per
pipeline stage, 40 to 48 KB, so four stages needed up to 192 KB: Triton
refuses to compile that on 99 KB GPUs (sm_86, sm_89: A10G, A40, L4, L40S,
RTX 3090/4090) and the 17-32 row tile on an A100 (163 KB), and the
Qwen3-TTS depth loop runs only through cp_linear.
The tile choice moves into _cp_linear_tile, which keeps the H200 choice
where it fits and otherwise keeps fewer stages, then shorter K steps,
until the stages fit the device's max_shared_mem. Every H200 tile is
unchanged. The stale harness names in the old comment are gone.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: on Base, batch each codec chunk width apart (R1.3)
Codec steps batch chunks of any width, each padded to the widest. An
in-context clone's first chunk leads with the reference tail, 73 frames
wide by default, so every row batched with it paid a 73-frame decode,
about three of a full chunk. The width-mixing A/B that motivated the mix
ran CustomVoice, which has no clone chunk.
Base, the one variant that clones, opts back into
split_batches_by_capture_key, as the codec did before mixing; the other
variants keep mixing widths. The docs said the scheduler batched every
width apart, which it no longer did; they now say which variant does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS 1.7B config: default pools, bf16 codec scoped to H200 (R1.10, R1.16, M.2, M.4, M.6, part of R1.11)
configs/qwen3tts_1p7b.yaml, the default for `mstar serve qwen3_tts_1p7b`,
allocated 4096 Talker KV pages (~56 GiB) and 512 codec slots at startup:
it could not start on a 40-48 GB GPU. It now keeps the default pools; the
overload-sized ones stay in the YAML as a commented override for GPUs with
80 GB or more, with their measurement dated before admission started
queueing requests for KV pages and codec slots.
The bf16 codec claim is scoped to what was measured: an H200. The B300
result in the docs was main's windowed codec, and the docs said "measure
before switching" next to a config that had already switched; the YAML
and docs now say to serve float32 on other GPUs until it is measured.
The codec_dtype comment no longer claims float32 is bit-exact (the
parity test allows one PCM16 step), and the comments say RMSNorm and
RoPE are computed in fp32 from bf16 weights and frequencies rather than
kept in fp32: the bf16 gate measured these numerics, so they stay.
The 1.7B YAML headers said only the checkpoint differed.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: OmniVoice split and Qwen3-TTS codec text match the code (R1.11)
- OmniVoice's split config and its docs said code2wav's worker does the
silence trim and post-processing; that runs on the API side now, and the
split's measured gain was taken before the move, which they now say.
- The Qwen3-TTS benchmark said the codec filters EOS frames; the Talker
withholds them from the codec stream.
- A Qwen3-TTS stream's audio matches a whole-utterance decode to within
one PCM16 step, as its parity test checks, rather than equalling it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: Base takes four new prompts into a mixed step, as the other variants do (R1.19)
max_rows_in_combined split MIXED_MAX_PREFILL_ROWS over the two prefill
walks when the checkpoint supports reference audio, against a step that
mixes both. No variant runs both: Base requires one reference clip, so
every Base prompt runs the clone prefill, and the others never do. Base
was capped at two new prompts a step where the docs say four.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Recurrent state pool: refuse a max_slots override that leaves no slot (R1.20)
RecurrentStateConfig refused max_slots below 2 (with the sink) when it
was built, but the YAML override set it afterwards unchecked: a
deployment's `codec_state: {max_slots: 1}` left no slot past the sink,
and every request waited in admission forever. The override now runs the
same check, which also covers a sinkless pool of zero slots.
The Qwen3-TTS docs say what N slots hold: N - 1 streams, the sink taking
one.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: refuse codec chunk knobs that are not frame counts (R1.17)
codec_chunk_schedule and codec_chunk_frames went through int(): [2.5, 8]
ran as (2, 8), 12.5 as 12, [true] as 1, and a scalar schedule failed with
a bare TypeError while the model was built. apply_codec_chunking now
takes only positive integers and raises ValueError otherwise, as
Qwen3-Omni's same-named knobs do. An empty ramp, refused before because
the clone width read its first entry, starts at the steady chunk, as on
Qwen3-Omni.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT frontend: pad the last convolution launch to the smaller of a power of two and a multiple of 32 (R3.1)
The multiple-of-32 rule (R2.6) ran a last launch of 9-16 chunks as 32,
where the power-of-two rule had run it as 16, so a single 9-16 s
Qwen3-ASR clip convolved 2-3.5x its chunks. Taking the smaller of the
two roundings (still at least 8, at most conv_chunksize) pads no launch
more than either rule did: 9-16 -> 16, 17-64 as before, 65-96 -> 96,
and at most 31 empty chunks, for 18 conv shapes instead of 17.
GPU-unverified: the encoder's time on Qwen3-ASR at c=1 with 10 s and
14 s clips, multiple of 32 vs this rule (CUDA events around
AuTEncoder.frontend, or worker_phases on the encoder node).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* ARNodeSubmodule: narrow the PREFILL_TOKEN_BUCKETS comment to what holds (R3.2)
The comment said every submodule that captures prefill names its ladder
PREFILL_TOKEN_BUCKETS. BAGEL's LLM (PREFILL_TEXT_TOKEN_BUCKETS), the
Qwen3-Omni Talker (TALKER_PREFILL_TOKEN_BUCKETS) and Cosmos3's DiT (an
instance ladder with an env-var override) don't, and an override passed
to one of them would be ignored. The comment now states the obligation
instead: a submodule that takes `prefill_token_buckets` must read its
capture ladder off `self.PREFILL_TOKEN_BUCKETS`, as Qwen3.5, the Omni
Thinker and GLM-5.2 do.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: fold the depth weights without autograd; drop dead hooks, caches and keys (R1.15)
fold_depth_weights multiplied each projection weight by its norm weight
with grad enabled, so every folded weight kept a grad_fn and its fp32
intermediates alive: about 0.34 GB a Talker on the 1.7B. It now runs
under no_grad.
Also gone, none with a reader:
- the lm_head_weight buffer, a third copy of the LM heads beside the
modules and the folded heads (its comment still cited the piecewise
loop); the heads are stacked only to fold them;
- both submodules' can_use_accelerator_graphs overrides (the engine never
calls it);
- the Talker's per-tuple id tensor cache, which only fed _fixed_embed's
builds, already cached per key;
- the Codec's codec_chunk_frames/codec_chunk_schedule step_metadata.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor kernels: cp_embed_accumulate says how many partials it wrote (R1.22)
cp_linear's norm reads n_part row partials that the producer wrote, and
the Qwen3-TTS depth loop worked out how many cp_embed_accumulate wrote by
hand (embed width // EMBED_BLOCK); a wrong count normalizes silently
wrong. cp_embed_accumulate now returns the count with the input, and
refuses a batch wider than the partials' rows instead of relying on
cp_linear's check running first. Its docstring no longer cites a
constant that doesn't exist.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS codec, code predictor kernels: ValueError for limits a caller can hit (R1.23)
The incremental decoder raised TypeError for a codec whose conv geometry
it can't decode, and cp_linear checked its row, width and partials
limits (tied to this checkpoint's shapes) with bare asserts. Both now
raise ValueError with the limit in the message; the stride checks stay
asserts.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS codec: drop a clamp that never applies (M.5)
prepare_inputs clamped a chunk's context items to its frame count, but
StreamBuffer.pop_chunk already caps context_items at the chunk's items,
and every item is one frame.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: cover cp_linear's two-row mode; drop Qwen3-TTS codec duplicates and pins (R1.24)
- cp_linear's two-row tile was never run; the parity test now covers it.
- The codec's empty-final-flush test, and the empty-chunk half of its
padding test, repeated cases of the EOS-withholding stream test
(num_frames 4 and 0); the state-pool test re-checked the block layout
that CodecSubmodule already refuses to build without.
- The codec graph-config test pinned CAPTURE_BATCH_SIZES and asserted
can_use_accelerator_graphs, a hook the engine never calls.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS codec bench: time and check the served decode, nothing else (R1.14)
The bench's "base" arm called the window decode that no checkout in the
stack has (_trim, a three-argument _decode, codec_window_contexts), its
"nwc" arm timed a 180-line prototype of the decoder, its "ncw" arm (the
served decoder) was labelled the other way round, and a SlotPool plus a
whole-step compile stood in for the state pool.
The base arm, the prototype and SlotPool are gone. ServedStep builds
the codec's state pool from the model's own codec_state declaration and
drives each step as the engine does (declare, admit, plan, the codec's
own _decode with its compiled decode, commit); timing and kernels replay
a planned step per chunk width, and parity streams real codes through it
in the model's chunk ramp against a whole float32 decode.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: run the Qwen3-TTS codec parity on the served CUDA path; fix the real-weights codec test (R1.18)
The streamed-codec parity test ran only eager fp32 on the CPU, while the
server runs the decode compiled (Inductor) over channels-last cuDNN
convs. It now also runs on CUDA, where CodecSubmodule compiles its
decode, with TF32 off so both decodes are fp32.
The real-weights codec test called _decode with the signature two
codecs ago. It now runs a served step (the model's declared state pool,
admitted and planned) and checks the bf16 codec against the float32 one
on the same codes.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Name the TP x SP head-sharding test for the kinds it covers; reflow a docs paragraph
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* AuT encoder: say a step holds up to 4096 tokens; import its loader at the top
The token budget applies to Qwen3-ASR as well as Qwen3-Omni, so the ASR
docs and the node's docstrings now say a larger step runs its clips one
by one. Both models import load_aut_encoder_submodule beside
aut_attention_spec, from the module they already load at the top.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 9B DP2 config: give the GDN pool prefix snapshot slots (RM.1)
The dp2 config was written before #877 and has no `gdn_state.prefix_slots`,
so the recurrent pool can't serve a cached prefix and the engine keeps the
LLM node's KV prefix cache closed with it. Each replica runs at world size 1,
so it takes the same 128 snapshot slots as qwen3_5_9b.yaml.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests, docs: run the Qwen3-TTS codec parity with TF32 off for matmuls too; scope its claim (R2.1)
Importing mstar.engine sets float32_matmul_precision('high'), so turning off
only cuDNN's TF32 left the codec's GEMM-lowered transposed and kernel-1 convs
in TF32 on the CUDA cases. Both decodes now run under
fp32_precision_scope(Fp32Precision("highest", cudnn_allow_tf32=False)). The
one-PCM16-step claim in the model's comment and docs/models.rst is scoped to
float32 with TF32 off (as on CPU); the served CUDA codec runs its GEMMs in TF32.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS codec bench: drop --tf32, which could not turn TF32 off (R2.5)
The flag was applied before the mstar imports set float32_matmul_precision
('high'), so every run used TF32 matmuls, as the served decode does, and a
run without it was mislabelled as plain float32.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: drop the cp_linear tile test that restates the shared-memory formula (R2.3)
It asserted _cp_linear_tile's own stage_bytes expression, so it could not
catch a formula that undercounts what Triton allocates. The real check is
test_cp_linear_matches_the_unfused_ops compiling on a 99 KB (sm_86/sm_89) GPU.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: drop Talker asserts on can_use_accelerator_graphs, a hook no engine path calls (R2.4)
R1.15 deleted the Talker's override, so these three asserts only tested the
base class's default; their comment described the deleted behaviour.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: say that Base's four-prompt mixed cap can spill long clones to eager (R2.2)
The cap counts prompts, not tokens, so four clones with references of about
18 s or more outgrow the 1024-token top mixed bucket and that step runs eager.
A token cap would be a new scheduler concept for a rare case; the comment
states the limit instead.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice, code predictor kernels: say what the knob check and the partials limit cover (R2.6)
unmask_knobs' finite-number check stops non-finite knobs only; finite
extremes that overflow the reveal scores still fail in postprocess, so the
docstring and the test comment no longer claim it stops that case.
cp_embed_accumulate's ValueError names the partials-count limit as well as
the row limit, since either trips it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan 2.2: keep the dit's heads whole over SP with shard_over_sp=False
The node-level splits_heads flag is gone; the dit's one ragged attention
declares the same thing on the resource. Under SP each rank still runs
every head (the ranks split CFG branches), and TP still splits them.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: bring precision and bucket comments in line with R2 fixes (R3.1-R3.3)
R3.1: the mixed-step class comment no longer says the prompt cap always keeps
a step in the captured buckets; it points at max_rows_in_combined, where the
long-reference eager spill is described.
R3.2: the served fp32 codec runs its cuDNN convs in TF32 as well as its
GEMMs; the CUDA test case is the served code path at fp32, not served
precision.
R3.3: codec_bench parity compares against a whole-sequence decode at the
served precision (TF32), not a float32 or reference decode.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: run a guided window as its own walk on one node too
A single-node (or LLM-disaggregated) deployment ran a guided request as one
image_gen walk and read each step's branch count off the loop counter against
a CFGWindow in step metadata, with a guard for a counter that disagreed with
the seeded latents. The CFG-parallel layouts already ran the window as its own
walk with a fixed count. Every layout now does: a guided request runs
image_gen_head (if the window starts late), then image_gen_guided2 or
image_gen_guided3 on LLM alone, then image_gen (or image_decode when the window
reaches the last step). Each step runs the same branches as before.
DenoiseWalk gains `split` (a node per branch or not), ImageGenPhase loses its
`branches`, and the per-step count, its guard, the CFG_BRANCHES kwarg and the
cfg_window / image_gen_branches step metadata are deleted. The disaggregated
configs and the docs' examples list the new walks for the image LLM.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: fix stale codec parity and mixed-step wording
R1.1, R1.2 (tail review): the codec_bench parity reference is always a
float32 decode (TF32 GEMMs and convs), not "the served precision" under a
bf16 config; and the model doc now notes that four long-reference clones
can outgrow the top mixed token bucket and run eager, as the class comment
already does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Say the api_switch_interval_s numbers predate #849's output-path rewrite
The row still said the knob would be re-measured once #849 lands. It has
landed and changed the message-thread and data-worker path the knob
tunes, so the numbers are from before that change and the re-measure is
still pending.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the head-sharding paragraph from the submodule attributes section
It was there to document the removed splits_heads attribute; what is
left repeats the RaggedAttentionConfig row.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the TP x SP head-sharding test that main's sizing already passes
Head sizing now equals main's, so the six cases pinned nothing this
branch changes. The shard_over_sp=False test stays; the helper loses its
recurrent branch, which only the deleted test used.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: hold guidance KV only for the branches a request's window runs
A request is guided when its CFG window has a step, not when a scale is above
1: a window with no step (a cfg_interval that excludes every timestep) runs no
guidance, so it no longer prefills or decodes cfg_img or forks cfg_text.
`_requires_cfg` is gone; process_prompt resolves the window before the layout
and the conductor reads it from the stamped window.
The layout reserves a branch label's denoise scratch only where the window
runs that branch: cfg_img's only with image guidance, and neither branch's
for an unguided request, which no layout denoises on them any more. A guided
text-to-image request no longer holds a latent-sized cfg_img scratch it never
writes.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the shard_over_sp sizing test, which main's sizing already passes
Main already sizes a shard_over_sp=False ragged attention over the TP
group, so the test pinned nothing this branch changes. The file's
coverage goes back to main's plus the manager registry test.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: comments state the contract instead of the code they replaced
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL docs: give the measured output shift of guided text-to-image
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test_kv_admission: name the CFG walk BAGEL defines
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Cosmos3 and LingBot decode through wan_vae.FastDecoder and decode_dtype
wan_vae.FastDecoder and decode_dtype had one user, Wan2.2, while Cosmos3's
two decoder nodes and LingBot-World kept their own copies of the cuDNN
>= 9.16 dtype gate, of the channels-last + DecoderStep setup and of the
autocast-off decode_causal call. All of them now use the shared helpers,
and the copies (_vae_decode_dtype, _fast_step, the streaming node's _step,
LingBot's inline gate and _decode_step) are deleted.
Numerics and switches are unchanged: Cosmos3 keeps COSMOS3_VAE_DECODE_FP32,
its compile switch, its cast before the layout switch, the streaming
step's non-deterministic compile (built only with the warm cache), its
latents per call and cudnn_benchmark scope; LingBot keeps no override
(decode_dtype's variable is now optional), and bf16 and channels-last only
off reference numerics. Every node keeps the layout it set up, so
FastDecoder's per-call layout check is a no-op for them.
FastDecoder's autocast-off scope now follows the latents' device instead of
naming cuda.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan 2.2: drop the FFN GELU epilogue, which has no measurement of its own
MLP(gelu_epilogue=) and Wan22DiT's fuse_ffn_gelu ran the FFN GELU in the
up-projection's cuBLASLt epilogue through the private
torch._addmm_activation. Nothing measured it apart from the ragged
attention change it shipped with, so it goes, as unmeasured performance
paths do: MLP is back to its plain form and the DiT builds without the
flag.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3 text encoder: refuse a row without text_len instead of never replaying
The shared encoder node picks its captured (masked) graph from each row's
text_len step metadata. A model that reused the node without setting it
got every batch run eager, with no error. cg_key_info now raises
ValueError for such a row; FLUX.2-klein and Z-Image set it at intake.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Image output knobs: an explicit null is the default, and only integers pass intake
image_output_format let png_compress_level: null through, and encode_image
then failed on int(None) after the denoise and decode. encode_image now
reads null as the default level, as it already did for output_compression.
Intake also refused nothing that int() accepted, so True passed as 1 and
50.9 was truncated to 50; it now takes only integers (not bools) in range,
as BAGEL's width and height check does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* V-JEPA2: say the B=1 guard routes eager steps only, and drop "as before" comments
A step that leases the AC rollout capture replays forward_batched at any
batch size, B=1 included, whatever can_batch returns, so can_batch's B=1
guard covers eager steps only; its comment now says so. Two comments
described the code by comparison with an earlier version ("as before the
capture", "same as before") and now state what the code does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* wan_vae.decode_dtype: say that an override of 0 forces bf16
The override variable is two-way, as Wan2.2's and Cosmos3's were before it
was shared (and as the Wan fast-decode test uses it): set, it forces fp32,
or bf16 when it is 0/false/no/off. The docstring read as if 0 meant "not
forced".
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Higgs Audio: declare text_output so inline decode tokens skip the per-row postprocess
Higgs's decode token rides inline, but without a declared text_output
Model.postprocess_inline built a tensor and called postprocess for every
row on every step. Declaring TextOutput(self._detokenizer), with no stop
ids as before, gives it the memoized id-to-bytes path Qwen3-ASR and
Whisper use; postprocess hands text to the base class. Bytes are
unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5: say the API_SWITCH_INTERVAL_S gain predates #849
docs/serving.rst says the number was measured before #849 and awaits a
re-measure; the comment pointing at it now says the same.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS codec bench: unambiguous float32 reference wording (tail R2 P3)
The parenthetical read as if the server runs float32; the shipped config
runs a bfloat16 codec. Say the reference matches a float32-codec server.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Conductor: build its instance table from WorkerGraph.replica_ranks
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Image output knobs: serve the whole numbers main served
Intake took only Python ints, so a whole-number float (90.0) or a number in
a JSON string ("90"), which main read through int() and served, became a
400. It now accepts any whole number int() reads in range, and still refuses
bools, fractions such as 50.9, infinities and out-of-range values.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: build only the window walks a deployment can run
A guided window runs on the branch nodes on CFG-parallel deployments and on
LLM alone elsewhere, never both, yet CFG-parallel builds still placed the
unsplit image_gen_guided2/3 walks and listed them in every request's
needed-label map. LLM_GRAPH_WALKS is now keyed by the layout, and the walk
graphs and the label map both read it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: say why a split LLM still publishes both branch labels
The comment dated from when a remote image_gen walk's branch count varied
per step. Each walk now fixes its count; what keeps cfg_img published on a
split deployment is a remote decode reading it and the single-node denoise
walks naming their labels per request rather than per walk.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL and Cosmos3: drop loop-counter leftovers from comments and names
denoise_loop's docstring still named the step's loop counter, which BAGEL
no longer reads; declare_step's walk-branch narrowing kept a requires_cfg
guard that active_labels already implies, and its guided-row set shadowed
the module's guided(). Cosmos3's decode-dtype comment now gives both
directions of its override.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Image output: check only the chosen format's compression knob
Intake refused a request whose other-format knob was out of range (a JPEG
request with png_compress_level 10, a PNG request with output_compression
150), which main served because encode_image never reads it. Check only
the knob the format applies.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL: say LLM-disaggregated, not split, in the branch-label comment
"Split" means CFG-parallel in the neighbouring code (DenoiseWalk.split,
LLM_GRAPH_WALKS).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* TP read-failure test: derive the loop-back hooks from device_loopback = None
The stand-in stubbed the deleted _device_loopback and restated the three
hooks DeviceLoopback now derives; clearing the declaration gives the same
answers through the base class.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* TP read-failure test: drop the stale "emit" from the stand-in's comment
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Merge-up fixes: DeviceLoopback in the withhold check, chunk-sink stand-in poll, postprocess_s on inline failure
- _refuse_unroutable_withholds read TokenStop.loopback_signal, which main
moved to ARNodeSubmodule.device_loopback; it now reads the signal from
the submodule's DeviceLoopback, and its test declares one.
- test_chunk_sink: its request profile stub lacked output_postprocess_s,
so every pass raised before relaying a chunk; it uses a real
RequestProfile, its transport's poll ends the pass (the loop no longer
sleeps), and an error in the pass is raised instead of retried.
- An inline value whose batched postprocess_inline call failed is charged
its share of that call as well as its own retry.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Withhold check: the loop-back signal is exempt only on its own walk
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Withhold check test: the loop-back exemption is scoped to its walk
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-ASR/Omni AuT node: pin the replay's windows and packed-clip parity on the shared node
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor: pin the unprojected depth path, drop the unused DepthAttention.forward
The depth-loop parity test now runs at both Talker widths: 2048 through
the input projection (Qwen3-TTS 1.7B) and 1024 without one (Qwen3-Omni,
Qwen3-TTS 0.6B), where each step's input is the raw embedding and the
first norm reads cp_embed_accumulate's per-chunk partials. The
cp_embed_accumulate test checks those partials for both the raw and the
projected input.
DepthAttention.forward and its attend reference were the unfused layer,
and nothing runs them: run_depth_layers reads the weights and attends
through the predictor's DepthKV. The class stays as the layers' parameter
container that takes none of the node's resources, which the binding
test now checks directly. run_depth_layers' eager note keeps only the
verified reason.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* cp_attend: say why the multi-token path stays
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor: keep the compiled depth step free of per-position guards
depth_step took the LM head index as an int and passed it to the disabled
run_depth_layers, so under a static-int compile Dynamo compiled depth_step
once per head (guards "head is None", "head == 0", ...): one frame per depth
position, up to the recompile limit. depth_step itself is now the eager
boundary, with run_depth_layers folded into it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* cp_attend: name the row cap a multi-token opening would need raised
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor: name every eager reader of the depth ints; count the test's groups
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS codec: compile the incremental decode through compile_regions
The codec compiled its own decode with torch.compile(dynamic=False), which
recompiled once per captured shape (4 chunk widths x 14 batch buckets = 55
compiles) and skipped the engine's inductor overrides and recompile
accounting. On an empty cache that was ~16 of the ~19 minutes of a 1.7B cold
start. The decode is now a compile region: the engine compiles it once, with
dynamic shapes, before capture.
H100, slice-6 head (empty Inductor/Triton cache, 26 CPUs): cold start
1119 -> 444 s, warm 540 -> 333 s. Closed loop audio-s/s c=1/16/32/64, 2
rounds each: push4a 18.4-18.5 / 137-145 / 209-212 / 281-284 -> 18.4 /
139-146 / 208-212 / 289-301; push4c (server on 2 CPUs) 18.6/159/237/370 ->
18.4-18.5 / 160 / 237-251 / 371-379. Pinning the chunk width static (one
compile per width, batch dynamic) measured no faster and started in 765 s.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS codec: test and bench compile the decode as the engine does
The codec's decode is compiled by the engine's compile_regions pass, so
a CodecSubmodule built directly ran it eagerly. The per-node body of
Engine._compile_regions becomes compile_node_regions, which the engine
loop, the streamed-codec test's CUDA case and codec_bench's load_codec
all call. The compile_regions docstring and models.rst now say what
dynamic=True gives: a few graphs (batch 1 and width 1 specialise) for
every chunk width and batch bucket.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Codec and compile_node_regions docstrings: the measured cold start, no dynamic=None aside
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS 1.7B: start the codec stream with a 2-frame chunk
The deployment config's ramp began with a 4-frame chunk, so a stream's
first audio waited for four Talker steps. [2, 4, 8, 16] sends it after two,
and the 160 ms first chunk still covers the time until the next one arrives
under load. [1, 2, 5, 16] is documented as the lower-latency opt-in.
H100 (push4a and push4c), closed loop, c=64: TTFA p50 84-95 -> 63-78 ms
(c=1 33 -> 25 ms), no playback stalls either way, audio-s/s 379-424 vs
363-416 over 4-6 rounds each (same within noise). Open-loop SLO capacity
(TTFA <= 200 ms, stall < 10 ms): Seed-TTS 80 @90/99 and passages 28 @90
either way. WER (Whisper large-v3, 128 sentences): 2.09% -> 1.73%.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* qwen3tts_1p7b config: the header names the 2-frame start; date the TTFA numbers
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS 1.7B: the bf16 codec is measured on H100 too
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS 1.7B bf16 codec: state what the H100 runs measured
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS bf16 codec: tie each H100 WER to its chunk schedule
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Attention: select_last_hidden is one eager boundary, so prefill buckets don't recompile its frame
The disabled `qo_indptr_buf` call broke the graph inside `select_last_hidden`,
leaving a resume frame that holds both the bucket's static token count
(`hidden`) and its row count (`qo_indptr`). Captured forwards compile with
dynamic=False, so that frame took one cache entry per (tokens, rows) bucket:
Qwen3.5's 9 token buckets by 7 prefill/mixed row counts made 63, and the
13-step prefill ladder made 91, past the recompile limit of 84.
Disabling the whole method moves the break to its call site, where the
caller already breaks; no compiled frame holds the pair. Same for the MLA
and XPU managers, which share the shape.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL takes its prefill's last rows through select_last_hidden; test docstrings
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* GDN prefill run: state the zeroed-tail precondition the mixed wrapper also meets
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice: key the backbone's step outputs by the engine's row ids
The packed backbone keyed its outputs by CanvasItem.request_id, which
prepare_inputs now fills from fwd_info.request_id, the wire request id. The
engine looks each row's outputs up by its worker-local handle, found none,
and every OmniVoice request failed in postprocess with KeyError 'c_logits'.
forward_batched now keys rows by engine_inputs.request_ids, as the other
batched submodules do, and CanvasItem carries no request id at all, so the
canvas can't be keyed by the wrong one again. The unread request_id kwarg
goes with it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni 2-GPU config: keep the 8192 Thinker ladder off the 80 GB default
qwen3omni_2gpu.yaml, the CLI default and the config sized for 80 GB GPUs,
ran the Thinker's prefill ladder to 8192. On an H100 its 4096-8192 captures
don't fit beside the 1,024-page KV pool (59 of 86 Thinker buckets captured),
so those steps run eagerly and a deployment that requires accelerator graphs
refuses to start. The default goes back to the model's 128-2048 ladder, as on
the base. The 8192 ladder, measured on 2xH200, stays in
qwen3omni_2gpu_large_kv.yaml, which is already the H200 variant; its header,
the default's, the submodule comment and the Qwen3-Omni notes say so.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS 1.7B configs: back to the base's codec, chunk ramp and pools (owned by #978)
Kirill's #978 sets these himself (the [4, 8, 16] ramp, the bf16 codec,
the codec_state slots, the GC knob and their comments), so this PR leaves
both 1.7B configs as its base has them. This reverts the config hunks of
fce78c050, f63eb7994, aba669bac, 9ef1d1e66, 08c8b2350, e4450ea53 and
5428cca16, including the [2, 4, 8, 16] first-chunk schedule.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: drop the incremental codec, mixed walk and fused depth loop (owned by #978)
Kirill's #978 makes the same Qwen3-TTS changes (from the same
stacked-head work): the incremental codec with its codec_state pool and
chunk knobs, the combined prefill+decode walk and its captures, the fixed
prefill embeddings, and a fused depth loop (cp_linear with folded norms,
cp_attend, fold_depth_weights) beside the codec's own compile. This PR
leaves all of it to him: Qwen3-TTS, the shared code predictor kernels and
Qwen3-Omni's two call sites go back to the base's code, and so do their
tests, the real-weights codec test, the Qwen3-TTS and Qwen3-Omni
code-predictor paragraphs in docs/models.rst, and the codec bench, which
timed the removed decoder.
That reverts fce78c050 and the fixes made to its code since: 59d1b3854,
f5b0422ff, 2f6333a90, e59639f21, 251d13e34, add26fbc7, f7887b71a,
4d9f54584, e6b2a2fef, a2331d014, ce5693c3d, 1158293d0, 8ebf63986,
613160070, b13ab78e4, 9a588f978, 26ef71b6e, 15413c499, 5fea4023e,
ab3350e18, 4f935c1f6, a8c537e7d, c4ed93f0c, 88b5b68de, 5bfeeac60,
a5687bdd7, the code-predictor half of 4791cb3ca and the Qwen3-TTS halves
of f63eb7994, 7906d0902 and 13b35b17a.
Kept: db8b241c7, a decode row reads its text token by its loop iteration
rather than a counter a re-prepared step advances twice. The counter is
the base's (and main's) code, which #978 leaves as it is.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Engine: drop compile_node_regions, which only the removed codec test and bench called
13b35b17a and 7a8be7397 lifted Engine._compile_regions' body into a
module-level compile_node_regions so the Qwen3-TTS codec test and bench
could compile the decode as the engine does. Both are gone with the
incremental codec (owned by #978), so _compile_regions takes its body
back, as on the base. No behaviour change.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-Omni 2-GPU config: don't claim the remaining captures fit 80 GB
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Pin video decode threads to the CPUs the plan leaves under pinning
A case with a pinned CPU fails if video_decode_threads goes back to
usable_cpus(); the MSTAR_STEP_THREAD_CORE row names FFmpeg's per-decode
threads among what the pins shrink.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Ragged attention: a wrapper skips replanning the layout it last planned
The dense manager kept the (cu_q, cu_k, causal) each of its wrappers was
last planned with and skipped an unchanged plan; the ragged self- and
cross-attention managers replanned every step, so a denoise loop paid
FlashInfer's plan (a kernel and a stream sync) once per step for the same
segments. The skip now lives in the wrapper both use
(_RaggedPrefillBase._plan): each wrapper is one (label, slot) or one
(bucket, slot, label), so the schedule it keeps is the one the step reading
it planned. The dense manager's own table goes.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan2.2 docs: date the attention and CFG-parallel timings as before #963
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Conductor: leave the idle wait to its own change
The conductor's 1 ms nap becomes a wait on the poll in a change of its
own, which waits until the next timed check instead of a fixed 1 ms. This
branch's capped poll, its failed-pass back-off, the conductor stand-ins'
poll_for_messages and their test go, so those files are main's again.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server tests: the failed-wait test pins its back-off; stand-in waits as main has them
The message loop always passes a computed timeout, so the 1 ms defaults
left by the fixed poll go: the chunk-sink stand-in takes main's again and
the failed-wait stand-in takes it as required. That test now asserts the
10 ms back-off, without which a persistent wait error would spin.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Preprocess pool: at most the CPUs the plan leaves unpinned
Every other API-server pool is sized from the CPU plan; the preprocess
pool took its declared or configured size uncapped, so under
step_thread_core a model declaring 8 workers could run them, each with its
own FFmpeg threads, on fewer free CPUs.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Data worker: drain routes before an input only when it prepares here
The drain keeps a call's held items from waiting behind a preprocess that
may take seconds. An input prepared on the pool only stores and sends when
it is handed on, and each pass already drains routes at its top, so it no
longer takes a drain of its own.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: say what a RecurrentStateSpec max_slots override must leave
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 vision: compile the block loop, so a step past the captured buckets runs compiled
disable_torch_compile left the block loop uncompiled whenever it did not
replay: a long video (over 16384 patches) ran the blocks' fp32 rope, norms
and activations as eager kernels, slower than the whole-forward compile it
replaced. The block loop is now a compile region like the embed and the
merger, dynamic in the patch count; it compiles in its capture's warm-up.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server: drop the api_switch_interval_s knob
Its gain was measured before #849 and #939 changed the message-thread and
data-worker path, and the re-measure it was waiting on never ran. The
serving-config key, its validation, docs rows and test go; workers keep
main's MSTAR_PY_SWITCH_INTERVAL_SEC unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the GIL switch-interval knob's model defaults (the knob is removed in #858)
Reverts #859's additions on top of #858's api_switch_interval_s: the
Model.API_SWITCH_INTERVAL_S class default, Qwen3.5's 0.001 override,
the model argument to entrypoint._set_switch_interval, the doc wording
and the test fakes that carried the attribute. Every file #858 owns is
back to #858's text, so #858's removal of the knob merges in cleanly.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: the TP2 publish test's KV managers hold one rank's head
KVManager shards its config in build, so the test's direct constructor
shards it the same way and each rank moves its own shard.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* API server: drop warmup_requests' concurrency option
No config sets it, so its parsing, docs and tests covered a dead path. Each
entry still goes out once per data-parallel replica, now as
warmup-<spec>.<copy>.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 vision: a request too large for a replay gets an encoder step of its own
Up to 8 requests share an encoder step, and a step that cannot replay as
one runs them back to back (can_batch). For requests that are each past the
captured buckets (long videos) that only held every request's embeddings,
and so their prefills, until the last one was encoded: under load the
prefills bunched into waves (TTFT +30-50% on 60 s video at c=16/32, decode
almost never interleaved with a prefill). The vision node now splits its
steps by capture key, a request too large for a replay being a group of its
own, and the engine lets a node whose only capture is a piecewise region
split that way.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Tests: Canary's decode cap leaves the prefill's token in the budget
#794's cap-versus-stop-check agreement tests for Canary and Canary-Qwen assumed
a stop check that counts max_tokens; TokenStop leaves the budget to the loop's
cap, so they pin the cap instead.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: last_prefill_flag for a prompt prefilled over several passes; GLM-5.2 test comment names the loop cap
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 vision split: review cleanups
- Engine.capture_group asks Engine.splits_by_capture_key; the module helper
existed only for a test fake, which now builds a real Engine.
- The split docstrings say a piecewise region is not tied to a walk, and that
on a piecewise-only node the key only groups the scheduler's rows.
- Stale 'eager' comments on the vision tower: past the top bucket the block
loop runs compiled, uncaptured.
- docs/models.rst: drop the one-boot measurement of a different overlay.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Chatterbox, Higgs Audio tests: a stop id at prefill ends the partition
#858 moved TokenStop's prefill_walks onto END_PARTITION; the step closing
a reply that ended at prefill no longer exists.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Orpheus test: the stop check's fake info names its walk
TokenStop now reads each stopping row's walk to choose between its loop
and END_PARTITION.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5 9B DP2 config: the replica pick follows #920's router
The conductor now weighs a predicted prefix hit beside the live requests.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Review wording: the vision split's capture-only scope, Qwen3.5 configs' start-up compile, Qwen3-Omni's encoder header
- The Qwen3.5 vision split applies while the block loop is captured, and
requests that fit alone but not together still run one after another in
one step; cg_key_info's and capture_group's docstrings say a
piecewise-only node's key is a scheduler group, not a bucket.
- The Qwen3.5 configs' warmup comments: the vision regions compile at
start-up.
- Qwen3-Omni's submodule header takes main's line for its audio encoder.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan2.2 docs: date the timings by the DiT step they measured, not a PR number
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan2.2 docs: an earlier version of the DiT block, not an earlier DiT step
"DiT step" reads as a denoising step in this section.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Engine: a capped deployment holds cuDNN's capture search to its share
capture_cudnn_benchmark's search runs in capture sizing since the sized
warm-ups took the node's scope, and it tries workspaces as large as the
GPU has free: Kokoro reserved ~140 GB of an H200 for ~25 s at startup,
past the gpu_memory_fraction a co-located worker (or an 80 GB-capped
comparison) plans around. With a fraction set, the allocator is held to
it while such a node is sized, and an over-cap try is skipped.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Docs: the preprocess pool's CPU cap does not count a prepare's FFmpeg threads
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Chunk policy: leave the longer primed first chunk to #977
#977 makes the same change (ScheduledLeftContextChunkPolicy.prime past
left_context, and its test), and its consumer, the incremental Qwen3-TTS
codec, is no longer in this stack: every primer here passes at most
left_context (Qwen3-TTS's lead is min(ref_frames, left_context_frames)).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* ViT patch embed: the conv timing in its comment is Qwen3.5's
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Engine: _search_cap's docstring separates a skipped search try from a failed sizing
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5-9B DP2 config: say the vision regions compile at start-up
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Engine: capture_cudnn_benchmark's search runs in capture sizing, held to a capped deployment's share
Sizing runs a node's warm-up forwards; the capture after it warms up with
no forward, so the scope around warmup_and_capture covered no conv on CUDA
and Kokoro's graphs kept cuDNN's heuristic picks. The scope now wraps
sizing.
The search tries workspaces as large as the GPU has free, and the allocator
keeps them until sizing's empty_cache: Kokoro reserved ~140 GB of an H200
for ~25 s at startup, past the gpu_memory_fraction a co-located worker (or
an 80 GB-capped comparison) plans around. With a fraction set, the
allocator is held to it while such a node is sized and restored to its
previous value after; an over-cap search try is skipped, and the node's own
sizing fails the boot when its eager step does not fit in the share.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Fish S2 Pro: capture the decode buckets uncompiled, as on main
Compiling the seven decode buckets took a cold start from 180 s to 1372 s
on H100, with no measured throughput to show for it. The 64-row bucket and
the packed prefill captures stay.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Fish S2 Pro: TF32 codec decode is the model's default
The shipped config set codec_exact_fp32: false while the model defaulted to
the exact decode, so other Fish configs decoded differently. The default now
lives in the model; codec_exact_fp32: true still opts into the exact decode.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor: open the depth loop with position 0 alone
Both Talkers opened with one-token passes on main; the two-token opening
changed greedy codes for both, and a one-token pass keeps every pass at bs
rows. The opening is now main's.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3-TTS: drop the 24- and 48-row decode buckets
They took the 1.7B's cold start from 204 s to 307 s on H100, and a step at
64 rows costs about what it does at 32, so padding into 32 and 64 is nearly
free.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* run_layers: the residual stream starts from its own copy
The stream was the caller's hidden_states, overwritten in place from layer
0, so a caller had to know to pass a copy (the Qwen3-TTS Talker did for its
packed replay's static input). One [tokens, hidden] copy per forward, made
inside the forward so a capture records it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* OmniVoice silence test: pydub is the oracle, installed with dev
The comparison imported the reference package, which CI does not install,
so it always skipped there. It now restates the reference's trim over
pydub.silence (raw-sample segments, no ffmpeg), with pydub in the dev
extras, and covers the rates and edge cases the self-referential trailing
trim test stood in for; that test is gone.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Code predictor: each Talker keeps its own embedding-sum dtype
The shared depth loop summed a frame's codec embeddings in fp32 for both
Talkers. Qwen3-TTS accumulated in the model dtype, rounding after every add,
so its greedy codes could differ at near-ties. ResidualCodePredictor takes
embed_sum_dtype from its Talker: None (the embeddings' dtype) for Qwen3-TTS,
fp32 for Qwen3-Omni, whose fp32 sum is rounded once to the model dtype as
before. The accumulate kernel adds in fp32 and rounds to the sum's dtype, as
torch's add_ does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Review wording: Qwen3-TTS's embedding sum, the cuDNN search in capture sizing; audioop-lts for pydub on 3.13
- Qwen3-TTS's per-add rounding is the sum this Talker had in M*; the
reference sums in fp32 and rounds once.
- adding_models.rst and Kokoro say the cuDNN search runs in capture sizing.
- pydub 0.25.1 imports audioop, which Python 3.13 removed: the dev extras
add audioop-lts there, so the silence oracle does not skip.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* TokenStop: withholding outputs and the runtime check refuse each other
The graph runtime's stop check never calls check_stop, which is what
withholds a stop step's outputs, so a TokenStop with both would drop the
withhold silently. It is a frozen dataclass now, and that combination
raises at construction.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test_knobs: take main's version back now that ingraph_decode_rope stays
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan2.2 CFG-parallel: the scaffold keeps ragged_for, docs say a rank runs its branch over the batch
Since the dit moved onto the scaffold's denoise loop, DenoiseLoopSubmodule is
RaggedAttentionMixin's only user, so the mixin folds back into it as on main.
The declare_step test keeps only what the per-rank segment test does not
cover.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* BAGEL docs: quote the guided text-to-image PSNR and the configuration of each number
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* wan22 extra: install flashinfer-python for the default ragged attention
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Wan2.2: leave CFG-parallel to the cfg_size axis
Drop wan22_cfg_parallel_2gpu.yaml, my_rows and the dit's sp-group plumbing,
which used the SP slot to mean a CFG split. The cfg_size axis carries it.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Z-Image-Turbo: capture batch sizes 1, 2, 4, 8; every size up to 8 is a commented opt-in
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* V-JEPA2-AC test: a rollout request's stamped clip shape keys the encoder capture
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5-9B configs: keep main's 2176 bucket for a mixed step
The 9B ladders replace the default, so they lost the slack bucket that
lets a mixed step of four 512-token prompts and its decode rows replay at
2176 rather than pad to 4096. Adding it keeps the compile count at 82,
strictly under Dynamo's limit of 84; the 9B ladders keep 3072-7168 in place
of main's 4352 and 8704. The mixed-step bucket and the recompile limit are
now checked on every ladder a Qwen3.5 config ships, not only the default.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Sampler RNG offset test: both captured kernels, each against a draw it can match
At V=64 the captured step takes fused_sample, whose Triton RNG never agrees
with the eager step's FlashInfer draw, so the test passed whether or not the
two steps shared an offset. The eager case now draws its captured step through
FlashInfer too, and a second case runs two captured steps on fused_sample,
which reads the same offset buffer the replay advances.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Data worker: refuse a bad preprocess_workers before the thread starts
The thread reads preprocess_workers as it starts, so a value it cannot take
("auto") ended it while the server came up, and every request then waited
for its timeout; 0 or a negative quietly meant inline, and true meant 1.
PreprocessWorker now refuses anything but a positive integer, on the
caller's thread.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* TokenStepNode: copy only the outputs the per-request stop check reads
The per-request fallback aliased text_inputs as new_token with a dict splat,
which reads every output name, so the lazy host outputs copied all of them
to the host. A ChainMap over the outputs reads only what the check reads.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Check every shipped prefill ladder, not only Qwen3.5's
The ladder tests read only configs/qwen3_5*.yaml, so the Qwen3-Omni large_kv
config's ladder was never checked. Every config that sets
model_kwargs.prefill_token_buckets is now built into its own model's prefill
capture; the Qwen3.5-specific checks keep to Qwen3.5 configs.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Ladder test: pin the config-to-model plumbing, drop the Qwen3.5 repeat
Every shipped ladder's model must take prefill_token_buckets as its own
argument (it would otherwise vanish into **kwargs); the Thinker's capture
check stays for Qwen3-Omni, whose ladder no other test builds.
Co-Authored-By: Claude Opus 5…
merceod
pushed a commit
that referenced
this pull request
Nov 3, 2026
* Qwen3.5 9B DP2 config: give the GDN pool prefix snapshot slots (RM.1) The dp2 config was written before #877 and has no `gdn_state.prefix_slots`, so the recurrent pool can't serve a cached prefix and the engine keeps the LLM node's KV prefix cache closed with it. Each replica runs at world size 1, so it takes the same 128 snapshot slots as qwen3_5_9b.yaml. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests, docs: run the Qwen3-TTS codec parity with TF32 off for matmuls too; scope its claim (R2.1) Importing mstar.engine sets float32_matmul_precision('high'), so turning off only cuDNN's TF32 left the codec's GEMM-lowered transposed and kernel-1 convs in TF32 on the CUDA cases. Both decodes now run under fp32_precision_scope(Fp32Precision("highest", cudnn_allow_tf32=False)). The one-PCM16-step claim in the model's comment and docs/models.rst is scoped to float32 with TF32 off (as on CPU); the served CUDA codec runs its GEMMs in TF32. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS codec bench: drop --tf32, which could not turn TF32 off (R2.5) The flag was applied before the mstar imports set float32_matmul_precision ('high'), so every run used TF32 matmuls, as the served decode does, and a run without it was mislabelled as plain float32. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: drop the cp_linear tile test that restates the shared-memory formula (R2.3) It asserted _cp_linear_tile's own stage_bytes expression, so it could not catch a formula that undercounts what Triton allocates. The real check is test_cp_linear_matches_the_unfused_ops compiling on a 99 KB (sm_86/sm_89) GPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: drop Talker asserts on can_use_accelerator_graphs, a hook no engine path calls (R2.4) R1.15 deleted the Talker's override, so these three asserts only tested the base class's default; their comment described the deleted behaviour. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: say that Base's four-prompt mixed cap can spill long clones to eager (R2.2) The cap counts prompts, not tokens, so four clones with references of about 18 s or more outgrow the 1024-token top mixed bucket and that step runs eager. A token cap would be a new scheduler concept for a rare case; the comment states the limit instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice, code predictor kernels: say what the knob check and the partials limit cover (R2.6) unmask_knobs' finite-number check stops non-finite knobs only; finite extremes that overflow the reveal scores still fail in postprocess, so the docstring and the test comment no longer claim it stops that case. cp_embed_accumulate's ValueError names the partials-count limit as well as the row limit, since either trips it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan 2.2: keep the dit's heads whole over SP with shard_over_sp=False The node-level splits_heads flag is gone; the dit's one ragged attention declares the same thing on the resource. Under SP each rank still runs every head (the ranks split CFG branches), and TP still splits them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: bring precision and bucket comments in line with R2 fixes (R3.1-R3.3) R3.1: the mixed-step class comment no longer says the prompt cap always keeps a step in the captured buckets; it points at max_rows_in_combined, where the long-reference eager spill is described. R3.2: the served fp32 codec runs its cuDNN convs in TF32 as well as its GEMMs; the CUDA test case is the served code path at fp32, not served precision. R3.3: codec_bench parity compares against a whole-sequence decode at the served precision (TF32), not a float32 or reference decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * compile_utils: docstring covers the VAE decode side and group_norm (pre-push review P3-1) The module docstring named only the flow transformers and decode_utils as callers; group_norm also asks kept_exact for the VAE's fused GroupNorm+SiLU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: run a guided window as its own walk on one node too A single-node (or LLM-disaggregated) deployment ran a guided request as one image_gen walk and read each step's branch count off the loop counter against a CFGWindow in step metadata, with a guard for a counter that disagreed with the seeded latents. The CFG-parallel layouts already ran the window as its own walk with a fixed count. Every layout now does: a guided request runs image_gen_head (if the window starts late), then image_gen_guided2 or image_gen_guided3 on LLM alone, then image_gen (or image_decode when the window reaches the last step). Each step runs the same branches as before. DenoiseWalk gains `split` (a node per branch or not), ImageGenPhase loses its `branches`, and the per-step count, its guard, the CFG_BRANCHES kwarg and the cfg_window / image_gen_branches step metadata are deleted. The disaggregated configs and the docs' examples list the new walks for the image LLM. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: fix stale codec parity and mixed-step wording R1.1, R1.2 (tail review): the codec_bench parity reference is always a float32 decode (TF32 GEMMs and convs), not "the served precision" under a bf16 config; and the model doc now notes that four long-reference clones can outgrow the top mixed token bucket and run eager, as the class comment already does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Say the api_switch_interval_s numbers predate #849's output-path rewrite The row still said the knob would be re-measured once #849 lands. It has landed and changed the message-thread and data-worker path the knob tunes, so the numbers are from before that change and the re-measure is still pending. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the head-sharding paragraph from the submodule attributes section It was there to document the removed splits_heads attribute; what is left repeats the RaggedAttentionConfig row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the TP x SP head-sharding test that main's sizing already passes Head sizing now equals main's, so the six cases pinned nothing this branch changes. The shard_over_sp=False test stays; the helper loses its recurrent branch, which only the deleted test used. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: hold guidance KV only for the branches a request's window runs A request is guided when its CFG window has a step, not when a scale is above 1: a window with no step (a cfg_interval that excludes every timestep) runs no guidance, so it no longer prefills or decodes cfg_img or forks cfg_text. `_requires_cfg` is gone; process_prompt resolves the window before the layout and the conductor reads it from the stamped window. The layout reserves a branch label's denoise scratch only where the window runs that branch: cfg_img's only with image guidance, and neither branch's for an unguided request, which no layout denoises on them any more. A guided text-to-image request no longer holds a latent-sized cfg_img scratch it never writes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the shard_over_sp sizing test, which main's sizing already passes Main already sizes a shard_over_sp=False ragged attention over the TP group, so the test pinned nothing this branch changes. The file's coverage goes back to main's plus the manager registry test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: comments state the contract instead of the code they replaced Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL docs: give the measured output shift of guided text-to-image Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test_kv_admission: name the CFG walk BAGEL defines Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Cosmos3 and LingBot decode through wan_vae.FastDecoder and decode_dtype wan_vae.FastDecoder and decode_dtype had one user, Wan2.2, while Cosmos3's two decoder nodes and LingBot-World kept their own copies of the cuDNN >= 9.16 dtype gate, of the channels-last + DecoderStep setup and of the autocast-off decode_causal call. All of them now use the shared helpers, and the copies (_vae_decode_dtype, _fast_step, the streaming node's _step, LingBot's inline gate and _decode_step) are deleted. Numerics and switches are unchanged: Cosmos3 keeps COSMOS3_VAE_DECODE_FP32, its compile switch, its cast before the layout switch, the streaming step's non-deterministic compile (built only with the warm cache), its latents per call and cudnn_benchmark scope; LingBot keeps no override (decode_dtype's variable is now optional), and bf16 and channels-last only off reference numerics. Every node keeps the layout it set up, so FastDecoder's per-call layout check is a no-op for them. FastDecoder's autocast-off scope now follows the latents' device instead of naming cuda. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan 2.2: drop the FFN GELU epilogue, which has no measurement of its own MLP(gelu_epilogue=) and Wan22DiT's fuse_ffn_gelu ran the FFN GELU in the up-projection's cuBLASLt epilogue through the private torch._addmm_activation. Nothing measured it apart from the ragged attention change it shipped with, so it goes, as unmeasured performance paths do: MLP is back to its plain form and the DiT builds without the flag. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3 text encoder: refuse a row without text_len instead of never replaying The shared encoder node picks its captured (masked) graph from each row's text_len step metadata. A model that reused the node without setting it got every batch run eager, with no error. cg_key_info now raises ValueError for such a row; FLUX.2-klein and Z-Image set it at intake. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Image output knobs: an explicit null is the default, and only integers pass intake image_output_format let png_compress_level: null through, and encode_image then failed on int(None) after the denoise and decode. encode_image now reads null as the default level, as it already did for output_compression. Intake also refused nothing that int() accepted, so True passed as 1 and 50.9 was truncated to 50; it now takes only integers (not bools) in range, as BAGEL's width and height check does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2: say the B=1 guard routes eager steps only, and drop "as before" comments A step that leases the AC rollout capture replays forward_batched at any batch size, B=1 included, whatever can_batch returns, so can_batch's B=1 guard covers eager steps only; its comment now says so. Two comments described the code by comparison with an earlier version ("as before the capture", "same as before") and now state what the code does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * wan_vae.decode_dtype: say that an override of 0 forces bf16 The override variable is two-way, as Wan2.2's and Cosmos3's were before it was shared (and as the Wan fast-decode test uses it): set, it forces fp32, or bf16 when it is 0/false/no/off. The docstring read as if 0 meant "not forced". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Higgs Audio: declare text_output so inline decode tokens skip the per-row postprocess Higgs's decode token rides inline, but without a declared text_output Model.postprocess_inline built a tensor and called postprocess for every row on every step. Declaring TextOutput(self._detokenizer), with no stop ids as before, gives it the memoized id-to-bytes path Qwen3-ASR and Whisper use; postprocess hands text to the base class. Bytes are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: say the API_SWITCH_INTERVAL_S gain predates #849 docs/serving.rst says the number was measured before #849 and awaits a re-measure; the comment pointing at it now says the same. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS codec bench: unambiguous float32 reference wording (tail R2 P3) The parenthetical read as if the server runs float32; the shipped config runs a bfloat16 codec. Say the reference matches a float32-codec server. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Conductor: build its instance table from WorkerGraph.replica_ranks Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Image output knobs: serve the whole numbers main served Intake took only Python ints, so a whole-number float (90.0) or a number in a JSON string ("90"), which main read through int() and served, became a 400. It now accepts any whole number int() reads in range, and still refuses bools, fractions such as 50.9, infinities and out-of-range values. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: build only the window walks a deployment can run A guided window runs on the branch nodes on CFG-parallel deployments and on LLM alone elsewhere, never both, yet CFG-parallel builds still placed the unsplit image_gen_guided2/3 walks and listed them in every request's needed-label map. LLM_GRAPH_WALKS is now keyed by the layout, and the walk graphs and the label map both read it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: say why a split LLM still publishes both branch labels The comment dated from when a remote image_gen walk's branch count varied per step. Each walk now fixes its count; what keeps cfg_img published on a split deployment is a remote decode reading it and the single-node denoise walks naming their labels per request rather than per walk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL and Cosmos3: drop loop-counter leftovers from comments and names denoise_loop's docstring still named the step's loop counter, which BAGEL no longer reads; declare_step's walk-branch narrowing kept a requires_cfg guard that active_labels already implies, and its guided-row set shadowed the module's guided(). Cosmos3's decode-dtype comment now gives both directions of its override. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Image output: check only the chosen format's compression knob Intake refused a request whose other-format knob was out of range (a JPEG request with png_compress_level 10, a PNG request with output_compression 150), which main served because encode_image never reads it. Check only the knob the format applies. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: say LLM-disaggregated, not split, in the branch-label comment "Split" means CFG-parallel in the neighbouring code (DenoiseWalk.split, LLM_GRAPH_WALKS). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TP read-failure test: derive the loop-back hooks from device_loopback = None The stand-in stubbed the deleted _device_loopback and restated the three hooks DeviceLoopback now derives; clearing the declaration gives the same answers through the base class. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TP read-failure test: drop the stale "emit" from the stand-in's comment Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Merge-up fixes: DeviceLoopback in the withhold check, chunk-sink stand-in poll, postprocess_s on inline failure - _refuse_unroutable_withholds read TokenStop.loopback_signal, which main moved to ARNodeSubmodule.device_loopback; it now reads the signal from the submodule's DeviceLoopback, and its test declares one. - test_chunk_sink: its request profile stub lacked output_postprocess_s, so every pass raised before relaying a chunk; it uses a real RequestProfile, its transport's poll ends the pass (the loop no longer sleeps), and an error in the pass is raised instead of retried. - An inline value whose batched postprocess_inline call failed is charged its share of that call as well as its own retry. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Withhold check: the loop-back signal is exempt only on its own walk Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Withhold check test: the loop-back exemption is scoped to its walk Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-ASR/Omni AuT node: pin the replay's windows and packed-clip parity on the shared node Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: pin the unprojected depth path, drop the unused DepthAttention.forward The depth-loop parity test now runs at both Talker widths: 2048 through the input projection (Qwen3-TTS 1.7B) and 1024 without one (Qwen3-Omni, Qwen3-TTS 0.6B), where each step's input is the raw embedding and the first norm reads cp_embed_accumulate's per-chunk partials. The cp_embed_accumulate test checks those partials for both the raw and the projected input. DepthAttention.forward and its attend reference were the unfused layer, and nothing runs them: run_depth_layers reads the weights and attends through the predictor's DepthKV. The class stays as the layers' parameter container that takes none of the node's resources, which the binding test now checks directly. run_depth_layers' eager note keeps only the verified reason. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * FLUX.2 klein: one shared kept-module forward, so blocks don't recompile per instance exclude_from_compile set a functools.partial forward on each kept 1-D RMSNorm. Dynamo ID-guards a partial set on an instance, so the single-stream block frame klein's blocks share failed its guards once per block until the recompile limit, and the remaining blocks ran eagerly: -6 to -10% img/s at c=1-4 on klein 4B and 9B against main. A kept module now changes class instead, to a subclass made once per (class, path) whose forward is the class's: the exact RMSNorm op for a plain 1-D nn.RMSNorm, torch.compiler.disable of the base forward for every other kept module (the same graph break as before). compile_kept_exact is a class attribute of that subclass, so kept_exact (Z-Image's fused q/k norm + rotary, the VAE's norm_act) still answers. Outputs are unchanged: the same exact_rms_norm call, the same disabled base forward. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * compile_utils: state the guard that recompiled per block Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * cp_attend: say why the multi-token path stays Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: keep the compiled depth step free of per-position guards depth_step took the LM head index as an int and passed it to the disabled run_depth_layers, so under a static-int compile Dynamo compiled depth_step once per head (guards "head is None", "head == 0", ...): one frame per depth position, up to the recompile limit. depth_step itself is now the eager boundary, with run_depth_layers folded into it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * compile_utils: every block passes the same guards, not one check Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * cp_attend: name the row cap a multi-token opening would need raised Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: name every eager reader of the depth ints; count the test's groups Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS codec: compile the incremental decode through compile_regions The codec compiled its own decode with torch.compile(dynamic=False), which recompiled once per captured shape (4 chunk widths x 14 batch buckets = 55 compiles) and skipped the engine's inductor overrides and recompile accounting. On an empty cache that was ~16 of the ~19 minutes of a 1.7B cold start. The decode is now a compile region: the engine compiles it once, with dynamic shapes, before capture. H100, slice-6 head (empty Inductor/Triton cache, 26 CPUs): cold start 1119 -> 444 s, warm 540 -> 333 s. Closed loop audio-s/s c=1/16/32/64, 2 rounds each: push4a 18.4-18.5 / 137-145 / 209-212 / 281-284 -> 18.4 / 139-146 / 208-212 / 289-301; push4c (server on 2 CPUs) 18.6/159/237/370 -> 18.4-18.5 / 160 / 237-251 / 371-379. Pinning the chunk width static (one compile per width, batch dynamic) measured no faster and started in 765 s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS codec: test and bench compile the decode as the engine does The codec's decode is compiled by the engine's compile_regions pass, so a CodecSubmodule built directly ran it eagerly. The per-node body of Engine._compile_regions becomes compile_node_regions, which the engine loop, the streamed-codec test's CUDA case and codec_bench's load_codec all call. The compile_regions docstring and models.rst now say what dynamic=True gives: a few graphs (batch 1 and width 1 specialise) for every chunk width and batch bucket. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Codec and compile_node_regions docstrings: the measured cold start, no dynamic=None aside Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B: start the codec stream with a 2-frame chunk The deployment config's ramp began with a 4-frame chunk, so a stream's first audio waited for four Talker steps. [2, 4, 8, 16] sends it after two, and the 160 ms first chunk still covers the time until the next one arrives under load. [1, 2, 5, 16] is documented as the lower-latency opt-in. H100 (push4a and push4c), closed loop, c=64: TTFA p50 84-95 -> 63-78 ms (c=1 33 -> 25 ms), no playback stalls either way, audio-s/s 379-424 vs 363-416 over 4-6 rounds each (same within noise). Open-loop SLO capacity (TTFA <= 200 ms, stall < 10 ms): Seed-TTS 80 @90/99 and passages 28 @90 either way. WER (Whisper large-v3, 128 sentences): 2.09% -> 1.73%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * qwen3tts_1p7b config: the header names the 2-frame start; date the TTFA numbers Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B: the bf16 codec is measured on H100 too Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: klein's exact q/k norm, measured with the shared-forward fix Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B bf16 codec: state what the H100 runs measured Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: klein's shared kept-norm forward, measured against a forward per block Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS bf16 codec: tie each H100 WER to its chunk schedule Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: kept norms share a forward per class Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Attention: select_last_hidden is one eager boundary, so prefill buckets don't recompile its frame The disabled `qo_indptr_buf` call broke the graph inside `select_last_hidden`, leaving a resume frame that holds both the bucket's static token count (`hidden`) and its row count (`qo_indptr`). Captured forwards compile with dynamic=False, so that frame took one cache entry per (tokens, rows) bucket: Qwen3.5's 9 token buckets by 7 prefill/mixed row counts made 63, and the 13-step prefill ladder made 91, past the recompile limit of 84. Disabling the whole method moves the break to its call site, where the caller already breaks; no compiled frame holds the pair. Same for the MLA and XPU managers, which share the shape. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL takes its prefill's last rows through select_last_hidden; test docstrings Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * GDN prefill run: state the zeroed-tail precondition the mixed wrapper also meets Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice: key the backbone's step outputs by the engine's row ids The packed backbone keyed its outputs by CanvasItem.request_id, which prepare_inputs now fills from fwd_info.request_id, the wire request id. The engine looks each row's outputs up by its worker-local handle, found none, and every OmniVoice request failed in postprocess with KeyError 'c_logits'. forward_batched now keys rows by engine_inputs.request_ids, as the other batched submodules do, and CanvasItem carries no request id at all, so the canvas can't be keyed by the wrong one again. The unread request_id kwarg goes with it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni 2-GPU config: keep the 8192 Thinker ladder off the 80 GB default qwen3omni_2gpu.yaml, the CLI default and the config sized for 80 GB GPUs, ran the Thinker's prefill ladder to 8192. On an H100 its 4096-8192 captures don't fit beside the 1,024-page KV pool (59 of 86 Thinker buckets captured), so those steps run eagerly and a deployment that requires accelerator graphs refuses to start. The default goes back to the model's 128-2048 ladder, as on the base. The 8192 ladder, measured on 2xH200, stays in qwen3omni_2gpu_large_kv.yaml, which is already the H200 variant; its header, the default's, the submodule comment and the Qwen3-Omni notes say so. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B configs: back to the base's codec, chunk ramp and pools (owned by #978) Kirill's #978 sets these himself (the [4, 8, 16] ramp, the bf16 codec, the codec_state slots, the GC knob and their comments), so this PR leaves both 1.7B configs as its base has them. This reverts the config hunks of fce78c0, f63eb79, aba669b, 9ef1d1e, 08c8b23, e4450ea and 5428cca, including the [2, 4, 8, 16] first-chunk schedule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: drop the incremental codec, mixed walk and fused depth loop (owned by #978) Kirill's #978 makes the same Qwen3-TTS changes (from the same stacked-head work): the incremental codec with its codec_state pool and chunk knobs, the combined prefill+decode walk and its captures, the fixed prefill embeddings, and a fused depth loop (cp_linear with folded norms, cp_attend, fold_depth_weights) beside the codec's own compile. This PR leaves all of it to him: Qwen3-TTS, the shared code predictor kernels and Qwen3-Omni's two call sites go back to the base's code, and so do their tests, the real-weights codec test, the Qwen3-TTS and Qwen3-Omni code-predictor paragraphs in docs/models.rst, and the codec bench, which timed the removed decoder. That reverts fce78c0 and the fixes made to its code since: 59d1b38, f5b0422, 2f6333a, e59639f, 251d13e, add26fb, f7887b7, 4d9f545, e6b2a2f, a2331d0, ce5693c, 1158293, 8ebf639, 6131600, b13ab78, 9a588f9, 26ef71b, 15413c4, 5fea402, ab3350e, 4f935c1, a8c537e, c4ed93f, 88b5b68, 5bfeeac, a5687bd, the code-predictor half of 4791cb3 and the Qwen3-TTS halves of f63eb79, 7906d09 and 13b35b1. Kept: db8b241, a decode row reads its text token by its loop iteration rather than a counter a re-prepared step advances twice. The counter is the base's (and main's) code, which #978 leaves as it is. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: drop compile_node_regions, which only the removed codec test and bench called 13b35b1 and 7a8be73 lifted Engine._compile_regions' body into a module-level compile_node_regions so the Qwen3-TTS codec test and bench could compile the decode as the engine does. Both are gone with the incremental codec (owned by #978), so _compile_regions takes its body back, as on the base. No behaviour change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni 2-GPU config: don't claim the remaining captures fit 80 GB Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Pin video decode threads to the CPUs the plan leaves under pinning A case with a pinned CPU fails if video_decode_threads goes back to usable_cpus(); the MSTAR_STEP_THREAD_CORE row names FFmpeg's per-decode threads among what the pins shrink. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ragged attention: a wrapper skips replanning the layout it last planned The dense manager kept the (cu_q, cu_k, causal) each of its wrappers was last planned with and skipped an unchanged plan; the ragged self- and cross-attention managers replanned every step, so a denoise loop paid FlashInfer's plan (a kernel and a stream sync) once per step for the same segments. The skip now lives in the wrapper both use (_RaggedPrefillBase._plan): each wrapper is one (label, slot) or one (bucket, slot, label), so the schedule it keeps is the one the step reading it planned. The dense manager's own table goes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 docs: date the attention and CFG-parallel timings as before #963 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Conductor: leave the idle wait to its own change The conductor's 1 ms nap becomes a wait on the poll in a change of its own, which waits until the next timed check instead of a fixed 1 ms. This branch's capped poll, its failed-pass back-off, the conductor stand-ins' poll_for_messages and their test go, so those files are main's again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * API server tests: the failed-wait test pins its back-off; stand-in waits as main has them The message loop always passes a computed timeout, so the 1 ms defaults left by the fixed poll go: the chunk-sink stand-in takes main's again and the failed-wait stand-in takes it as required. That test now asserts the 10 ms back-off, without which a persistent wait error would spin. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Preprocess pool: at most the CPUs the plan leaves unpinned Every other API-server pool is sized from the CPU plan; the preprocess pool took its declared or configured size uncapped, so under step_thread_core a model declaring 8 workers could run them, each with its own FFmpeg threads, on fewer free CPUs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Data worker: drain routes before an input only when it prepares here The drain keeps a call's held items from waiting behind a preprocess that may take seconds. An input prepared on the pool only stores and sends when it is handed on, and each pass already drains routes at its top, so it no longer takes a drain of its own. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: say what a RecurrentStateSpec max_slots override must leave Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision: compile the block loop, so a step past the captured buckets runs compiled disable_torch_compile left the block loop uncompiled whenever it did not replay: a long video (over 16384 patches) ran the blocks' fp32 rope, norms and activations as eager kernels, slower than the whole-forward compile it replaced. The block loop is now a compile region like the embed and the merger, dynamic in the patch count; it compiles in its capture's warm-up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * API server: drop the api_switch_interval_s knob Its gain was measured before #849 and #939 changed the message-thread and data-worker path, and the re-measure it was waiting on never ran. The serving-config key, its validation, docs rows and test go; workers keep main's MSTAR_PY_SWITCH_INTERVAL_SEC unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the GIL switch-interval knob's model defaults (the knob is removed in #858) Reverts #859's additions on top of #858's api_switch_interval_s: the Model.API_SWITCH_INTERVAL_S class default, Qwen3.5's 0.001 override, the model argument to entrypoint._set_switch_interval, the doc wording and the test fakes that carried the attribute. Every file #858 owns is back to #858's text, so #858's removal of the knob merges in cleanly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the TP2 publish test's KV managers hold one rank's head KVManager shards its config in build, so the test's direct constructor shards it the same way and each rank moves its own shard. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * API server: drop warmup_requests' concurrency option No config sets it, so its parsing, docs and tests covered a dead path. Each entry still goes out once per data-parallel replica, now as warmup-<spec>.<copy>. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision: a request too large for a replay gets an encoder step of its own Up to 8 requests share an encoder step, and a step that cannot replay as one runs them back to back (can_batch). For requests that are each past the captured buckets (long videos) that only held every request's embeddings, and so their prefills, until the last one was encoded: under load the prefills bunched into waves (TTFT +30-50% on 60 s video at c=16/32, decode almost never interleaved with a prefill). The vision node now splits its steps by capture key, a request too large for a replay being a group of its own, and the engine lets a node whose only capture is a piecewise region split that way. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: Canary's decode cap leaves the prefill's token in the budget #794's cap-versus-stop-check agreement tests for Canary and Canary-Qwen assumed a stop check that counts max_tokens; TokenStop leaves the budget to the loop's cap, so they pin the cap instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: last_prefill_flag for a prompt prefilled over several passes; GLM-5.2 test comment names the loop cap Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision split: review cleanups - Engine.capture_group asks Engine.splits_by_capture_key; the module helper existed only for a test fake, which now builds a real Engine. - The split docstrings say a piecewise region is not tied to a walk, and that on a piecewise-only node the key only groups the scheduler's rows. - Stale 'eager' comments on the vision tower: past the top bucket the block loop runs compiled, uncaptured. - docs/models.rst: drop the one-boot measurement of a different overlay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chatterbox, Higgs Audio tests: a stop id at prefill ends the partition #858 moved TokenStop's prefill_walks onto END_PARTITION; the step closing a reply that ended at prefill no longer exists. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus test: the stop check's fake info names its walk TokenStop now reads each stopping row's walk to choose between its loop and END_PARTITION. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 9B DP2 config: the replica pick follows #920's router The conductor now weighs a predicted prefix hit beside the live requests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Review wording: the vision split's capture-only scope, Qwen3.5 configs' start-up compile, Qwen3-Omni's encoder header - The Qwen3.5 vision split applies while the block loop is captured, and requests that fit alone but not together still run one after another in one step; cg_key_info's and capture_group's docstrings say a piecewise-only node's key is a scheduler group, not a bucket. - The Qwen3.5 configs' warmup comments: the vision regions compile at start-up. - Qwen3-Omni's submodule header takes main's line for its audio encoder. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 docs: date the timings by the DiT step they measured, not a PR number Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 docs: an earlier version of the DiT block, not an earlier DiT step "DiT step" reads as a denoising step in this section. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: a capped deployment holds cuDNN's capture search to its share capture_cudnn_benchmark's search runs in capture sizing since the sized warm-ups took the node's scope, and it tries workspaces as large as the GPU has free: Kokoro reserved ~140 GB of an H200 for ~25 s at startup, past the gpu_memory_fraction a co-located worker (or an 80 GB-capped comparison) plans around. With a fraction set, the allocator is held to it while such a node is sized, and an over-cap try is skipped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: the preprocess pool's CPU cap does not count a prepare's FFmpeg threads Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunk policy: leave the longer primed first chunk to #977 #977 makes the same change (ScheduledLeftContextChunkPolicy.prime past left_context, and its test), and its consumer, the incremental Qwen3-TTS codec, is no longer in this stack: every primer here passes at most left_context (Qwen3-TTS's lead is min(ref_frames, left_context_frames)). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * ViT patch embed: the conv timing in its comment is Qwen3.5's Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: _search_cap's docstring separates a skipped search try from a failed sizing Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5-9B DP2 config: say the vision regions compile at start-up Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: capture_cudnn_benchmark's search runs in capture sizing, held to a capped deployment's share Sizing runs a node's warm-up forwards; the capture after it warms up with no forward, so the scope around warmup_and_capture covered no conv on CUDA and Kokoro's graphs kept cuDNN's heuristic picks. The scope now wraps sizing. The search tries workspaces as large as the GPU has free, and the allocator keeps them until sizing's empty_cache: Kokoro reserved ~140 GB of an H200 for ~25 s at startup, past the gpu_memory_fraction a co-located worker (or an 80 GB-capped comparison) plans around. With a fraction set, the allocator is held to it while such a node is sized and restored to its previous value after; an over-cap search try is skipped, and the node's own sizing fails the boot when its eager step does not fit in the share. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish S2 Pro: capture the decode buckets uncompiled, as on main Compiling the seven decode buckets took a cold start from 180 s to 1372 s on H100, with no measured throughput to show for it. The 64-row bucket and the packed prefill captures stay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish S2 Pro: TF32 codec decode is the model's default The shipped config set codec_exact_fp32: false while the model defaulted to the exact decode, so other Fish configs decoded differently. The default now lives in the model; codec_exact_fp32: true still opts into the exact decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: open the depth loop with position 0 alone Both Talkers opened with one-token passes on main; the two-token opening changed greedy codes for both, and a one-token pass keeps every pass at bs rows. The opening is now main's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: drop the 24- and 48-row decode buckets They took the 1.7B's cold start from 204 s to 307 s on H100, and a step at 64 rows costs about what it does at 32, so padding into 32 and 64 is nearly free. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * run_layers: the residual stream starts from its own copy The stream was the caller's hidden_states, overwritten in place from layer 0, so a caller had to know to pass a copy (the Qwen3-TTS Talker did for its packed replay's static input). One [tokens, hidden] copy per forward, made inside the forward so a capture records it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice silence test: pydub is the oracle, installed with dev The comparison imported the reference package, which CI does not install, so it always skipped there. It now restates the reference's trim over pydub.silence (raw-sample segments, no ffmpeg), with pydub in the dev extras, and covers the rates and edge cases the self-referential trailing trim test stood in for; that test is gone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: each Talker keeps its own embedding-sum dtype The shared depth loop summed a frame's codec embeddings in fp32 for both Talkers. Qwen3-TTS accumulated in the model dtype, rounding after every add, so its greedy codes could differ at near-ties. ResidualCodePredictor takes embed_sum_dtype from its Talker: None (the embeddings' dtype) for Qwen3-TTS, fp32 for Qwen3-Omni, whose fp32 sum is rounded once to the model dtype as before. The accumulate kernel adds in fp32 and rounds to the sum's dtype, as torch's add_ does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Review wording: Qwen3-TTS's embedding sum, the cuDNN search in capture sizing; audioop-lts for pydub on 3.13 - Qwen3-TTS's per-add rounding is the sum this Talker had in M*; the reference sums in fp32 and rounds once. - adding_models.rst and Kokoro say the cuDNN search runs in capture sizing. - pydub 0.25.1 imports audioop, which Python 3.13 removed: the dev extras add audioop-lts there, so the silence oracle does not skip. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStop: withholding outputs and the runtime check refuse each other The graph runtime's stop check never calls check_stop, which is what withholds a stop step's outputs, so a TokenStop with both would drop the withhold silently. It is a frozen dataclass now, and that combination raises at construction. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test_knobs: take main's version back now that ingraph_decode_rope stays Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 CFG-parallel: the scaffold keeps ragged_for, docs say a rank runs its branch over the batch Since the dit moved onto the scaffold's denoise loop, DenoiseLoopSubmodule is RaggedAttentionMixin's only user, so the mixin folds back into it as on main. The declare_step test keeps only what the per-rank segment test does not cover. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL docs: quote the guided text-to-image PSNR and the configuration of each number Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * wan22 extra: install flashinfer-python for the default ragged attention Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2: leave CFG-parallel to the cfg_size axis Drop wan22_cfg_parallel_2gpu.yaml, my_rows and the dit's sp-group plumbing, which used the SP slot to mean a CFG split. The cfg_size axis carries it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Z-Image-Turbo: capture batch sizes 1, 2, 4, 8; every size up to 8 is a commented opt-in Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2-AC test: a rollout request's stamped clip shape keys the encoder capture Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5-9B configs: keep main's 2176 bucket for a mixed step The 9B ladders replace the default, so they lost the slack bucket that lets a mixed step of four 512-token prompts and its decode rows replay at 2176 rather than pad to 4096. Adding it keeps the compile count at 82, strictly under Dynamo's limit of 84; the 9B ladders keep 3072-7168 in place of main's 4352 and 8704. The mixed-step bucket and the recompile limit are now checked on every ladder a Qwen3.5 config ships, not only the default. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2-AC capture-key test: decode the sampled frames as the request does process_prompt now decodes the frames load_video sampled, by index, so the test hands it sampled indices and stands in for the decoder rather than a decoded clip, which a request no longer carries. The Z-Image capture-size figures say they predate the fused q/k norm + rotary. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Sampler RNG offset test: both captured kernels, each against a draw it can match At V=64 the captured step takes fused_sample, whose Triton RNG never agrees with the eager step's FlashInfer draw, so the test passed whether or not the two steps shared an offset. The eager case now draws its captured step through FlashInfer too, and a second case runs two captured steps on fused_sample, which reads the same offset buffer the replay advances. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Data worker: refuse a bad preprocess_workers before the thread starts The thread reads preprocess_workers as it starts, so a value it cannot take ("auto") ended it while the server came up, and every request then waited for its timeout; 0 or a negative quietly meant inline, and true meant 1. PreprocessWorker now refuses anything but a positive integer, on the caller's thread. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2 pool sizes: say what they cost and when to change them The configs' preprocess_workers is taken as given, so each says it was measured on an H200 host, how many intra-op threads the pool then runs, and to lower it (or drop the key for the data worker's CPU-aware default) on a smaller host. The model comments name CpuPlan.preprocess_threads, and the pool docs name V-JEPA2 beside Qwen3-Omni. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStepNode: copy only the outputs the per-request stop check reads The per-request fallback aliased text_inputs as new_token with a dict splat, which reads every output name, so the lazy host outputs copied all of them to the host. A ChainMap over the outputs reads only what the check reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Check every shipped prefill ladder, not only Qwen3.5's The ladder tests read only configs/qwen3_5*.yaml, so the Qwen3-Omni large_kv config's ladder was never checked. Every config that sets model_kwargs.prefill_token_buckets is now built into its own model's prefill capture; the Qwen3.5-specific checks keep to Qwen3.5 configs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: pin the config-to-model plumbing, drop the Qwen3.5 repeat Every shipped ladder's model must take prefill_token_buckets as its own argument (it would otherwise vanish into **kwargs); the Thinker's capture check stays for Qwen3-Omni, whose ladder no other test builds. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: say what the plumbing check covers, and why it fails Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: read the Thinker's prefill capture by walk Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: dispatch on the model class, and name what it checks Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Address #971 review: num_step on the data worker, slot per label, postprocess on CPU CI - OmniVoice: process_prompt parses num_step too, so a non-integer is a 400 rather than a raise at the conductor that hangs the client. - RecurrentStateConfig: refuse a pool whose usable slots cannot hold one request's labels (two labels, one usable slot never admits). - test_omnivoice_silence: stand in for omnivoice.utils.audio's fade and pad where the package is missing, so the postprocess test (NaN refusal, clone reference level, PCM16) runs in CPU CI. - code2wav: masked_fill instead of a second waveform-sized buffer. - Qwen3-TTS: reword the loop-counter comment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * AuT encoder: stage the trim indices with to_device_nonblocking The shared helper pins and copies without a host stall, in place of a hand-rolled pin_memory + non_blocking copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Configs: drop prefill rungs no step reaches, and the restated speculation default Qwen3.5-9B's 2048-token step budget caps prefill and mixed steps, so its ladder stops at 2176; the Thinker's is 1024, so qwen3omni_2gpu_large_kv.yaml goes back to Omni's default ladder. speculation_mixes_walks {Thinker: true} restated the default (the Thinker's mixed walk has a token budget). The docs drop the claims measured on those rungs, and the DP2 note keeps to video, with the TP-async-on numbers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Thinker: start the residual stream from a copy of the input embeddings The fused add + norms update the stream in place, so started from input_embeds they wrote into a captured graph's static input. The copy moves from layer_0_embed to the stream (same count), and _stage's docstring goes back to main's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: warn when MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS is still set It is no longer read; the warning names model_kwargs.prefill_token_buckets. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Address #972 review: say why V-JEPA2 opts out of batching, pin torchcodec's keyframe index - vjepa2 can_batch: lead with the throughput reason, and leave a TODO for warming the encoder's batch sizes if batching is measured to pay. - test_media_decode_limits: the installed torchcodec still has the private _get_key_frame_indices the GOP split reads (skips without torchcodec). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * exact_rms_norm: a first call under capture takes torch's kernel Review: the replica's one-time check does a torch.equal, a host read. Warm-up runs it before any capture today, but a first call that lands in a capture now records torch's kernel (the same result) and leaves the check for the next eager call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2: drop the 16-thread preprocess override Review asked whether 16 intra-op threads beat the CpuPlan default. On an H100 host, 64-frame 720p uploads: 3.34 vs 3.30 req/s at c=1 and 14.19 vs 14.34 at c=16 (16 vs default), within noise, so the override goes and both V-JEPA2 models take the default. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * vjepa2_ac.yaml: re-wrap the pool comment Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
merceod
added a commit
that referenced
this pull request
Nov 3, 2026
…oice and others (#1140) * BAGEL: comments state the contract instead of the code they replaced Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL docs: give the measured output shift of guided text-to-image Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test_kv_admission: name the CFG walk BAGEL defines Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Cosmos3 and LingBot decode through wan_vae.FastDecoder and decode_dtype wan_vae.FastDecoder and decode_dtype had one user, Wan2.2, while Cosmos3's two decoder nodes and LingBot-World kept their own copies of the cuDNN >= 9.16 dtype gate, of the channels-last + DecoderStep setup and of the autocast-off decode_causal call. All of them now use the shared helpers, and the copies (_vae_decode_dtype, _fast_step, the streaming node's _step, LingBot's inline gate and _decode_step) are deleted. Numerics and switches are unchanged: Cosmos3 keeps COSMOS3_VAE_DECODE_FP32, its compile switch, its cast before the layout switch, the streaming step's non-deterministic compile (built only with the warm cache), its latents per call and cudnn_benchmark scope; LingBot keeps no override (decode_dtype's variable is now optional), and bf16 and channels-last only off reference numerics. Every node keeps the layout it set up, so FastDecoder's per-call layout check is a no-op for them. FastDecoder's autocast-off scope now follows the latents' device instead of naming cuda. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan 2.2: drop the FFN GELU epilogue, which has no measurement of its own MLP(gelu_epilogue=) and Wan22DiT's fuse_ffn_gelu ran the FFN GELU in the up-projection's cuBLASLt epilogue through the private torch._addmm_activation. Nothing measured it apart from the ragged attention change it shipped with, so it goes, as unmeasured performance paths do: MLP is back to its plain form and the DiT builds without the flag. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3 text encoder: refuse a row without text_len instead of never replaying The shared encoder node picks its captured (masked) graph from each row's text_len step metadata. A model that reused the node without setting it got every batch run eager, with no error. cg_key_info now raises ValueError for such a row; FLUX.2-klein and Z-Image set it at intake. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Image output knobs: an explicit null is the default, and only integers pass intake image_output_format let png_compress_level: null through, and encode_image then failed on int(None) after the denoise and decode. encode_image now reads null as the default level, as it already did for output_compression. Intake also refused nothing that int() accepted, so True passed as 1 and 50.9 was truncated to 50; it now takes only integers (not bools) in range, as BAGEL's width and height check does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2: say the B=1 guard routes eager steps only, and drop "as before" comments A step that leases the AC rollout capture replays forward_batched at any batch size, B=1 included, whatever can_batch returns, so can_batch's B=1 guard covers eager steps only; its comment now says so. Two comments described the code by comparison with an earlier version ("as before the capture", "same as before") and now state what the code does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * wan_vae.decode_dtype: say that an override of 0 forces bf16 The override variable is two-way, as Wan2.2's and Cosmos3's were before it was shared (and as the Wan fast-decode test uses it): set, it forces fp32, or bf16 when it is 0/false/no/off. The docstring read as if 0 meant "not forced". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Higgs Audio: declare text_output so inline decode tokens skip the per-row postprocess Higgs's decode token rides inline, but without a declared text_output Model.postprocess_inline built a tensor and called postprocess for every row on every step. Declaring TextOutput(self._detokenizer), with no stop ids as before, gives it the memoized id-to-bytes path Qwen3-ASR and Whisper use; postprocess hands text to the base class. Bytes are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: say the API_SWITCH_INTERVAL_S gain predates #849 docs/serving.rst says the number was measured before #849 and awaits a re-measure; the comment pointing at it now says the same. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS codec bench: unambiguous float32 reference wording (tail R2 P3) The parenthetical read as if the server runs float32; the shipped config runs a bfloat16 codec. Say the reference matches a float32-codec server. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Conductor: build its instance table from WorkerGraph.replica_ranks Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Image output knobs: serve the whole numbers main served Intake took only Python ints, so a whole-number float (90.0) or a number in a JSON string ("90"), which main read through int() and served, became a 400. It now accepts any whole number int() reads in range, and still refuses bools, fractions such as 50.9, infinities and out-of-range values. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: build only the window walks a deployment can run A guided window runs on the branch nodes on CFG-parallel deployments and on LLM alone elsewhere, never both, yet CFG-parallel builds still placed the unsplit image_gen_guided2/3 walks and listed them in every request's needed-label map. LLM_GRAPH_WALKS is now keyed by the layout, and the walk graphs and the label map both read it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: say why a split LLM still publishes both branch labels The comment dated from when a remote image_gen walk's branch count varied per step. Each walk now fixes its count; what keeps cfg_img published on a split deployment is a remote decode reading it and the single-node denoise walks naming their labels per request rather than per walk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL and Cosmos3: drop loop-counter leftovers from comments and names denoise_loop's docstring still named the step's loop counter, which BAGEL no longer reads; declare_step's walk-branch narrowing kept a requires_cfg guard that active_labels already implies, and its guided-row set shadowed the module's guided(). Cosmos3's decode-dtype comment now gives both directions of its override. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Image output: check only the chosen format's compression knob Intake refused a request whose other-format knob was out of range (a JPEG request with png_compress_level 10, a PNG request with output_compression 150), which main served because encode_image never reads it. Check only the knob the format applies. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: say LLM-disaggregated, not split, in the branch-label comment "Split" means CFG-parallel in the neighbouring code (DenoiseWalk.split, LLM_GRAPH_WALKS). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TP read-failure test: derive the loop-back hooks from device_loopback = None The stand-in stubbed the deleted _device_loopback and restated the three hooks DeviceLoopback now derives; clearing the declaration gives the same answers through the base class. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TP read-failure test: drop the stale "emit" from the stand-in's comment Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Merge-up fixes: DeviceLoopback in the withhold check, chunk-sink stand-in poll, postprocess_s on inline failure - _refuse_unroutable_withholds read TokenStop.loopback_signal, which main moved to ARNodeSubmodule.device_loopback; it now reads the signal from the submodule's DeviceLoopback, and its test declares one. - test_chunk_sink: its request profile stub lacked output_postprocess_s, so every pass raised before relaying a chunk; it uses a real RequestProfile, its transport's poll ends the pass (the loop no longer sleeps), and an error in the pass is raised instead of retried. - An inline value whose batched postprocess_inline call failed is charged its share of that call as well as its own retry. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Withhold check: the loop-back signal is exempt only on its own walk Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Withhold check test: the loop-back exemption is scoped to its walk Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-ASR/Omni AuT node: pin the replay's windows and packed-clip parity on the shared node Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: pin the unprojected depth path, drop the unused DepthAttention.forward The depth-loop parity test now runs at both Talker widths: 2048 through the input projection (Qwen3-TTS 1.7B) and 1024 without one (Qwen3-Omni, Qwen3-TTS 0.6B), where each step's input is the raw embedding and the first norm reads cp_embed_accumulate's per-chunk partials. The cp_embed_accumulate test checks those partials for both the raw and the projected input. DepthAttention.forward and its attend reference were the unfused layer, and nothing runs them: run_depth_layers reads the weights and attends through the predictor's DepthKV. The class stays as the layers' parameter container that takes none of the node's resources, which the binding test now checks directly. run_depth_layers' eager note keeps only the verified reason. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * FLUX.2 klein: one shared kept-module forward, so blocks don't recompile per instance exclude_from_compile set a functools.partial forward on each kept 1-D RMSNorm. Dynamo ID-guards a partial set on an instance, so the single-stream block frame klein's blocks share failed its guards once per block until the recompile limit, and the remaining blocks ran eagerly: -6 to -10% img/s at c=1-4 on klein 4B and 9B against main. A kept module now changes class instead, to a subclass made once per (class, path) whose forward is the class's: the exact RMSNorm op for a plain 1-D nn.RMSNorm, torch.compiler.disable of the base forward for every other kept module (the same graph break as before). compile_kept_exact is a class attribute of that subclass, so kept_exact (Z-Image's fused q/k norm + rotary, the VAE's norm_act) still answers. Outputs are unchanged: the same exact_rms_norm call, the same disabled base forward. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * compile_utils: state the guard that recompiled per block Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * cp_attend: say why the multi-token path stays Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: keep the compiled depth step free of per-position guards depth_step took the LM head index as an int and passed it to the disabled run_depth_layers, so under a static-int compile Dynamo compiled depth_step once per head (guards "head is None", "head == 0", ...): one frame per depth position, up to the recompile limit. depth_step itself is now the eager boundary, with run_depth_layers folded into it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * compile_utils: every block passes the same guards, not one check Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * cp_attend: name the row cap a multi-token opening would need raised Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: name every eager reader of the depth ints; count the test's groups Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS codec: compile the incremental decode through compile_regions The codec compiled its own decode with torch.compile(dynamic=False), which recompiled once per captured shape (4 chunk widths x 14 batch buckets = 55 compiles) and skipped the engine's inductor overrides and recompile accounting. On an empty cache that was ~16 of the ~19 minutes of a 1.7B cold start. The decode is now a compile region: the engine compiles it once, with dynamic shapes, before capture. H100, slice-6 head (empty Inductor/Triton cache, 26 CPUs): cold start 1119 -> 444 s, warm 540 -> 333 s. Closed loop audio-s/s c=1/16/32/64, 2 rounds each: push4a 18.4-18.5 / 137-145 / 209-212 / 281-284 -> 18.4 / 139-146 / 208-212 / 289-301; push4c (server on 2 CPUs) 18.6/159/237/370 -> 18.4-18.5 / 160 / 237-251 / 371-379. Pinning the chunk width static (one compile per width, batch dynamic) measured no faster and started in 765 s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS codec: test and bench compile the decode as the engine does The codec's decode is compiled by the engine's compile_regions pass, so a CodecSubmodule built directly ran it eagerly. The per-node body of Engine._compile_regions becomes compile_node_regions, which the engine loop, the streamed-codec test's CUDA case and codec_bench's load_codec all call. The compile_regions docstring and models.rst now say what dynamic=True gives: a few graphs (batch 1 and width 1 specialise) for every chunk width and batch bucket. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Codec and compile_node_regions docstrings: the measured cold start, no dynamic=None aside Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B: start the codec stream with a 2-frame chunk The deployment config's ramp began with a 4-frame chunk, so a stream's first audio waited for four Talker steps. [2, 4, 8, 16] sends it after two, and the 160 ms first chunk still covers the time until the next one arrives under load. [1, 2, 5, 16] is documented as the lower-latency opt-in. H100 (push4a and push4c), closed loop, c=64: TTFA p50 84-95 -> 63-78 ms (c=1 33 -> 25 ms), no playback stalls either way, audio-s/s 379-424 vs 363-416 over 4-6 rounds each (same within noise). Open-loop SLO capacity (TTFA <= 200 ms, stall < 10 ms): Seed-TTS 80 @90/99 and passages 28 @90 either way. WER (Whisper large-v3, 128 sentences): 2.09% -> 1.73%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * qwen3tts_1p7b config: the header names the 2-frame start; date the TTFA numbers Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B: the bf16 codec is measured on H100 too Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: klein's exact q/k norm, measured with the shared-forward fix Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B bf16 codec: state what the H100 runs measured Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: klein's shared kept-norm forward, measured against a forward per block Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS bf16 codec: tie each H100 WER to its chunk schedule Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: kept norms share a forward per class Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Attention: select_last_hidden is one eager boundary, so prefill buckets don't recompile its frame The disabled `qo_indptr_buf` call broke the graph inside `select_last_hidden`, leaving a resume frame that holds both the bucket's static token count (`hidden`) and its row count (`qo_indptr`). Captured forwards compile with dynamic=False, so that frame took one cache entry per (tokens, rows) bucket: Qwen3.5's 9 token buckets by 7 prefill/mixed row counts made 63, and the 13-step prefill ladder made 91, past the recompile limit of 84. Disabling the whole method moves the break to its call site, where the caller already breaks; no compiled frame holds the pair. Same for the MLA and XPU managers, which share the shape. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL takes its prefill's last rows through select_last_hidden; test docstrings Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * GDN prefill run: state the zeroed-tail precondition the mixed wrapper also meets Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice: key the backbone's step outputs by the engine's row ids The packed backbone keyed its outputs by CanvasItem.request_id, which prepare_inputs now fills from fwd_info.request_id, the wire request id. The engine looks each row's outputs up by its worker-local handle, found none, and every OmniVoice request failed in postprocess with KeyError 'c_logits'. forward_batched now keys rows by engine_inputs.request_ids, as the other batched submodules do, and CanvasItem carries no request id at all, so the canvas can't be keyed by the wrong one again. The unread request_id kwarg goes with it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni 2-GPU config: keep the 8192 Thinker ladder off the 80 GB default qwen3omni_2gpu.yaml, the CLI default and the config sized for 80 GB GPUs, ran the Thinker's prefill ladder to 8192. On an H100 its 4096-8192 captures don't fit beside the 1,024-page KV pool (59 of 86 Thinker buckets captured), so those steps run eagerly and a deployment that requires accelerator graphs refuses to start. The default goes back to the model's 128-2048 ladder, as on the base. The 8192 ladder, measured on 2xH200, stays in qwen3omni_2gpu_large_kv.yaml, which is already the H200 variant; its header, the default's, the submodule comment and the Qwen3-Omni notes say so. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS 1.7B configs: back to the base's codec, chunk ramp and pools (owned by #978) Kirill's #978 sets these himself (the [4, 8, 16] ramp, the bf16 codec, the codec_state slots, the GC knob and their comments), so this PR leaves both 1.7B configs as its base has them. This reverts the config hunks of fce78c050, f63eb7994, aba669bac, 9ef1d1e66, 08c8b2350, e4450ea53 and 5428cca16, including the [2, 4, 8, 16] first-chunk schedule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: drop the incremental codec, mixed walk and fused depth loop (owned by #978) Kirill's #978 makes the same Qwen3-TTS changes (from the same stacked-head work): the incremental codec with its codec_state pool and chunk knobs, the combined prefill+decode walk and its captures, the fixed prefill embeddings, and a fused depth loop (cp_linear with folded norms, cp_attend, fold_depth_weights) beside the codec's own compile. This PR leaves all of it to him: Qwen3-TTS, the shared code predictor kernels and Qwen3-Omni's two call sites go back to the base's code, and so do their tests, the real-weights codec test, the Qwen3-TTS and Qwen3-Omni code-predictor paragraphs in docs/models.rst, and the codec bench, which timed the removed decoder. That reverts fce78c050 and the fixes made to its code since: 59d1b3854, f5b0422ff, 2f6333a90, e59639f21, 251d13e34, add26fbc7, f7887b71a, 4d9f54584, e6b2a2fef, a2331d014, ce5693c3d, 1158293d0, 8ebf63986, 613160070, b13ab78e4, 9a588f978, 26ef71b6e, 15413c499, 5fea4023e, ab3350e18, 4f935c1f6, a8c537e7d, c4ed93f0c, 88b5b68de, 5bfeeac60, a5687bdd7, the code-predictor half of 4791cb3ca and the Qwen3-TTS halves of f63eb7994, 7906d0902 and 13b35b17a. Kept: db8b241c7, a decode row reads its text token by its loop iteration rather than a counter a re-prepared step advances twice. The counter is the base's (and main's) code, which #978 leaves as it is. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: drop compile_node_regions, which only the removed codec test and bench called 13b35b17a and 7a8be7397 lifted Engine._compile_regions' body into a module-level compile_node_regions so the Qwen3-TTS codec test and bench could compile the decode as the engine does. Both are gone with the incremental codec (owned by #978), so _compile_regions takes its body back, as on the base. No behaviour change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni 2-GPU config: don't claim the remaining captures fit 80 GB Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Pin video decode threads to the CPUs the plan leaves under pinning A case with a pinned CPU fails if video_decode_threads goes back to usable_cpus(); the MSTAR_STEP_THREAD_CORE row names FFmpeg's per-decode threads among what the pins shrink. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ragged attention: a wrapper skips replanning the layout it last planned The dense manager kept the (cu_q, cu_k, causal) each of its wrappers was last planned with and skipped an unchanged plan; the ragged self- and cross-attention managers replanned every step, so a denoise loop paid FlashInfer's plan (a kernel and a stream sync) once per step for the same segments. The skip now lives in the wrapper both use (_RaggedPrefillBase._plan): each wrapper is one (label, slot) or one (bucket, slot, label), so the schedule it keeps is the one the step reading it planned. The dense manager's own table goes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 docs: date the attention and CFG-parallel timings as before #963 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Conductor: leave the idle wait to its own change The conductor's 1 ms nap becomes a wait on the poll in a change of its own, which waits until the next timed check instead of a fixed 1 ms. This branch's capped poll, its failed-pass back-off, the conductor stand-ins' poll_for_messages and their test go, so those files are main's again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * API server tests: the failed-wait test pins its back-off; stand-in waits as main has them The message loop always passes a computed timeout, so the 1 ms defaults left by the fixed poll go: the chunk-sink stand-in takes main's again and the failed-wait stand-in takes it as required. That test now asserts the 10 ms back-off, without which a persistent wait error would spin. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Preprocess pool: at most the CPUs the plan leaves unpinned Every other API-server pool is sized from the CPU plan; the preprocess pool took its declared or configured size uncapped, so under step_thread_core a model declaring 8 workers could run them, each with its own FFmpeg threads, on fewer free CPUs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Data worker: drain routes before an input only when it prepares here The drain keeps a call's held items from waiting behind a preprocess that may take seconds. An input prepared on the pool only stores and sends when it is handed on, and each pass already drains routes at its top, so it no longer takes a drain of its own. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: say what a RecurrentStateSpec max_slots override must leave Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision: compile the block loop, so a step past the captured buckets runs compiled disable_torch_compile left the block loop uncompiled whenever it did not replay: a long video (over 16384 patches) ran the blocks' fp32 rope, norms and activations as eager kernels, slower than the whole-forward compile it replaced. The block loop is now a compile region like the embed and the merger, dynamic in the patch count; it compiles in its capture's warm-up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * API server: drop the api_switch_interval_s knob Its gain was measured before #849 and #939 changed the message-thread and data-worker path, and the re-measure it was waiting on never ran. The serving-config key, its validation, docs rows and test go; workers keep main's MSTAR_PY_SWITCH_INTERVAL_SEC unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the GIL switch-interval knob's model defaults (the knob is removed in #858) Reverts #859's additions on top of #858's api_switch_interval_s: the Model.API_SWITCH_INTERVAL_S class default, Qwen3.5's 0.001 override, the model argument to entrypoint._set_switch_interval, the doc wording and the test fakes that carried the attribute. Every file #858 owns is back to #858's text, so #858's removal of the knob merges in cleanly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: the TP2 publish test's KV managers hold one rank's head KVManager shards its config in build, so the test's direct constructor shards it the same way and each rank moves its own shard. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * API server: drop warmup_requests' concurrency option No config sets it, so its parsing, docs and tests covered a dead path. Each entry still goes out once per data-parallel replica, now as warmup-<spec>.<copy>. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision: a request too large for a replay gets an encoder step of its own Up to 8 requests share an encoder step, and a step that cannot replay as one runs them back to back (can_batch). For requests that are each past the captured buckets (long videos) that only held every request's embeddings, and so their prefills, until the last one was encoded: under load the prefills bunched into waves (TTFT +30-50% on 60 s video at c=16/32, decode almost never interleaved with a prefill). The vision node now splits its steps by capture key, a request too large for a replay being a group of its own, and the engine lets a node whose only capture is a piecewise region split that way. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: Canary's decode cap leaves the prefill's token in the budget #794's cap-versus-stop-check agreement tests for Canary and Canary-Qwen assumed a stop check that counts max_tokens; TokenStop leaves the budget to the loop's cap, so they pin the cap instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: last_prefill_flag for a prompt prefilled over several passes; GLM-5.2 test comment names the loop cap Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision split: review cleanups - Engine.capture_group asks Engine.splits_by_capture_key; the module helper existed only for a test fake, which now builds a real Engine. - The split docstrings say a piecewise region is not tied to a walk, and that on a piecewise-only node the key only groups the scheduler's rows. - Stale 'eager' comments on the vision tower: past the top bucket the block loop runs compiled, uncaptured. - docs/models.rst: drop the one-boot measurement of a different overlay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chatterbox, Higgs Audio tests: a stop id at prefill ends the partition #858 moved TokenStop's prefill_walks onto END_PARTITION; the step closing a reply that ended at prefill no longer exists. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus test: the stop check's fake info names its walk TokenStop now reads each stopping row's walk to choose between its loop and END_PARTITION. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 9B DP2 config: the replica pick follows #920's router The conductor now weighs a predicted prefix hit beside the live requests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Review wording: the vision split's capture-only scope, Qwen3.5 configs' start-up compile, Qwen3-Omni's encoder header - The Qwen3.5 vision split applies while the block loop is captured, and requests that fit alone but not together still run one after another in one step; cg_key_info's and capture_group's docstrings say a piecewise-only node's key is a scheduler group, not a bucket. - The Qwen3.5 configs' warmup comments: the vision regions compile at start-up. - Qwen3-Omni's submodule header takes main's line for its audio encoder. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 docs: date the timings by the DiT step they measured, not a PR number Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 docs: an earlier version of the DiT block, not an earlier DiT step "DiT step" reads as a denoising step in this section. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: a capped deployment holds cuDNN's capture search to its share capture_cudnn_benchmark's search runs in capture sizing since the sized warm-ups took the node's scope, and it tries workspaces as large as the GPU has free: Kokoro reserved ~140 GB of an H200 for ~25 s at startup, past the gpu_memory_fraction a co-located worker (or an 80 GB-capped comparison) plans around. With a fraction set, the allocator is held to it while such a node is sized, and an over-cap try is skipped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: the preprocess pool's CPU cap does not count a prepare's FFmpeg threads Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunk policy: leave the longer primed first chunk to #977 #977 makes the same change (ScheduledLeftContextChunkPolicy.prime past left_context, and its test), and its consumer, the incremental Qwen3-TTS codec, is no longer in this stack: every primer here passes at most left_context (Qwen3-TTS's lead is min(ref_frames, left_context_frames)). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * ViT patch embed: the conv timing in its comment is Qwen3.5's Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: _search_cap's docstring separates a skipped search try from a failed sizing Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5-9B DP2 config: say the vision regions compile at start-up Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: capture_cudnn_benchmark's search runs in capture sizing, held to a capped deployment's share Sizing runs a node's warm-up forwards; the capture after it warms up with no forward, so the scope around warmup_and_capture covered no conv on CUDA and Kokoro's graphs kept cuDNN's heuristic picks. The scope now wraps sizing. The search tries workspaces as large as the GPU has free, and the allocator keeps them until sizing's empty_cache: Kokoro reserved ~140 GB of an H200 for ~25 s at startup, past the gpu_memory_fraction a co-located worker (or an 80 GB-capped comparison) plans around. With a fraction set, the allocator is held to it while such a node is sized and restored to its previous value after; an over-cap search try is skipped, and the node's own sizing fails the boot when its eager step does not fit in the share. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish S2 Pro: capture the decode buckets uncompiled, as on main Compiling the seven decode buckets took a cold start from 180 s to 1372 s on H100, with no measured throughput to show for it. The 64-row bucket and the packed prefill captures stay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish S2 Pro: TF32 codec decode is the model's default The shipped config set codec_exact_fp32: false while the model defaulted to the exact decode, so other Fish configs decoded differently. The default now lives in the model; codec_exact_fp32: true still opts into the exact decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: open the depth loop with position 0 alone Both Talkers opened with one-token passes on main; the two-token opening changed greedy codes for both, and a one-token pass keeps every pass at bs rows. The opening is now main's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-TTS: drop the 24- and 48-row decode buckets They took the 1.7B's cold start from 204 s to 307 s on H100, and a step at 64 rows costs about what it does at 32, so padding into 32 and 64 is nearly free. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * run_layers: the residual stream starts from its own copy The stream was the caller's hidden_states, overwritten in place from layer 0, so a caller had to know to pass a copy (the Qwen3-TTS Talker did for its packed replay's static input). One [tokens, hidden] copy per forward, made inside the forward so a capture records it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice silence test: pydub is the oracle, installed with dev The comparison imported the reference package, which CI does not install, so it always skipped there. It now restates the reference's trim over pydub.silence (raw-sample segments, no ffmpeg), with pydub in the dev extras, and covers the rates and edge cases the self-referential trailing trim test stood in for; that test is gone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Code predictor: each Talker keeps its own embedding-sum dtype The shared depth loop summed a frame's codec embeddings in fp32 for both Talkers. Qwen3-TTS accumulated in the model dtype, rounding after every add, so its greedy codes could differ at near-ties. ResidualCodePredictor takes embed_sum_dtype from its Talker: None (the embeddings' dtype) for Qwen3-TTS, fp32 for Qwen3-Omni, whose fp32 sum is rounded once to the model dtype as before. The accumulate kernel adds in fp32 and rounds to the sum's dtype, as torch's add_ does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Review wording: Qwen3-TTS's embedding sum, the cuDNN search in capture sizing; audioop-lts for pydub on 3.13 - Qwen3-TTS's per-add rounding is the sum this Talker had in M*; the reference sums in fp32 and rounds once. - adding_models.rst and Kokoro say the cuDNN search runs in capture sizing. - pydub 0.25.1 imports audioop, which Python 3.13 removed: the dev extras add audioop-lts there, so the silence oracle does not skip. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStop: withholding outputs and the runtime check refuse each other The graph runtime's stop check never calls check_stop, which is what withholds a stop step's outputs, so a TokenStop with both would drop the withhold silently. It is a frozen dataclass now, and that combination raises at construction. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test_knobs: take main's version back now that ingraph_decode_rope stays Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2 CFG-parallel: the scaffold keeps ragged_for, docs say a rank runs its branch over the batch Since the dit moved onto the scaffold's denoise loop, DenoiseLoopSubmodule is RaggedAttentionMixin's only user, so the mixin folds back into it as on main. The declare_step test keeps only what the per-rank segment test does not cover. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL docs: quote the guided text-to-image PSNR and the configuration of each number Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * wan22 extra: install flashinfer-python for the default ragged attention Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Wan2.2: leave CFG-parallel to the cfg_size axis Drop wan22_cfg_parallel_2gpu.yaml, my_rows and the dit's sp-group plumbing, which used the SP slot to mean a CFG split. The cfg_size axis carries it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Z-Image-Turbo: capture batch sizes 1, 2, 4, 8; every size up to 8 is a commented opt-in Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2-AC test: a rollout request's stamped clip shape keys the encoder capture Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5-9B configs: keep main's 2176 bucket for a mixed step The 9B ladders replace the default, so they lost the slack bucket that lets a mixed step of four 512-token prompts and its decode rows replay at 2176 rather than pad to 4096. Adding it keeps the compile count at 82, strictly under Dynamo's limit of 84; the 9B ladders keep 3072-7168 in place of main's 4352 and 8704. The mixed-step bucket and the recompile limit are now checked on every ladder a Qwen3.5 config ships, not only the default. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2-AC capture-key test: decode the sampled frames as the request does process_prompt now decodes the frames load_video sampled, by index, so the test hands it sampled indices and stands in for the decoder rather than a decoded clip, which a request no longer carries. The Z-Image capture-size figures say they predate the fused q/k norm + rotary. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Sampler RNG offset test: both captured kernels, each against a draw it can match At V=64 the captured step takes fused_sample, whose Triton RNG never agrees with the eager step's FlashInfer draw, so the test passed whether or not the two steps shared an offset. The eager case now draws its captured step through FlashInfer too, and a second case runs two captured steps on fused_sample, which reads the same offset buffer the replay advances. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Data worker: refuse a bad preprocess_workers before the thread starts The thread reads preprocess_workers as it starts, so a value it cannot take ("auto") ended it while the server came up, and every request then waited for its timeout; 0 or a negative quietly meant inline, and true meant 1. PreprocessWorker now refuses anything but a positive integer, on the caller's thread. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2 pool sizes: say what they cost and when to change them The configs' preprocess_workers is taken as given, so each says it was measured on an H200 host, how many intra-op threads the pool then runs, and to lower it (or drop the key for the data worker's CPU-aware default) on a smaller host. The model comments name CpuPlan.preprocess_threads, and the pool docs name V-JEPA2 beside Qwen3-Omni. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Attention: MLA and XPU share one select_last_hidden on the base class The FlashInfer MLA and XPU managers carried the same select_last_hidden body with its own torch.compiler.disable. One method on AttentionManager reads the backend's qo_indptr_buf and is eager as a whole, so a captured prefill forward keeps no compiled frame holding both the bucket's token count and its row count, and a new paged backend gets it without remembering the decorator. The base class declares qo_indptr_buf too, raising an error that names a backend which plans no query offsets. FlashInfer keeps an override of its own: a label planned as a decode has no query offsets there (every token is a row's last), and a label with no plan is an error. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish S2-Pro: decode each codec frame once from per-stream state; withhold <|im_end|> with TokenStop The codec re-decoded every window's left context. It now keeps each stream's causal conv and transformer state in a codec_state recurrent pool and decodes only a chunk's new frames (IncrementalDACDecoder, with the streaming conv steps and a time-major Snake in components/streaming_codec.py), with the transposed-conv weights laid out channels-last for their GEMM. The AR declares its stop as a TokenStop that withholds the <|im_end|> frame, so the codec stream carries only frames to decode and the codec needs no device read to drop it; the frame budget is the loop's cap, as for every TokenStop adopter, and batched AR outputs share one stop-check copy per step. The TF32 figures in the docs and the model's TF32-default comment are marked as measured on the earlier windowed codec. Tests: incremental decode stitches into the whole-utterance decode (eager and captured, any chunking), the conv weight layout, a CUDA-only compiled-vs-eager check, the TokenStop stop check, and the frameless final chunk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice: score and reveal the whole step's canvases in one batched pass The unmask step scored and revealed each request's canvas in postprocess, one request at a time. forward_batched now combines CFG, draws the Gumbel noise and reveals every canvas of the step in one pass (unmask_step), keyed by the engine's row ids, with the CFG combine's dead buffers freed before the sampler. The backbone captures up to 8192 packed tokens, and the codec's acoustic decoder compiles as a region and is warmed at start-up for every canvas-length range. Tests: the parity tests drive the served unmask step, including a mixed-temperature batch; codec warm-up coverage; a CUDA-only compiled-vs-eager check of the codec decoder. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: a cheaper ViT step and a GC threshold Vision encoder: the shared ViT's host grid helpers (components/qwen3_vit.py) build position-table taps once per frame shape and tile them over a video's frames through numpy, on the host only; Qwen3-Omni's vision encoder, their other caller, builds them there too and copies them without a sync, as Qwen3.5's does. The block loop gets a 12288-patch capture bucket (a 30 s video at 2 fps, which 16384 padded by 39%), and each ViT block compiles as its own region, so Dynamo traces one block and the rest reuse the graph instead of tracing the unrolled loop at every start; the start-up examples warm the embed and merger regions as before. Configs: 2B, 4B, 9B TP2 and 27B TP4 raise the gen-0 GC threshold. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Attention: fused q/k norm, RoPE and KV write helpers on AttentionCallable The Qwen3 LM (Qwen3-ASR, Canary-Qwen, Fish's slow AR) called its fused q/k RMSNorm + rotary + KV-write prologue through its own glue. AttentionCallable now owns it, with one helper per rotation: norm_rope_attend rotates by positions and a theta, and norm_rope_table_attend by a precomputed cos/sin table (the multimodal-RoPE case). The table variant runs on a new Triton kernel, qk_norm_rope_table_kv_write, in fp32 and rounded once. Both refuse a Gemma-style norm, whose (1 + w) scale the kernels don't apply. Tests: the table kernel against an fp64 reference and the unfused norm and rotation, and both helpers' Gemma-norm refusal. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: an H100 config, and the Thinker's q/k norm, mRoPE and KV write in one kernel configs/qwen3omni_2gpu_h100.yaml serves 80 GB GPUs: the encoders run on the Talker's GPU so an encoder step doesn't wait behind a Thinker step, and the Thinker KV pool is 512 pages. That size was measured on an H100 while the encoders shared the Thinker's GPU and its text ladder ran to 8192: beside 1,024 pages 27 Thinker captures, all among those of 4096+ tokens (text and vision), ran eagerly, and 640 pages still left two vision buckets out. Chunked prefill now caps a Thinker step at 1,024 tokens, so those large buckets no longer replay, but they are still captured and still take that memory; whether 1,024 pages fit with today's ladders and encoder placement is unmeasured. The other 2-GPU configs and the Qwen3-Omni notes say which GPU each is for, and when a full Thinker pool needs max_concurrent_requests. The Thinker's q/k RMSNorm, interleaved multimodal RoPE and KV write run in one kernel through AttentionCallable.norm_rope_table_attend, in fp32 and rounded once where the unfused ops round to bf16 at each step. docs/models.rst, the fused-kernel section of adding_models.rst and the kernel's test module name the Thinker as its caller. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fused_experts: per-pair GEMV experts for small bf16 steps; H100 tile tables A bf16/fp16 step of up to 7 tokens with no shared expert runs one block per (token, expert) pair (routed_experts) instead of the sorted grouped GEMM, on NVIDIA GPUs, for shapes whose tiles are whole and unclamped (routed_fits). H100 tile tables for Qwen3-Omni's Thinker and Talker experts; the per-pair path was swept on H100 only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * FLUX.2-klein, Z-Image: config and compile docs keep the reasons, not the narratives The klein and Z-Image configs and compile_transformer_forward point to compile_exact_ops in docs/models.rst for measured PSNR and cost instead of repeating the numbers; klein ships [norms], Z-Image the plain compile, both on sdpa. The docs give SGLang's PSNR from klein's output as the yardstick. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Cosmos3, Orpheus, Pi0.5: per-model serving fixes - Cosmos3: a generation step past max_batched_vision_tokens (a model kwarg, 16384 by default) runs its requests one at a time. - Orpheus: SNAC windows past the decoder's left reach pad to the widest window, whose default left context comes from that reach. - Pi0.5: camera images travel to the worker as uint8, and the model refuses to start without a tokenizer (the docs say it needs the gated PaliGemma one). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * TokenStepNode: copy only the outputs the per-request stop check reads The per-request fallback aliased text_inputs as new_token with a dict splat, which reads every output name, so the lazy host outputs copied all of them to the host. A ChainMap over the outputs reads only what the check reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Check every shipped prefill ladder, not only Qwen3.5's The ladder tests read only configs/qwen3_5*.yaml, so the Qwen3-Omni large_kv config's ladder was never checked. Every config that sets model_kwargs.prefill_token_buckets is now built into its own model's prefill capture; the Qwen3.5-specific checks keep to Qwen3.5 configs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: pin the config-to-model plumbing, drop the Qwen3.5 repeat Every shipped ladder's model must take prefill_token_buckets as its own argument (it would otherwise vanish into **kwargs); the Thinker's capture check stays for Qwen3-Omni, whose ladder no other test builds. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: say what the plumbing check covers, and why it fails Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: read the Thinker's prefill capture by walk Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Ladder test: dispatch on the model class, and name what it checks Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Address #971 review: num_step on the data worker, slot per label, postprocess on CPU CI - OmniVoice: process_prompt parses num_step too, so a non-integer is a 400 rather than a raise at the conductor that hangs the client. - RecurrentStateConfig: refuse a pool whose usable slots cannot hold one request's labels (two labels, one usable slot never admits). - test_omnivoice_silence: stand in for omnivoice.utils.audio's fade and pad where the package is missing, so the postprocess test (NaN refusal, clone reference level, PCM16) runs in CPU CI. - code2wav: masked_fill instead of a second waveform-sized buffer. - Qwen3-TTS: reword the loop-counter comment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * AuT encoder: stage the trim indices with to_device_nonblocking The shared helper pins and copies without a host stall, in place of a hand-rolled pin_memory + non_blocking copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Configs: drop prefill rungs no step reaches, and the restated speculation default Qwen3.5-9B's 2048-token step budget caps prefill and mixed steps, so its ladder stops at 2176; the Thinker's is 1024, so qwen3omni_2gpu_large_kv.yaml goes back to Omni's default ladder. speculation_mixes_walks {Thinker: true} restated the default (the Thinker's mixed walk has a token budget). The docs drop the claims measured on those rungs, and the DP2 note keeps to video, with the TP-async-on numbers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Thinker: start the residual stream from a copy of the input embeddings The fused add + norms update the stream in place, so started from input_embeds they wrote into a captured graph's static input. The copy moves from layer_0_embed to the stream (same count), and _stage's docstring goes back to main's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: warn when MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS is still set It is no longer read; the warning names model_kwargs.prefill_token_buckets. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Address #1140 review: size Fish's codec pool to the AR batch, warm OmniVoice's short canvases - Fish S2-Pro: the codec state pool holds 64 streams (+ sink), as many as the AR decodes a step, not 256: about 0.55 GiB instead of 2.1 GiB. - OmniVoice: warm the 1- and 2-frame codec canvases too, so a very short reply does not compile on the GPU thread. - test_model_kv_layouts: Fish joins the decode-cap case that ties get_loop_caps to where TokenStop's decode ends. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish: cap concurrent requests at the codec streams; tighten the cap test Review of the previous commit: a request keeps its codec stream until it ends, so past 64 requests a request waited a whole request for its first audio while its AR rows slowed the streams. The shipped yaml now sets max_concurrent_requests: 64 so later requests queue before the AR, and validate_config_yaml refuses a cap the codec pool cannot hold. The decode-cap case keys on TokenStop's own loop; the OmniVoice warm-up note says why one and two frames are warmed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Address #972 review: say why V-JEPA2 opts out of batching, pin torchcodec's keyframe index - vjepa2 can_batch: lead with the throughput reason, and leave a TODO for warming the encoder's batch sizes if batching is measured to pay. - test_media_decode_limits: the installed torchcodec still has the private _get_key_frame_indices the GOP split reads (skips without torchcodec). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Cosmos3: leave the batched vision-token cap off by default Review: 16384 was measured on H100 only, and at 832x480 two 121-frame Edge clips (about 12k vision tokens each) never batched under it. It is a per-GPU number, so it is now off unless a deployment sets max_batched_vision_tokens in its model_kwargs (16384 on an H100). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * exact_rms_norm: a first call under capture takes torch's kernel Review: the replica's one-time check does a torch.equal, a host read. Warm-up runs it before any capture today, but a first call that lands in a capture now records torch's kernel (the same result) and leaves the check for the next eager call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish: a null max_slots override means the default; drop a stale large_kv note Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * V-JEPA2: drop the 16-thread preprocess override Review asked whether 16 intra-op threads beat the CpuPlan default. On an H100 host, 64-frame 720p uploads: 3.34 vs 3.30 req/s at c=1 and 14.19 vs 14.34 at c=16 (16 vs default), within noise, so the override goes and both V-JEPA2 models take the default. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * vjepa2_ac.yaml: re-wrap the pool comment Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Fish: pin that a null max_slots override validates; re-wrap a docs line Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: re-wrap the Qwen3-Omni H100 line Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * OmniVoice graph test: the fake node carries lora_pools --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: merceod <atj10@stanford.edu> Co-authored-by: merceod <atindra@mstarlabs.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #944: this branch carries #944's earlier commit (6ac1bc6, the warmup fix) and is merged with main at 7eac23c. #944's own branch has moved since (56b49e1) and conflicts with that commit in
cosmos3/submodules.pyandtest_denoise_compile.py, so this needs a merge once #944 lands.Wan2.2-5B's denoise step on B300 spent ~17 ms of a 72 ms step in non-matmul kernels and 7.5 ms with the GPU idle. This takes the compiled c=1 step at 832x480x33 from 72.5 to 62.4 ms (−13.9%) and c=2 from 29.8 to 33.1 videos/min, and stops Cosmos3-Edge replanning FlashInfer's schedule every step (−2.8% c=1 step, byte-identical output). Wan's output is not main's: the per-frame time embedding can round differently from the reference's per-token one, which puts it on a different trajectory (see Correctness).
What changed
Per-frame modulation and a one-pass adaLN (
components/diffusion/modulation.py,wan22/components/dit.py). Every token of a latent frame shares its timestep (I2V zeroes frame 0 only), so the DiT takes its timestep per frame,[B, frames], and computes the time embedding and the six modulation vectors once per frame instead of once per token. Each block's norms go throughada_layer_norm: an optional gated residual add, LayerNorm, and per-group shift/scale. Compiled, that is a Triton kernel that holds the row in registers, so the stream is read once. Inductor lowers a LayerNorm wider than 1024 to a looped Welford that reads the row twice, and before this change each block also read six[7020, 3072]modulation chunks. Eager calls run the expressions the block used to write inline.q/k RMSNorm + RoPE (
components/diffusion/qk_norm_rope.py). The reference rope writesout[..., 0::2]/out[..., 1::2], and Inductor lowers those slice writes to a kernel that stores eight fp32 intermediates per block, plus two kernels that read them back. Eager now uses the shared stacked rotation (rope.apply_rotary_interleaved), which gives the slice-write rotation's bits on Wan's tables (test_eager_rms_norm_rope_is_the_slice_write_rotation_bitwise). Compiled, one kernel normalizes q and k across heads, rounds where the eager norm rounds, and rotates in fp32.No host syncs in the step (
wan22/submodules.py,wan22/components/unipc.py).forwardstarted withtime_index.item(), moved the UniPC sigma/timestep tables to the device from pageable memory (a stream sync each), and the corrector'slinalg.solvechecked its info on the host. So each step drained the GPU, and the next step's prep, Dynamo's guard check (2.3 ms per call on this host, torch.profiler's "TorchDynamo Cache Lookup") and the solver's ~100 small kernels ran with the GPU idle. Nowprepare_inputsandforwardtakekfrom the engine's loop counter (the per-batch snapshot inForwardPassInfoWrapper, so it cannot move under a queued step) and thetime_indexedge is gone, so a caller with no counter gets aKeyErrorinstead of the sync. The tables go up throughh2d.to_device_nonblocking, and the solve issolve_ex. The strict read depends on the runtime reporting the counter at iteration 0, which I checked on a live server on both graph runtimes (MSTAR_RUST_GRAPHauto and0).Dense attention: don't replan an unchanged layout (
engine/resources/attn/dense.py). Each slot's wrapper keeps its last schedule until it is planned again, andrunonly reads it. So when a step's(cu_q, cu_k, causal)matches what that slot's wrapper was last planned with, the plan is skipped. For Cosmos3 that skips FlashInfer'splan_kernel(1.7 ms) and the stream sync after it on every step after the first two. This commit is engine-side and independent of the rest; I'm happy to split it out.The two components switch on
torch.compiler.is_compiling(): inside the compiled transformer they replace what Inductor would emit, and eager calls keep the formulation the port's bit-exact tests hold. Compiled output already differed from eager by Inductor's own rounding; these agree with eager to bf16 rounding (test_diffusion_dit_kernels.py).Performance
Gate on 1x B300 (sm_103, 1100 W cap), torch 2.12.1+cu132,
-O2,MSTAR_SHM_ARENA=1, server pinned to 8 cores and client to 4, fresh servers alternated, one shared warm Inductor cache, no recompiles. Main is 7eac23c plus #944's commit, without which the stock Cosmos3 configs fail at warmup. Wan runsconfigs/wan22.yamlwithcompile_dit_sizes: ["832x480x33"]added, since the stock config compiles only 1280x704x81 and runs this size eager, without either kernel. 20 steps, CFG 5, VAE decode as configured. Cosmos3-Edge is the stock config at 832x480x121. c=1 is the p50 of 10 requests after 2 warmups, c=2 is 20 requests after one wave. "Step" is DiT ms per step: (compute span − text encode − VAE decode) / 20 from the server stats, since the server's own per-step figure double-counts once steps queue behind each other. Cosmos has no text-encode phase and its steps never overlap, so it uses the server's per-step mean.Per change, on this branch's earlier base (f7cb5ca + #944, main 72.8 / 72.1 / 71.4 ms, mean 72.1), before main was merged; each arm is that base plus the one change:
nsys, one warm c=1 Wan request on GPU 1, same earlier base (median of the middle 18 steps): main 72.2 ms GPU span = GEMM 39.3, norm 10.1, attention 8.0, elementwise 3.9, rotary 2.9, gaps 7.5. This PR: 61.4 ms = GEMM 41.1, attention 8.5, elementwise 5.2, norm 4.4, rotary 2.0, gaps 0.18. The GEMMs are the same kernels and take 5-8% longer per call. The step runs at the power cap, and with fewer low-power memory-bound kernels between them the GEMMs run at a lower clock (~1520 MHz at c=1 vs ~1760 on main). That is a hypothesis from the clock samples; I didn't profile power per kernel.
Not included: the FFN's GELU in cuBLASLt's bias+GELU epilogue (
_addmm_activation). It took 2.1% off the step alone and 3.9% on top of the three changes here, so it's a candidate follow-up rather than part of this PR.Correctness
test_wan22_native_dit_equivalence.py): weights, RoPE tables, text stream, patchify, all 30 blocks and the head are bitwise, with both modules given the same time embedding and modulation rows. The rows themselves are held to kernel rounding:temb2e-3, modulation 3.2e-2.tembdiffers by up to 9.8e-4, because cuBLAS picks another kernel for the second Linear at 18 rows than at 7020, and the served videos on the stock config sit 25 dB from the base's, against 43 dB between two base servers. The base's own compiled and eager outputs are 21 dB apart, so this is a different trajectory, not a worse one. (Gaurav's run, 1x H100; I have no H100 to repeat it on.)NVIDIA_TF32_OVERRIDE=0the two row counts differ:tembby 3.1e-5, modulation by 3.9e-3. Switching TF32 off in the time embedder alone moves them by 9.8e-4 and 1.6e-2, the sametembfigure as H100, so my hypothesis is that H100 picks TF32 at one row count and not the other. The native test passes in all three setups, with every block at max_abs 0.test_wan22_reference_equivalence.py,NVIDIA_TF32_OVERRIDE=0,CUBLAS_WORKSPACE_CONFIG=:4096:8), max_abs over the run. Running the embedder at the reference's row count takes the branch back to 0.0, so the time embedding is the only departure; the sequential case stays in the suite as a 0.0 gate. The perturbed column adds ±1.1e-4 totemb's fp32 rows, the size of the TF32 kernel change, and is what the new bounds are set from (sequential 4.5, no-CFG 1.0).Tests
ruff check .: clean.python -m pytest test/modular test/distributed test/chatterbox test/kokoro test/omnivoice test/sampling_test mstar/model/*/tests benchmark/worker_phases, CUDA hidden,-p xdist -n 8): main + Cosmos3: one builder for the batched step's rows, so the warmup can't break boot #944 50 failed / 8779 passed / 785 skipped; this branch 50 failed / 8783 passed / 794 skipped, the same 50 by test id (32test_loop_output_ownership, 5test_send_registration_failure, 4test_wan_vae_decode, 3test_z_image_native_equivalence, 3test_eval_compare_multigpu, 3 others).test_piecewise_config_signatureandtest_tp_async_lockstep_simran serially: 19 passed on both, counted above.test_flex_attention_resource,test_waypoint_componentsandtest_waypoint_ditcrash in native code on this box's CPU on both.test_step_prep_host_syncs.py::test_wan22_forward_takes_no_host_sync(CUDA, T2V and I2V, stub transformer underset_sync_debug_mode("error"); fails onlinalg.solve, a host read of a device tensor, and a pageable copy of either table, each checked by mutating the source),test_wan22_step_prep_fails_without_the_loop_counter,test_sequential_cfg_trajectory_is_bitexact_on_per_token_rows, the rope test above (fails on a flipped rotation and on rotating in bf16),test_dense_attention.py::test_an_unchanged_layout_is_not_replanned, andtest_diffusion_dit_kernels.py, which holds the compiled kernels to the eager result within one bf16 step everywhere and equal on ≥99.9% (adaLN) / ≥99% (q/k rope) of elements.test_wan22_reference_equivalence.pywith oracles recorded on this box: 11 passed.test_wan_vae_decode.py: 22 passed, 4 failed, the same 4 on main.Not in this PR
_RaggedPrefillBase._plan, so the FLUX.2-klein, Z-Image and LTX-2.5 managers get it: its own PR.DenoiseLoopSubmodulecarries, for Wan and Cosmos3's four generation paths. Cosmos3's step syncs the same way Wan's did (forward_batchedreadstime_index.item(), and diffusers'UniPCMultistepScheduler.stepmoves sigmas from pageable memory and checkslinalg.solve's info). That's the follow-up after this lands.ada_layer_norm. FLUX.2's LayerNorm runs in bf16 per its reference, so it would need a variant that rounds the same way. Cosmos3's generation tower has no adaLN.