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

[ishovkun]

Dev Engineer Review

• Adds nvidia/Cosmos3-Edge support for text-to-image, text-to-video, and image-to-video at native 480p.
• Adds Nemotron-dense architecture recipes, normalization, RoPE resolution, generator-only K-norm handling, strict checkpoint coverage, tokenizer loading, and native flow scheduling.
• Adds checkpoint-family and mode-specific defaults, prompt metadata handling, negative prompts, scheduler caching, and envelope advisories.
• Corrects the generator normalization key to use_und_k_norm_for_gen.
• Preserves existing Cosmos3 behavior through regression coverage.
• Documents supported and unsupported Edge capabilities.
• Uses consistent model identifiers and valid test-list entries.
• The stricter checkpoint loader and mode-dependent defaults require continued checkpoint validation.

QA Engineer Review

• Adds unit coverage for architecture recipes, RoPE settings, Nemotron normalization, transformer structure, generator-only K-norm behavior, native flow scheduling, strict loading, defaults, sampling validation, envelope warnings, tokenizer behavior, and Diffusers parity.
• Adds CUDA-gated Edge checkpoint tests for text-to-image, text-to-video, and image-to-video.
• Adds TP and Ulysses parity tests for the Edge backbone.
• Adds pipeline regression tests for metadata, negative prompts, scheduler behavior, default propagation, and conditioning-anchor validation.
• Adds the cosmos3_edge_i2v_example test and Edge text-to-image, text-to-video, and image-to-video LPIPS tests.
• Adds golden LPIPS references for all three Edge modes.
• Covers the Edge integration tests in tests/integration/test_lists/test-db/l0_b200.yml.
• The standalone unit, parity, and multi-GPU tests are not individually listed in the supplied test-list change.
• Verdict: sufficient.

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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f6b9904a-a02e-4f11-ab08-70c46d053fb2

📥 Commits

Reviewing files that changed from the base of the PR and between 82510a5 and fe85ad8.

📒 Files selected for processing (2)

• tests/unittest/_torch/visual_gen/test_cosmos3_edge.py
• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

---

Walkthrough

