[None][feat] VisualGen: Attention2D + Ulysses & Multi-GPU LPIPS Evals - #13944
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR introduces ring-overlapped column parallelism for Attention2D and unifies sequence sharding across visual generation models through a restructured VisualGenMapping. Key changes include a shared RingComm helper, Attention2D forward_overlapped path with online LSE merging, consolidated parallel wrapping in attention modules, and updated CLI tools to report all parallelism dimensions simultaneously. ChangesRing Attention and Unified Sequence Parallelism
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tensorrt_llm/_torch/visual_gen/pipeline.py (1)
1-4:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required SPDX/NVIDIA header.
This modified Python source still starts directly with imports, so it misses the repository-mandated copyright/SPDX header for 2026.
As per coding guidelines
**/*.{cpp,h,hpp,cu,cuh,py}: All C++, Python, and other source files must contain NVIDIA copyright header with current modification year.🤖 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/pipeline.py` around lines 1 - 4, This file (tensorrt_llm/_torch/visual_gen/pipeline.py) is missing the required NVIDIA/SPDX header for 2026; add the repository-mandated copyright/SPDX block as the very first lines of the file (above the existing imports such as itertools, os, time and the typing imports) so the header precedes all code and imports.tensorrt_llm/_torch/visual_gen/modules/attention.py (1)
1-3:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required SPDX/NVIDIA header.
This modified Python source file is missing the NVIDIA copyright/SPDX header with the current modification year.
As per coding guidelines
**/*.{cpp,h,hpp,cu,cuh,py}: All C++, Python, and other source files must contain NVIDIA copyright header with current modification year.🤖 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/modules/attention.py` around lines 1 - 3, This file lacks the required NVIDIA copyright/SPDX header; add the standard NVIDIA copyright and SPDX-License-Identifier header block at the very top of tensorrt_llm/_torch/visual_gen/modules/attention.py (above the imports) with the current modification year, following the project's header format used for other .py files so tools and license checks recognize it.tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py (1)
1-3:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required SPDX/NVIDIA header.
This modified Python source file is missing the NVIDIA copyright/SPDX header with the current modification year.
As per coding guidelines
**/*.{cpp,h,hpp,cu,cuh,py}: All C++, Python, and other source files must contain NVIDIA copyright header with current modification year.🤖 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/wan/transformer_wan.py` around lines 1 - 3, Add the required NVIDIA copyright/SPDX header at the top of tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py by inserting the project's standard multi-line comment block (including the current modification year and the SPDX-License-Identifier) as the very first lines of the file; ensure the header matches the project's canonical NVIDIA header template used across .py files so it includes the copyright holder, current year, and the SPDX tag.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py (1)
505-523: 💤 Low valueVerify deadlock avoidance with odd world sizes.
The even/odd send-first pattern avoids deadlock for power-of-2 ring sizes. For odd world sizes (e.g., 3 GPUs), the pattern may still work but is less common in practice. Consider adding a comment noting the expected world size constraints, or verify the pattern holds for odd sizes.
🤖 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/attention_backend/parallel.py` around lines 505 - 523, The neighbor exchange may deadlock for some non-power-of-two world sizes; update _ring_send_recv/_ring_wait to either assert or document acceptable world sizes and ensure deadlock avoidance: add a check using self.pg.size() (or equivalent) and either raise/handle when world_size is 1 or when an unsupported odd world size is detected, or add a concise comment explaining that the send-first pattern relies on even world sizes (or power-of-two) and that _send_first alternation is intended to prevent deadlock; reference the _ring_send_recv method, the _send_first flag, self.pg (process group) and the _p2p_reqs list when making the change.
🤖 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/visual_gen_flux.py`:
- Around line 366-370: parallel_str currently references undefined CLI fields
via getattr(args, 'cfg_size', 1) and getattr(args, 'ring_size', 1), so update
the code to use the real configured values instead of hardcoded defaults: after
parse_args() call and after build_diffusion_args(args) (or wherever
diffusion_args is created), read cfg_size and ring_size from
diffusion_args.parallel (or the appropriate fields on the object returned by
build_diffusion_args) and use those values in the parallel_str construction;
alternatively, if you prefer CLI options, add --cfg_size and --ring_size to
parse_args() and use args.cfg_size and args.ring_size when building
parallel_str. Ensure you reference the symbols parallel_str, parse_args,
build_diffusion_args, and diffusion_args.parallel so the change is applied in
the correct location.
In `@examples/visual_gen/visual_gen_ltx2.py`:
- Around line 386-389: The log string uses getattr(args, 'ring_size', 1) but
--ring_size is not defined, so either add a CLI argument or stop referencing it:
locate the parallel_str construction (symbol parallel_str) and either add an
argparse argument like add_argument('--ring_size', type=int, default=1) to the
args setup (symbol args) so args.ring_size exists, or replace/remove the
Ring(size=...) segment from parallel_str so it doesn't reference a non-existent
option; ensure whichever path you choose keeps the string formatting consistent
with the other size fields (CFG, Attention2D, Ulysses).
In `@examples/visual_gen/visual_gen_wan_i2v.py`:
- Around line 313-316: The log string uses getattr(args, 'ring_size', 1) but
parse_args() in this file does not define --ring_size, so add a proper CLI
argument to the parser to make the logged value accurate: update the argument
parser in parse_args() (or where arguments are defined) to include
parser.add_argument('--ring_size', type=int, default=1, help='ring size for
Ulysses/Ring component') and then use args.ring_size in the parallel_str
construction (the existing parallel_str variable can remain the same but replace
getattr with args.ring_size) so behavior matches visual_gen_wan_t2v.py and the
log reflects the configured value.
In `@examples/visual_gen/visual_gen_wan_t2v.py`:
- Around line 306-312: The log currently hard-codes num_heads = 40 which can be
incorrect for different Ulysses variants; instead derive the head count
dynamically and use that value in the logger call — e.g., obtain the number of
heads from the model/config (or from a mapping keyed by checkpoint/model size)
and replace the literal 40 with that derived variable before computing heads per
GPU; update the block using args.ulysses_size, num_heads, and logger.info so the
message reflects the actual num_heads used at runtime.
- Around line 314-317: The banner now advertises the combined topology built
into parallel_str, but CLI help still claims those combinations are
unimplemented; either update the argparse help text for --ulysses_size,
--attn2d_row_size, and --attn2d_col_size to state the combination (Attention2D +
Ulysses + Ring) is supported, or reintroduce validation in the argument parsing
logic (the code that reads args and builds parallel_str) to detect disallowed
combinations and raise an error; modify the help strings on the
parser.add_argument calls for ulysses_size/attn2d_* or add a guard that checks
args.cfg_size/args.attn2d_row_size/args.attn2d_col_size/args.ulysses_size and
exits with a clear message to keep CLI messaging consistent with parallel_str.
In `@tensorrt_llm/_torch/visual_gen/mapping.py`:
- Around line 17-27: This file is missing the required NVIDIA SPDX header; add a
header block at the very top of tensorrt_llm/_torch/visual_gen/mapping.py
containing the current modification year and the NVIDIA copyright notice plus
the SPDX license identifier (e.g., SPDX-FileCopyrightText: Copyright 2026 NVIDIA
CORPORATION and SPDX-License-Identifier: <license>) so the file begins with the
required SPDX copyright/license lines before the imports and definitions like
DeviceMeshTopologyImpl, SingleProcessGroup, _DEVICE_MESH_DIM_ORDER_LEGACY, and
_DEVICE_MESH_DIM_ORDER_ATTN2D.
- Around line 210-237: The cached-mesh path skips initializing the derived
seq_mesh and Attn2D groups causing stale topology for new VisualGenMapping
instances; update build_mesh (and the early-return branch that checks
DeviceMeshTopologyImpl.device_mesh) to still compute and assign
VisualGenMapping.seq_mesh (using the existing cp/cp_row/cp_col/ulysses logic)
and call _attach_attn2d_groups_from_device_mesh() when the device_mesh is reused
so SequenceSharder.from_vgm and collectives see correct groups; ensure you
reference and reuse the same checks and index logic currently around
VisualGenMapping.seq_mesh and self._use_attn2d_plane so the code path is
executed both for fresh and cached device_mesh scenarios.
In `@tensorrt_llm/_torch/visual_gen/utils.py`:
- Around line 3-11: The file tensorrt_llm._torch.visual_gen.utils is missing the
required NVIDIA SPDX header; add the standard two-line header at the very top of
the module (before any imports) including the current modification year (2026)
and the SPDX license identifier (e.g., "SPDX-FileCopyrightText: Copyright (c)
2026 NVIDIA CORPORATION & AFFILIATES" and "SPDX-License-Identifier:
Apache-2.0"), so the header appears above the existing symbols such as the
module imports and the VisualGenMapping import.
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py`:
- Around line 15-17: This test file
(tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py) is
missing the required NVIDIA SPDX/copyright header; add the appropriate
SPDX-License-Identifier and the NVIDIA copyright header with the current
modification year at the very top of the file (above imports and any shebangs),
following the project's header style so the file includes the required license
lines before the VisualGenMapping import and any other code.
In `@tests/unittest/_torch/visual_gen/test_utils.py`:
- Around line 126-135: The test fails early because SequenceSharder.from_vgm
rejects seq_size>1 when seq_group is None; update the fake vgm in
test_from_vgm_head_divisibility to set seq_group to a non-None sentinel (e.g., 0
or an object) so the code bypasses the seq_group guard and reaches the
num_attention_heads divisibility check invoked by SequenceSharder.from_vgm; keep
other fields (seq_size, seq_rank, ulysses_size) unchanged since gather()
behavior is not being tested here.
---
Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py`:
- Around line 1-3: Add the required NVIDIA copyright/SPDX header at the top of
tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py by inserting the
project's standard multi-line comment block (including the current modification
year and the SPDX-License-Identifier) as the very first lines of the file;
ensure the header matches the project's canonical NVIDIA header template used
across .py files so it includes the copyright holder, current year, and the SPDX
tag.
In `@tensorrt_llm/_torch/visual_gen/modules/attention.py`:
- Around line 1-3: This file lacks the required NVIDIA copyright/SPDX header;
add the standard NVIDIA copyright and SPDX-License-Identifier header block at
the very top of tensorrt_llm/_torch/visual_gen/modules/attention.py (above the
imports) with the current modification year, following the project's header
format used for other .py files so tools and license checks recognize it.
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 1-4: This file (tensorrt_llm/_torch/visual_gen/pipeline.py) is
missing the required NVIDIA/SPDX header for 2026; add the repository-mandated
copyright/SPDX block as the very first lines of the file (above the existing
imports such as itertools, os, time and the typing imports) so the header
precedes all code and imports.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py`:
- Around line 505-523: The neighbor exchange may deadlock for some
non-power-of-two world sizes; update _ring_send_recv/_ring_wait to either assert
or document acceptable world sizes and ensure deadlock avoidance: add a check
using self.pg.size() (or equivalent) and either raise/handle when world_size is
1 or when an unsupported odd world size is detected, or add a concise comment
explaining that the send-first pattern relies on even world sizes (or
power-of-two) and that _send_first alternation is intended to prevent deadlock;
reference the _ring_send_recv method, the _send_first flag, self.pg (process
group) and the _p2p_reqs list when making the change.
🪄 Autofix (Beta)
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: 26822c0a-a923-4af2-bd1b-6067214ea819
📒 Files selected for processing (23)
examples/visual_gen/visual_gen_flux.pyexamples/visual_gen/visual_gen_ltx2.pyexamples/visual_gen/visual_gen_wan_i2v.pyexamples/visual_gen/visual_gen_wan_t2v.pytensorrt_llm/_torch/visual_gen/attention_backend/__init__.pytensorrt_llm/_torch/visual_gen/attention_backend/interface.pytensorrt_llm/_torch/visual_gen/attention_backend/parallel.pytensorrt_llm/_torch/visual_gen/config.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.pytensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.pytensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.pytensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.pytensorrt_llm/_torch/visual_gen/modules/attention.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytensorrt_llm/_torch/visual_gen/utils.pytests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.pytests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.pytests/unittest/_torch/visual_gen/multi_gpu/test_wan_ring.pytests/unittest/_torch/visual_gen/test_utils.pytests/unittest/_torch/visual_gen/test_visual_gen_args.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/visual_gen/pipeline_loader.py
NVShreyas
force-pushed
the
user/shreyasm/attn2d-ulysses
branch
2 times, most recently
from
May 19, 2026 19:47
dee0527 to
06f06ab
Compare
NVShreyas
reviewed
May 19, 2026
NVShreyas
force-pushed
the
user/shreyasm/attn2d-ulysses
branch
from
May 21, 2026 01:06
65176f7 to
6b0bf1f
Compare
NVShreyas
marked this pull request as ready for review
May 21, 2026 01:06
NVShreyas
changed the title
Collaborator
|
@@coderabbitai summary |
Collaborator
|
@coderabbitai full review |
Contributor
✅ Actions performedFull review triggered. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py (1)
398-465:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRefresh the inner-backend kwargs after the gathers.
Both Attention2D paths change the effective Q/KV sequence lengths before calling
inner_backend.forward_with_lse(), but they still forward the pre-gatherseq_len/seq_len_kvthat came fromAttention._attn_impl(). The non-overlapped path also normalizesattention_maskand then drops it entirely. Any backend that reshapes, validates, or selects kernels from those kwargs will see inconsistent inputs here.🩹 Suggested fix
seq_len = q.shape[1] + kv_seq_len = k.shape[1] + + kwargs["batch_size"] = B + kwargs["seq_len"] = seq_len + kwargs["seq_len_kv"] = kv_seq_len if self._inner_layout == AttentionTensorLayout.HND: q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) - output, lse = self.inner_backend.forward_with_lse(q=q, k=k, v=v, **kwargs) + output, lse = self.inner_backend.forward_with_lse( + q=q, + k=k, + v=v, + attention_mask=attention_mask, + **kwargs, + ) @@ self._ensure_buffers(q, k) + seq_len = q.shape[1] + kv_seq_len = k.shape[1] + kwargs["batch_size"] = B + kwargs["seq_len"] = seq_len + kwargs["seq_len_kv"] = kv_seq_len kv_bufs = self._kv_ring_bufsAlso applies to: 486-545
🤖 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/attention_backend/parallel.py` around lines 398 - 465, The non-overlapped and overlapped forward paths mutate q/k/v (changing effective seq lengths) but continue to pass the original kwargs (seq_len/seq_len_kv and possibly an attention_mask) into inner_backend.forward_with_lse; update/refresh the kwargs immediately after the all-gather/reshape steps (in Attention2DAttention.forward and forward_overlapped) so that seq_len and seq_len_kv reflect q.shape[1] and k.shape[1] respectively, and ensure attention_mask is normalized or removed/reshaped to match the new per-process sequence lengths before calling inner_backend.forward_with_lse; adjust the dict/object you pass to inner_backend.forward_with_lse (the keyword args variable) rather than relying on the pre-gather values.examples/visual_gen/visual_gen_ltx2.py (1)
381-392:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winFix SyntaxError and remove obsolete Attn2D+Ulysses rejection in main()
main()hasraise ValueError(starting at line 383 with no closing)(SyntaxError: '(' was never closed), so the script won’t run.- The stale guard
if attn2d_size > 1 and args.ulysses_size > 1:still rejects combining--attn2d_row_size/--attn2d_col_sizewith--ulysses_size.Suggested fix
- if attn2d_size > 1 and args.ulysses_size > 1: - raise ValueError( - "Combining --ulysses_size with --attn2d_row_size/--attn2d_col_size is not yet implemented." - if bool(args.spatial_upsampler_path) != bool(args.distilled_lora_path): missing = ( "--distilled_lora_path" if args.spatial_upsampler_path else "--spatial_upsampler_path" ) raise ValueError( f"Two-stage pipeline requires both --spatial_upsampler_path and " f"--distilled_lora_path, but {missing} was not provided." )🤖 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 `@examples/visual_gen/visual_gen_ltx2.py` around lines 381 - 392, There are two problems in main(): the ValueError for attn2d_size is missing a closing parenthesis causing a SyntaxError, and the stale guard that rejects combining attn2d_size with ulysses_size is no longer desired. Fix by closing the raise ValueError(...) block for the attn2d_size check (ensure the string literal and trailing parenthesis are present) and remove (or comment out) the entire if block that checks "if attn2d_size > 1 and args.ulysses_size > 1:" so combining --attn2d_row_size/--attn2d_col_size with --ulysses_size is allowed; keep the subsequent two-stage pipeline check that validates spatial_upsampler_path vs distilled_lora_path as-is.
♻️ Duplicate comments (1)
tensorrt_llm/_torch/visual_gen/mapping.py (1)
208-219:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRebuild per-instance Attention2D groups on cache hits.
After the cache-hit
returnon Line 219, the new instance never executes_attach_attn2d_groups_from_device_mesh(), soattn2d_row_groupandattn2d_col_groupstayNoneeven though the shareddevice_meshis valid. Once cache reuse is fixed, the second Attention2D mapping in a process will still break row/col collectives unless this branch reattaches the per-instance groups before returning.Also applies to: 258-259
🤖 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/mapping.py` around lines 208 - 219, The cached-device_mesh early-return skips per-instance reconstruction of Attention2D groups causing attn2d_row_group/attn2d_col_group to remain None; before returning from the cache-hit branch in VisualGenMapping.build_mesh (when cls.device_mesh is reused), call self._attach_attn2d_groups_from_device_mesh() to reattach the per-instance groups so row/col collectives are valid; apply the same fix to the other cache-return branch that similarly reuses cls.device_mesh (the branch referenced around the second return) so every instance reconstructs its attn2d groups even on cache hits.
🧹 Nitpick comments (1)
tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py (1)
518-520: No QA test-list update is needed for this change set.These updates are in
tests/unittest/...scope and do not add/alter integration tests undertests/integration/defs/.As per coding guidelines, “If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional.”
🤖 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_visual_gen_mapping.py` around lines 518 - 520, Add an explicit note in the PR description or commit message stating that QA test-list updates are unnecessary since this change only touches unit tests under tests/unittest (e.g., the new/modified test function test_attn2d_ulysses_seq_rank_matches_global_rank and its call to _run_multi_gpu) and does not modify integration tests under tests/integration/defs; update the PR body to include a single clear sentence like "No QA test-list update required: changes are limited to tests/unittest only."
🤖 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/visual_gen_wan_i2v.py`:
- Around line 365-368: The banner string built in parallel_str claims an
"Attention2D + Ulysses" combo but the code still blocks that combination with
the guard using attn2d_size and args.ulysses_size; update parallel_str
construction to reflect the actual allowed topology by conditionally including
the Attention2D and Ulysses parts only when the corresponding guard passes
(e.g., check attn2d_size > 1 and args.ulysses_size > 1 before appending the
f"Attention2D(...)+ Ulysses(...)" segment) or omit/replace the Ulysses portion
when the guard disallows it so the printed banner always matches the runtime
constraints.
In `@tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py`:
- Around line 403-415: The forward and forward_overlapped methods currently
reject None even though the class docstring treats None as the FULL-attention
sentinel; change the validation to accept either PredefinedAttentionMask.FULL or
None as equivalent full-mask values. Specifically, update the check in forward
(method forward(q, k, v, attention_mask: PredefinedAttentionMask =
PredefinedAttentionMask.FULL, ...)) and the analogous check in
forward_overlapped so they do not raise when attention_mask is None (i.e., allow
attention_mask in (PredefinedAttentionMask.FULL, None)). Ensure any error
message still reports unexpected values and reference
PredefinedAttentionMask.FULL, attention_mask, forward, and forward_overlapped
when making the change.
In `@tensorrt_llm/_torch/visual_gen/mapping.py`:
- Around line 160-178: The code removed setting self._order causing
_validate_vae_ranks_share_cfg_group() to raise AttributeError; restore a
replacement by assigning self._order = dim_order in both branches (when
self._use_attn2d_plane is True and False) before any validation runs so
downstream methods like _validate_vae_ranks_share_cfg_group() can read the
expected order string; ensure this assignment happens just after computing
dim_order and before setting self._dim_names/self._dim_sizes.
- Around line 161-163: Normalize _dim_names to a tuple when assigned (replace
dim_order.split("-") list with a tuple) so cached_dim_names comparison in
VisualGenMapping.build_mesh succeeds; when a cached mesh is reused, call or
inline the logic from _attach_attn2d_groups_from_device_mesh to set
self._attn2d_row_group and self._attn2d_col_group on cache-hit paths as well;
and update the error message in _validate_vae_ranks_share_cfg_group to reference
an existing attribute (e.g. self._dim_names) instead of the non-existent
self._order so raising NotImplementedError yields the intended message.
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py`:
- Around line 304-305: The docstring contains the Unicode multiplication sign
"×" which triggers RUF002 lint warnings; edit the docstring line that mentions
Mesh ``cfg-tp-cp_row-cp_col-ulysses`` and the attn2d size "2×2" to use ASCII "x"
instead (i.e., change "2×2" to "2x2" and any other occurrences of "×" in that
docstring), preserving the surrounding text and formatting so tests and
descriptions remain the same.
---
Outside diff comments:
In `@examples/visual_gen/visual_gen_ltx2.py`:
- Around line 381-392: There are two problems in main(): the ValueError for
attn2d_size is missing a closing parenthesis causing a SyntaxError, and the
stale guard that rejects combining attn2d_size with ulysses_size is no longer
desired. Fix by closing the raise ValueError(...) block for the attn2d_size
check (ensure the string literal and trailing parenthesis are present) and
remove (or comment out) the entire if block that checks "if attn2d_size > 1 and
args.ulysses_size > 1:" so combining --attn2d_row_size/--attn2d_col_size with
--ulysses_size is allowed; keep the subsequent two-stage pipeline check that
validates spatial_upsampler_path vs distilled_lora_path as-is.
In `@tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py`:
- Around line 398-465: The non-overlapped and overlapped forward paths mutate
q/k/v (changing effective seq lengths) but continue to pass the original kwargs
(seq_len/seq_len_kv and possibly an attention_mask) into
inner_backend.forward_with_lse; update/refresh the kwargs immediately after the
all-gather/reshape steps (in Attention2DAttention.forward and
forward_overlapped) so that seq_len and seq_len_kv reflect q.shape[1] and
k.shape[1] respectively, and ensure attention_mask is normalized or
removed/reshaped to match the new per-process sequence lengths before calling
inner_backend.forward_with_lse; adjust the dict/object you pass to
inner_backend.forward_with_lse (the keyword args variable) rather than relying
on the pre-gather values.
---
Duplicate comments:
In `@tensorrt_llm/_torch/visual_gen/mapping.py`:
- Around line 208-219: The cached-device_mesh early-return skips per-instance
reconstruction of Attention2D groups causing attn2d_row_group/attn2d_col_group
to remain None; before returning from the cache-hit branch in
VisualGenMapping.build_mesh (when cls.device_mesh is reused), call
self._attach_attn2d_groups_from_device_mesh() to reattach the per-instance
groups so row/col collectives are valid; apply the same fix to the other
cache-return branch that similarly reuses cls.device_mesh (the branch referenced
around the second return) so every instance reconstructs its attn2d groups even
on cache hits.
---
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py`:
- Around line 518-520: Add an explicit note in the PR description or commit
message stating that QA test-list updates are unnecessary since this change only
touches unit tests under tests/unittest (e.g., the new/modified test function
test_attn2d_ulysses_seq_rank_matches_global_rank and its call to _run_multi_gpu)
and does not modify integration tests under tests/integration/defs; update the
PR body to include a single clear sentence like "No QA test-list update
required: changes are limited to tests/unittest only."
🪄 Autofix (Beta)
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: 7b203e82-74d8-402d-85fc-d93ee65357e1
📒 Files selected for processing (13)
examples/visual_gen/visual_gen_flux.pyexamples/visual_gen/visual_gen_ltx2.pyexamples/visual_gen/visual_gen_wan_i2v.pyexamples/visual_gen/visual_gen_wan_t2v.pytensorrt_llm/_torch/visual_gen/attention_backend/parallel.pytensorrt_llm/_torch/visual_gen/config.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.pytensorrt_llm/_torch/visual_gen/modules/attention.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/utils.pytests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.pytests/unittest/_torch/visual_gen/test_visual_gen_args.py
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49532 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49532 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49716 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49716 [ run ] completed with state
|
Collaborator
|
PR_Github #51399 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51504 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51504 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51568 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51568 [ run ] completed with state
|
Collaborator
CI triage: the bulk of the failures are a test-env leak in this PR (not infra)I dug into the red CI on this PR (latest run L0_MergeRequest_PR #40909, commit
This is PR-specific, not a fleet/infra problem. Contemporaneous L0 builds from other PRs in the same window are clean (0 Ray failures, several full SUCCESS runs), while this exact cluster reproduces on every build of this PR (#40688, #40810, #40909) and survived re-triggering the same commit. Root cause #1 — module-scope env mutation leaks across the pytest workerThree new test files set
import os
os.environ["TLLM_DISABLE_MPI"] = "1" # module scope → runs during pytest collectionBecause pytest imports every test module during collection, this pollutes the shared worker environment. Every later test that constructs Suggested fix: don't mutate global env at module scope. Set Root cause #2 — three genuine VisualGen test failures (independent of the leak)
Fixing #1 should clear the ~3k cascade and let the real signal through; #2 are small, localized fixes in the new VisualGen tests / 🤖 Automated CI triage |
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51623 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51623 [ run ] completed with state
|
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
/bot kill |
Collaborator
|
PR_Github #51688 [ run ] triggered by Bot. Commit: |
Collaborator
|
/bot kill |
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51695 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51688 [ run ] completed with state |
Collaborator
|
PR_Github #51695 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #51782 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51782 [ run ] completed with state |
This was referenced Jun 9, 2026
luyiyun1021
added a commit
to luyiyun1021/TensorRT-LLM
that referenced
this pull request
Jun 10, 2026
…kwarg The enable_ulysses kwarg was a rebase artifact in PR NVIDIA#13978. Before that PR merged, base Attention had only enable_sequence_parallel as the single sequence-parallel master switch (Attention2D PR NVIDIA#13944 had already consolidated the kwarg surface). NVIDIA#13978's rebase resolution kept the older enable_ulysses kwarg in a union with enable_sequence_parallel, re-introducing the redundant gate. Only LTX2Attention explicitly forwarded enable_ulysses=use_ulysses, and it always co-varied with enable_sequence_parallel (line 591 self-attn: both True; line 605 cross-attn: both False). No model uses the SP-on-but-Ulysses-off escape hatch the kwarg notionally provided. Changes: - Drop enable_ulysses kwarg from base Attention.__init__. - Simplify the use_ulysses local: gate on (ulysses_size > 1 and enable_sequence_parallel and (qkv_mode != SEPARATE_QKV or async_ulysses)). - Drop the enable_ulysses=use_ulysses forwarding from LTX2Attention.__init__. LTX2Attention's own use_ulysses kwarg is preserved -- it still drives async_ulysses gating (line 109) and the audio dual-attn (_has_dual_attn) setup (line 176) independent of the base. Verified locally: 6/6 test_ltx2_attention.py + 14/14 test_ltx2_transformer.py pass on the parent k_pe-drop branch which carries the same Attention surface. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes
New Features
RingAttentionimplementation for distributed attention computation.Improvements
Description
Test Coverage
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)
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.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.