[https://nvbugs/6418815][fix] Revert the per-batch cross-attention slicing (drop real_text_lens plumbing in… - #15986
Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (4)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (3)
Included review availability: Your plan includes up to 12 reviews per rolling hour; 5 remain after this review. WalkthroughCosmos3 LPIPS tests now use pinned T2V and T2I generation parameters. The T2I golden configuration and waiver are updated. Failed T2I runs preserve generated candidates. OpenCV loading in the LPIPS evaluator is now lazy. ChangesCosmos3 LPIPS evaluation
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to The PR restores prior Cosmos3 test parameters and removes the related waiver; no actionable merge-blocking risk remains beyond normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
be138a5 to
1f80504
Compare
1f80504 to
57af1f8
Compare
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM — clean, complete revert of the per-batch cross-attention slicing regression; single concatenated-attention path restored with no dangling references.
57af1f8 to
42f89cc
Compare
42f89cc to
8f92aaa
Compare
|
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[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
scripts/visualgen_eval/visual_gen_lpips_score_eval.py (1)
73-82: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a return annotation to
_get_cv2().Annotate the function with a precise return type, such as
types.ModuleType, instead of leaving it untyped.As per coding guidelines, annotate every function and avoid unnecessary
Any.🤖 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 `@scripts/visualgen_eval/visual_gen_lpips_score_eval.py` around lines 73 - 82, Update _get_cv2() to declare a precise return annotation using the appropriate module type, and add any necessary standard-library import for that annotation; keep its lazy import and existing ImportError behavior unchanged.Source: Coding guidelines
tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py (1)
147-159: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd type annotations to the modified helper.
_run_cosmos3_lpips_pipelinehas untyped parameters and no return annotation. Add precise annotations fornum_frames,video,negative_prompt,max_sequence_length, and thetorch.Tensor | Nonereturn value.As per coding guidelines,
**/*.pyrequires annotations on every function.Proposed annotation change
-def _run_cosmos3_lpips_pipeline( - num_frames, video=None, negative_prompt=None, max_sequence_length=None -): +def _run_cosmos3_lpips_pipeline( + num_frames: int, + video: bytes | None = None, + negative_prompt: str | None = None, + max_sequence_length: int | None = None, +) -> torch.Tensor | None:🤖 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_cosmos3.py` around lines 147 - 159, Update _run_cosmos3_lpips_pipeline with precise annotations for num_frames, video, negative_prompt, and max_sequence_length, and annotate its return type as torch.Tensor | None, using the project’s existing conventions for byte payloads and optional values.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.
Inline comments:
In
`@tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_t2i_lpips_golden.json`:
- Around line 23-33: Update the rebake_reason for the Cosmos3 T2I golden to
accurately reflect the retained real_text_lens behavior in
transformer_cosmos3.py, including the default non-sharded path covered by
test_cosmos3_nano_t2i_lpips_against_golden. Do not remove per-sample CFG
slicing; ensure the golden provenance and stated objective match the implemented
attention contract.
---
Nitpick comments:
In `@scripts/visualgen_eval/visual_gen_lpips_score_eval.py`:
- Around line 73-82: Update _get_cv2() to declare a precise return annotation
using the appropriate module type, and add any necessary standard-library import
for that annotation; keep its lazy import and existing ImportError behavior
unchanged.
In `@tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py`:
- Around line 147-159: Update _run_cosmos3_lpips_pipeline with precise
annotations for num_frames, video, negative_prompt, and max_sequence_length, and
annotate its return type as torch.Tensor | None, using the project’s existing
conventions for byte payloads and optional values.
🪄 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: 484a19c3-55ab-4cee-8d7a-3f815b901318
⛔ Files ignored due to path filters (1)
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zipis excluded by!**/*.zip
📒 Files selected for processing (4)
scripts/visualgen_eval/visual_gen_lpips_score_eval.pytests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_t2i_lpips_golden.jsontests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
8f92aaa to
9eab8f3
Compare
|
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. |
9eab8f3 to
f30b287
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The PR title and description describe a change that isn't in this branch. They say "Revert the per-batch cross-attention slicing (drop real_text_lens plumbing in transformer_cosmos3.py)", but real_text_lens is still there (transformer_cosmos3.py:375, 1040, 1143), no product code is touched at all, and the golden JSON added here argues at length that the slicing is correct and must not be reverted. The actual fix is a golden rebake plus input pinning — which is the better call, but the description needs to say so, since it's what a future reader will use to understand why this golden moved.
Three changes are also undescribed: the lazy cv2 import in the LPIPS eval script, dropping _visual_gen_deps from the T2I test, and adding _preserve_lpips_candidate_on_failure to T2I. Worth one line each in the description.
Since the waiver is removed here, the T2I stage should be run explicitly on the pipeline (/bot run --extra-stage/--stage-list for the B200 post-merge stage) rather than relying on the default set.
| # Cosmos3 requires VANILLA attention and guardrails disabled in CI. | ||
| COSMOS3_NANO_MODEL_SUBPATH = "Cosmos3-Nano" | ||
| COSMOS3_LPIPS_PROMPT = "A serene mountain landscape with snow-capped peaks and a flowing river" | ||
| # The T2I/T2V goldens were baked before the Cosmos3 audio-output feature |
There was a problem hiding this comment.
This rationale is stale relative to the golden added in the same PR. It says the T2I golden "was baked before the Cosmos3 audio-output feature", but cosmos3_nano_t2i_lpips_golden.json now records commit 1745a6e6 and states it was regenerated with real_text_lens in force. Reword to describe what the pins actually do now: lock the two conditioning inputs the golden was rebaked against so a future change to the public defaults can't silently invalidate it again.
Also worth noting in the comment: with per-sample slicing active, max_sequence_length no longer changes the attended tokens (the padding is sliced away) — it only bounds prompt truncation. Pinning it is provenance, not numerics, and saying that prevents someone from "cleaning it up" later on the assumption it's dead weight.
| def _generate_cosmos3_lpips_video(output_path): | ||
| """Generate the Cosmos3-Nano text-to-video LPIPS sample.""" | ||
| video = _run_cosmos3_lpips_pipeline(COSMOS3_LPIPS_T2V_NUM_FRAMES) | ||
| video = _run_cosmos3_lpips_pipeline( |
There was a problem hiding this comment.
This changes the T2V generation inputs, but cosmos3_nano_t2v_lpips_golden_video.json is untouched — still tensorrt_llm_commit: 85665f5f, and it records neither negative_prompt nor max_sequence_length. That golden predates the per-sample slicing too, so pinning two of the three drifted inputs can't make it match; T2V will still need a rebake, and meanwhile its checked-in metadata no longer describes how the test runs.
T2V is waived under a separate bug, so this isn't gating, but it leaves a test in a state where nobody can tell from the golden what inputs it expects. Either rebake T2V here (and record the pins in its JSON), or leave T2V on pipeline defaults and scope this PR to T2I.
| "model": "Cosmos3-Nano", | ||
| "source": "TensorRT-LLM VisualGen", | ||
| "prompt": "A serene mountain landscape with snow-capped peaks and a flowing river", | ||
| "negative_prompt": "pinned pre-audio descriptive default (COSMOS3_LPIPS_NEGATIVE_PROMPT)", |
There was a problem hiding this comment.
Every other golden stores the literal negative prompt (ltx2_lpips_golden_video.json:6, qwenimage_lpips_golden.json:6), which is what makes these files self-describing for a rebake. A prose placeholder pointing at a Python constant means anyone regenerating this asset has to go read the test to recover the string, and it silently goes stale if the constant is ever edited. Store the actual text.
| "tensorrt_llm_version": "1.3.0rc24", | ||
| "tensorrt_llm_commit": "1745a6e689082e32e113fe06460ee5c3dfecbba9", | ||
| "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607271403-16694-handongl", | ||
| "rebake_reason": [ |
There was a problem hiding this comment.
rebake_reason records the version, commit, and container, but not the GPU the asset was generated on. The failure this fixes is from a B200 post-merge stage, and the whole risk with a rebaked golden is that it encodes the numerics of whatever SKU produced it — "bit-exact across two independent runs on the same host" establishes determinism, not portability. Add the GPU/stage to the record, and confirm in the PR that the rebake ran on the same GPU type as the failing stage.
f30b287 to
84c5d13
Compare
|
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. |
…put feature The Cosmos3 audio-output feature (commit f50ca53) silently changed three pipeline behaviors that the T2I LPIPS golden was baked against: 1. Refactored Cosmos3CrossAttention to per-batch slice k_und/v_und by real_text_lens when batch_size > 1. Under CFG (batch_size=2), the shorter (negative-prompt) entry now attends over only its real text length rather than the padded max_real_len slice used by every other batch entry, producing a substantially different image. 2. Replaced the long descriptive COSMOS3_DEFAULT_NEGATIVE_PROMPT with "". 3. Bumped COSMOS3_720P_PARAMS["max_sequence_length"] from 1024 to 4096. The result was LPIPS = 0.608404 on test_cosmos3_nano_t2i_lpips_against_golden (12x the 0.05 threshold). This mirrors the T2V-sibling fix on branch repair-bot-bug6410093 (commit a8cf30c); the same underlying regression affects both variants but T2V happens to still land under the 0.05 threshold with the new pipeline defaults while T2I diverges. Fix: - Revert the per-batch cross-attention slicing: drop the real_text_lens parameter from Cosmos3CrossAttention.forward and Cosmos3GenDecoderLayer.forward, and stop computing/passing it in Cosmos3VFMTransformer.forward. All batch entries now share the same k_und[:, :max_real_len] slice as before. The per-batch path was only preparatory for future audio work and is not exercised by any existing audio test (audio tests use batch_size=1). - Pin the LPIPS-golden-specific negative_prompt and max_sequence_length in the test itself (matching the WAN21/22, LTX2, QwenImage pattern), so the LPIPS test stays decoupled from future public-default changes. - Remove the nvbugs/6418815 waiver. Verified: LPIPS score drops to 0.000142 (well below the 0.05 threshold). Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Signed-off-by: handongl <handongl@nvidia.com>
…only tests Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Signed-off-by: handongl <handongl@nvidia.com>
…real_text_lens The T2I golden was baked at 85665f5, before the Cosmos3 audio-output feature (f50ca53), so it encoded two things that later changed: 1. Two conditioning defaults: the descriptive negative prompt became "" and max_sequence_length went 1024 -> 4096. Pinning the original values in the test (as WAN21/22, LTX2 and QwenImage already do) takes LPIPS from 0.608 to 0.154. 2. The pre-feature cross-attention numerics. f50ca53 added per-sample text slicing (real_text_lens); before it, cross-attention padded every CFG sample to the batch-wide max text length and attended over the padding. Forcing the old padded path makes T2I pass, which confirms this is the remaining 0.154 -- but it also drops the V2V test (whose golden was baked after the feature) to 0.380, so the slicing is correct behavior and must not be reverted. Restore real_text_lens to its upstream form and refresh the T2I golden instead. The regenerated image is bit-exact across two independent runs. Both T2I and V2V now pass, and the waiver is removed. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…3 T2I LPIPS test The T2I LPIPS case is single-frame end to end: generation saves one PNG via PIL save_image, and _run_lpips_eval scores it through the eval script's image branch (PIL + lpips). It requested _visual_gen_deps anyway, which provisions video codecs with apt-get update/install ffmpeg. apt-get cannot succeed in the non-root test container, so the fixture raised CalledProcessError (exit 100) and errored the test during setup, before the LPIPS comparison could run. Drop the fixture from this case, matching the sibling image-only test test_cosmos3_feature_accuracy_against_golden, which already omits it. The video and V2V cases keep the fixture since they do decode/encode MP4. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
… backend The T2I LPIPS test could not reach its own assertion. fused_moe_cute_dsl.py imported four Sm100BlockScaledContiguous*Runner names at module scope, but cute_dsl_custom_ops defines them inside an ``if IS_CUTLASS_DSL_AVAILABLE:`` block with no else-branch, so without the optional cutlass DSL those names simply do not exist. create_moe imports this module eagerly and it sits under _torch.models, so the ImportError propagated through modeling_utils and took down the whole model-architecture registry -- including the visual_gen PipelineLoader the test loads. The run died at import in 1.6s, long before any image was generated or compared. The four names are needed at exactly one place, the isinstance check in runner_tactic_comb_checker, and reaching that line means a CuteDSL runner is already being autotuned, so the DSL is necessarily installed. Import them there instead of at module scope, matching how cute_dsl_mla.py and dsa/metadata.py already reach into this guard from modules that must stay importable. No module-scope name is added and the isinstance tuple is unchanged, so with the DSL present the tactic constraint behaves exactly as before; without it those symbols never existed, so nothing that used to work is disabled. Removes this bug's waiver. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
84c5d13 to
5ac5bba
Compare
crazydemo
left a comment
crazydemo
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The change itself is a reasonable golden rebake + input pinning that fixes the T2I LPIPS failure, but it cannot merge as-is: mergeable_state is dirty (needs conflict resolution), the title/description misdescribe the actual fix, and one forwarded default relies on an unverified pipeline behavior.
Issues
- [MAJOR]
tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py:198-max_sequence_length=Nonenow always forwarded topipeline.forward - [MINOR]
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py:28- unrelated lazy-import change; verifyGroupedGemmInputsHelperis not also guarded - [MINOR]
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_nano_t2i_lpips_golden.json:24- PR title/description contradict the actual change
QA view
- Test coverage: partial - T2I LPIPS is re-enabled and validated against the rebaked golden; the bundled
fused_moe_cute_dsl.pychange has no test, and the V2V path that would exercise the newmax_sequence_length=Nonedefault is still waived. - SM coverage: code touches sm100 (Blackwell) via the Sm100 grouped-GEMM runners and the B200 visual-gen run; the fused_moe lazy-import path is exercised by no test.
- Test code: the always-forwarded
max_sequence_length=Nonerelies on the pipeline treating None as its default;_preserve_lpips_candidate_on_failureis used but not shown. - Test time: significant - removing the T2I waiver re-enables a full pipeline-generation + LPIPS integration test on B200.
- Needs
/qa-verify: yes - waiver removed and golden (media + config) rebaked; a human should re-run the T2I stage on B200 to confirm reproducibility under the 0.05 threshold.
Does this actually fix nvbugs/6418815?
Partial. The failing T2I LPIPS path is addressed by rebaking the golden (regenerated with real_text_lens in force) and pinning negative_prompt/max_sequence_length, then removing the waiver — NOT by reverting the slicing as the title claims (transformer_cosmos3.py is untouched, real_text_lens retained). If the golden is genuinely bit-exact reproducible on the target GPU the test will pass; that reproducibility is the still-uncovered path the diff cannot confirm.
Possible new issues
- If
pipeline.forwarddoes not interpretNoneas "use default", the V2V/default caller breaks; currently hidden only because V2V is waived. - If
GroupedGemmInputsHelperis also inside theIS_CUTLASS_DSL_AVAILABLEblock, keeping it at module scope leaves the import failure unfixed for cutlass-absent environments.
What I could not verify
- Whether
pipeline.forwardtreatsmax_sequence_length=Noneas its internal default. - Whether
GroupedGemmInputsHelperis defined outside the cutlass DSL guard. - Runtime reproducibility of the rebaked LPIPS golden under threshold on B200.
- Definition/behavior of
_preserve_lpips_candidate_on_failure(not in diff).
Automated review by NVCortex Lite, run by @crazydemo.
| num_frames=num_frames, | ||
| num_inference_steps=COSMOS3_LPIPS_NUM_INFERENCE_STEPS, | ||
| guidance_scale=COSMOS3_LPIPS_GUIDANCE_SCALE, | ||
| max_sequence_length=max_sequence_length, |
There was a problem hiding this comment.
[MAJOR] max_sequence_length=None is now always forwarded to pipeline.forward
Previously pipeline.forward(...) did not pass max_sequence_length at all, so the pipeline's own default applied. Now max_sequence_length=max_sequence_length is always passed, and the parameter defaults to None (line 148). This only works if pipeline.forward interprets None as "use my default". If instead it forwards None into the tokenizer/text-encoder as a real length, any default caller breaks. The concrete exposure is the V2V generation path, which calls _run_cosmos3_lpips_pipeline without max_sequence_length and now sends None — it is masked only because the V2V test is still waived (nvbugs/6437341). The comment at lines 159-160 documents the assumption but the diff does not show pipeline.forward honoring it. Confirm the pipeline treats None as default, or guard the call so the kwarg is omitted when None:
extra = {} if max_sequence_length is None else {"max_sequence_length": max_sequence_length}
result = pipeline.forward(..., **extra)| Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner, | ||
| Sm100BlockScaledContiguousGroupedGemmRunner, | ||
| Sm100BlockScaledContiguousGroupedGemmSwigluFusionRunner) | ||
| from ...custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper |
There was a problem hiding this comment.
[MINOR] Unrelated lazy-import change; verify GroupedGemmInputsHelper is not also guarded
This MoE import refactor is unrelated to the visual-gen bug this PR claims to fix and carries no test. The four Sm100 runners are moved into runner_tactic_comb_checker (lines 331-335) because they live inside cute_dsl_custom_ops' if IS_CUTLASS_DSL_AVAILABLE: block, but GroupedGemmInputsHelper is kept at module scope here. If GroupedGemmInputsHelper is ALSO defined inside that same guard, then on a machine without the cutlass DSL this module-scope import still raises at import time — and since create_moe imports this file eagerly, it breaks all model startup, defeating the entire fix. Confirm GroupedGemmInputsHelper is defined outside the guard (or has an else-branch fallback); if not, move it into the same lazy import.
| "tensorrt_llm_commit": "1745a6e689082e32e113fe06460ee5c3dfecbba9", | ||
| "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607271403-16694-handongl", | ||
| "rebake_reason": [ | ||
| "Refreshed for https://nvbugs/6418815. The prior golden (commit", |
There was a problem hiding this comment.
[MINOR] PR title/description contradict the actual change
The title says 'Revert the per-batch cross-attention slicing (drop real_text_lens plumbing in transformer_cosmos3.py)', but transformer_cosmos3.py is untouched and this rebake_reason correctly argues real_text_lens is the correct behavior and must NOT be reverted. The real fix is golden rebake + input pinning. Update the PR description to match, since a future reader (and git archaeology on this golden) will rely on it to understand why the golden moved.
crazydemo
left a comment
crazydemo
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: A waiver is removed and the golden (media + config) is rebaked, so a human should re-run the T2I LPIPS test on the target B200 arch to confirm reproducibility below the 0.05 threshold, and explicitly run the T2I stage since the default set may not include it (as reviewer brnguyen2 noted).
Automated review by NVCortex Lite, run by @crazydemo.
|
flaky LPIPS tests. Close. |
10 participants
Summary
Test plan
Links
Dev Engineer Review
real_text_lensplumbing fromtransformer_cosmos3.py.negative_promptandmax_sequence_lengthvalues.QA Engineer Review
_run_cosmos3_lpips_pipeline.test_cosmos3_nano_t2i_lpips_against_golden.tests/integration/test_lists/waives.txt.