[https://nvbugs/6418815][fix] Revert the per-batch cross-attention slicing (drop real_text_lens plumbing in…

[trtllm-agent]

Summary

• Root cause: Commit f50ca53 (Cosmos3 Audio Output Support) silently changed three inputs the T2I LPIPS golden was baked against: (1) added CFG per-batch cross-attention slicing keyed on real_text_lens, (2) replaced the rich COSMOS3_DEFAULT_NEGATIVE_PROMPT with "", (3) bumped default max_sequence_length 1024→4096.
• Fix: Revert the per-batch cross-attention slicing (drop real_text_lens plumbing in transformer_cosmos3.py) and pin negative_prompt + max_sequence_length in the test to their pre-f50ca53dae values (matching the WAN/LTX2/QwenImage LPIPS-test pattern); remove the T2I waiver.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6418815

Summary by CodeRabbit

• Bug Fixes

  • Improved Cosmos3 generation behavior by using a consistent attention path, which should make outputs more stable.
  • Updated the Cosmos3 LPIPS test setup to use fixed prompt settings and sequence length for more reliable golden comparisons.

• Tests

  • Re-enabled a previously waived Cosmos3 Nano LPIPS integration test.

[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: a36becbf-77ee-4c0b-a974-fcc68c3a8505

📥 Commits

Reviewing files that changed from the base of the PR and between 011e849 and e904942.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py
• tests/integration/defs/examples/visual_gen/test_visual_gen.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

---

📝 Walkthrough

Walkthrough

Cosmos3CrossAttention and Cosmos3GenDecoderLayer no longer accept real_text_lens, replacing per-batch conditional slicing of UND keys/values with a full-attention concatenation path. Cosmos3VFMTransformer.forward updated accordingly. Integration test adds pinned negative prompt/max sequence length constants for Cosmos3 LPIPS, and removes a corresponding test waiver.

Changes

Cosmos3 Cross-Attention Simplification and Test Updates

Layer / File(s)	Summary
Remove real_text_lens from cross-attention and gen decoder layer tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py	Cosmos3CrossAttention.forward and Cosmos3GenDecoderLayer.forward drop the real_text_lens parameter; attention always concatenates full UND K/V with generation K/V using a FULL attention mask instead of per-example slicing.
Propagate removal into transformer forward tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py	Cosmos3VFMTransformer.forward stops computing real_text_lens from text_mask and no longer passes it to gen-layer calls.
Cosmos3 LPIPS test config and waiver update tests/integration/defs/examples/visual_gen/test_visual_gen.py, tests/integration/test_lists/waives.txt	Adds COSMOS3_LPIPS_NEGATIVE_PROMPT and COSMOS3_LPIPS_MAX_SEQUENCE_LENGTH constants, passes them into the Cosmos3 LPIPS pipeline call, and removes the corresponding test waiver entry.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15959: Both PRs remove SKIP/waiver entries for examples/visual_gen tests in tests/integration/test_lists/waives.txt.
• NVIDIA/TensorRT-LLM#15968: Both PRs modify the waiver for test_cosmos3_nano_t2i_lpips_against_golden in the same waives file.

Suggested reviewers: yibinl-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required ticket/type format and summarizes the main Cosmos3 cross-attention and test-pin fix.
Description check	✅ Passed	The description includes Summary, Test plan, and Links sections and clearly explains the bug, fix, validation, and bug reference.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[BowenFu]

LGTM — clean, complete revert of the per-batch cross-attention slicing regression; single concatenated-attention path restored with no dangling references.