Skip to content

Cosmos3: one builder for the batched step's rows, so the warmup can't break boot #944

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Gaurav-Shah05 merged 3 commits into from
Oct 30, 2026
Merged

Cosmos3: one builder for the batched step's rows, so the warmup can't break boot #944

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bada56b
Cosmos3: the batched-step warmup hands the step the rotary table
merceod Oct 29, 2026
3a009a1
merge main
merceod Oct 30, 2026
56b49e1
Cosmos3: one builder for the batched step's rows
merceod Oct 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions mstar/model/cosmos3/submodules.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,8 +398,9 @@ def warmup_compiles(self) -> None:
dtype=self.transformer.proj_in.weight.dtype,
),
"vision_timesteps": torch.zeros(static["num_noisy_vision_tokens"], device=device),
"position_ids_cond": static["vision_mrope_ids"],
"position_ids_uncond": static["vision_mrope_ids"],
# the step takes the mRoPE table the statics built, not the position ids
"rotary_cond": static["rotary"],

@Gaurav-Shah05 Gaurav-Shah05 Oct 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 merceod Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

"rotary_uncond": static["rotary"],
"vision_token_shapes": static["vision_token_shapes"],
"vision_noisy_frame_indexes": static["vision_noisy_frame_indexes"],
"vision_mse_loss_indexes": static["mse_gen_indexes"],
Expand Down
12 changes: 11 additions & 1 deletion mstar/model/cosmos3/tests/test_denoise_compile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,14 @@ def __init__(self):
self.proj_in = torch.nn.Linear(2, 2)
self.batched_calls = []

# what the real step indexes per request; a warmup row missing one would raise here as it did on a server
REQ_KEYS = ("latents", "vision_timesteps", "rotary_cond", "rotary_uncond", "vision_token_shapes",

@Gaurav-Shah05 Gaurav-Shah05 Oct 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 merceod Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

"vision_noisy_frame_indexes", "vision_mse_loss_indexes")

def denoise_step_batched(self, requests, label, attn):
for req in requests:
for key in self.REQ_KEYS:
req[key]
self.batched_calls.append((requests, label, attn))
return [(None, None)] * len(requests)

Expand All @@ -103,11 +110,12 @@ def test_warmup_compiles_runs_one_two_request_step_at_the_default_size(monkeypat
t = _RecordingTransformer()
sub = Cosmos3DiTSubmodule(t, cfg)
built = []
rotary = (torch.zeros(4, 2), torch.zeros(4, 2))

def fake_static(ids, height, width, num_frames, fps, has_image_condition, device):
built.append((height, width, num_frames, has_image_condition))
return {
"num_noisy_vision_tokens": 4, "vision_mrope_ids": torch.zeros(3, 4),
"num_noisy_vision_tokens": 4, "rotary": rotary,
"vision_token_shapes": [(1, 2, 2)], "vision_noisy_frame_indexes": [torch.zeros(1)],
"mse_gen_indexes": torch.arange(4),
}
Expand All @@ -121,6 +129,8 @@ def fake_static(ids, height, width, num_frames, fps, has_image_condition, device
[(requests, label, attn)] = t.batched_calls
assert len(requests) == 2 and label == submodules.CFG_BATCHED_LABEL
assert requests[0]["latents"].shape == sub._latent_shape(480, 832, 121)
# both branches take the statics' table, as a served window's rows do
assert requests[0]["rotary_cond"] is rotary and requests[1]["rotary_uncond"] is rotary
# the stand-in reads no plan and answers with the query's shape
q = torch.randn(5, 2, 3)
assert not attn.requires_kv_write and torch.equal(attn.run(q, k=q, v=q), torch.zeros_like(q))
Expand Down
Loading