[TRTLLM-11410][feat] MoT World Model Support#14012
Conversation
📝 WalkthroughWalkthroughThis PR integrates a Cosmos3 visual-to-video diffusion pipeline into TensorRT-LLM, introducing text/video guardrails, a dual-pathway mRoPE-based transformer, and a complete generation flow supporting both text-to-video and image-to-video modes. It also refines Ulysses sequence parallelism control across pipelines. ChangesCosmos3 Pipeline & Guardrails
🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 16
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
tensorrt_llm/_torch/visual_gen/config.py (1)
1-3:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA SPDX/copyright header.
This modified Python source file currently has no header block at the top.
As per coding guidelines:
All C++, Python, and other source files must contain NVIDIA copyright header with current modification year.🤖 Prompt for 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. In `@tensorrt_llm/_torch/visual_gen/config.py` around lines 1 - 3, This file is missing the required NVIDIA SPDX/copyright header; add the standard NVIDIA header block (including the SPDX-License-Identifier and the copyright line with the current modification year) at the very top of the file before any imports (i.e., before the existing import json / from enum import Enum / from pathlib import Path lines), preserving encoding and formatting conventions used in other Python sources in the repo.tensorrt_llm/_torch/visual_gen/modules/attention.py (1)
1-3:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA SPDX/copyright header.
This modified Python source file currently has no header block at the top.
As per coding guidelines:
All C++, Python, and other source files must contain NVIDIA copyright header with current modification year.🤖 Prompt for 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. In `@tensorrt_llm/_torch/visual_gen/modules/attention.py` around lines 1 - 3, Add the required NVIDIA copyright/SPDX header at the very top of the file tensorrt_llm/_torch/visual_gen/modules/attention.py (above the existing imports), replacing the missing header; include the current modification year, the NVIDIA copyright owner wording, and the SPDX-License-Identifier line as required by project policy so the header appears before the lines containing Enum and typing imports.tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py (1)
1-3:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA SPDX/copyright header.
This modified file is missing the header block required for Python source files.
As per coding guidelines:
All C++, Python, and other source files must contain NVIDIA copyright header with current modification year.🤖 Prompt for 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. In `@tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py` around lines 1 - 3, Add the required NVIDIA copyright/SPDX header to the top of transformer_wan.py (above the existing imports like "import math" and "from typing import Tuple"); include the current modification year, the NVIDIA Corporation copyright statement and the SPDX-License-Identifier line exactly as required by the project's header policy so the file complies with the C++/Python source file header guidelines.tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py (1)
1-2:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUpdate the copyright year on this modified file.
Line 1 still reflects 2025, but this file is modified in 2026.
As per coding guidelines:
Include NVIDIA copyright header on all new files; update year on modified files.🤖 Prompt for 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. In `@tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py` around lines 1 - 2, Update the SPDX header year in the file tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py by changing the copyright year on the top-of-file header (the lines starting with "# SPDX-FileCopyrightText" and "# SPDX-License-Identifier") from 2025 to 2026 so the modified file reflects the current year.
🤖 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 `@requirements.txt`:
- Line 92: The requirement line "retinaface @
git+https://github.com/NVShreyas/retinaface.git@main" uses a moving branch; pin
it to an immutable commit by replacing "@main" with a specific commit SHA (e.g.
"@<commit-hash>") from the retinaface repo, pick the desired stable commit via
the repo's commits page or git ls-remote, update the line in requirements.txt to
"retinaface @ git+https://github.com/NVShreyas/retinaface.git@<commit-hash>",
then regenerate/verify your lock or dependency install (pip-compile / pip
install -r requirements.txt) to ensure reproducible builds.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py`:
- Around line 1-3: This new module lacks the required NVIDIA copyright header;
add the repository’s standard NVIDIA copyright header block at the very top of
tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py before any imports so
the file complies with the policy; keep the existing import of
Cosmos3OmniMoTPipeline and the __all__ export unchanged (refer to the module
name Cosmos3OmniMoTPipeline and the file __init__.py when making the edit).
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py`:
- Around line 15-19: Module docstring incorrectly references "Wan" pipelines;
update the module-level docstring in defaults.py to reference "Cosmos3" (e.g.,
"Per-model default generation parameters for Cosmos3 pipelines.") and replace or
remove any mentions of WanPipeline and WanImageToVideoPipeline in that
docstring—if there are Cosmos3 pipeline class names (e.g., Cosmos3Pipeline /
Cosmos3ImageToVideoPipeline) use those exact names instead.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py`:
- Around line 1-13: This file is missing the required NVIDIA SPDX/copyright
header; add the repository-standard header comment block (including
SPDX-License-Identifier and NVIDIA copyright line with the correct year) at the
very top of tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py before
the existing from __future__ import, ensuring the header format matches other
files in the repo and update the year if this is a modified file.
- Around line 126-233: Optional guardrail initialization may raise other runtime
exceptions (OSError, RuntimeError, etc.) beyond ImportError/FileNotFoundError;
wrap the Qwen guardrail and the video safety filter loading in broader exception
handlers so failures disable the guardrail instead of hard-failing.
Specifically, for the Qwen block that defines qwen_tokenizer, qwen_model and
_qwen_check (used by text_guardrail), replace the bare except ImportError with
except Exception as e and log the exception detail and skip adding the checker;
likewise in build_video_guardrail wrap the SiglipModel/SiglipProcessor and
classifier loading (siglip_model, siglip_processor, classifier, and the
_safety_check closure) with except Exception as e (not just
ImportError/FileNotFoundError), log the error including e, and leave
safety_checker as None so the pipeline continues. Ensure logs include the
exception message to aid debugging.
- Around line 361-370: The output-rank restoration is inverted: record the
original input rank (e.g., orig_dim = video_tensor.dim()) before converting to
numpy, then after creating result = torch.from_numpy(frames_np) restore to the
original shape—if orig_dim == 5 then result = result.unsqueeze(0) to return a 5D
tensor, and if orig_dim == 4 leave result as 4D; update the code around v,
video_tensor, v.dim(), torch.from_numpy and video_guardrail in
check_video_safety accordingly so the unsqueeze matches the original input rank.
- Line 206: Remove the unnecessary nonlocal declarations for siglip_model and
classifier in the guardrails closure(s): locate the lines containing "nonlocal
siglip_model, classifier" and any other "nonlocal classifier" inside the nested
function scopes in guardrails.py and delete those nonlocal statements (these
variables are only read, not reassigned). Keep all other logic unchanged so the
nested functions continue to reference the outer-scope siglip_model and
classifier as read-only. Run lint/tests to confirm the F824 error is resolved.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 61-63: Normalize skip_components at the start of
load_standard_components (e.g., set to an empty list or set if None) so
membership checks like "PipelineComponent.VAE not in skip_components" are safe,
and update the logic around self.vae_scale_factor_spatial to only reference it
when the VAE is actually being loaded (i.e., when PipelineComponent.VAE is not
in skip_components and/or when the VAE has been initialized), avoiding
unconditional access that breaks skip_components=[PipelineComponent.VAE].
- Around line 439-445: Rank 0 currently checks text_guardrail (and similarly the
video-guardrail later) and returns early, but other ranks continue causing
deadlocks; change this so rank 0 computes a single boolean "blocked" (e.g.,
iterate prompts and set blocked=True if any text_guardrail returns false) and
then synchronize that decision across ranks using a torch.distributed collective
(broadcast or all_reduce) before any return; update the block that uses
self.rank, use_guardrails, text_guardrail, logger.warning, timer.mark_end,
timer.fill, PipelineOutput to have rank 0 set the blocked flag and an optional
message, broadcast them to all ranks, and have every rank perform
timer.mark_end() and return timer.fill(PipelineOutput()) when blocked; apply the
same pattern to the video-guardrail branch around the code at the later block
(lines ~575-580) so all ranks make the same early-exit decision.
- Around line 1-29: Add the standard NVIDIA copyright header at the top of the
new module pipeline_cosmos3.py before the first import; ensure the exact header
text used across the repo (including current year) is inserted as a block
comment so licensing scanners pick it up. Locate the file by the module name or
imports such as AutoencoderKLWan, Cosmos3VFMTransformer, and add the header
above those imports, keeping existing imports and code unchanged.
- Around line 90-104: Current logic gates both guardrails together; update the
block so text and video guardrails are loaded independently: check
TRTLLM_DISABLE_COSMOS3_GUARDRAILS and then separately if
PipelineComponent.TEXT_GUARDRAIL not in skip_components to initialize
self.text_guardrail (use model_config.extra_attrs['guardrail_checkpoint_dir'] or
download_guardrail_checkpoint() once) via build_text_guardrail(), and separately
if PipelineComponent.VIDEO_GUARDRAIL not in skip_components to initialize
self.video_guardrail via build_video_guardrail(); ensure the checkpoint
directory is resolved only once and reused for both builds and respect the
global TRTLLM_DISABLE_COSMOS3_GUARDRAILS flag.
- Around line 221-229: The code truncates token_ids before appending the
mandatory EOS and "<|vision_start|>" tokens, which can exceed
max_sequence_length; update the logic in the block handling token_ids and
max_sequence_length so you first ensure there is room for two suffix tokens by
truncating to max_sequence_length - 2 (or 0 min) before appending
self.tokenizer.eos_token_id and
self.tokenizer.convert_tokens_to_ids("<|vision_start|>"), then compute seq_len,
pad_len, attention_mask and pad with self.tokenizer.pad_token_id (or 0) as
before; apply this change around the token_ids manipulation to prevent exceeding
the model positional limit.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`:
- Around line 1-19: This new module transformer_cosmos3.py is missing the
required NVIDIA copyright header at the top of the file; add the standard NVIDIA
license/header block as the very first lines (before the imports such as import
math, import torch, etc.), using the canonical header used across the repo and
updating the year if needed, so files like transformer_cosmos3.py (containing
symbols like TimestepEmbedding, DiffusionModelConfig, Attention, GatedMLP,
DynamicLinearWeightLoader) conform to the repository's copyright guidelines.
- Around line 662-671: The Ulysses divisibility check currently uses a logical
AND so it only raises when both self.num_attention_heads and self.num_kv_heads
are non-divisible by ulysses_size; change the condition in the block that
references use_ulysses, self.num_attention_heads, self.num_kv_heads, and
ulysses_size to use OR so the validation raises if either head count is not
divisible by ulysses_size (keep the ValueError message but ensure it triggers
when one or the other fails).
- Around line 85-98: The fps-division branch uses fps even when effective_fps
can be None (T<=1); update the enable_fps_modulation block (symbols:
enable_fps_modulation, fps, temporal_compression_factor, base_fps,
frame_indices, t_index, temporal_offset, grid_t, grid_h, grid_w) to first guard
on fps being not None — if fps is None, fall back to the non-modulation logic
that builds integer frame indices (like the else branch) or compute using
base_fps only; ensure you avoid doing fps / temporal_compression_factor when fps
is None and preserve the same tensor shapes/expansions for t_index.
- Around line 928-948: The current seq-parallel padding logic miscomputes shards
and pads the wrong axis: when S_gen % seq_parallel_size != 0 you must compute
pad and apply it to hidden_gen and to the sequence axis of cached_freqs_gen,
then compute S_shard from the padded sequence length (not the original S_gen)
and slice both hidden_gen and the padded cos/sin using seq_parallel_rank *
S_shard : (seq_parallel_rank+1) * S_shard; update cached_freqs_gen to the padded
versions and set freqs_gen from those padded, sliced cos/sin so no tail tokens
are dropped and shapes match. Ensure you reference hidden_gen,
self.cached_freqs_gen, S_gen, pad, S_shard, self.seq_parallel_size,
self.seq_parallel_rank, and freqs_gen when making the changes.
---
Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py`:
- Around line 1-2: Update the SPDX header year in the file
tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py by changing the
copyright year on the top-of-file header (the lines starting with "#
SPDX-FileCopyrightText" and "# SPDX-License-Identifier") from 2025 to 2026 so
the modified file reflects the current year.
In `@tensorrt_llm/_torch/visual_gen/config.py`:
- Around line 1-3: This file is missing the required NVIDIA SPDX/copyright
header; add the standard NVIDIA header block (including the
SPDX-License-Identifier and the copyright line with the current modification
year) at the very top of the file before any imports (i.e., before the existing
import json / from enum import Enum / from pathlib import Path lines),
preserving encoding and formatting conventions used in other Python sources in
the repo.
In `@tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py`:
- Around line 1-3: Add the required NVIDIA copyright/SPDX header to the top of
transformer_wan.py (above the existing imports like "import math" and "from
typing import Tuple"); include the current modification year, the NVIDIA
Corporation copyright statement and the SPDX-License-Identifier line exactly as
required by the project's header policy so the file complies with the C++/Python
source file header guidelines.
In `@tensorrt_llm/_torch/visual_gen/modules/attention.py`:
- Around line 1-3: Add the required NVIDIA copyright/SPDX header at the very top
of the file tensorrt_llm/_torch/visual_gen/modules/attention.py (above the
existing imports), replacing the missing header; include the current
modification year, the NVIDIA copyright owner wording, and the
SPDX-License-Identifier line as required by project policy so the header appears
before the lines containing Enum and typing imports.
🪄 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: abfe7a97-dcbb-477f-9bc9-a2f4dc9f9a7a
📒 Files selected for processing (13)
requirements.txttensorrt_llm/_torch/visual_gen/attention_backend/vanilla.pytensorrt_llm/_torch/visual_gen/config.pytensorrt_llm/_torch/visual_gen/models/__init__.pytensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.pytensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.pytensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.pytensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.pytensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.pytensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.pytensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.pytensorrt_llm/_torch/visual_gen/modules/attention.pytensorrt_llm/_torch/visual_gen/pipeline_registry.py
ff3d46e to
f38b227
Compare
chang-l
left a comment
There was a problem hiding this comment.
Can we add some unit/integration/accuracy tests later?
|
/bot run --disable-fail-fast |
|
PR_Github #49232 [ run ] triggered by Bot. Commit: |
|
PR_Github #49232 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49715 [ run ] triggered by Bot. Commit: |
|
PR_Github #49715 [ run ] completed with state |
ca2640d to
015c08f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #50389 [ run ] triggered by Bot. Commit: |
|
PR_Github #50389 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #50395 [ run ] triggered by Bot. Commit: |
|
PR_Github #50395 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #50547 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
PR_Github #50552 [ kill ] triggered by Bot. Commit: |
|
PR_Github #50557 [ run ] completed with state
|
4e79009 to
46e618f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #50602 [ run ] triggered by Bot. Commit: |
|
PR_Github #50602 [ run ] completed with state
|
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Mista <shreyasm@nvidia.com>
Signed-off-by: Shreyas Mista <shreyasm@nvidia.com>
Signed-off-by: Shreyas Mista <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
46e618f to
06b027e
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #50655 [ run ] triggered by Bot. Commit: |
|
PR_Github #50655 [ run ] completed with state |
…class Reviewer suggested using ``enable_ulysses`` (PR NVIDIA#14012) to simplify our PR's hand-rolled v2a Ulysses wrapper. Two surgical changes that together delete the manual wrapper construction + the parallel ``use_ulysses_cross`` ctor flag. 1. Base ``Attention.__init__`` wrap dispatch (modules/attention.py): The combined gate ``(ring_size > 1 or ulysses_size > 1) and qkv_mode != SEPARATE_QKV`` excluded cross-attn from both Ring and Ulysses wrapping. Split it into two single-purpose gates: * Ring stays gated on ``qkv_mode != SEPARATE_QKV`` — Ring shards the seq dim and requires ``S_q == S_kv``. * Ulysses now keys on ``use_ulysses`` only — Ulysses shards the head dim and is value-preserving for ``S_q != S_kv``, so cross-attn is valid. Also fixes a latent bug where ``enable_ulysses=False`` would still build the wrapper while making the inner backend use the unsharded head count. 2. LTX2Attention (models/ltx2/transformer_ltx2.py): * Drop the ``use_ulysses_cross`` ctor flag. Caller passes a single ``use_ulysses=True`` for both self-attn and cross-attn modules. Internally, cross-attn carries an additional ``cp_size == 1`` gate (Ring/Attention2D + cross-attn Ulysses is unvalidated). * Delete the 30-line ``_ulysses_cross_attn = UlyssesAttention(...)`` hand wrap that bypassed the base class. With Step 1 the base class now builds the wrap for cross-attn too. * Collapse the cross-attn dual-attn pair into the existing self-attn dual-attn path (single ``_ulysses_attn`` / ``_plain_attn`` pair, single ``_has_dual_attn`` flag, single branch in ``set_ulysses_active`` / ``is_ulysses_active``). * Rename the locally derived flag from ``wants_ulysses`` to ``enable_ulysses`` so it matches the kwarg passed to the base class. Net: 94→58 lines in transformer_ltx2.py ctor + 30 fewer lines in toggle methods. ``use_ulysses_cross`` call site at ``video_to_audio_attn = LTX2Attention(..., use_ulysses=True)``. Verified: all 4 PR-touched test files PASS (26/26 in 6m57s) — e2e LTXModel parity (VANILLA + FA4, no_pad + audio_pad), multi-rank UlyssesAttention parity, FA4/Vanilla key_padding_mask single-rank. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
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)
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.Summary by CodeRabbit
Release Notes