[merceod]

Main can't start Cosmos3 with the stock Edge or Nano AR configs. #927's warmup_compiles builds its two warmup rows with position_ids_cond/position_ids_uncond, but since #896 the batched step reads rotary_cond/rotary_uncond; #927 was written against the step before #896 and the merge didn't conflict. The engine's warmup raises KeyError: 'rotary_cond' and the conductor shuts the deployment down. A config with compile_denoise_batched: false skips the warmup and boots, which is how the video benchmarks kept running on main.

That key mismatch was possible because four places built the step's input by hand. Now one does.

Change

Cosmos3OmniTransformer.denoise_step_row(layout, latents, vision_timesteps, rotary_cond, rotary_uncond, noisy_token_mask=None, sound_latents=None, sound_timesteps=None) sits right above denoise_step_batched and is the only place a row is built. A new input to the step changes this signature, so every caller has to pass it, the warmup included. Where each caller's inputs come from:

caller	layout	latents	vision_timesteps	rotary_cond / rotary_uncond	extra
warmup_compiles	the statics it builds at the default size	zeros	zeros, one per noisy token	static["rotary"] for both
forward_batched	st["cond"]	the request's latents	t per noisy token	st["cond"]["rotary"] / st["uncond"]["rotary"]
_forward_batched_sound	st["cond"], which also has the sound band's token layout	the request's latents	t per noisy token	st["cond"]["rotary"] / st["uncond"]["rotary"]	sound_latents, sound_timesteps (t per st["num_sound"])
forward_captured	_capture_layout[latent size]	latents[i]	vision_timesteps[i]	(cos_cond[i], sin_cond[i]) / (cos_uncond[i], sin_uncond[i])	noisy_token_mask[i]

The rows are what they were. I fingerprinted every row from all four callers on the tree before and after this change (keys in order, dtype, size, bytes, which statics or input object each value is), and the dumps are identical. The builder adds about 40 ns a row over the dict literal (202 vs 161 ns), which is why it isn't benchmarked on its own; the served numbers are below.

The callers reach it through self.transformer: importing the tower at the top of submodules.py would pull about 400 more modules into every process that loads cosmos3_model, which imports it lazily today.

denoise_step_action_batched is a different step and isn't touched: its rows are built in one place (_forward_batched_action) and it has no warmup row.

Tests

test_denoise_compile.py no longer has a stand-in that lists the step's keys. test_warmup_compiles_runs_the_real_step_and_a_served_row_reuses_its_compile runs warmup_compiles through the real step on test_rotary_table._tiny() (two layers, CPU) with compile_denoise on, then runs a row of another size against the attention a served step binds under fail_on_recompile. About 20 s, nearly all of it the compile. With this test file on main it fails in 3.7 s with the server's error:

mstar/model/cosmos3/components/transformer.py:1341: in denoise_step_batched
    cc, sc = req["rotary_cond"]
E   KeyError: 'rotary_cond'

test_served_batched_paths_run_the_real_step runs forward_batched, the sound path and a captured batch through the real step and checks the eager and captured batches agree to 1e-4. Nothing on CPU ran the batched sound row before.

What those tests catch, each edit tried on a scratch copy:

edit	result
the step renames rotary_cond/rotary_uncond, the builder follows	the warmup test passes: the warmup follows by construction
the step renames a key, the builder doesn't	warmup and served-path tests fail with KeyError
the step reads a new key no caller passes	warmup and served-path tests fail with KeyError
the builder drops sound_mse_gen_indexes	served-path test fails with KeyError
the layer stack compiles static instead of dynamic	warmup test fails: fail_on_recompile names the cos size guard

