[None][fix] Fix passing scaled timestep to time_embedder in Cosmos3

[bastefaniak]

Description

This MR fixes regression introduced in MR14687 for Cosmos3 model, after changes we are passing scaled timestep to Cosmos3VFMTransformer.time_embedder which makes inference produce noise, the fix is to pass raw timestep to the time_embedder and scaled timestep elsewhere. This can be tested by running example examples/visual_gen/models/cosmos3_ti2v.py

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

A new attention_timestep parameter (normalized by num_train_timesteps) is added to Cosmos3VFMTransformer.forward. Internally, both the UND language_model call and each GEN decoder layer call are switched to use attention_timestep instead of timestep. The pipeline's forward_fn computes and passes this normalized value alongside the raw timestep.

Changes

Cosmos3 attention_timestep separation

Layer / File(s)	Summary
attention_timestep parameter and internal usage in transformer tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py	Cosmos3VFMTransformer.forward adds an optional attention_timestep: Optional[torch.Tensor] parameter with docstring. The language_model(...) call and each GEN decoder layer call switch their timestep= argument from timestep to attention_timestep.
Pipeline: compute and pass attention_timestep tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py	Inside the denoising loop's forward_fn, attention_timestep is computed as timestep / self.scheduler.config.num_train_timesteps and forwarded to self.transformer(...).

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies a specific fix to the Cosmos3 model regarding timestep passing to time_embedder, matching the main change described in the pull request.
Description check	✅ Passed	The PR description clearly identifies the regression, references the causing PR, explains the fix approach, and provides testing instructions.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[chang-l]

/bot run

[tensorrt-cicd]

PR_Github #55280 [ run ] triggered by Bot. Commit: 460b989 Link to invocation