[merceod]

The Cosmos3 policy's compiled denoise_step (dynamic=False) recompiled for every new prompt length: 2.2-3.4 s on that request with a cold Inductor cache, 0.5-0.6 s warm, and a busy server walking toward the 84-recompile limit. TORCH_LOGS=recompiles shows each new prompt failing one guard, the stride of position_ids_cond: the action and sound paths pass position_ids[:, und_len:], a view whose row stride is the prompt plus generation length, while the text K/V only enters after the attention graph break.

The static inputs now carry gen_position_ids, a contiguous copy made once per request in packing.py, used at every call site. The video paths already passed their own tensor. No padding or bucketing needed.

Numbers, 2x B300, 8 droid_100 observations with 8 distinct prompt lengths (6-22 tokens), seeds 1000+i:

• Recompiles per log: main 23 (one per new prompt), this branch 15, all at startup.
• First request on a new prompt minus its repeat: main +2190 to +3405 ms cold, +518 to +608 warm; this branch -91 to +47 cold, -49 to +38 warm.
• Eager main vs eager branch: bit-identical 8/8. Compiled vs eager: 1.4% / 2.2% median / max rel. L2 (main 1.6% / 2.5%).

Tests: a CPU test that the generation band has the same size, stride and offset for 4, 9 and 23 prompt tokens (fails on main). mstar/model/cosmos3/tests, test_cosmos3_edge_model.py, test_cosmos3_graph_switch.py: 217 passed, 21 skipped (main 216 / 21). e2e --tier full cosmos3_edge_droid and cosmos3_nano_ar pass on B300 with the flashinfer configs. ruff clean.

[mstar-review]

Automated review

🟡 No findings on 0556690, light review

Description

• The latency and recompile numbers don't give the command or harness that produced them, so they can't be reproduced. Missing: how the server was launched, how requests were sent, and how 'first request minus its repeat' was measured.
• AGENTS.md asks for the torch/CUDA build alongside the GPU hardware for the B300 e2e and parity runs, and the description doesn't give it.

What the review checked

• Both static-input builders now set gen_position_ids: build_static_inputs, which also covers build_t2i_static_inputs through its wrapper, and build_action_static_inputs. In submodules.py, the request statics come only from _build_static and _build_action_static, so every dict that reaches a denoise call site carries the key.
• No position_ids[:, und_len:] slice is left in mstar/model/cosmos3/ outside test helpers, so every action, sound, batched-CFG and batched-request call site uses the contiguous copy.
• position_ids is never reassigned or changed in place after packing, so gen_position_ids can't fall out of sync with it. The other position_ids writes in submodules.py belong to the reasoner's own inputs, not the DiT statics.
• Cond and uncond prompts have different lengths, so their slices used to have different strides; both now pass contiguous tensors with the same layout.
• .contiguous() runs once per request at packing time, not per denoise step, so it adds no per-step CPU or GPU cost (AGENTS.md invariant 6).
• The new test in mstar/model/cosmos3/tests/test_action.py is CPU-only, runs under the CPU Core paths and checks that the values equal the old slice.

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

[NSagan271]

LGTM, pending some live testing that I'm running

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Numbers come from one allocation, so treat small deltas as provisional.

Live test: #922 Cosmos3 generation position ids

Head 05566900 vs its base b47e35b3 (no Cosmos3 changes on main since). cosmos3_edge_droid, 1× H100 80GB, SHM, Rust runtime, TORCH_LOGS=recompiles. 15 DROID prompts of distinct lengths (3–47 words), each sent once and then repeated, one droid_100 frame per request, seeds 1000+i. Builds alternated; the Inductor cache was cold on 3 boots per build and warm on 2–3. Video+sound used cosmos3_nano.

Verdict: ✅ A new prompt length no longer recompiles the denoise step, and eager outputs are bit-identical to base.

Claim	Result	Evidence
A new prompt length doesn't recompile	✅	Recompiles after startup: base 12 on every one of 6 boots, all from position_ids_cond stride mismatch at index 0 in denoise_step_batched_cfg; PR 0 on all 6. Totals including startup: base 27, PR 15
No first-request penalty	✅	First request minus its repeat, median per boot. Cold cache: base +1,604 / +1,869 / +1,886 ms (max +2,359), PR −16 / +9 / +35 ms. Warm cache: base +388 / +421 ms (max +838), PR −10 / +3 / +16 ms
Eager output unchanged	✅	TORCH_COMPILE_DISABLE=1: PR vs base 15/15 bit-identical, the same as base vs base and PR vs PR
Batched action path (_forward_batched_action)	✅	4 concurrent policy requests, 15 lengths × 2: recompiles after startup, base 8 / 5, PR 0 / 0. 30/30 returned 200 on every boot
Video + sound path (_forward_video_sound_gen)	✅	cosmos3_nano, generate_sound, 5 prompts of distinct lengths × 2, 29 frames at 256×256, 4 steps. Base hits the same position_ids_cond stride guard 10 times per boot, PR 0. A new prompt length's first request minus its repeat: base +399 to +450 ms, PR +39 to +87 ms. Eager: video and audio bit-identical to base, 10/10 each. Compiled: identical within each build across boots, different between builds
Compiled close to eager (1.4% / 2.2% on B300)	➖ not reproduced on H100 for either build	Compiled vs eager, median / max relative L2 over the 15 actions. PR: 5.1% / 10.7% on all 4 caches (identical across them). Base: 1.5% / 31%, 7.7% / 23%, 15.5% / 21%, 15.5% / 32%, depending on the cache. Base compiles a separate graph per prompt length, and its error varies with which kernels each compile picked; the PR's single graph is stable and its worst case is lower

Other

• Latency once warm is unchanged: repeat median 162–214 ms on both builds. Every request returned 200, and no server logged an error.
• Base's error on H100 (up to 15% median, 32% max, compiled vs eager) is larger than the description's B300 figures and comes from the existing compile, not this PR. It may be worth its own look.
• Not tested: B300/B200; the batched-sound path (_forward_batched_sound), which these runs, at concurrency 1, didn't reach.