• ruff check . clean.
• python -m pytest mstar/model/cosmos3/tests, CPU: 192 passed, 21 skipped (main: 191 passed, 21 skipped).
• The same directory on a B300 with the Nano and Edge checkpoints: 190 passed, 8 failed, 15 skipped. The real-weight batched checks pass (test_cross_request_batch_matches_individual, test_batched_cfg_matches_sequential, test_action_cross_request_batch_matches_individual). The 8 failures are Nano video parity (test_engine_cache x4, test_sound_engine_path_flashinfer, test_t2i_parity_vs_diffusers, test_t2v_parity_vs_diffusers, test_i2v_parity_vs_diffusers) and fail the same way without this change, with the same PSNRs to two decimals. The 15 skips need the DROID checkpoint or Edge reference dumps, or are marked skipped in the file.
• CPU Core selection (CI's shard, -p xdist, no GPU): 8857 passed, 41 failed, 787 skipped; main 8854 passed, 43 failed, 787 skipped. All 41 fail on main too (test_loop_output_ownership x32, test_send_registration_failure x5, test_eval_compare_multigpu x2, test_client_error_text, test_serving_metrics); main's other two are test_eval_compare and test_eval_compare_multigpu cases.

Boot check

2x B300, torch 2.12.1+cu132, stock configs, -O2, server on 8 cores, MSTAR_SHM_ARENA=1, TORCH_LOGS=recompiles.

config	main (7eac23c)	this PR
configs/cosmos3_edge.yaml	KeyError: 'rotary_cond' at warmup	starts in 63-64 s, t2v 832x480x121 served
configs/cosmos3_nano_ar.yaml	KeyError: 'rotary_cond' at warmup	starts in 245-324 s (weights on a network volume), chained t2v 832x480x121 served

Bytes, same seeds and pins: Edge's c=1 mp4s (3 seeds at 832x480, 1 at 480x832) are identical on bada56b, on its merge with main and on this head, across four servers, and Nano's chained mp4 is identical across three. c=1 doesn't run the batched step, so I checked c=2 too: Edge's four 832x480 mp4s match across the three trees, but at c=2 this head differs from itself on 2 of 4 between its two servers. I read that as which steps happen to batch (the hashes repeat across trees), not traced. Nano's second run on this head matches its own first run on 6 of 12 c=2 mp4s and the tree before this change on 8 of 12, and all 12 match one or the other.

Recompiling function lines after startup, the same counts and functions on the tree before this change:

cell	Edge	Nano AR
c=1 t2v	0	50 (this config has no warmup_requests, so it's the first use)
first c=2 at 832x480	0	29
first c=2 at 480x832	24	24
second c=2 at 480x832	0	0

None of those is the batched step: nothing names gen_layers_batched, and no guard sees a two-request pack (48360 rows at 121 frames), while the worker log shows bs=2.00 in these cells. What does recompile is the VAE decoder at a new latent size and the one-request denoise_step_batched_cfg, which compiles per size and runs while a request steps alone (its guards read 24180 rows or the new latent size). Both are per-size compiles that predate this PR.

Results

GPU 0, server on 8 cores, client on 4 cores of another NUMA node, two servers per arm alternated. Main is 7eac23c with compile_denoise_batched: false, the only way it boots.

cell	main, batched step eager (s1 / s2)	this PR (s1 / s2)
Edge t2v 832x480x121 c=1 e2e p50	3.37 / 3.36 s	3.35 / 3.36 s
Edge t2v c=4 videos/min	13.9 / 14.0	19.2 / 19.4
Edge t2v c=4 e2e p50	16.8 / 16.2 s	11.6 / 11.7 s
Nano chained t2v 832x480x121 c=1 TTFC / chunk gap	2.23 / 1.95, 2.22 / 1.95 s	2.22 / 1.96, 2.21 / 1.95 s
Nano chained t2v c=4 TTFC / chunk gap	11.6 / 10.5, 11.2 / 10.4 s	8.45 / 7.77, 8.93 / 7.77 s
Nano chained t2v c=4 aggregate frames/s	8.46 / 8.46	11.0 / 11.0

Review measured the Edge cells on one H100 too: c=4 7.5 → 9.6 videos/min, c=1 unchanged.

Follow-ups

• The 8 Nano parity failures above, which aren't from this PR.
• A portrait request's first step still compiles the one-request step and the VAE decoder at that size; warmup_requests only covers the default video sizes.

[mstar-review]

Automated review

🟡 No findings on bada56b, light review

Description

• The Nano AR boot check, which the first paragraph names as broken on main, hasn't been reported yet and is promised in a later comment, so the fix is only shown working on configs/cosmos3_edge.yaml so far.

What the review checked

• denoise_step_batched (mstar/model/cosmos3/components/transformer.py:1316-1373) reads rotary_cond/rotary_uncond as (cos, sin) tuples plus latents, vision_timesteps, token shapes, noisy frame indexes and mse indexes; the warmup row now provides every unconditional key, and noisy_token_mask/sound_* are read with .get or only when sound latents are present.
• static["rotary"] is always set when the warmup runs: _build_static calls _attach_rotary (submodules.py:541-551), which sets it whenever a transformer exists, and warmup_compiles returns early when it doesn't (submodules.py:386).
• Sharing the cond table for the uncond branch at warmup matches the shape the served path passes (submodules.py:2101-2102, 2154-2155), and the compiled gen_layers_batched is dynamic in pack length, so the table's contents don't affect the compiled graph.
• The test stand-in's REQ_KEYS matches the unconditional keys the real batched step reads, and the fake _build_static now returns rotary the way the real one does.

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

[merceod]

Boot check is in: both stock configs start on a B300 (Edge 98 s, Nano 70 s), and with the batched step compiled again Edge t2v c=4 goes from 13.8 to 19.4 videos/min and Nano c=4 TTFC from 11.4 to 8.6 s. Table in the description.

[Gaurav-Shah05]

The two keys fix the boot, but the warmup still builds the step's input by hand, so the next change to what denoise_step_batched reads can break every Cosmos3 boot again.

On one H100, both stock configs fail on main with KeyError: 'rotary_cond' at warmup and start on this head. The first batched requests after the warmup, at 832x480 and at 480x832, don't recompile the step. Edge t2v at 832x480x121 and 20 steps, against main with compile_denoise_batched: false, at least two server sessions per arm:

	main, batched step eager	this head	Δ
c=1 e2e p50	7.1 s	7.0 s	1.0x
c=4 e2e p50	29.2 s	22.1 s	1.32x ↓
c=4 videos/min	7.5	9.6	1.29x ↑

[Gaurav-Shah05]

This row is the fourth hand-built copy of denoise_step_batched's input, beside the ones in forward_batched, _forward_batched_sound and forward_captured, so every change to what the step reads has to find all four. One builder next to denoise_step_batched in transformer.py, called by all four with their statics or capture layout, should build these rows, so the warmup follows the step by construction.

[merceod]

Done: Cosmos3OmniTransformer.denoise_step_row sits right above the step and warmup_compiles, forward_batched, _forward_batched_sound and forward_captured all build their rows with it, so a new input to the step changes its signature and every caller has to pass it. I fingerprinted each caller's rows before and after (keys in order, bytes, which statics object each value is) and they're identical, and the description says where each caller's inputs come from.

[Gaurav-Shah05]

REQ_KEYS is a second copy of what denoise_step_batched reads: when I renamed the step's rotary keys and the served rows but not the warmup, this test still passed. The test should run warmup_compiles through the real step on test_rotary_table._tiny() (a two-layer tower on CPU) with compile_denoise turned on, and on main that fails in about 4 s with the server's KeyError: 'rotary_cond'.

[merceod]

Done: REQ_KEYS and the recording transformer are gone, and the test runs warmup_compiles through the real step on test_rotary_table._tiny() with compile_denoise on. With this test file on main it fails in 3.7 s with the server's KeyError: 'rotary_cond', and it also fails if the step renames a key and the builder doesn't; rename both, as you tried, and it passes because the warmup follows.

[merceod]

Agreed: one row builder next to denoise_step_batched in transformer.py, called by forward_batched, _forward_batched_sound, forward_captured and the warmup, and the test runs warmup_compiles through the real step on test_rotary_table._tiny() with compile_denoise on, so it fails on main the way the server did. Rework coming.

[merceod]

Pushed 3a009a1 (merge of main) and 56b49e1: one row builder next to the step, the four callers on it, the warmup test through the real step, and a served-path test that also runs the batched sound row on CPU. Description rewritten with the B300 boot check, recompile counts and the c=4 numbers; the 8 Nano parity failures are pre-existing and listed as a follow-up. Re-requesting.

[Gaurav-Shah05]

LGTM