[TRTLLM-11410][feat] Cosmos3 Support by NVShreyas · Pull Request #14824 · NVIDIA/TensorRT-LLM

NVShreyas · 2026-06-01T14:06:13Z

Summary by CodeRabbit

New Features
- Added support for NVIDIA Cosmos3-Nano and Cosmos3-Super diffusion models with FP8 dynamic quantization
- Added example scripts and multi-GPU configurations for Cosmos3 text-to-video and text+image-to-video generation
Documentation
- Updated supported models documentation with Cosmos3 feature matrix and capability specifications

Description

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 · 2026-06-01T14:10:31Z

📝 Walkthrough

Walkthrough

This PR adds support for NVIDIA's Cosmos3 visual generation models (Nano and Super variants) to TensorRT-LLM. Changes include pipeline registration with Hugging Face model IDs, YAML configuration templates for single and multi-GPU deployments, a CLI example script for text-to-video/image-to-video generation, and documentation updates listing the new supported models and their capabilities.

Changes

Cosmos3 Visual Generation Feature

Layer / File(s)	Summary
Pipeline registration and metadata `tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`	`Cosmos3OmniMoTPipeline` decorator is updated to include explicit Hugging Face model IDs and a documentation string, removing the prior TODO and completing the pipeline metadata.
Configuration templates for different hardware setups `examples/visual_gen/configs/cosmos3-nano-1gpu.yaml`, `examples/visual_gen/configs/cosmos3-super-4gpu.yaml`	Two new YAML configs provide quantization (FP8 dynamic), attention backend (VANILLA), and parallelism parameters tailored for 1-GPU and 4-GPU deployments.
CLI example script and documentation updates `examples/visual_gen/models/cosmos3_ti2v.py`, `docs/source/models/visual-generation.md`	A new CLI script demonstrates text-to-video and image-to-video generation with argparse options for model selection, config loading, and output saving. Documentation tables are updated to list Cosmos3-Nano and Cosmos3-Super as supported models with their feature capabilities.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Suggested labels

VisualGen

Suggested reviewers

chang-l
arysef

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is completely empty of actual content—only the repository template is present with no filled-in sections for description, test coverage, or checklist items.	Fill in the Description section explaining what Cosmos3 support includes, provide Test Coverage details listing relevant tests, and complete the PR Checklist items.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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 and specifically identifies Cosmos3 support as the main feature being added across multiple files and configurations.

