[None][feat] Add Cosmos3-Edge (Nemotron-dense) support

[ishovkun]

Dev Engineer Review

• Added nvidia/Cosmos3-Edge support for text-to-image, text-to-video, and image-to-video at native 480p.
• Added Nemotron-dense architecture recipes, NemotronRMSNorm, MLP variants, RoPE resolution, generator-only key normalization, and strict checkpoint coverage checks.
• Added family- and mode-specific defaults, native-flow scheduling, tokenizer loading, prompt-conditioning fixes, scheduler caching, and I2V device/dtype validation.
• Registered Edge in documentation, CLI help, and Cosmos3 pipeline configuration.
• Preserved Nano, Super, distilled, and audio-capable checkpoint behavior through compatibility handling and regression tests.
• Corrected the generator normalization key to use_und_k_norm_for_gen.
• Added the Cosmos3 video negative-prompt specification.
• Updated _merge_defaults and VisualGen.default_params to preserve caller-versus-pipeline default state across serialization.
• Configuration and test-list updates are consistent with the supported Edge modes. The B200 list adds one unit test, one I2V integration test, and three LPIPS tests.
• Edge action generation, audio, V2V, and transfer remain unsupported.
• Review follow-up should verify CI results for the latest pipeline because one reported run was force-killed. Earlier successful runs validated unit, distributed, checkpoint-loading, tensor-parallel, and Ulysses coverage.

QA Engineer Review

Test functions and coverage added or modified:

• tests/unittest/_torch/visual_gen/test_cosmos3_edge.py
  • Added architecture, RoPE, normalization, transformer, generator K-norm, native-flow, strict-loading, defaults, sampling, scheduler-cache, envelope, parity, tokenizer, pipeline, and checkpoint tests.
  • Covered by l0_b200.yml through the Cosmos3 Edge unit entry.
• tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py
  • Added standalone Diffusers/TRT-LLM per-step parity validation.
  • Not directly listed in l0_b200.yml; verify invocation through the Edge unit test.
• tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py
  • Added default-provenance, mode-switching, serialization, system-prompt, and conditioning-anchor tests.
  • Not directly listed in the modified test list.
• tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py
  • Added prompt metadata, negative-prompt, stale-field, JSON, aspect-ratio, and V2V scheduler tests.
  • Not directly listed in the modified test list.
• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py
  • Added optional rope_scaling and has_action_weights regression coverage.
  • Not directly listed in the modified test list.
• tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py
  • Added Edge TP2 and Ulysses2 parity tests.
  • Not directly listed in the modified test list.
• tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py
  • Added Edge I2V example coverage and T2V, I2V, and T2I LPIPS golden tests.
  • Covered by l0_b200.yml through the Edge I2V and three LPIPS entries.
• Added three Edge LPIPS golden reference files.
• Added Edge coverage to tests/integration/test_lists/test-db/l0_b200.yml.

Verdict: needs follow-up. The primary Edge unit and integration tests are listed for CI. Several modified unit and distributed test modules are not explicitly represented in the changed test list, and CBTS coverage data is unavailable.

Description

