fix(packing): gate timestep-embedding fast path on timestep values, not tensor shape

[umairjavaid]

Fixes #235.

Summary

uses_single_timestep gates the one-row-and-broadcast fast path in _embed_packed_timesteps. Its docstring describes a property of timestep values ("all noised tokens share one input timestep scalar", sequence.py:42) but it was computed as input_timesteps.numel() == 1, a property of the tensor's shape. The fast path was therefore skipped whenever the timesteps were uniform but the tensor held more than one element.

This makes the gate implement its own documented contract, and keeps it honest for the one code path that rewrites timesteps after packing.

Timestep tensor	Uniform values?	Fast path before	Fast path after
[0.5] — one sample	yes	yes	yes
[0.5, 0.5] — batch sharing a timestep	yes	no	yes
zeros(2B) — use_batched_cfg	yes	no	yes
full((1, T), 0.5) — uniform diffusion forcing	yes	no	yes
[0.25, 0.75] — per-sample sigmas	no	no	no
[[0.1, 0.2, 0.3, 0.4]] — per-frame sigmas	no	no	no

The three regressed rows include every step of the template inference path, where _copy_timestep_to_template (omni_mot_model.py:2734) broadcasts a single scalar over all noisy tokens, so uniformity is guaranteed by construction.

The fix, in two parts

1. Derive the flag from the values (packers.py). Extracted as a small helper because the condition is now needed in more than one place. pack_input_sequence already rejects CUDA input (packers.py:193), so this reads a CPU tensor and adds no device synchronization:

def uses_single_timestep(input_timesteps: torch.Tensor) -> bool:
    if input_timesteps.numel() == 0:
        return False
    flat = input_timesteps.reshape(-1)
    return bool((flat == flat[0]).all())

Empty input returns False, byte-identical to the previous behaviour. NaN never compares equal, so it declines the fast path.

2. Keep the flag honest when timesteps are rewritten after packing (omni_mot_model.py). This is what makes a values check safe rather than merely correct-looking. Under independent_action_schedule / independent_sound_schedule, training_step replaces the packed action and sound timesteps with one sigma per sample after the packer has already set the flag from the vision timesteps. A uniform vision batch with differing action sigmas would otherwise leave the flag true and let _embed_packed_timesteps broadcast sample 0's action sigma across the whole batch. The old numel() == 1 derivation masked this, since a single-element tensor implies batch size 1, which makes every modality trivially uniform. Both rewrite sites now AND in the uniformity of what they wrote:

packed_sequence.uses_single_timestep &= uses_single_timestep(sample_ts)

To be explicit about reachability: training sigmas are independent continuous draws per sample (rectified_flow.py), so a uniform vision batch at B > 1 is essentially unreachable with today's sampler. Part 2 is defensive — it makes the flag's contract hold for any caller (a fixed-sigma evaluation, an external packer) rather than by accident of the sampler.

I audited every writer to a packed .timesteps field. _copy_timestep_to_template (expand_as), omni_mot_causal_model.py:2885 (timestep.flatten()[0].repeat(n)) and omni_mot_causal_model.py:3090 (torch.zeros) are uniform by construction; temporal_causal.py:284 derives each frame's value from the same per-sample input_timestep the packer-level check already covers; teacher_forcing.py:30 writes an empty tensor, which the consumer's numel() > 1 guard already handles. The two independent-schedule sites were the only ones that could produce non-uniform values.

Tests

New cosmos_framework/data/generator/sequence_packing/packers_test.py — the first coverage for this flag. Construction mirrors create_packed_sequence in context_parallel_test.py. Restricted to the tests that use only the public API:

	Result
On main	3 failed, 5 passed
With this change	8 passed

The five that already passed are the negative cases, which pin that the fix does not over-enable the fast path. The full file is 17 tests, covering the helper directly, the packer end to end, and the post-packing composition from part 2. CPU-only, ~3 s:

uv run pytest cosmos_framework/data/generator/sequence_packing/packers_test.py

One test asserts the underlying invariant rather than the flag: when the flag is true, every packed timestep really does equal timesteps[:1], which is what the consumer broadcasts. That holds regardless of how the flag is derived.

Numerics

The fast path is not bit-identical to the general path, because an M=1 GEMM selects a different kernel than M=N and accumulation order differs:

Setting	Max abs difference
CPU fp32	3e-07 – 5e-07
A100, TF32 enabled (equivalent TimestepEmbedding module, measured on the diffusers Cosmos3-Edge path)	~5e-05
A100, TF32 disabled (same)	~3.5e-07

This is pre-existing behaviour: the fast path is already taken at batch size 1 on main, so this change extends an already-accepted approximation to more shapes rather than introducing a new class of difference. Flagging it explicitly because outputs for B > 1 will no longer be bit-identical to before this patch.

For contrast, forcing the flag true on genuinely mixed timesteps (a synthetic check) produces a max abs difference of 0.103 — four orders of magnitude larger, which is the wrong-conditioning error part 2 prevents.

Performance

The saving is confined to the conditioning path: the fp32 MLP runs on 1 row instead of one per noisy token. Large factor on the module in isolation, but well under one percent of a full generator forward in bf16, so I would frame this as making the gate match its documented contract, with the use_batched_cfg interaction as the concrete motivation, rather than as a speedup.

For scale, the equivalent module at equivalent shapes — measured on the diffusers Cosmos3-Edge path at 480p/61 frames, not in this repo — evaluates 6,240 rows where 1 suffices.

Notes

• No new device synchronization; the check runs on a CPU tensor by existing contract.
• torch.compile: the value check runs in the packer, outside any compiled region, so it cannot introduce a graph break. The flag is consumed inside _encode_vision / _encode_action, which apply_compile compiles individually (parallelize_vfm_network.py:37-41); there it is a Python bool that Dynamo already guards on today. In inference the flag is now constant True across batch sizes, so this removes a guard variant rather than adding one.
• No behaviour change for empty timestep tensors.
• ruff check and ruff format --check produce identical output to main for both edited files (the two pre-existing I001 findings are left untouched to keep the diff focused); the new test file is clean under both.
• Commit is signed off per CONTRIBUTING.md.

Follow-up, deliberately not in this PR

uses_single_timestep lives on the PackedSequence but is consumed per modality — vision, action and sound each get their own _embed_packed_timesteps call while sharing one flag. That is why part 2 needs an AND rather than a per-modality answer, and it means a uniform-vision batch with non-uniform action sigmas now conservatively loses the fast path for vision too. Moving the flag onto the modality would fix that and delete the AND. Happy to do it as a separate change if you'd like it.

[Copilot]

🟢 Approval recommended

The change aligns the flag with its documented semantics, preserves safety for post-pack timestep rewrites, and adds focused tests that reproduce the prior failures and validate the corrected behavior.

Pull request overview

This PR fixes the PackedSequence.uses_single_timestep gate so it matches its documented contract (“all noised tokens share one timestep scalar”) by checking timestep values rather than the timestep tensor’s shape, ensuring the timestep-embedding broadcast fast path is used for uniform batches across more real inference/training shapes.

Changes:

• Add uses_single_timestep(input_timesteps) helper that returns True iff all timestep values are identical (and False for empty tensors / NaNs).
• Switch pack_input_sequence to derive PackedSequenceBuilder(uses_single_timestep=...) from the helper rather than numel() == 1.
• Ensure training_step clears (ANDs) the flag when action/sound timesteps are overwritten post-pack under independent schedules, and add a new targeted test suite covering positive/negative cases and the “post-pack rewrite” composition.

File summaries

File	Description
cosmos_framework/model/generator/omni_mot_model.py	ANDs uses_single_timestep with rewritten per-sample action/sound schedules so the embedding fast path remains correct after post-pack timestep overrides.
cosmos_framework/data/generator/sequence_packing/packers.py	Introduces value-based uses_single_timestep helper and uses it to initialize the packed-sequence flag during packing.
cosmos_framework/data/generator/sequence_packing/packers_test.py	Adds regression tests covering uniform vs non-uniform timestep inputs and the independent-modality post-pack invalidation behavior.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[umairjavaid]

Kindly tag reviewers for this PR

[pengcuo]

Merged — thank you for this contribution!

Really appreciated the thoroughness here: the audit of every post-pack writer of .timesteps, the explicit AND at the independent action/sound schedule sites, and the honest note that the fast path is not bit-identical to the general path. That made the review straightforward. The new packers_test.py is also the first coverage this flag has had, which is great to have.