_{✏️ 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 and usage tips.}

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/visual_gen/models/cosmos3_ti2v.py`:
- Line 78: Update the function signature for main by adding an explicit return
type annotation; change the definition of the top-level function main (def
main()) to def main() -> None: so it conforms to the project's typing rule,
leaving the body unchanged and ensuring no implicit return values are returned.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 84a37da3-55d0-4db8-81e8-df09db376217

📥 Commits

Reviewing files that changed from the base of the PR and between 06456e1 and 685b570.

📒 Files selected for processing (5)

docs/source/models/visual-generation.md
examples/visual_gen/configs/cosmos3-nano-1gpu.yaml
examples/visual_gen/configs/cosmos3-super-4gpu.yaml
examples/visual_gen/models/cosmos3_ti2v.py
tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py

NVShreyas · 2026-06-01T23:41:23Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-01T23:48:56Z

PR_Github #51443 [ run ] triggered by Bot. Commit: b0c02da Link to invocation

tensorrt-cicd · 2026-06-02T06:33:30Z

PR_Github #51443 [ run ] completed with state SUCCESS. Commit: b0c02da
/LLM/main/L0_MergeRequest_PR pipeline #40852 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

Link to invocation

chang-l

Can you please also update this doc: https://nvidia.github.io/TensorRT-LLM/1.3.0rc17/models/supported-models.html?

NVShreyas · 2026-06-02T22:39:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-02T22:45:39Z

PR_Github #51686 [ run ] triggered by Bot. Commit: a66fdc9 Link to invocation

tensorrt-cicd · 2026-06-03T06:36:29Z

PR_Github #51686 [ run ] completed with state SUCCESS. Commit: a66fdc9
/LLM/main/L0_MergeRequest_PR pipeline #41066 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

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

NVShreyas · 2026-06-03T14:39:23Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-03T14:45:18Z

PR_Github #51861 [ run ] triggered by Bot. Commit: d9d6e8b Link to invocation

tensorrt-cicd · 2026-06-04T14:46:03Z

PR_Github #51861 [ run ] completed with state ABORTED. Commit: d9d6e8b

Link to invocation

NVShreyas · 2026-06-04T15:12:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-04T15:19:31Z

PR_Github #52098 [ run ] triggered by Bot. Commit: d9d6e8b Link to invocation

tensorrt-cicd · 2026-06-04T21:17:13Z

PR_Github #52098 [ run ] completed with state SUCCESS. Commit: d9d6e8b
/LLM/main/L0_MergeRequest_PR pipeline #41428 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

NVShreyas · 2026-06-04T21:46:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-04T21:52:25Z

PR_Github #52165 [ run ] triggered by Bot. Commit: d9d6e8b Link to invocation

tensorrt-cicd · 2026-06-05T02:30:59Z

PR_Github #52165 [ run ] completed with state SUCCESS. Commit: d9d6e8b
/LLM/main/L0_MergeRequest_PR pipeline #41486 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

NVShreyas · 2026-06-05T04:15:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-05T04:21:10Z

PR_Github #52235 [ run ] triggered by Bot. Commit: d9d6e8b Link to invocation

tensorrt-cicd · 2026-06-05T08:51:20Z

PR_Github #52235 [ run ] completed with state SUCCESS. Commit: d9d6e8b
/LLM/main/L0_MergeRequest_PR pipeline #41551 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

NVShreyas · 2026-06-05T14:44:39Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-05T14:50:58Z

PR_Github #52352 [ run ] triggered by Bot. Commit: d9d6e8b Link to invocation

NVShreyas · 2026-06-05T14:53:33Z

/bot kill

tensorrt-cicd · 2026-06-05T14:59:54Z

PR_Github #52354 [ kill ] triggered by Bot. Commit: d9d6e8b Link to invocation

NVShreyas · 2026-06-05T15:03:18Z

/bot skip --comment "PR only adds docs and examples, tests are run on B200 which pass. failures are GB200 which are unrelated"

tensorrt-cicd · 2026-06-05T15:03:47Z

PR_Github #52352 [ run ] completed with state ABORTED. Commit: d9d6e8b

Link to invocation

tensorrt-cicd · 2026-06-05T15:03:58Z

PR_Github #52354 [ kill ] completed with state SUCCESS. Commit: d9d6e8b
Successfully killed previous jobs for commit d9d6e8b

Link to invocation

tensorrt-cicd · 2026-06-05T15:09:03Z

PR_Github #52356 [ skip ] triggered by Bot. Commit: d9d6e8b Link to invocation

tensorrt-cicd · 2026-06-05T15:15:51Z

PR_Github #52356 [ skip ] completed with state SUCCESS. Commit: d9d6e8b
Skipping testing for commit d9d6e8b

Link to invocation

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com> Signed-off-by: NVFB <186336021+NVFB@users.noreply.github.com>

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

NVShreyas requested review from a team as code owners June 1, 2026 14:06

NVShreyas requested review from chang-l and kaiyux June 1, 2026 14:06

github-actions Bot assigned NVShreyas Jun 1, 2026

coderabbitai Bot reviewed Jun 1, 2026

View reviewed changes

Comment thread examples/visual_gen/models/cosmos3_ti2v.py

NVShreyas force-pushed the user/shreyasm/cosmos3 branch 2 times, most recently from c028dfd to cf040bb Compare June 1, 2026 16:58

NVShreyas changed the title ~~[None][feat] Cosmos3 Support~~ [TRTLLM-11410][feat] Cosmos3 Support Jun 1, 2026

chang-l approved these changes Jun 2, 2026

View reviewed changes

Comment thread docs/source/models/visual-generation.md Outdated

Comment thread tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

NVShreyas added 7 commits June 3, 2026 07:37

cosmos3 examples

e522452

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

update example

0aca7f4

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

smoke tests

5f872e1

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

User sequence sharder; multi-gpu test

2ba2dd8

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

enable tp for cosmos3

71cac19

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

multi-gpu unit tests

764be26

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

update docs

d9d6e8b

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

NVShreyas force-pushed the user/shreyasm/cosmos3 branch from a66fdc9 to d9d6e8b Compare June 3, 2026 14:39

NVShreyas merged commit 73c824c into NVIDIA:main Jun 5, 2026
7 checks passed

fbxai pushed a commit to fbxai/TensorRT-LLM that referenced this pull request Jun 5, 2026

[TRTLLM-11410][feat] Cosmos3 Support (NVIDIA#14824)

91041ec

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com> Signed-off-by: NVFB <186336021+NVFB@users.noreply.github.com>

2ez4bz pushed a commit to 2ez4bz/TensorRT-LLM that referenced this pull request Jun 8, 2026

[TRTLLM-11410][feat] Cosmos3 Support (NVIDIA#14824)

e588c04

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

tensorrt-cicd mentioned this pull request Jun 11, 2026

[https://nvbugs/6290965][fix] Rename model_config → pipeline_config on the two Cosmos3 test asserts and… #15263

Closed

2 tasks

coderabbitai Bot mentioned this pull request Jul 9, 2026

[None][feat] Cosmos3 video-to-video (V2V) generation #16155

Open

1 task

Uh oh!

Conversation

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

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

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

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

NVShreyas commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

chang-l left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

NVShreyas commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

NVShreyas commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

NVShreyas commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

NVShreyas commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

NVShreyas commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

NVShreyas commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

NVShreyas commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

NVShreyas commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

NVShreyas commented Jun 1, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jun 1, 2026 •

edited

Loading