Cosmos3 Edge support adds Nemotron-dense architecture handling, checkpoint-aware generation defaults and scheduling, prompt processing, CLI and documentation updates, and unit, parity, integration, and LPIPS validation for T2I, T2V, and I2V modes.

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, MRoPE validation, recipe-specific normalization and MLPs, generation-key handling, temporal compression settings, and strict checkpoint coverage checks.
Checkpoint-aware sampling and pipeline execution tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py, tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py, tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py, tensorrt_llm/_torch/visual_gen/models/cosmos3/negative_prompt.py, tensorrt_llm/_torch/visual_gen/executor.py, tensorrt_llm/visual_gen/visual_gen.py, tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py, tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py	Adds Edge generation tables and envelopes, native flow scheduling, cached video and audio schedulers, mode-aware parameter resolution, negative prompts, metadata formatting, conditioning checks, and caller-default tracking.
Unit, parity, and distributed validation tests/unittest/_torch/visual_gen/test_cosmos3_edge.py, tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py, tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py	Adds architecture, scheduler, loading, tokenizer, pipeline, Diffusers parity, and multi-GPU validation.
Examples and public model documentation examples/visual_gen/models/cosmos3/*, docs/source/models/*, .pre-commit-config.yaml	Documents Cosmos3 Edge capabilities and usage, updates CLI model help, and adds LOD to the codespell ignore list.
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 golden references, end-to-end generation tests, LPIPS comparisons, failure artifacts, and B200 test entries.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#16155: The PR directly builds on Cosmos3 V2V changes in the same pipeline, defaults, and test areas.
• NVIDIA/TensorRT-LLM#17325: The PR adds action-generation functionality to the same Cosmos3 pipeline, defaults, transformer, and test components.

Suggested reviewers: bowenfu, qijune, arysef

🚥 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, scheduling, pipeline, documentation, and integration changes beyond the naming fix requested by [#5239].	Split the broader Cosmos3-Edge feature work into separate linked issues, or link issues that explicitly cover the additional architecture, pipeline, documentation, and test changes.
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 Cosmos3-Edge feature and follows the repository's required [ticket][type] format.
Description check	✅ Passed	The description explains the implementation, scope, tests, limitations, and checklist status in the required sections.
Linked Issues check	✅ Passed	The PR corrects the generator key naming and adds tests validating the use_und_k_norm_for_gen behavior required by [#5239].

✨ 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

[tburt-nv]

This is a lot of test time, can we move some if it to post-merge?

[yibinl-nvidia]

+1, 1 cosmos3 example test is enough for pre-merge test.

[zhaoyangwang-nvidia]

The new paragraph says None means "the mode's default", but this line — unchanged — still promises default_params lets you inspect what the model will use. After this PR Cosmos3Pipeline.default_generation_params forces height/width/num_inference_steps/guidance_scale to None unconditionally, so that no longer holds — and note this hits existing Cosmos3-Nano/Super users too, not just Edge: they previously got concrete 720p values here.

[ishovkun]

We don't want to frustrate the Nano/Super users :-)
I reverted the docstring to the original state from before my nested commits in 891c529.

Nano/Super users get their 720p now. The key was to check whether the user actually set the value.

Note that the default value is slightly misleading as in Cosmos3 the defaults are mode dependent, and, therefore, the users are getting the defaults for the video mode. That's a tricky model 😄 .

[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

[tensorrt-cicd]

PR_Github #63592 [ run ] triggered by Bot. Commit: 28c430b Link to invocation

[tensorrt-cicd]

PR_Github #63592 [ run ] completed with state FAILURE. Commit: 28c430b

Link to invocation

[ishovkun]

/bot run

[tensorrt-cicd]

PR_Github #63609 [ run ] triggered by Bot. Commit: ce42d30 Link to invocation

[tensorrt-cicd]

PR_Github #63609 [ run ] completed with state FAILURE. Commit: ce42d30
/LLM/main/L0_MergeRequest_PR pipeline #51570 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 #63627 [ run ] triggered by Bot. Commit: ce42d30 Link to invocation

[tensorrt-cicd]

PR_Github #63627 [ run ] completed with state FAILURE. Commit: ce42d30
/LLM/main/L0_MergeRequest_PR pipeline #51582 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 #63685 [ run ] triggered by Bot. Commit: ce42d30 Link to invocation

[tensorrt-cicd]

PR_Github #63685 [ run ] completed with state SUCCESS. Commit: ce42d30
/LLM/main/L0_MergeRequest_PR pipeline #51637 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (17)

tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py (1)

410-415: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use an explicit optional configuration type for both parameters.

pretrained_dict: dict = None appears in _forward and _build_ref_and_parallel. Define a precise configuration type and annotate both parameters as ConfigType | None to remove the implicit Optional and unparameterized dict.

Test coverage: nine tests cover TP, Ulysses, CFG, audio, and Attention2D. The file is listed in tests/integration/test_lists/test-db/l0_dgx_b200.yml; coverage is sufficient.

🤖 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/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py`
around lines 410 - 415, Define or reuse the precise ConfigType configuration
type, then update the pretrained_dict parameters in both _forward and
_build_ref_and_parallel to use ConfigType | None with an appropriate default.
Remove the unparameterized dict and implicit Optional annotations while
preserving existing behavior.

Sources: Coding guidelines, Learnings, Linters/SAST tools

tests/unittest/_torch/visual_gen/test_cosmos3_edge.py (3)

1146-1353: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add negative coverage for the modes Edge does not support.

The PR states that action generation, audio, V2V, and transfer stay unsupported for Edge. This class covers only the supported paths. No test asserts that an Edge pipeline rejects an unsupported request with a clear error. A future change that silently accepts one of these requests would pass this suite.

Add one parametrized test that calls forward for each unsupported mode and asserts the documented exception. Keep it in the checkpoint-gated class, or use _bare_pipeline if the rejection happens before weights are needed.

I can draft the parametrized negative test if you confirm the exception type and the request fields for each unsupported mode.

🤖 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/unittest/_torch/visual_gen/test_cosmos3_edge.py` around lines 1146 -
1353, Add a parametrized negative test to TestEdgeCheckpoint covering action
generation, audio, V2V, and transfer requests. For each unsupported mode, call
edge_pipeline.forward with the corresponding request field and assert the
documented exception type and clear error message; use _bare_pipeline if
validation occurs before model loading. Keep the test checkpoint-gated and
aligned with the existing request-field conventions.

---

260-321: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Move the recipe-validation tests into TestArchRecipe.

test_qwen3_flag_contradiction_raises (line 311) and test_edge_latent_geometry_invariants (line 316) exercise resolve_arch_recipe, not resolve_rope_axes_dim. TestArchRecipe already groups that behavior. Moving them keeps the coverage map readable.

Also import resolve_rope_axes_dim at module scope with the other transformer_cosmos3 symbols (lines 38-46), then delete the _resolve helper's local import.

🤖 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/unittest/_torch/visual_gen/test_cosmos3_edge.py` around lines 260 -
321, Move test_qwen3_flag_contradiction_raises and
test_edge_latent_geometry_invariants from TestRopeAxesResolution into the
existing TestArchRecipe class, preserving their assertions and parametrization.
Add resolve_rope_axes_dim to the module-level transformer_cosmos3 imports, then
remove the local import from TestRopeAxesResolution._resolve.

---

147-216: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Tighten the helper signatures and drop the redundant alias.

_synthetic_state_dict and _edge_state_dict return dict without parameters. The coding guidelines ask for precise types instead of bare dict. _edge_state_dict also adds no behavior over _synthetic_state_dict.

♻️ Proposed refactor

-def _synthetic_state_dict(cfg: SimpleNamespace) -> dict:
+def _synthetic_state_dict(cfg: SimpleNamespace) -> dict[str, torch.Tensor]:
     """A complete synthetic checkpoint for either recipe (diffusers key layout)."""

-def _edge_state_dict(cfg: SimpleNamespace) -> dict:
-    return _synthetic_state_dict(cfg)

Then call _synthetic_state_dict directly at the _edge_state_dict call sites (lines 603, 621, 627, 640, 680, 699).

🤖 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/unittest/_torch/visual_gen/test_cosmos3_edge.py` around lines 147 -
216, Update _synthetic_state_dict to use a precise state-dict type annotation
instead of bare dict, remove the redundant _edge_state_dict wrapper, and replace
every _edge_state_dict call at the indicated call sites with direct
_synthetic_state_dict calls.

Source: Coding guidelines

tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py (3)

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

Tighten the return annotation to list[int].

resolve_rope_axes_dim returns MRoPE axis sizes. Annotate the element type so callers get precise typing.

♻️ Proposed annotation change

-def resolve_rope_axes_dim(pretrained_config) -> list:
+def resolve_rope_axes_dim(pretrained_config) -> list[int]:

As per coding guidelines: "use precise types instead of dict/object/Any" and "prefer built-in generic types".

🤖 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/cosmos3/transformer_cosmos3.py` at line
65, Update the return annotation of resolve_rope_axes_dim to list[int],
preserving its existing behavior and implementation.

Source: Coding guidelines

---

633-634: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add the return annotation to _layer_norm_cls.

The coding guidelines require an annotation on every function. Declare the returned class type.

♻️ Proposed annotation change

-def _layer_norm_cls(recipe: Cosmos3ArchRecipe):
+def _layer_norm_cls(recipe: Cosmos3ArchRecipe) -> type[Qwen3VLTextRMSNorm]:
     return NemotronRMSNorm if recipe.nemotron_norms else Qwen3VLTextRMSNorm

As per coding guidelines: "Annotate every function, use None for procedures".

🤖 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/cosmos3/transformer_cosmos3.py` around
lines 633 - 634, Update _layer_norm_cls to include a return type annotation
describing the RMS normalization class it returns, covering both NemotronRMSNorm
and Qwen3VLTextRMSNorm branches while preserving the existing selection logic.

Source: Coding guidelines

---

460-473: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Refactor the rotary helper to support key-only rotation.

At line 470, qwen3_apply_rotary_pos_emb recomputes the already rotated q and discards it. Rotate k_for_gen without recomputing q.

🤖 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/cosmos3/transformer_cosmos3.py` around
lines 460 - 473, Update qwen3_apply_rotary_pos_emb and the k_for_gen branch in
the transformer forward flow to support key-only rotary rotation, so k_for_gen
is rotated without recomputing or discarding q. Preserve the existing q/k
rotation behavior for the primary path and assign only the rotated key for the
generation path.

tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py (4)

302-306: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider emitting the envelope advisory on rank 0 only.

forward guards its generation-dimension info log with if self.rank == 0. This warning has no rank guard, so a tensor-parallel or Ulysses run emits one identical warning per rank for every out-of-envelope request.

♻️ Proposed rank guard

-        if outside:
+        if outside and self.rank == 0:
             logger.warning(
                 "Request is outside the model-card validated envelope "
                 f"({'; '.join(outside)}); generation proceeds but quality may degrade."
             )

🤖 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/cosmos3/pipeline_cosmos3.py` around
lines 302 - 306, Update the out-of-envelope warning in forward to emit only when
self.rank == 0, matching the existing generation-dimension logging guard;
preserve the current warning message and outside detection for all ranks.

---

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

Annotate the sampling and transformer parameters.

Both module-level validators accept an unannotated object. _validate_sampling_recipe reads sampling.is_distilled, and _validate_temporal_compression reads transformer.temporal_compression_factor. Declare the types so the expected interface is explicit.

Use Cosmos3SamplingPolicy for sampling. For transformer, use Cosmos3VFMTransformer; both are already imported in this module.

As per coding guidelines: "Annotate every function, use None for procedures".

Also applies to: 160-160

🤖 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/cosmos3/pipeline_cosmos3.py` at line
120, Annotate the sampling parameter of _validate_sampling_recipe with
Cosmos3SamplingPolicy, and annotate the transformer parameter of
_validate_temporal_compression with Cosmos3VFMTransformer. Keep both validators’
existing return annotation as None and do not change their behavior.

Source: Coding guidelines

---

537-547: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Note the field-coverage gap between infer and forward.

infer re-resolves four fields against the request's own mode: height, width, num_inference_steps, and guidance_scale. forward re-resolves seven, adding num_frames, max_sequence_length, and frame_rate.

For the three extra fields, infer forwards req.params.* as concrete values that the executor merged from the video table. forward then sees non-None and keeps them, so the mode table never applies. This is harmless today: the image tables omit frame_rate and fall back to the video table, both Edge tables declare the same max_sequence_length, and the is_t2i branch overwrites num_frames with 1.

The gap becomes a wrong-default bug the moment an image table declares its own max_sequence_length or frame_rate. Route the same field set through as_given so both entry points agree.

♻️ Proposed alignment

         resolved = self._resolve_generation_params(
             "image" if is_t2i else "video",
             height=as_given("height"),
             width=as_given("width"),
             num_inference_steps=as_given("num_inference_steps"),
             guidance_scale=as_given("guidance_scale"),
+            num_frames=as_given("num_frames"),
+            max_sequence_length=as_given("max_sequence_length"),
+            frame_rate=as_given("frame_rate"),
         )
         height = resolved["height"]
         width = resolved["width"]
         num_inference_steps = resolved["num_inference_steps"]
         guidance_scale = resolved["guidance_scale"]

Then pass the resolved num_frames, max_sequence_length, and frame_rate to forward instead of the raw req.params values.

🤖 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/cosmos3/pipeline_cosmos3.py` around
lines 537 - 547, Align infer’s parameter resolution with forward by routing
num_frames, max_sequence_length, and frame_rate through as_given when calling
_resolve_generation_params, alongside the existing four fields. Pass these
resolved values to forward instead of the raw req.params values, while
preserving the existing mode-specific handling.

---

226-229: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Initialize _base_scheduler and _base_audio_scheduler in the constructor.

__init__ initializes use_native_flow_schedule, family, and _scheduler_cache, but not _base_scheduler or _base_audio_scheduler. load_standard_components creates them, and it skips that block when PipelineComponent.SCHEDULER is in skip_components.

_scheduler_for compensates for the gap at lines 509-515 with getattr(self, "_base_scheduler", None) or self.scheduler plus a lazy re-creation of _scheduler_cache. Initializing all three attributes here lets those three guards become direct attribute reads.

The coding guidelines require initializing externally visible class members in the constructor.

♻️ Proposed initialization

         self.use_native_flow_schedule = False
         self.family = resolve_arch_recipe(primary_pretrained_config).name
         # Schedulers are identified by their resolved (flow_shift, karras) pair
         self._scheduler_cache: dict = {}
+        # Untouched per-stream bases for _scheduler_for(); load_standard_components
+        # replaces them with the checkpoint's scheduler instances.
+        self._base_scheduler = None
+        self._base_audio_scheduler = None

As per coding guidelines: "initialize externally visible class members in the constructor".

🤖 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/cosmos3/pipeline_cosmos3.py` around
lines 226 - 229, Initialize _base_scheduler, _base_audio_scheduler, and
_scheduler_cache in the constructor alongside the existing scheduler-related
members, using appropriate empty/default values. Then update _scheduler_for to
use direct attribute reads and remove the getattr fallback and lazy
_scheduler_cache initialization, while preserving its existing
scheduler-selection behavior.

Source: Coding guidelines

tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py (2)

161-170: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Note the coupling between these keys and the recipe names.

The family keys are literal strings. The pipeline resolves self.family from resolve_arch_recipe(...).name in tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py, then indexes this table. A rename of Cosmos3ArchRecipe.name breaks the lookup.

The failure is loud (a KeyError at pipeline construction, not a wrong default), so this is safe today. Tests already reference QWEN3_RECIPE.name and NEMOTRON_DENSE_RECIPE.name for the same values. Consider a short comment here that names Cosmos3ArchRecipe.name as the source of these strings, so the coupling is visible at the table.

🤖 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/cosmos3/defaults.py` around lines 161 -
170, Add a concise comment immediately above COSMOS3_GENERATION_DEFAULTS
documenting that its family keys must match the Cosmos3ArchRecipe.name values
returned by resolve_arch_recipe, specifically the QWEN3_RECIPE.name and
NEMOTRON_DENSE_RECIPE.name identifiers. Keep the existing mapping unchanged.

---

165-175: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Annotate both tables with their key and value types.

COSMOS3_GENERATION_DEFAULTS is keyed by a (family, mode) pair and COSMOS3_ENVELOPES by family name. A bare Dict hides both. Declare the concrete types so a wrong lookup key fails type checking rather than at request time.

♻️ Proposed annotation change

-COSMOS3_GENERATION_DEFAULTS: Dict = {
+COSMOS3_GENERATION_DEFAULTS: dict[tuple[str, str], dict[str, Any]] = {
     ("qwen3", "video"): COSMOS3_720P_PARAMS,
     ("qwen3", "image"): COSMOS3_T2I_PARAMS,
     ("nemotron_dense", "video"): COSMOS3_EDGE_VIDEO_PARAMS,
     ("nemotron_dense", "image"): COSMOS3_EDGE_T2I_PARAMS,
 }
 
 # Families without an entry get no envelope advisory.
-COSMOS3_ENVELOPES: Dict = {
+COSMOS3_ENVELOPES: dict[str, dict[str, Any]] = {
     "nemotron_dense": COSMOS3_EDGE_ENVELOPE,
 }

As per coding guidelines: "use precise types instead of dict/object/Any" and "prefer built-in generic types".

🤖 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/cosmos3/defaults.py` around lines 165 -
175, Annotate COSMOS3_GENERATION_DEFAULTS with a tuple-of-strings key and the
appropriate generation-parameters value type, and annotate COSMOS3_ENVELOPES
with string keys and the envelope value type. Replace both bare Dict annotations
with precise built-in generic types so invalid family, mode, or envelope lookups
are caught by type checking.

Source: Coding guidelines

tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py (3)

272-283: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Apply the native-flow karras default before the early return.

Line 272 returns the scheduler unchanged when target_shift and use_karras_sigmas are both None. Lines 282-283 then apply the native-flow implicit use_karras_sigmas=False — but only for calls that got past that return.

So set_flow_shift(scheduler, None) on a native-flow policy returns the checkpoint scheduler with karras sigmas still enabled. That is the wrong-trajectory case this flag exists to prevent, and it fails silently.

No current caller reaches it: forward always resolves a concrete target_shift from the mode table, checkpoint_flow_shift, or the V2V constant. The invariant therefore rests on every present and future caller passing a shift. Move the implicit assignment above the early return so the policy enforces it itself.

♻️ Proposed reordering

         if self.unipc_base_config is None:
             return scheduler
+        # The native flow schedule is defined on explicit linear sigmas, which
+        # UniPC's karras branch discards; treat it as an implicit request for
+        # the uniform grid rather than a separate code path. Applied before the
+        # no-knob early return so a caller that passes neither knob still gets
+        # the uniform grid.
+        if use_karras_sigmas is None and self.native_flow_schedule:
+            use_karras_sigmas = False
         if target_shift is None and use_karras_sigmas is None:
             return scheduler
 
         current_shift = float(_config_get(scheduler.config, "flow_shift", 1.0) or 1.0)
         target_shift = self.checkpoint_flow_shift if target_shift is None else float(target_shift)
         current_karras = bool(_config_get(scheduler.config, "use_karras_sigmas", False))
         base_karras = bool(_config_get(self.unipc_base_config, "use_karras_sigmas", False))
-        # The native flow schedule is defined on explicit linear sigmas, which
-        # UniPC's karras branch discards; treat it as an implicit request for
-        # the uniform grid rather than a separate code path.
-        if use_karras_sigmas is None and self.native_flow_schedule:
-            use_karras_sigmas = False
         target_karras = base_karras if use_karras_sigmas is None else bool(use_karras_sigmas)

🤖 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/cosmos3/sampling.py` around lines 272 -
283, Move the native-flow default assignment for use_karras_sigmas above the
early return in set_flow_shift. Ensure native-flow policies set
use_karras_sigmas to false before evaluating whether target_shift and
use_karras_sigmas are both None, while preserving the existing return behavior
for other policies.

---

110-117: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider stating why the distilled branch drops native_flow_schedule.

The UniPC branch forwards native_flow_schedule to the constructed policy. The distilled branch at line 164 omits it, so the policy reports False regardless of the caller's argument.

This is safe today. _validate_sampling_recipe in pipeline_cosmos3.py rejects both Edge-plus-distilled and non-Edge-plus-native, so a distilled policy carrying the flag is unreachable. The omission reads as an oversight, though. A short comment on the distilled return would record that the flag is deliberately not applicable there.

Also applies to: 139-143

🤖 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/cosmos3/sampling.py` around lines 110 -
117, Clarify the distilled branch in Cosmos3SamplingPolicy.from_scheduler by
adding a short comment at its return explaining that native_flow_schedule is
intentionally not applicable there. Leave the existing behavior unchanged,
including forwarding the flag in the UniPC branch.

---

221-227: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low value

Correct the interval in the inline comment. np.linspace(... )[:-1] includes 1 - 1/T and excludes 0, so describe the base grid as [1 - 1/T, 0). Keep the NumPy array because diffusers 0.39.0 accepts both arguments, but its flow-shift arithmetic does not support a Python list.

🤖 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/cosmos3/sampling.py` around lines 221 -
227, Update the inline comment in the native_flow_schedule branch to describe
the generated sigma grid as [1 - 1/T, 0), reflecting that np.linspace(... )[:-1]
includes the upper bound and excludes zero. Keep the existing NumPy array
implementation and scheduler.set_timesteps call unchanged.

tensorrt_llm/_torch/visual_gen/executor.py (1)

420-424: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Guard model_fields_set mutation behind one helper. The current Pydantic implementation returns a live set, but the public API does not guarantee that the returned set remains a mutable alias. If it returns a copy, both discard calls silently fail and image requests can retain video defaults. Add a helper beside VisualGenParams that removes fields and asserts each removal, then use it at both call sites.

🤖 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/executor.py` around lines 420 - 424, Add a
helper beside VisualGenParams that removes a named field from model_fields_set
and asserts the removal succeeded, then replace the direct discard at both
affected sites in tensorrt_llm/_torch/visual_gen/executor.py (lines 420-424) and
tensorrt_llm/visual_gen/visual_gen.py (lines 319-325) with that helper; preserve
the existing pipeline-default behavior.

🤖 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/README.md`:
- Line 19: Add the required NVIDIA copyright header at the beginning of the
README, before the existing heading, preserving all current model documentation
below it.

In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 513-521: Update the scheduler-cache flow around _scheduler_cache
and set_flow_shift to prevent unbounded growth from arbitrary target_shift
values: validate or constrain flow_shift to an allowed range before caching, or
bound the cache with an eviction/size limit. Preserve distinct valid scheduler
configurations and ensure invalid requests cannot create cache entries.

In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`:
- Around line 814-819: Update the rope_type initialization near
resolve_rope_axes_dim to read rope_scaling tolerantly, treating a missing or
None value as an empty mapping before calling get("rope_type", "default"). Keep
the existing fallback and ensure it matches the access behavior used by
resolve_rope_axes_dim.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/executor.py`:
- Around line 420-424: Add a helper beside VisualGenParams that removes a named
field from model_fields_set and asserts the removal succeeded, then replace the
direct discard at both affected sites in
tensorrt_llm/_torch/visual_gen/executor.py (lines 420-424) and
tensorrt_llm/visual_gen/visual_gen.py (lines 319-325) with that helper; preserve
the existing pipeline-default behavior.

In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py`:
- Around line 161-170: Add a concise comment immediately above
COSMOS3_GENERATION_DEFAULTS documenting that its family keys must match the
Cosmos3ArchRecipe.name values returned by resolve_arch_recipe, specifically the
QWEN3_RECIPE.name and NEMOTRON_DENSE_RECIPE.name identifiers. Keep the existing
mapping unchanged.
- Around line 165-175: Annotate COSMOS3_GENERATION_DEFAULTS with a
tuple-of-strings key and the appropriate generation-parameters value type, and
annotate COSMOS3_ENVELOPES with string keys and the envelope value type. Replace
both bare Dict annotations with precise built-in generic types so invalid
family, mode, or envelope lookups are caught by type checking.

In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 302-306: Update the out-of-envelope warning in forward to emit
only when self.rank == 0, matching the existing generation-dimension logging
guard; preserve the current warning message and outside detection for all ranks.
- Line 120: Annotate the sampling parameter of _validate_sampling_recipe with
Cosmos3SamplingPolicy, and annotate the transformer parameter of
_validate_temporal_compression with Cosmos3VFMTransformer. Keep both validators’
existing return annotation as None and do not change their behavior.
- Around line 537-547: Align infer’s parameter resolution with forward by
routing num_frames, max_sequence_length, and frame_rate through as_given when
calling _resolve_generation_params, alongside the existing four fields. Pass
these resolved values to forward instead of the raw req.params values, while
preserving the existing mode-specific handling.
- Around line 226-229: Initialize _base_scheduler, _base_audio_scheduler, and
_scheduler_cache in the constructor alongside the existing scheduler-related
members, using appropriate empty/default values. Then update _scheduler_for to
use direct attribute reads and remove the getattr fallback and lazy
_scheduler_cache initialization, while preserving its existing
scheduler-selection behavior.

In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py`:
- Around line 272-283: Move the native-flow default assignment for
use_karras_sigmas above the early return in set_flow_shift. Ensure native-flow
policies set use_karras_sigmas to false before evaluating whether target_shift
and use_karras_sigmas are both None, while preserving the existing return
behavior for other policies.
- Around line 110-117: Clarify the distilled branch in
Cosmos3SamplingPolicy.from_scheduler by adding a short comment at its return
explaining that native_flow_schedule is intentionally not applicable there.
Leave the existing behavior unchanged, including forwarding the flag in the
UniPC branch.
- Around line 221-227: Update the inline comment in the native_flow_schedule
branch to describe the generated sigma grid as [1 - 1/T, 0), reflecting that
np.linspace(... )[:-1] includes the upper bound and excludes zero. Keep the
existing NumPy array implementation and scheduler.set_timesteps call unchanged.

In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`:
- Line 65: Update the return annotation of resolve_rope_axes_dim to list[int],
preserving its existing behavior and implementation.
- Around line 633-634: Update _layer_norm_cls to include a return type
annotation describing the RMS normalization class it returns, covering both
NemotronRMSNorm and Qwen3VLTextRMSNorm branches while preserving the existing
selection logic.
- Around line 460-473: Update qwen3_apply_rotary_pos_emb and the k_for_gen
branch in the transformer forward flow to support key-only rotary rotation, so
k_for_gen is rotated without recomputing or discarding q. Preserve the existing
q/k rotation behavior for the primary path and assign only the rotated key for
the generation path.

In
`@tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py`:
- Around line 410-415: Define or reuse the precise ConfigType configuration
type, then update the pretrained_dict parameters in both _forward and
_build_ref_and_parallel to use ConfigType | None with an appropriate default.
Remove the unparameterized dict and implicit Optional annotations while
preserving existing behavior.

In `@tests/unittest/_torch/visual_gen/test_cosmos3_edge.py`:
- Around line 1146-1353: Add a parametrized negative test to TestEdgeCheckpoint
covering action generation, audio, V2V, and transfer requests. For each
unsupported mode, call edge_pipeline.forward with the corresponding request
field and assert the documented exception type and clear error message; use
_bare_pipeline if validation occurs before model loading. Keep the test
checkpoint-gated and aligned with the existing request-field conventions.
- Around line 260-321: Move test_qwen3_flag_contradiction_raises and
test_edge_latent_geometry_invariants from TestRopeAxesResolution into the
existing TestArchRecipe class, preserving their assertions and parametrization.
Add resolve_rope_axes_dim to the module-level transformer_cosmos3 imports, then
remove the local import from TestRopeAxesResolution._resolve.
- Around line 147-216: Update _synthetic_state_dict to use a precise state-dict
type annotation instead of bare dict, remove the redundant _edge_state_dict
wrapper, and replace every _edge_state_dict call at the indicated call sites
with direct _synthetic_state_dict calls.

🪄 Autofix

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: 8ac5caca-3a81-4add-a895-23a45a23ae27

📥 Commits

Reviewing files that changed from the base of the PR and between 5533f66 and de434cf.

⛔ 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 (23)

• .pre-commit-config.yaml
• docs/source/models/supported-models.md
• docs/source/models/visual-generation.md
• examples/visual_gen/models/cosmos3/README.md
• examples/visual_gen/models/cosmos3/cosmos3.py
• tensorrt_llm/_torch/visual_gen/executor.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/negative_prompt.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/visual_gen/visual_gen.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/test_visual_gen.py
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py
• tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.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

🚧 Files skipped from review as they are similar to previous changes (7)

• 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/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py
• docs/source/models/supported-models.md
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

[karljang]

LGTM, but just two small comments.

[ishovkun]

/bot run --disable-fail-fast

[coderabbitai]

Actionable comments posted: 3

🤖 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 `@tests/unittest/_torch/visual_gen/test_cosmos3_edge.py`:
- Around line 968-973: Update test_cacheable_set_is_derived_from_the_mode_tables
to patch the checkpoint_flow_shift, video-mode, and image-mode table entries
with distinct values, then assert _cacheable_flow_shifts() contains all three
patched shifts plus COSMOS3_V2V_FLOW_SHIFT. Preserve the test’s focus on
deriving cacheable shifts from each independent source rather than relying on
current Edge mode-table values.
- Around line 955-999: Annotate
tests/unittest/_torch/visual_gen/test_cosmos3_edge.py lines 955-999: add
parameter and return annotations to _pipeline, and add -> None to
test_cacheable_set_is_derived_from_the_mode_tables,
test_arbitrary_shifts_never_enter_the_cache,
test_declared_shifts_are_memoized_and_shared, and
test_release_drops_retained_solver_state. In
tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py lines 510-528,
annotate rope_scaling and add -> None to its function.

In `@tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py`:
- Around line 509-528: Add a separate test case in
test_rope_type_tolerates_missing_rope_scaling using a SimpleNamespace that omits
the rope_scaling attribute entirely, rather than assigning it None. Construct
Qwen3VLTextRotaryEmbedding with this config and assert it still resolves
rope_type to "default" and preserves mrope_section from rope_axes_dim; retain
the existing None and empty-dict cases.

🪄 Autofix

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: 04f4ec2f-913b-4048-85bb-0b3e497dab5e

📥 Commits

Reviewing files that changed from the base of the PR and between de434cf and 82510a5.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py
• tests/unittest/_torch/visual_gen/test_cosmos3_edge.py
• tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py
• tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py

[tensorrt-cicd]

PR_Github #64119 [ run ] triggered by Bot. Commit: 82510a5 Link to invocation

[ishovkun]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64138 [ run ] triggered by Bot. Commit: fe85ad8 Link to invocation

[tensorrt-cicd]

PR_Github/16773-82510a5 #64119 was force-killed by a newer pipeline run.
L0 job information not available (job may not have been triggered yet).

Link to superseding invocation