Adds inference support for nvidia/Cosmos3-Edge, the 4B Cosmos3 variant whose
reasoner and generator towers use the Nemotron-dense architecture instead of
Nano/Super's Qwen3. Supported tasks: T2I, T2V, I2V (480p-native).
Builds on the sampling-policy and checkpoint-declared-defaults plumbing
introduced by #16690 (distilled I2V-4Step), now merged into main.
Reference implementations: diffusers Cosmos3OmniTransformer/
Cosmos3OmniPipeline (huggingface/diffusers#14181 + #14246) for the
transformer, and cosmos-framework's PyTorch backend for the sampling schedule.

Architecture recipes (the core change). backbone_type in the transformer
config selects a complete, validated architecture recipe rather than
individual flags steering construction — a mixed config cannot half-apply and
fails at startup with the offending key named. The Nemotron-dense recipe:
non-gated squared-ReLU MLP (reusing _torch/modules/mlp.py::MLP with the
existing relu2 activation, as modeling_nemotron.py does — no Edge MLP
class), no reasoner Q/K norm, and a new per-layer k_norm_und_for_gen
applied to the reasoner's keys only where the generator consumes them: in the
cached-KV design, forward_with_kv attends with rope(k_raw) internally but
caches rope(k_norm_und_for_gen(k_raw)) for the generator layers. This
distinction is subtle enough that both public references shipped it wrong
initially (diffusers normed the reasoner pathway too, fixed in #14246;
vllm-omni read a transposed config key and skipped the norm entirely, fixed
in vllm-project/vllm-omni#5239) — it is pinned here by a dedicated
regression test. Edge norms use the Nemotron flavor (weight multiply in fp32
before downcast); F.rms_norm is bit-exact to that flavor, so
NemotronRMSNorm is a one-line fused wrapper, and the attention Q/K norms of
both recipes now route through norm modules instead of a hardcoded
functional — byte-identical for Nano by construction (pinned by test).

Native flow schedule. Edge's model_index.json declares
use_native_flow_schedule: true; the intended schedule (per the model card
and cosmos-framework's PyTorch backend) is linear flow sigmas with a runtime
shift of 3.0. Stock diffusers does not reproduce it — its UniPC karras branch
discards the provided sigmas (the checkpoint config ships
use_karras_sigmas: true), which is why the schedule is validated against
cosmos-framework directly. The implementation configures diffusers' own
Apache-2.0 UniPCMultistepScheduler (karras off, flow_shift from the
per-mode table, explicit linspace sigmas): timesteps are bit-identical to
cosmos-framework's FlowUniPCMultistepScheduler and full multi-step step()
trajectories agree to ≤1.6e-7 relative across three (shift, steps) fixtures
recorded in the tests — no OpenMDW code is copied.

A scheduler's identity is its resolved (flow_shift, use_karras_sigmas)
pair, not the request mode, so forward() resolves that pair — mode table,
then V2V's stronger shift, then a caller-supplied flow_shift — and takes
the instance from a cache keyed on it. Each configuration is constructed at
most once and never rebuilt; a mode that is never served never builds one,
and modes that resolve to the same pair share an instance (Edge's video and
image tables both declare shift 3.0, so Edge builds one scheduler, not two).
Streams key separately so audio gets its own instance at the same
configuration — schedulers mutate state on every step(), and the two
streams denoise in lockstep. This replaces the per-request rebuild that
#16155 introduced for V2V, which reconstructed a scheduler on every mode
alternation.

Strict weight loading (both recipes benefit). load_weights previously
collected skipped modules and dropped the list on return — a missing tensor
silently kept its init values. Coverage is now default-fail in both
directions at parameter granularity: any parameter of a constructed module
(root parameters included) without a checkpoint tensor raises with names, a
partially covered fused QKV raises, and mapped tensors that land on no
constructed parameter warn with names. The explicit skip list (lm_head, the
reasoner's final norm — text-output path VisualGen never computes — and the
action heads) is logged, never silent.

Pipeline. The tokenizer loads via AutoTokenizer (Edge ships a Nemotron
PreTrainedTokenizerFast; eos = pad = 11, <|vision_start|> = 20).
Generation defaults become a (family, mode) table — family resolved from
the same recipe function, mode from the request's output type, never from the
checkpoint name — with Edge's model-card values (video 832×480 × 121 frames,
50 steps, guidance 5.0, shift 3.0; T2I 640×640, guidance 4.0). forward()'s
numeric parameters are now None-defaulted and resolve through the same
tables, so direct callers get checkpoint-appropriate values (this also lets
the sampling policy fill fixed distilled steps/guidance for direct callers).
Requests outside the model-card envelope (256p/480p, 50–150 frames, 12–30
fps) log one advisory line and proceed — the reference runtime accepts a
wider range, so the envelope is documented support, not enforced validation.
The checkpoint's action weights are skipped loudly and the init log says
action generation is not supported; enable_safety_checker: false in the
model_index is deliberately not honored (guardrail policy unchanged).

Verified on 1×B200: per-step transformer velocity parity vs diffusers main is
0.7–1.6% relative (bf16, both CFG branches, masked I2V path included); E2E
LPIPS vs reference goldens: T2V 0.045, I2V 0.078, T2I 0.006.

Out of scope for Edge: action generation (lands on the in-flight cosmos3
action work), audio (the checkpoint has no audio tower — requesting it
raises), and V2V/transfer — not an Edge capability, since the generator's
inputs are text, image, and action trajectories only. V2V itself is
supported for Nano/Super via #16155; this PR merges that work and keeps its
request path intact, it just does not extend it to Edge.

Reference-parity fixes for prompt conditioning (follow-up commits). Three
divergences from cosmos-framework in how prompt text is built, found while
investigating an artifact report on Edge. Each is verified by comparing the text
reaching the tokenizer against the reference on the real assets, and each is
independently corroborated by vllm-omni, which matches cosmos-framework on every
point.

1. The negative prompt took the JSON field-injection branch. forward() routed it
   through _format_prompt_with_metadata, which for a JSON prompt injects
   duration/fps/resolution/aspect_ratio as object fields. cosmos-framework
   applies its plain-text formatter to the negative unconditionally
   (inference.py:541) and reserves field injection for the positive; vllm-omni does
   the same (pipeline_cosmos3.py:2095). 125 of ~2900 tokens differed, all at the tail
   — immediately before the handoff to the video tokens. Now byte-identical at 14936
   chars / 2910 tokens.

2. JSON metadata values. duration "5.0s" → "5s", fps 24 → 24.0,
   resolution {"W","H"} → {"H","W"}, aspect_ratio from the reduced ratio
   "15,26" to the nearest W,H bucket "16,9", and json.dumps back to the default
   ensure_ascii=True. Stills now pop duration/fps instead of merely skipping them.
   Two of these were not cosmetic: the checkpoint's own example_i2v_prompt.json
   already carries "aspect_ratio": "16,9", which the reduced-ratio computation
   overwrote with a string outside the bucket vocabulary the model saw in training; and
   a T2I request kept the source prompt's stale "duration": "7s".
   _aspect_ratio_bucket round-trips all 25 entries of cosmos-framework's
   VIDEO_RES_SIZE_INFO.

3. Default negative prompt for video modes. The pipeline defaulted every mode to
   empty, matching diffusers and vllm-omni but not cosmos-framework, whose per-mode
   defaults wire neg_prompts.json into text2video, image2video, video2video and
   audio_image2video while leaving text2image/image2image unset. Video modes now default
   to it, carried as a Python literal in negative_prompt.py so it ships in every
   install mode without a package_data entry. Image modes keep the empty default — the
   prompt is almost entirely temporal vocabulary a still cannot exhibit, and the model
   card advertises nothing for image generation.

Blast radius. (3) changes default output for every Cosmos3 video request that
does not pass a negative prompt — Nano and Super included, not only Edge. The
LPIPS goldens were generated against an empty uncond branch, so their three helpers
now pin negative_prompt="" explicitly rather than inheriting a default they never
recorded. .pre-commit-config.yaml gains one codespell word (LOD); the vendored
text cannot be edited without breaking parity.

These are parity fixes, not artifact fixes. They make the conditioning match the
reference; they do not reduce generation artifacts. Measured across 5 seeds on the
model-card config, pre- vs post-fix moves the artifact rate from 11.82x to 10.87x
median — improved or tied on every seed, but p=0.71, i.e. indistinguishable.

Test Coverage

• Unit (tests/unittest/_torch/visual_gen/test_cosmos3_edge.py, 78 tests, no
  checkpoint required for the unit tier): recipe selection/validation with
  rejection cases (unknown backbone_type, contradicting flags, latent-geometry
  invariants, inconsistent patch_latent_dim); Nemotron RMSNorm bit-exactness
  vs the Qwen flavor; the generator-only K-norm regression (perturbing
  k_norm_und_for_gen leaves reasoner attention bit-identical) plus a
  numeric identity check; Nano Q/K byte-identical pin; native-flow schedule
  parity against three recorded cosmos-framework fixtures (bit-equal
  timesteps, full step() trajectories to 1e-9); strict-loading coverage in
  both directions (missing generic/architecture-specific/root tensors raise,
  partial fused raises, skip list and unconsumed-tensor warnings asserted by
  content); per-family defaults/warmup/advisory resolution and the
  executor-defaults shape.
• Unit, checkpoint-gated: tokenizer specials + chat template; cond token ids
  vs a diffusers-main golden and the uncond CFG branch vs a recorded
  self-golden (keep-metadata semantics documented); model-index detection;
  recipe/scheduler wiring (both mode schedulers at shift 3.0, karras off);
  full 549-tensor load + forward; direct forward(None) resolution reaching
  the generation path with Edge values; T2V/I2V/T2I sanity generations with
  shape and non-black checks.
• Parity (env-gated on DIFFUSERS_MAIN_PATH, subprocess because the pinned
  diffusers predates the Edge classes): per-step velocity comparison vs
  diffusers main on the real checkpoint, threshold 5% (observed 0.7–1.6%).
• Integration smoke (test_cosmos3_edge_i2v_example, l0_b200.yml): the
  documented example invocation at the deployed 480p × 121-frame, 50-step
  shape; asserts a non-empty MP4.
• Integration quality gates (test_cosmos3_edge_{t2v,i2v,t2i}_lpips_against_golden,
  l0_b200.yml): LPIPS vs goldens produced by the reference implementation
  (diffusers main with the scheduler patched to the native flow schedule —
  equivalent to cosmos-framework at fp32-ulp), with matched initial noise and
  matched cond/uncond prompt texts, so the gates check the denoising
  trajectory against the reference rather than regression against a past
  TRT-LLM run. Thresholds 0.10 / 0.13 / 0.05 (measured 0.0447 / 0.0778 /
  0.0056 at creation; a wrong-seed I2V run measures 0.858). The I2V gate runs
  10 steps — cross-stack drift accumulates per step and the deployed 50-step
  shape is covered by the example test. Full provenance in
  golden/visual_gen_lpips/cosmos3_edge_*.json.
• Pre-existing cosmos3 suite (115 tests) passes unchanged; Nano behavior is
  pinned byte-identical through the norm-routing and loader refactors.
• Prompt-conditioning parity (test_cosmos3_pipeline.py): the negative prompt keeps
  its serialized JSON and gains sentences rather than injected fields, checked
  byte-for-byte against the reference algorithm; the positive branch still injects,
  so the two paths stay deliberately distinct; aspect-ratio bucket mapping over seven
  resolutions; non-ASCII escaping; stills dropping stale duration/fps; and the
  video/image split of the default negative prompt including its serialization and
  cache identity.
• Note for reviewers: Cosmos3-Edge (8.7 GB) must be staged in the CI
  llm-models/ storage, or the checkpoint-gated tests skip.

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.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

Cosmos3-Edge support adds Nemotron-dense architecture handling, checkpoint-aware generation defaults, native-flow sampling, prompt and scheduler updates, CLI and documentation entries, and unit, parity, distributed, and LPIPS integration coverage.

Changes

Cosmos3 visual generation support

Layer / File(s)	Summary
Architecture recipes and checkpoint loading tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py, tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py	Adds Qwen3 and Nemotron-dense architecture recipes, RoPE resolution, recipe-specific normalization and MLPs, generation-key normalization, temporal-compression handling, and strict checkpoint coverage validation.
Checkpoint-aware sampling and pipeline execution tensorrt_llm/_torch/visual_gen/models/cosmos3/*, tensorrt_llm/_torch/visual_gen/executor.py, tensorrt_llm/visual_gen/visual_gen.py, tests/unittest/_torch/visual_gen/test_cosmos3_*.py	Adds Edge defaults and envelopes, native-flow schedules, family and mode resolution, cached video and audio schedulers, negative-prompt handling, metadata formatting, default provenance tracking, and I2V latent validation.
Unit, parity, and distributed validation tests/unittest/_torch/visual_gen/*, tests/unittest/_torch/visual_gen/multi_gpu/*	Adds Edge architecture, scheduler, checkpoint, tokenizer, pipeline, Diffusers parity, and distributed TP/Ulysses validation.
Examples and public model documentation docs/source/models/*, examples/visual_gen/models/cosmos3/*, .pre-commit-config.yaml	Documents Cosmos3-Edge tasks, architecture, defaults, validation limits, CLI selection, and an I2V example.
Integration and LPIPS coverage tests/integration/defs/examples/visual_gen/*, tests/integration/test_lists/test-db/l0_b200.yml	Adds Edge T2I, T2V, and I2V example tests, golden metadata, LPIPS comparisons, and B200 test-list entries.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#16155: Overlaps in Cosmos3 pipeline, defaults, transformer, and V2V/action-generation support.
• NVIDIA/TensorRT-LLM#16394: Overlaps in Cosmos3 pipeline, defaults, and transformer code.
• NVIDIA/TensorRT-LLM#17325: Overlaps in Cosmos3 defaults, pipeline, transformer, and action-generation tests.

Suggested reviewers: bowenfu, arysef

Sequence Diagram(s)

sequenceDiagram
  participant User
  participant VisualGen
  participant Cosmos3OmniMoTPipeline
  participant Cosmos3VFMTransformer
  participant LPIPSTest
  User->>VisualGen: select nvidia/Cosmos3-Edge and output mode
  VisualGen->>Cosmos3OmniMoTPipeline: merge mode-aware defaults
  Cosmos3OmniMoTPipeline->>Cosmos3VFMTransformer: denoise Edge latents
  Cosmos3OmniMoTPipeline-->>User: return image or video
  LPIPSTest->>Cosmos3OmniMoTPipeline: run deterministic generation
  LPIPSTest->>LPIPSTest: compare output with golden reference

Loading

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The PR includes broad Edge architecture, pipeline, documentation, defaults, and integration changes beyond issue #5239's configuration-key rename.	Split general Cosmos3-Edge support into a separate linked issue or PR, and retain only the configuration-key fix here.
Docstring Coverage	⚠️ Warning	Docstring coverage is 35.44% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: adding Cosmos3-Edge Nemotron-dense support.
Description check	✅ Passed	The description is detailed, follows the required sections, and documents implementation scope and test coverage.
Linked Issues check	✅ Passed	The PR addresses issue #5239 by implementing and testing the corrected generator key-normalization configuration behavior.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (2)

tests/integration/defs/examples/visual_gen/test_visual_gen.py (1)

2116-2141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Edge LPIPS gates lack on-failure candidate preservation. The three Edge gates assert the threshold but never call _preserve_lpips_candidate_on_failure, so a cross-host failure leaves no archived media to refresh the golden — yet the Edge goldens' threshold_rationale explicitly banks on that helper for the ~0.04 cross-host headroom. The distilled I2V gate (Lines 2299-2305) already wires this up; mirror it here (each gate takes request and passes its generated path + golden basename).

• tests/integration/defs/examples/visual_gen/test_visual_gen.py#L2116-L2141: add request to the signature and call _preserve_lpips_candidate_on_failure before _assert_lpips_below_threshold with generated_path / cosmos3_edge_t2v_lpips_golden_video.mp4.
• tests/integration/defs/examples/visual_gen/test_visual_gen.py#L2144-L2172: same, with cosmos3_edge_i2v_lpips_golden_video.mp4.
• tests/integration/defs/examples/visual_gen/test_visual_gen.py#L2175-L2201: same, with cosmos3_edge_t2i_lpips_golden.png (also add the _visual_gen_deps/request fixtures already noted separately).

🤖 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 `@tests/integration/defs/examples/visual_gen/test_visual_gen.py` around lines
2116 - 2141, Update the three Edge LPIPS gates in
tests/integration/defs/examples/visual_gen/test_visual_gen.py at lines
2116-2141, 2144-2172, and 2175-2201: add the request fixture to each test
signature, call _preserve_lpips_candidate_on_failure with the generated path and
the corresponding golden basename before _assert_lpips_below_threshold, and add
the _visual_gen_deps fixture to the T2I gate as needed. Use
cosmos3_edge_t2v_lpips_golden_video.mp4,
cosmos3_edge_i2v_lpips_golden_video.mp4, and cosmos3_edge_t2i_lpips_golden.png
respectively.

.gitignore (1)

125-126: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

.plans/ looks unrelated to this PR's scope.

This ignore entry isn't tied to Cosmos3 Edge/distilled support and reads like a local tooling artifact. Consider dropping it here and adding it in a dedicated change (or your personal global gitignore) to keep the PR focused.

As per coding guidelines: "Keep each pull request focused on one concern and avoid scope creep; split unrelated changes into separate pull requests."

🤖 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 @.gitignore around lines 125 - 126, Remove the unrelated .plans/ entry from
.gitignore to keep this change focused on Cosmos3 Edge/distilled support; do not
replace it with another ignore rule.

Source: Coding guidelines

🤖 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.

Nitpick comments:
In @.gitignore:
- Around line 125-126: Remove the unrelated .plans/ entry from .gitignore to
keep this change focused on Cosmos3 Edge/distilled support; do not replace it
with another ignore rule.

In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py`:
- Around line 2116-2141: Update the three Edge LPIPS gates in
tests/integration/defs/examples/visual_gen/test_visual_gen.py at lines
2116-2141, 2144-2172, and 2175-2201: add the request fixture to each test
signature, call _preserve_lpips_candidate_on_failure with the generated path and
the corresponding golden basename before _assert_lpips_below_threshold, and add
the _visual_gen_deps fixture to the T2I gate as needed. Use
cosmos3_edge_t2v_lpips_golden_video.mp4,
cosmos3_edge_i2v_lpips_golden_video.mp4, and cosmos3_edge_t2i_lpips_golden.png
respectively.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 82c2a32f-30dc-4081-8a65-f056cb82ae4c

📥 Commits

Reviewing files that changed from the base of the PR and between 0a401b4 and 0165d4f.

⛔ Files ignored due to path filters (1)

• tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip is excluded by !**/*.zip

📒 Files selected for processing (24)

• .gitignore
• docs/source/models/supported-models.md
• docs/source/models/visual-generation.md
• examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml
• examples/visual_gen/models/cosmos3/README.md
• examples/visual_gen/models/cosmos3/cosmos3.py
• requirements.txt
• tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py
• tensorrt_llm/_torch/visual_gen/pipeline.py
• tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json
• tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json
• tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json
• tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json
• tests/integration/defs/examples/visual_gen/test_visual_gen.py
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/visual_gen/conftest.py
• tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py
• tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py
• tests/unittest/_torch/visual_gen/test_cosmos3_edge.py
• tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py
• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

[fredricz-20070104]

Review summary - CONCERNS

Verdict: No blocking code defect found, but this cannot merge yet — the branch is dirty (conflicts with main) and the strict weight-loading change is high-blast-radius and warrants QA on existing checkpoints before merge.

Concerns

1. [MAJOR] PR is not mergeable (mergeable_state=dirty)

   • What is wrong: the branch conflicts with base main.
   • How it fails: a merge now either fails or forces an unreviewed manual conflict resolution.
   • Suggested fix: rebase onto current main, resolve conflicts, re-run the l0_b200 visual_gen suite, and confirm the state goes clean.

2. [MAJOR] tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py:~1523 - load_weights is now default-fail for every Cosmos3 checkpoint

   • What is wrong: load_weights now raises when any constructed parameter has no checkpoint tensor, replacing the previous behaviour that silently kept init values. This path is shared by Nano, Super, audio-capable and distilled variants, not just Edge.
   • How it fails: an existing Super/audio checkpoint that previously loaded because a benign tensor was silently skipped (name drift, an optional/extra tensor, or a module the recipe no longer builds) will now hard-fail at load. In this PR only Nano is clearly exercised end-to-end (LPIPS + test_cosmos3_example); Super and audio-capable real checkpoints are covered only via reduced synthetic state dicts.
   • Suggested fix: load each currently-supported real checkpoint (Nano, Super, audio, both distilled) through the new strict loader in CI/QA before merge, or scope the strict-fail to the new recipe until the older variants are confirmed clean.

Minor notes (non-blocking)

• tests/integration/defs/examples/visual_gen/test_visual_gen.py:~2116 - the three Edge LPIPS gates don't call _preserve_lpips_candidate_on_failure; mirror the distilled I2V gate so cross-host failures archive a candidate to refresh the golden.
• tests/integration/defs/examples/visual_gen/test_visual_gen.py:~2175 - test_cosmos3_edge_t2i_lpips_against_golden omits the _visual_gen_deps fixture its siblings take.

QA view

• Test coverage: adequate for the new Edge/distilled paths (recipe validation, Nemotron norm bit-exactness, generator-only K-norm regression, native-flow schedule parity, strict-loading positive/negative cases, mode resolution, denoise-loop contract). Uncovered: strict load of real Super/audio checkpoints; native-flow validated against recorded fixtures rather than a live cross-run.
• SM coverage: architecture-independent (bf16, VANILLA attention, RMSNorm/MLP; no arch guards in the diff). Tests run on B200 (sm100). Numeric goldens captured on B200; other-SM parity relies on documented cross-host LPIPS headroom.
• Test code: Edge LPIPS gates lack candidate preservation; one gate lacks _visual_gen_deps; the diffusers-parity subprocess script is env/checkpoint-gated and silently skips when unavailable.
• Test time: significant - four new B200 LPIPS gates + three example smoke tests added to l0_b200.yml (TIMEOUT 30/45/30 and 20/20/15/15), a refreshed golden media zip, and ~2200 lines of new unit tests.
• Needs /qa-verify: yes - resolve the dirty merge and confirm existing Super/audio checkpoints still load under the now strict load_weights.

Possible new issues

• Strict load_weights could crash an existing Super/audio checkpoint that previously loaded with a silently-skipped tensor.
• forward() now writes the I2V conditioning frame in place without the prior .clone() / .to(device,dtype); correct only if image_latent already matches the loop output's device/dtype.

What I could not verify

• Whether every real Super/audio/distilled checkpoint's full parameter set is covered by the remap so the strict loader does not false-positive (callers and real checkpoint layouts are not in the diff).
• Runtime LPIPS/parity numbers and whether the goldens reproduce on the CI B200 host within the stated thresholds.

Automated review by NVCortex Lite, run by @fredricz-20070104.

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65129 [ run ] triggered by Bot. Commit: 9d3553a Link to invocation

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65129 [ run ] completed with state FAILURE. Commit: 9d3553a
/LLM/main/L0_MergeRequest_PR pipeline #52924 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

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65229 [ run ] triggered by Bot. Commit: 9d3553a Link to invocation

[tensorrt-cicd]

PR_Github #65229 [ run ] completed with state FAILURE. Commit: 9d3553a
/LLM/main/L0_MergeRequest_PR pipeline #53012 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

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65246 [ run ] triggered by Bot. Commit: 9d3553a Link to invocation

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65318 [ run ] triggered by Bot. Commit: 9d3553a Link to invocation

[tensorrt-cicd]

PR_Github #65246 [ run ] completed with state ABORTED. Commit: 9d3553a

Link to invocation

[tensorrt-cicd]

PR_Github #65318 [ run ] completed with state SUCCESS. Commit: 9d3553a
/LLM/main/L0_MergeRequest_PR pipeline #53093 completed with status: 'SUCCESS'

CI Report

Link to invocation

[github-actions]

✅ LFS objects already in storage (1 file) — no sync needed.

These LFS-tracked files are already present in this repository's LFS storage:

• tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip