[None][fix] Fix passing scaled timestep to time_embedder in Cosmos3 by bastefaniak · Pull Request #15545 · NVIDIA/TensorRT-LLM

bastefaniak · 2026-06-23T13:01:36Z

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.

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

coderabbitai · 2026-06-23T13:05:53Z

📝 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 · 2026-06-23T16:53:06Z

/bot run

tensorrt-cicd · 2026-06-23T16:58:50Z

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

tensorrt-cicd · 2026-06-24T01:21:32Z

PR_Github #55280 [ run ] completed with state FAILURE. Commit: 460b989
/LLM/main/L0_MergeRequest_PR pipeline #44235 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chang-l · 2026-06-24T01:30:00Z

/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-1"

tensorrt-cicd · 2026-06-24T01:35:57Z

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

tensorrt-cicd · 2026-06-24T04:13:50Z

PR_Github #55366 [ run ] completed with state FAILURE. Commit: 460b989
/LLM/main/L0_MergeRequest_PR pipeline #44316 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chang-l · 2026-06-24T04:17:51Z

/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-1"

chang-l · 2026-06-24T04:20:01Z

/bot run --reuse-test

tensorrt-cicd · 2026-06-24T04:24:09Z

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

tensorrt-cicd · 2026-06-24T04:25:33Z

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

tensorrt-cicd · 2026-06-24T04:25:42Z

PR_Github #55407 [ run ] completed with state ABORTED. Commit: 460b989

Link to invocation

tensorrt-cicd · 2026-06-24T05:48:55Z

PR_Github #55408 [ run ] completed with state SUCCESS. Commit: 460b989
/LLM/main/L0_MergeRequest_PR pipeline #44352 completed with status: 'SUCCESS'

CI Report

Link to invocation

…VIDIA#15545) Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>

…d Cosmos3-Nano Add default-setting single-GPU LPIPS golden tests for QwenImage and Cosmos3-Nano, and complete the VisualGen LPIPS CI protection introduced by the preceding changes. Refresh all eight golden media entries with the pinned staging main image at TRT-LLM commit 85665f5, which contains the Cosmos3 accuracy fix from NVIDIA#15545. Record the TRT-LLM commit, package versions, and container digest in every golden JSON while retaining the original LPIPS thresholds. Remove every remaining LPIPS waiver and keep all eight single-GPU and eight multi-GPU cases registered in the B200 test lists. Restore the Cosmos3 guardrail environment after generation and disable TRT-LLM MPI bootstrap before the torch-distributed multi-GPU harness imports its helpers. Validated on B200: all eight single-GPU and eight multi-GPU LPIPS cases passed with the original thresholds. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>

…ation Add default-setting single-GPU LPIPS golden tests for QwenImage and Cosmos3-Nano, completing single-GPU VisualGen LPIPS CI coverage across all supported models. Refresh all eight golden media entries with the pinned staging main image at TRT-LLM commit 85665f5, which contains the Cosmos3 accuracy fix from NVIDIA#15545. Record the TRT-LLM commit, package versions, container digest, compile-off mode, and deterministic-algorithm mode in every golden JSON while retaining the existing LPIPS thresholds. Explicitly disable pipeline torch compile for every single-GPU LPIPS generator. Unwaive and register all eight single-GPU cases while retaining waivers for all eight multi-GPU LPIPS cases. Validated on B200: all eight single-GPU LPIPS cases passed at 0.000000 against the refreshed goldens with the original thresholds. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>

…ation Add default-setting single-GPU LPIPS golden tests for QwenImage and Cosmos3-Nano, completing single-GPU VisualGen LPIPS CI coverage across all supported models. Refresh all eight golden media entries with the pinned staging main image at TRT-LLM commit 85665f5, which contains the Cosmos3 accuracy fix from NVIDIA#15545. Record the TRT-LLM commit, package versions, container digest, compile-off mode, and deterministic-algorithm mode in every golden JSON while retaining the existing LPIPS thresholds. Explicitly disable pipeline torch compile for every LPIPS generator. Unwaive all eight single-GPU cases, test_wan_t2v_example, and the passing attn2d_2x2_ulysses2 multi-GPU case. Retain the six pre-existing multi-GPU waivers and waive only the new cfg2_ulysses2_attn2d_2x1 case that measured 0.255208 above the unchanged 0.25 threshold. Validated on B200: all eight single-GPU LPIPS cases passed at 0.000000, test_wan_t2v_example passed, and attn2d_2x2_ulysses2 passed at 0.232348. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>

…VIDIA#15545) Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Fix timestep passed to time_embedder

8ac72ae

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

bastefaniak requested a review from a team as a code owner June 23, 2026 13:01

github-actions Bot assigned bastefaniak Jun 23, 2026

chang-l approved these changes Jun 23, 2026

View reviewed changes

bastefaniak changed the title ~~[None][Fix] Fix passing scaled timestep to time_embedder in Cosmos3~~ [None][fix] Fix passing scaled timestep to time_embedder in Cosmos3 Jun 23, 2026

Merge branch 'main' into bstefaniak/fix_cosmos3_visual_gen_timestep

460b989

chang-l enabled auto-merge (squash) June 23, 2026 19:24

bobboli approved these changes Jun 24, 2026

View reviewed changes

chang-l merged commit 9e9137f into NVIDIA:main Jun 24, 2026
8 checks passed

coderabbitai Bot mentioned this pull request Jun 25, 2026

[None][refactor] Refine Skip Softmax follow-ups #15417

Merged

bobboli mentioned this pull request Jun 25, 2026

[TRTLLM-13120][feat] Cosmos3 Audio Output Support #14827

Merged

1 task

chang-l mentioned this pull request Jun 29, 2026

[None][test] Enable single-GPU LPIPS CI protection for VisualGen models #15564

Merged

BrianLi23 pushed a commit to BrianLi23/TensorRT-LLM that referenced this pull request Jul 9, 2026

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

968281e

…VIDIA#15545) Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Uh oh!

Conversation

bastefaniak commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

chang-l commented Jun 23, 2026

Uh oh!

tensorrt-cicd commented Jun 23, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

chang-l commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

chang-l commented Jun 24, 2026

Uh oh!

chang-l commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bastefaniak commented Jun 23, 2026 •

edited

Loading

coderabbitai Bot commented Jun 23, 2026 •

edited

Loading