[https://nvbugs/6410093][fix] Revert the per-batch cross-attn slicing (drop real_text_lens plumbing through…

[trtllm-agent]

Summary

• Root cause: f50ca53 refactored Cosmos3CrossAttention to per-batch slice k_und/v_und by real_text_lens under CFG (batch_size=2), plus changed negative-prompt default to "" and max_sequence_length 1024→4096, all of which diverged from the LPIPS golden.
• Fix: Revert the per-batch cross-attn slicing (drop real_text_lens plumbing through Cosmos3CrossAttention/Cosmos3GenDecoderLayer/Cosmos3VFMTransformer); pin the pre-audio negative_prompt and max_sequence_length in the LPIPS test itself; remove the nvbugs/6410082 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/6410093

Summary by CodeRabbit

• Bug Fixes
  • Improved text/video generation consistency by simplifying how cross-attention combines inputs, which may change generated output quality and stability.
  • Updated a visual generation benchmark configuration with fixed generation settings to make evaluation results more repeatable.
• Tests
  • Removed an outdated test waiver so the related visual generation check will run again.

[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: 25365abd-80e2-4090-b61f-d238f8fae6e1

📥 Commits

Reviewing files that changed from the base of the PR and between ca9c439 and a8cf30c.

📒 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

Removes real_text_lens-based per-batch truncation from Cosmos3 cross-attention, decoder layer, and transformer forward paths, replacing it with unconditional concatenation and full attention. Updates the Cosmos3 LPIPS golden test with pinned negative prompt and max sequence length, and removes the corresponding skip waiver.

Changes

Cosmos3 cross-attention simplification and test re-enablement

Layer / File(s)	Summary
Remove real_text_lens truncation logic tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py	Cosmos3CrossAttention.forward drops the real_text_lens parameter and per-batch truncation loop, always concatenating k_und/v_und with k/v and running full attention.
Propagate removal through decoder layer and transformer tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py	Cosmos3GenDecoderLayer.forward stops passing real_text_lens to cross-attention; Cosmos3VFMTransformer.forward removes computation of real_text_lens and omits it from gen-layer calls.
Pin LPIPS golden generation parameters and re-enable test 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, passes them into the pipeline forward call, and removes the skip waiver for test_cosmos3_nano_t2v_lpips_against_golden.

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

Sequence Diagram(s)

sequenceDiagram
  participant Transformer as Cosmos3VFMTransformer
  participant Layer as Cosmos3GenDecoderLayer
  participant Attn as Cosmos3CrossAttention

  Transformer->>Layer: layer(k_und, v_und, ...)
  Layer->>Attn: cross_attention(k, v, k_und, v_und)
  Attn->>Attn: concatenate k_und/v_und with k/v
  Attn->>Attn: run _attn_impl with FULL mask
  Attn-->>Layer: attention output
  Layer-->>Transformer: layer output

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#15894: Both PRs modify tests/integration/test_lists/waives.txt for the same Cosmos3 LPIPS golden test, with opposite waiver actions.

Suggested reviewers: chzblych, brb-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change: reverting Cosmos3 cross-attention slicing and dropping real_text_lens plumbing.
Description check	✅ Passed	The description clearly states the root cause, fix, test plan, and bug link, so it mostly satisfies the template.
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.}