[TRTLLM-14502][feat] LTX-2 two-stage: dual-topology parallel Stage 2 (cfg folds into ulysses) - #16502
Merged
luyiyun1021 merged 31 commits intoJul 28, 2026
Merged
Conversation
…ement stage2_ulysses (pipeline_config) keeps Ulysses enabled through Stage 2 and runs _upsample_and_refine on every rank instead of rank 0 only. Requires parallel VAE (decode_latents already leaves Stage-1 video+audio latents on every vae_rank), cfg_size == 1 so the Ulysses group spans all ranks (Stage 2 has no CFG), and ulysses_size > 1; otherwise falls back to rank-0 Stage 2 with a warning. All ranks seed the refinement noise identically, so the collective output matches the rank-0 path and the pre-decode broadcast is skipped. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
_GpuTimeline records named cuda.Events at every two-stage phase boundary (stage1, upsample, lora_bind, stage2_denoise, lora_restore, video_decode, audio_decode) and rank 0 logs the consecutive deltas plus the stage-1 pre/denoise/post breakdown from CudaPhaseTimer. Events complete in stream order, so the deltas are true GPU-timeline durations even for phases the host only enqueues (CUDA-graph stage 2 denoise measured 0.02s by host wall). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…me attention type
Delete the rank-0-only serial Stage-2 path and all runtime ulysses-vs-plain
switching:
- set_ulysses_enabled / set_ulysses_active / is_ulysses_active() and the
{_ulysses_attn, _plain_attn} dual-attn pair are removed; every module's
attention TYPE is fixed at construction. audio_attn1 builds plain attention
under CONDITIONAL (default) and Ulysses under legacy FULL (env constant).
- The AV cross-attn dispatch reads a construction-time is_ulysses constant.
- Stage 2 runs collectively on every rank: the stage2_ulysses knob (registry,
extra_attrs whitelist, gate), the non-primary early-exit and the pre-decode
latent broadcast are deleted. The internal stage-1 latent handoff keeps
latents on every rank (_latents_on_all_ranks; the decode_latents rank gate
applies only to real VAE decode), so the stage boundary needs no collective
in either parallel-VAE mode. Non-parallel-VAE final decode stays rank-0-only.
- configure_audio_ulysses slims to pad computation.
- Absorb draft PR NVIDIA#16301: v2a cross-attention runs distributed
Attention2D + Ulysses when attn2d is active (use_ulysses no longer skips
SEPARATE_QKV; ring + SEPARATE_QKV remains rejected).
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…ned ulysses groups
Stage 1 runs the configured cfg x ulysses layout; at the Stage-2 boundary the
transformer switches to the stage-2 topology and back, under a try/finally.
- Topology rule: cfg is flattened into ulysses (stage-2 ulysses size =
cfg_size x ulysses_size) while every other mesh dim's groups are preserved
verbatim, so cp-side (ring/attn2d) communicators and peers are shared by
both topologies. VisualGenMapping.flatten_cfg_ranks() supplies the layout
arithmetic; the groups are built and owned by the two-stage pipeline at
load (raw dist.new_group, all-rank collective). At cp>1 a cp-major
seq-plane group is added; dist.new_group sorts its rank list, so the
stage-2 SequenceSharder carries the shard order via seq_rank/gather_index.
- {default, stage2} attention stacks are both built at construction on
attn1, video_to_audio_attn and FULL-mode audio_attn1 (fresh inner backend
for heads/stage2-group-size and a fresh cp wrapper over the same group);
set_ulysses_topology() atomically flips the active sharder and stacks.
Single-stage pipelines and other models pass no stage-2 groups and build
nothing extra; at cfg_size == 1 the pipeline skips the switch entirely.
- CUDA-graph isolation: the two-stage graph key gains a topology string next
to the LoRA state (matrix test covers all four combinations); a plain int
read in LTXModel.forward gives torch.compile a per-topology value guard.
- Audio is padded once to the stage-2 multiple (valid for both topologies,
no repad at the boundary); full-resolution latent token count is validated
against the stage-2 seq plane at request entry (fail fast, no fallback).
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…; tests Codex audit P0: the Stage-2 exit path flipped the graph-key topology back to default before the transformer stacks were restored — an exception inside the restore would leave stage-2 stacks live under the default key (the runtime-swap incident class). Restore now mirrors the entry order inversely, so the key reads stage2 at every instant the stage-2 stacks might be live. Tests: flatten_cfg_ranks layout cases; construction gates (no stage-2 groups -> no extra stacks at any cfg, switch raises); (lora x topology) graph-key matrix; 4-GPU dual-topology alternation with shard/gather round-trip and full-forward numerical equivalence vs the single-GPU reference in both topologies. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…al topology
4-GPU: with the AudioShardMode constant forced to FULL, audio_attn1 must be
CONSTRUCTED as ulysses (no runtime selector exists) and, under cfg2 with
stage-2 groups, gets its own distinct {default, stage2} stack pair; forwards
complete in both topologies across a switch cycle.
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…ve tests Round-0 review gaps: the topology-mismatch failure class was guarded by convention only. Make each guard mechanical and cover it with a negative: - TextCache gains a topology marker stamped by prepare_text_cache() and validated at LTXModel.forward() entry (PE shards/pads are topology-dependent); the CUDA-graph static clone propagates it. - _LTX2TwoStageCUDAGraphRunner cross-checks the graph key's topology element against the transformer's active topology before every capture and replay. - New negatives: stale-topology cache (single-GPU fabricated + 4-GPU prepare-before-switch), wrong-key capture/replay, warmup dual-topology precapture with zero post-warmup recapture, stage-2 token divisibility (request entry), stage-2 head divisibility (construction, 4-GPU). - Remove dead SequenceSharder.disable()/enable() (no callers since the serial Stage-2 deletion) and their _enabled flag. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…duction _run_warmup The warmup pre-capture test is now driven by the inherited LTX2Pipeline._run_warmup (forward stubbed to model the two-stage boundary), so a regression where startup warmup stops crossing the Stage-2 boundary fails the test; the manually-driven variant remains as runner coverage. _check_topology also raises on a graph key missing the topology element instead of returning silently. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…remove stale sharder test Remove the TextCache.topology field, its forward() ordering check, the _clone_value propagation, and both stale-cache negatives (user decision: a validation-only field on the dataclass is not worth the intrusion). Ordering correctness rests on the single ordered call site, the 4-GPU alternation test, and the CUDA-graph runner topology validator, which is independent of TextCache and unchanged. Also delete test_disable_enable_no_collectives, left calling the removed SequenceSharder.disable()/enable(). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…ive-stack is_ulysses, TP-local stage-2 heads
- Gate the tile-parallel Stage-2 decode to vgm.vae_ranks (mirrors
decode_latents ownership); with parallel_vae_size < world_size the
branch previously ran on every rank, raising on non-VAE ranks (pg=None)
while VAE ranks sat in collectives.
- Recompute is_ulysses in set_active_attn(): the {default, stage2} stacks
can differ in type (cfg2 x u1 -> stage-2 ulysses), and the stale
construction-time value made v2a all-gather K/V on top of the ulysses
wrapper's own gather. Alternation test asserts flag/stack consistency.
- Size the stage-2 inner backend and its divisibility check from TP-local
head counts, matching the default stack.
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…tage-2 seq groups stay per TP fiber - Cross-attn (v2a) under attn2d WITHOUT ulysses is no longer wrapped in Attention2D: the AV dispatch keys the seq-sharded K/V path off the Ulysses wrapper, so an attn2d-only wrapper received all-gathered full K/V on top of its own sharded expectation. attn2d cross-attn now requires ulysses on top; otherwise plain backend + all-gather. - VisualGenMapping.flatten_cfg_seq_ranks(): seq-plane rank lists per tp coordinate (cp-major, then cfg, then ulysses). The two-stage group builder previously concatenated fold lists ACROSS TP fibers into one world seq group, sharding activations across ranks holding different weight shards; it now builds one seq group per tp fiber with a group-local gather_index. tp=1 behavior is bit-identical. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…ver padding is possible The audio pad multiple is the stage-2 seq size when stage-2 groups exist, so a cfg-only fold (cfg2 x u1) can pad audio even at ulysses_size == 1 — and TRTLLM self-attn silently drops the key_padding_mask. The VANILLA downgrade now triggers when ulysses > 1 OR stage-2 groups are present. 2-GPU construction test asserts the downgrade (and that video attn1 keeps TRTLLM). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59728 [ run ] triggered by Bot. Commit: |
Contributor
📝 WalkthroughWalkthroughLTX2 now supports default and stage-2 CFG/Ulysses topologies, including topology-specific attention stacks, sequence sharding, collective refinement, CUDA graph isolation, rank-aware decoding, GPU timing, and distributed validation. ChangesLTX2 dual-topology execution
Estimated code review effort: 5 (Critical) | ~120 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Pipeline
participant Transformer
participant CUDAGraphs
participant Decoder
Pipeline->>Transformer: execute stage 1 with latents on all ranks
Pipeline->>Transformer: switch to stage2 topology
Pipeline->>CUDAGraphs: capture or replay topology-keyed graph
CUDAGraphs->>Transformer: execute stage2 denoising
Pipeline->>Decoder: decode on selected ranks
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py (1)
123-148: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winInclude stage-2-only Ulysses in construction eligibility.
For
cfg_size > 1with configuredulysses_size == 1, the stage-2 group is non-trivial, but these gates still disable async Ulysses and Attention2D cross-attention wrapping. Consequently, the stage-2 v2a stack aliases the fallback path instead of constructing the intended Attention2D → Ulysses stack.Derive a shared effective Ulysses flag from both the default size and
stage2_ulysses_group.Proposed fix
+ stage2_ulysses_size = ( + torch_dist.get_world_size(group=stage2_ulysses_group) + if stage2_ulysses_group is not None + else 1 + ) + has_ulysses = ulysses_size > 1 or stage2_ulysses_size > 1 self._use_async_ulysses = bool( - enable_sequence_parallel and async_ulysses and vgm is not None and vgm.ulysses_size > 1 + enable_sequence_parallel and async_ulysses and has_ulysses ) ... enable_sp = enable_sequence_parallel and ( - cp_size == 1 or (attn2d_active and ulysses_size > 1) + cp_size == 1 or (attn2d_active and has_ulysses) )🤖 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/ltx2/transformer_ltx2.py` around lines 123 - 148, The Ulysses eligibility gates only consider vgm.ulysses_size and miss non-trivial stage2_ulysses_group configurations. In the initialization logic around _use_async_ulysses and cross-attention enable_sp, derive one effective Ulysses-enabled flag from the configured Ulysses size and stage2_ulysses_group, then use it for async-Ulysses and Attention2D cross-attention eligibility so cfg_size > 1 stage-2 groups construct the intended Attention2D → Ulysses stack.
🧹 Nitpick comments (4)
tensorrt_llm/_torch/visual_gen/utils.py (1)
5-5: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the repository’s Python 3.10 typing style.
Declare this as
gather_index: list[int] | None = Noneand remove the newListimport.As per coding guidelines, “prefer built-in generic types and
|.”Also applies to: 68-74
🤖 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/utils.py` at line 5, Update the gather_index declaration to use Python 3.10 built-in typing as list[int] | None, and remove the now-unused List import from the typing imports. Apply the same typing style to the corresponding declaration around the referenced additional location.Source: Coding guidelines
tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py (1)
599-611: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType the new topology contract precisely across its producer, carrier, and consumer.
tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py#L599-L611: usetorch_dist.ProcessGroupandlist[int] | NoneinStage2Groups.tensorrt_llm/_torch/visual_gen/mapping.py#L396-L424: returnlist[list[int]]from both rank-layout helpers.tensorrt_llm/_torch/visual_gen/utils.py#L5-L74: uselist[int] | Noneforgather_indexand removeList.As per coding guidelines, avoid
Any, annotate functions precisely, and prefer built-in generics with|.🤖 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/ltx2/transformer_ltx2.py` around lines 599 - 611, Tighten the topology types across all affected sites: in transformer_ltx2.py, annotate Stage2Groups.ulysses_group and seq_group with torch_dist.ProcessGroup and use list[int] | None for gather_index; in mapping.py, annotate both rank-layout helpers as returning list[list[int]]; in utils.py, use list[int] | None for gather_index and remove the List import. Avoid Any and use built-in generics with | throughout.Source: Coding guidelines
tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py (1)
359-388: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise the production group builder instead of copying it.
_build_stage2_groups_for_testduplicatesLTX2TwoStagesPipeline._stage2_transformer_groups, so these distributed tests can pass while production wiring drifts. Invoke the production method on a minimal pipeline namespace instead.As per path instructions, coverage of the actual production group builder is currently insufficient.
🤖 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_ltx2_ulysses.py` around lines 359 - 388, Replace the duplicated group-construction logic in _build_stage2_groups_for_test with a call to LTX2TwoStagesPipeline._stage2_transformer_groups using a minimal pipeline namespace configured with the test’s vgm. Preserve the helper’s returned Stage2Groups behavior while ensuring the tests execute the production builder directly.Source: Path instructions
tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py (1)
1468-1468: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd direct coverage for the new rank-ownership paths.
These paths bypass the centralized
decode_latentsownership contract, but the supplied tests do not exercise their behavior by rank.
tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py#L1468-L1468: test the internal flag on a non-primary rank.tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py#L2076-L2088: verify local video/audio unpatchification versus default rank gating.tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py#L1409-L1445: verify serial and parallel-VAE decode ownership, including rank-0-only audio.As per path instructions, add these cases to
tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py; current coverage is insufficient.🤖 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/ltx2/pipeline_ltx2.py` at line 1468, Add direct rank-aware tests in tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py for all affected paths: cover _latents_on_all_ranks on a non-primary rank; verify local video/audio unpatchification versus default rank gating in pipeline_ltx2.py:2076-2088; and verify serial and parallel-VAE decode ownership, including rank-0-only audio, in pipeline_ltx2_two_stages.py:1409-1445. Ensure each test asserts the expected per-rank behavior.Source: Path instructions
🤖 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_ltx2_transformer.py`:
- Around line 290-298: Extend the distributed CFG/Ulysses test in the multi-GPU
test module to assert that the constructed block’s audio_attn1 remains unwrapped
in CONDITIONAL mode. Use the existing block variable and Ulysses-enabled setup
there, rather than relying on test_conditional_audio_attn1_is_plain with
ulysses_size=1.
---
Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py`:
- Around line 123-148: The Ulysses eligibility gates only consider
vgm.ulysses_size and miss non-trivial stage2_ulysses_group configurations. In
the initialization logic around _use_async_ulysses and cross-attention
enable_sp, derive one effective Ulysses-enabled flag from the configured Ulysses
size and stage2_ulysses_group, then use it for async-Ulysses and Attention2D
cross-attention eligibility so cfg_size > 1 stage-2 groups construct the
intended Attention2D → Ulysses stack.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py`:
- Line 1468: Add direct rank-aware tests in
tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py for all affected paths:
cover _latents_on_all_ranks on a non-primary rank; verify local video/audio
unpatchification versus default rank gating in pipeline_ltx2.py:2076-2088; and
verify serial and parallel-VAE decode ownership, including rank-0-only audio, in
pipeline_ltx2_two_stages.py:1409-1445. Ensure each test asserts the expected
per-rank behavior.
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py`:
- Around line 599-611: Tighten the topology types across all affected sites: in
transformer_ltx2.py, annotate Stage2Groups.ulysses_group and seq_group with
torch_dist.ProcessGroup and use list[int] | None for gather_index; in
mapping.py, annotate both rank-layout helpers as returning list[list[int]]; in
utils.py, use list[int] | None for gather_index and remove the List import.
Avoid Any and use built-in generics with | throughout.
In `@tensorrt_llm/_torch/visual_gen/utils.py`:
- Line 5: Update the gather_index declaration to use Python 3.10 built-in typing
as list[int] | None, and remove the now-unused List import from the typing
imports. Apply the same typing style to the corresponding declaration around the
referenced additional location.
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py`:
- Around line 359-388: Replace the duplicated group-construction logic in
_build_stage2_groups_for_test with a call to
LTX2TwoStagesPipeline._stage2_transformer_groups using a minimal pipeline
namespace configured with the test’s vgm. Preserve the helper’s returned
Stage2Groups behavior while ensuring the tests execute the production builder
directly.
🪄 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: 27cd08c2-919f-4e42-8b78-70e1a220bbc5
📒 Files selected for processing (13)
tensorrt_llm/_torch/visual_gen/attention_backend/parallel.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.pytensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.pytensorrt_llm/_torch/visual_gen/modules/attention.pytensorrt_llm/_torch/visual_gen/utils.pytests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.pytests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.pytests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.pytests/unittest/_torch/visual_gen/test_ltx2_pipeline.pytests/unittest/_torch/visual_gen/test_ltx2_transformer.pytests/unittest/_torch/visual_gen/test_utils.py
💤 Files with no reviewable changes (1)
- tests/unittest/_torch/visual_gen/test_utils.py
Collaborator
|
PR_Github #59728 [ run ] completed with state
|
luyiyun1021
marked this pull request as draft
July 16, 2026 18:13
…s tp_size == 1 LTX-2 already rejects TP at construction, so the per-TP-fiber seq-plane helper was dead generality (user decision). The two-stage group builder flattens the fold lists inline (the tp=1 order where cp-major global ranks double as group ranks) and fail-fasts on tp > 1. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…ps process groups The builder makes the DiT's stage-2 groups (the VAE has its own), so name it that; the group fields are ProcessGroup, not Any. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…; drop _latents_on_all_ranks The latent path does no VAE work — the "decode" is a local unpatchify and every rank holds the full latents after the denoise loop — so the vae_ranks/rank-0 decode gate does not apply to it (that gate exists to skip real VAE decode, and stays for pixel output, matching Wan/Flux). The two-stage handoff is now just the default latent behavior. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
The latent early-returns inside decode_video_fn/decode_audio_fn made the same condition appear three times; the tail branch now does the unpatchify itself and the decode fns are pure pixel decoders. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
luyiyun1021
requested review from
yibinl-nvidia and
zhenhuaw-me
and removed request for
kris1025,
o-stoner and
yunruis
July 20, 2026 02:15
…ment caller-agnostic Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
luyiyun1021
marked this pull request as ready for review
July 20, 2026 02:25
zhenhuaw-me
reviewed
Jul 21, 2026
…thod The async-Ulysses path pre-quantizes hidden states once and passes a shared Fp4QuantizedTensor to to_q/to_k/to_v. The gate was construction-time only (checkpoint NVFP4 + input_scale present), so in the two-stage Stage-2 refinement -- where the persistent LoRA cache merges FP4 params as BF16 and swaps the Linear to UnquantizedLinearMethod -- F.linear received an Fp4QuantizedTensor and raised TypeError during warmup. Skip the shared quantize when the active quant_method is unquantized; Stage 2 then feeds BF16 through, matching the sync path's semantics. Verified: two-stage async_ulysses warmup + generation at cfg2xuly2 and cfg2xuly4 on 8xB200. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…gy test The single-GPU CONDITIONAL-mode test builds with ulysses_size=1, where nothing is Ulysses-wrapped, so it cannot distinguish the CONDITIONAL gate from the trivial u1 case. Assert in the 4-GPU dual-topology loop -- where attn1 IS wrapped in both topologies -- that audio_attn1 remains plain. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
The decode log used a host wall clock started right after the refinement loop; the graphed stage-2 leaves the host far ahead of the GPU, so the clock absorbed stage-2's queued tail (logged 2.73s for a 0.94s decode at 1 GPU). Bracket the decode section with a _TwoStagePhaseTimer event pair instead -- stream-timeline distance, no host sync inserted. Verified at 1 GPU: decode 0.94s (matches the pre-refactor event-based GPU timeline), e2e unchanged. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
luyiyun1021
force-pushed
the
dev/ltx2-stage2-dual-topology
branch
2 times, most recently
from
July 22, 2026 13:11
6ac6e95 to
fc7f329
Compare
Rank 0 logs the default->stage2 seq-parallel width when the dual topology folds cfg into ulysses for the refinement loop. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
luyiyun1021
force-pushed
the
dev/ltx2-stage2-dual-topology
branch
from
July 22, 2026 13:14
fc7f329 to
ce26900
Compare
Reviewer feedback: the chained list-comprehension expansion was hard to read. Same layout arithmetic, order-preserving: one fold list per coordinate of every non-(cfg, ulysses) mesh dim via itertools.product, cfg-major x ulysses within each list. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
yibinl-nvidia
approved these changes
Jul 27, 2026
yibinl-nvidia
left a comment
yibinl-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
LGTM overall. Thanks for making two stage pipeline fully parallelizable. Left some comments.
…LTX2TwoStagesPipeline
Keep the base LTX2Pipeline free of any stage-2 surface: _init_transformer splices a generic _extra_transformer_kwargs() hook (base returns {}), and the two-stage subclass overrides it to inject stage2_groups. _init_transformer runs inside BasePipeline.__init__, so the injection has to be a polymorphic method rather than an attribute set afterward. Behavior unchanged.
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
TestLTX2TopologyConstruction only exercised construction internals: private-attr aliasing when no stage-2 groups, and a defensive assert in set_ulysses_topology that is unreachable in the real flow (every caller already guards on _has_stage2). The dual-topology switch, shard/gather round-trip, audio_attn1-plain gate, and numerical parity vs the single-GPU reference are covered with real forwards in multi_gpu/test_ltx2_ulysses.py. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61898 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61898 [ run ] completed with state
|
pengbowang-nv
approved these changes
Jul 28, 2026
pengbowang-nv
left a comment
pengbowang-nv
left a comment
Collaborator
There was a problem hiding this comment.
AttentionPart LGTM
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62099 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62099 [ run ] completed with state |
luyiyun1021
changed the title
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
5 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.
@coderabbitai summary
Description
LTX-2 two-stage generation ran Stage 2 (distilled-LoRA refinement at full resolution) either serially on rank 0 — leaving all other GPUs idle — or behind an opt-in
stage2_ulyssesknob limited to cfg=1. This PR deletes the serial path and makes Stage 2 unconditionally parallel via a dual topology: Stage 1 denoises under the configuredcfg × (ring|attn2d) × ulyssesgroups; at the Stage-2 boundary (CFG no longer applies) the transformer switches to a cfg-flattened topology with ulysses sizecfg × ulysses, and switches back afterwards. cp (ring/attn2d) communicators are preserved verbatim across the switch; undercfg_size=1the topologies collapse to one instance and behavior is unchanged.Performance (8× B200, 768×1280, 121 frames, 40 + 3 steps, NVFP4 checkpoint, torch.compile + CUDA graph, FA4, parallel_vae 8)
stage2_ulyssesknob (opt-in, forfeits CFG parallelism)¹¹ Knob-path row measured on the pre-PR branch at the same hardware/config; the other rows are from this branch's CUDA-event GPU timeline (3-run mean for the final row). Stage 2 matches the knob path exactly (both run ulysses-8) while Stage 1 additionally gains CFG parallelism — the knob's either/or tradeoff is gone.
Main changes:
set_ulysses_topology(is_stage2)swaps whole stacks. Serial Stage-2 path, runtime ulysses↔plain switching, non-primary early-exit, and thestage2_ulyssesknob are removed.key_padding_mask).vae_ranks; theoutput_type="latent"contract is unchanged.VisualGenMapping.flatten_cfg_ranks()/flatten_cfg_seq_ranks()helpers build the stage-2 ulysses and per-TP-fiber seq-plane groups (cp-major shard order carried via a group-localgather_index, sincedist.new_groupsorts rank lists).Output quality — LPIPS (AlexNet, decoded frames) vs a cfg1×u8 reference; all within the chaotic-divergence benign band (≈0.1–0.2), montage-verified:
Test Coverage
multi_gpu/test_ltx2_ulysses.py: 4-GPU dual-topology alternation with full-forward numerical equivalence vs a single-GPU reference in BOTH topologies (pad/no-pad);is_ulysses/stack consistency; FULL audio-mode construction/switch; stage-2 head-divisibility rejection; TRTLLM audio downgrade under a cfg-only fold (2-GPU).test_ltx2_pipeline.py: (lora × topology) CUDA-graph key matrix; wrong-key capture/replay negatives; warmup pre-captures both topology graphs via the production_run_warmupwith zero post-warmup recapture; stage-2 token-divisibility rejection.test_ltx2_transformer.py: construction gates (no stage-2 groups → aliased stacks; switch without groups raises; CONDITIONAL audio stays plain).multi_gpu/test_visual_gen_mapping.py:flatten_cfg_ranks/flatten_cfg_seq_rankslayout arithmetic (cfg×u, cfg×ring×u, cfg×attn2d, TP fibers).PR Checklist
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES.
Test cases are provided for new code paths.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if significant design change.
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.