[TRTLLM-12653][feat] LTX-2 Ulysses cross-attention for v2a with audio padding by luyiyun1021 · Pull Request #14044 · NVIDIA/TensorRT-LLM

luyiyun1021 · 2026-05-12T08:59:46Z

Description

This PR adds Ulysses sequence-parallelism support for the v2a (video-to-audio) cross-attention in LTX-2 by reusing the existing UlyssesAttention wrapper. It replaces the all-gather (AG) path on the heavy video K/V side with the Ulysses 4-collective (Q + K + V + output) all-to-all, which is significantly cheaper when T_q << T_kv. The audio padding mode, FA4 key_padding_mask support, audio self-attn backend fallback, and the AG-fallback gating fix are all in service of making the Ulysses v2a path safe across all supported LTX-2 configurations.

Why v2a, not a2v?

LTX-2 has two AV cross-attentions per block:

Module	Q	K/V
`audio_to_video_attn` (a2v)	video (`T_v ≈ 6300` at 768×1280 / 121 frames)	audio (`T_a ≈ 126`, padded to 128 under `audio_pad_for_ulysses=True`)
`video_to_audio_attn` (v2a)	audio (`T_a`)	video (`T_v`)

Per-rank communication volume (cross-attn K/V path), `U = ulysses_size`

Path	Formula (bf16, 2 B/elem)	What it ships
AG K/V (baseline)	`2 × (U-1)/U × T_kv × H × D × 2`	each rank receives full K and V across ranks
Ulysses 4-a2a (Q + K + V + output)	`(U-1)/U² × 2 × (T_q + T_kv) × H × D × 2`	each rank exchanges its 1/U head shard with the other U-1

Setting Ulysses < AG and simplifying gives the break-even:

T_q / T_kv  <  U − 1

Plug in LTX-2 with U = 4:

Module	`T_q / T_kv`	vs `U − 1 = 3`	Decision
v2a	128 / 6300 ≈ 0.02	≪ 3	Ulysses wins — switch
a2v	6300 / 128 ≈ 49	≫ 3	AG wins — keep AG path

Concrete byte numbers (LTX-2 v2a: H × D = 32 × 128, U = 4, bf16, per v2a layer per step, per rank):

Path	Per-rank bytes shipped	Net per layer per step per rank
Baseline v2a (2 AG: K, V on full `T_v = 6300`)	`2 × 0.75 × 6300 × 32 × 128 × 2 B`	≈ 147.6 MB
OPT v2a wrapper (4 a2a: Q + K + V + out)	`0.75 × (128 + 6300) × 32 × 128 × 2 B`	≈ 37.7 MB
OPT a2v new AGs (audio K/V, small)	`2 × 0.75 × 128 × 32 × 128 × 2 B`	≈ 1.5 MB
OPT total (v2a wrapper + a2v new AGs)		≈ 39.2 MB

OPT ships 27% of baseline for the v2a/a2v cross-attn comm budget — a theoretical −73% byte volume per layer per step. This is the upper bound on wall-time savings; per-kernel launch overhead (OPT has roughly 2× the kernel count) is what eats into the realized win.

Hence only v2a gets the wrapper; a2v stays on AG. Hardcoded in BasicAVTransformerBlock._init_av_cross_modules via use_ulysses_cross=True on video_to_audio_attn only.

Performance

NCCL kernel time, per rank, single denoise step (`step 5` capture, 4× B200, CFG=1 × Ulysses=4, 768×1280, 121 frames)

Captured with nsys --capture-range=cudaProfilerApi --cuda-graph-trace=node, TLLM_PROFILE_VISUAL_GEN_START_STOP=5-5:

Path	AllGather total	SendRecv total	Per-rank NCCL total
Baseline (`audio_pad_for_ulysses=false`, AG fallback)	26.7 ms	27.4 ms	54.0 ms
OPT (`audio_pad_for_ulysses=true`, v2a wrapper)	3.3 ms	41.4 ms	44.7 ms
Δ (OPT − Baseline)	−23.4 ms	+14.0 ms	−9.3 ms (−17 %)

AG time drops 23.4 ms because the dominant cost item (gathering full video K/V) is gone — the remaining AGs ship only the small audio K/V (a2v) and avg per-op cost falls 8× (138 µs → 17 µs). SR time rises 14 ms because the 4 a2a collectives in the v2a wrapper are realized as ncclDevKernel_SendRecv pairs (U-1 SR ops per a2a per rank). The net 9.3 ms / step / rank saved is well below the 73% theoretical byte savings — launch overhead from going from 385 to 770 kernel calls eats the rest.

End-to-end wall time (8× B200, CFG=2 × Ulysses=4, 768×1280, 121 frames, 40 steps)

	Run 0	Run 1	Avg	Δ vs baseline
Baseline (`audio_pad_for_ulysses=false`, AG fallback)	9.506 s	9.502 s	9.504 s	—
OPT (`audio_pad_for_ulysses=true`, v2a wrapper engaged)	9.113 s	9.121 s	9.117 s	−4.07 %

Same prompt / seed / steps; output MP4s are visually identical (sample artifacts: tmp/ulysses-a2v/ltx2_{baseline,opt}_refactor.mp4).

Implementation

LTX2Attention.__init__ builds a dual-attn pair on video_to_audio_attn only — both a UlyssesAttention(inner_backend=...) wrapper and the base-class plain attention — and toggles self.attn between them at runtime via set_ulysses_active(bool). Symmetric getter is_ulysses_active() lets the block forward gate the AG fallback correctly under Attention2D / ring (where _audio_is_sharded can be true without a Ulysses wrapper being built).
The dual-attn block only fires when use_ulysses_cross=True ∧ ulysses_size > 1 ∧ cp_size == 1 (pure Ulysses). Combined modes (ring+Ulysses, attn2d+Ulysses) shard the seq axis by seq_size = cp_size × ulysses_size, which this wrapper has not been validated against — they fall through to the existing AG path.
configure_audio_ulysses(audio_seq_len) runs once per request: when parallel.audio_pad_for_ulysses=True (default) it pads audio to a multiple of ulysses_size so the wrapper can always engage; otherwise it falls back to the natural-divisibility check.
LTXModel.forward pads audio (zero-pad latent, repeat-last positions/timesteps) on entry and strips the padded tail on exit. prepare_text_cache pre-pads audio PE so the denoise hot path sees consistent shapes. A [B, T_a_padded] validity mask travels through TransformerArgs.audio_padding_mask and is consumed by audio_attn1 and audio_to_video_attn.

Backend safety for `key_padding_mask`

Padded audio means valid Q tokens would attend to padded K positions unless the backend honors key_padding_mask:

Backend	Where it runs in LTX-2	Honors mask?	Status
VANILLA	any	yes (PR adds `key_padding_mask` → SDPA `attn_mask`)	✓
FA4	any	yes (PR adds `key_padding_mask` → cute `seqused_k`)	✓
TRTLLM	self-attn only (cross-attn auto-falls-back to VANILLA per existing rule)	no — silently drops `**kwargs`	`audio_attn1` overrides backend to VANILLA when `audio_pad_for_ulysses=True` (mirrors the cross-attn TRTLLM→VANILLA fallback in `modules/attention.py`; audio self-attn is small at `T_a ≈ 126` so the backend downgrade is negligible)

v2a itself never receives the mask — padded audio Q rows are stripped on LTXModel.forward exit, so their attention output is discarded regardless.

Test Coverage

tests/unittest/_torch/visual_gen/test_vanilla_key_padding_mask.py — single-rank parity for VanillaAttention.forward(key_padding_mask=...) vs unpadded SDPA.
tests/unittest/_torch/visual_gen/test_fa4_key_padding_mask.py — FlashAttn4Attention.forward(key_padding_mask=...) on a B200: self-attn parity, cross-attn parity, and pad-fill invariance.

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)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
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.

coderabbitai · 2026-05-13T06:15:53Z

📝 Walkthrough

Walkthrough

Adds UlyssesCrossAttention backend wrapper for sequence-parallel cross-attention (S_q ≠ S_kv), integrates key padding mask support into VanillaAttention, introduces audio padding configuration and pipeline in LTXModel to satisfy Ulysses divisibility constraints, wires masking through transformer blocks, and validates correctness with multi-GPU tests for collectives and RoPE equivariance.

Changes

Ulysses Cross-Attention & Audio Padding

Layer / File(s)	Summary
UlyssesCrossAttention backend wrapper `tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py`	New AttentionBackend subclass performing three distributed collectives (Q all-to-all, fused K\|V 5D all-to-all, output all-to-all) to parallelize cross-attention when sequence lengths differ; reports NHD layout preference and disables fused QKV.
Key padding mask support in VanillaAttention `tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py`	Adds optional key_padding_mask parameter to forward(), validated for non-causal paths, expanded to broadcastable shape for SDPA, enabling attention masking over padded key-value sequences.
Configuration and data structures `tensorrt_llm/_torch/visual_gen/config.py`, `tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/transformer_args.py`	ParallelConfig gains audio_pad_for_ulysses flag (default True); TransformerArgs gains audio_padding_mask field to propagate validity masks through transformer blocks.
LTX2Attention Ulysses cross-attention support `tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py` (LTX2Attention class)	Adds use_ulysses_cross constructor parameter; conditionally wraps cross-attention with UlyssesCrossAttention; expands set_ulysses_active() to toggle either self or cross dual backends; threads key_padding_mask parameter through forward() to backends via attn_kwargs.
LTXModel audio padding pipeline `tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py` (LTXModel class)	Introduces _audio_pad and _audio_pad_for_ulysses state; adds _pad_pe() and _pad_modality_audio() helpers to extend RoPE and pad audio sequences; updates prepare_text_cache() to extend audio PE tuples; updates forward() to pad audio on entry, construct audio_padding_mask, thread through TransformerArgs, and strip padded tail on exit.
Transformer block audio attention paths `tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py` (transformer block wiring)	Enables use_ulysses_cross=True for video_to_audio_attn; threads audio.audio_padding_mask through audio self-attention and A2V cross-attention; restructures V2A cross-attention for strict-Ulysses handling without key_padding_mask.
UlyssesCrossAttention multi-GPU tests `tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_cross_attention.py`	Test harness validating UlyssesCrossAttention wrapper: verifies head-count configuration and fused QKV support flags, validates forward output shapes under sequence sharding, confirms parity against full-tensor scaled_dot_product_attention reference, ensures world-size-1 fast-path bit-exactness.
RoPE A2A equivariance tests `tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_cross_rope_equivariance.py`	Multi-process equivariance tests validating RoPE commutes with all-to-all collectives for K-side and Q-side asymmetric cross-attention scenarios using synthetic INTERLEAVED positional encodings.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested reviewers

laikhtewari
kaiyux
PerkzZheng
zhenhuaw-me

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 43.48% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: adding Ulysses cross-attention for v2a (video-to-audio) with audio padding support in LTX-2.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description comprehensively explains the feature, performance justification, implementation details, test coverage, and all checklist items are addressed.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py (1)

107-118: ⚡ Quick win

Consider using ValueError instead of assertions for validation.

Lines 108 and 109-115 use assert statements for input validation. Per coding guidelines (Pydantic section), validation should raise ValueError instead of using assertions, as assertions can be disabled with python -O and are intended for debugging rather than input validation in production code.

🔒 Suggested change

 if key_padding_mask is not None:
-    assert not is_causal, "key_padding_mask is not supported with causal attention"
-    assert key_padding_mask.dim() == 2 and key_padding_mask.shape == (
-        q.shape[0],
-        k.shape[2],
-    ), (
-        f"Invalid key_padding_mask shape: expected [B={q.shape[0]}, "
-        f"S_kv={k.shape[2]}], got {tuple(key_padding_mask.shape)}"
-    )
+    if is_causal:
+        raise ValueError("key_padding_mask is not supported with causal attention")
+    if key_padding_mask.dim() != 2 or key_padding_mask.shape != (q.shape[0], k.shape[2]):
+        raise ValueError(
+            f"Invalid key_padding_mask shape: expected [B={q.shape[0]}, "
+            f"S_kv={k.shape[2]}], got {tuple(key_padding_mask.shape)}"
+        )
     # [B, S_kv] -> [B, 1, 1, S_kv] so SDPA broadcasts over H and S_q.
     attn_mask = key_padding_mask[:, None, None, :]
     return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=self.scale)

As per coding guidelines: "Use @field_validator and @model_validator in Pydantic instead of manual validate() methods; raise ValueError instead of assertions"

🤖 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/vanilla.py` around lines 107
- 118, Replace the two assert checks used for input validation with explicit
ValueError raises so validation cannot be disabled: instead of "assert not
is_causal, 'key_padding_mask is not supported with causal attention'" raise
ValueError("key_padding_mask is not supported with causal attention") and
replace the shape/assertion block that checks key_padding_mask.dim() and
key_padding_mask.shape with a single ValueError that includes the same detailed
message (use q.shape[0] and k.shape[2] in the message to show expected [B, S_kv]
and tuple(key_padding_mask.shape) for actual). Keep the subsequent construction
of attn_mask and the call to F.scaled_dot_product_attention unchanged; update
only the assertions to ValueError in the code paths referencing key_padding_mask
and is_causal.

🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py`:
- Around line 107-118: Replace the two assert checks used for input validation
with explicit ValueError raises so validation cannot be disabled: instead of
"assert not is_causal, 'key_padding_mask is not supported with causal
attention'" raise ValueError("key_padding_mask is not supported with causal
attention") and replace the shape/assertion block that checks
key_padding_mask.dim() and key_padding_mask.shape with a single ValueError that
includes the same detailed message (use q.shape[0] and k.shape[2] in the message
to show expected [B, S_kv] and tuple(key_padding_mask.shape) for actual). Keep
the subsequent construction of attn_mask and the call to
F.scaled_dot_product_attention unchanged; update only the assertions to
ValueError in the code paths referencing key_padding_mask and is_causal.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 28647f03-f28c-4881-9bdb-b26f6ff56b35

📥 Commits

Reviewing files that changed from the base of the PR and between 27c143d and 2d3f7cc.

📒 Files selected for processing (7)

tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py
tensorrt_llm/_torch/visual_gen/config.py
tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/transformer_args.py
tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py
tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_cross_attention.py
tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_cross_rope_equivariance.py

yibinl-nvidia

Left some comments

@yibinl-nvidia

…dio-sharded The v2a forward used ``_audio_is_sharded`` as a proxy for "Ulysses wrapper is active", which is wrong under Attention2D: ``configure_audio_ulysses`` sets ``_audio_is_sharded=True`` whenever ``use_seq_parallel + seq_parallel_size>1``, including the Attention2D case where ``ulysses_size==1`` and the v2a wrapper was never built. The branch then skipped the all-gather and passed seq-sharded K/V to the plain attention backend, which expects full K/V — silently producing wrong outputs for v2a under Attention2D + audio padding. Add ``LTX2Attention.is_ulysses_active()`` as the getter symmetric with ``set_ulysses_active``, and gate the v2a K/V path on it. Returns False whenever no Ulysses wrapper was built or the module is currently routed to the plain attn, so the AG fallback fires in every case where the plain backend will run. Reported by @yibinl-nvidia on NVIDIA#14044. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

@yibinl-nvidia

…dio-sharded The v2a forward used ``_audio_is_sharded`` as a proxy for "Ulysses wrapper is active", which is wrong under Attention2D: ``configure_audio_ulysses`` sets ``_audio_is_sharded=True`` whenever ``use_seq_parallel + seq_parallel_size>1``, including the Attention2D case where ``ulysses_size==1`` and the v2a wrapper was never built. The branch then skipped the all-gather and passed seq-sharded K/V to the plain attention backend, which expects full K/V — silently producing wrong outputs for v2a under Attention2D + audio padding. Add ``LTX2Attention.is_ulysses_active()`` as the getter symmetric with ``set_ulysses_active``, and gate the v2a K/V path on it. Returns False whenever no Ulysses wrapper was built or the module is currently routed to the plain attn, so the AG fallback fires in every case where the plain backend will run. Reported by @yibinl-nvidia on NVIDIA#14044. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-05-22T09:57:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-22T10:03:22Z

PR_Github #49918 [ run ] triggered by Bot. Commit: 3709672 Link to invocation

NVShreyas

LGTM from ulysses cross attention and attention backend perspective.
Might be good to have tests validating output with ulysses enabled for audio?

tensorrt-cicd · 2026-05-22T14:59:14Z

PR_Github #49918 [ run ] completed with state SUCCESS. Commit: 3709672
/LLM/main/L0_MergeRequest_PR pipeline #39496 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@yibinl-nvidia

…dio-sharded The v2a forward used ``_audio_is_sharded`` as a proxy for "Ulysses wrapper is active", which is wrong under Attention2D: ``configure_audio_ulysses`` sets ``_audio_is_sharded=True`` whenever ``use_seq_parallel + seq_parallel_size>1``, including the Attention2D case where ``ulysses_size==1`` and the v2a wrapper was never built. The branch then skipped the all-gather and passed seq-sharded K/V to the plain attention backend, which expects full K/V — silently producing wrong outputs for v2a under Attention2D + audio padding. Add ``LTX2Attention.is_ulysses_active()`` as the getter symmetric with ``set_ulysses_active``, and gate the v2a K/V path on it. Returns False whenever no Ulysses wrapper was built or the module is currently routed to the plain attn, so the AG fallback fires in every case where the plain backend will run. Reported by @yibinl-nvidia on NVIDIA#14044. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-05-24T11:25:43Z

LGTM from ulysses cross attention and attention backend perspective. Might be good to have tests validating output with ulysses enabled for audio?

@NVShreyas Do you mean e2e test testing the output quality or testing the padding correctness for fa4 and vanilla backend? I think for the former one, there is no test yet and for the latter one, I've added two tests.

luyiyun1021 · 2026-05-26T05:41:40Z

/bot run --disable-fail-fast

luyiyun1021 · 2026-05-26T05:42:51Z

Added test_ulysses_with_key_padding_mask_parity test to test key_padding_mask works properly with ulysses

NVShreyas · 2026-05-26T14:05:56Z

LGTM from ulysses cross attention and attention backend perspective. Might be good to have tests validating output with ulysses enabled for audio?

@NVShreyas Do you mean e2e test testing the output quality or testing the padding correctness for fa4 and vanilla backend? I think for the former one, there is no test yet and for the latter one, I've added two tests.

yes I meant e2e test with a small input size. It is fine as long as there is some test ensuring output of transformer is consistent with ulysses enabled, for both self and cross attention.

luyiyun1021 · 2026-05-27T06:31:16Z

added multi_gpu/test_ltx2_ulysses.py for e2e test

luyiyun1021 · 2026-05-27T06:38:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-28T11:39:04Z

PR_Github #50765 [ run ] triggered by Bot. Commit: 693348b Link to invocation

tensorrt-cicd · 2026-05-28T13:58:18Z

PR_Github #50765 [ run ] completed with state FAILURE. Commit: 693348b
/LLM/main/L0_MergeRequest_PR pipeline #40243 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

NVShreyas · 2026-05-28T14:34:02Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-28T14:41:27Z

PR_Github #50808 [ run ] triggered by Bot. Commit: 693348b Link to invocation

tensorrt-cicd · 2026-05-28T16:02:38Z

PR_Github #50808 [ run ] completed with state FAILURE. Commit: 693348b
/LLM/main/L0_MergeRequest_PR pipeline #40279 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Adds a Ulysses sequence-parallel wrapper for the video-to-audio (v2a) cross-attention in LTX-2. v2a is the asymmetric direction where the wrapper saves communication: Q is audio (small, T_a~126) and K/V are video (large, T_v=15360 for 768x1280). Replacing the per-block video K/V all-gather with a Q+(K|V fused)+output three-collective pattern cuts per-rank communication volume by ~4x. The opposite a2v direction (Q=video, K/V=audio) is left on the existing all-gather path because the wrapper would move two video-sized tensors (Q a2a + output a2a) in exchange for replacing a cheap audio-K/V all-gather — net 30x MORE bytes per rank at U=4. For LTX-2 at 768x1280 / 121f / U=4: T_v = 15360, T_a = 126. AllGather path (current main for both cross-attns): 2 * (U-1)/U * B * T_kv * D_a * 2 a2v (K/V = audio, T_kv = T_a = 126): ~0.78 MB per rank/block v2a (K/V = video, T_kv = T_v = 15360): 94 MB per rank/block UlyssesCrossAttention path (3 a2a): (U-1)/U^2 * B * (T_q + 2*T_kv + T_q) * D_a * 2 a2v (T_q = T_v, T_kv = T_a): 24 MB per rank/block (30x MORE) v2a (T_q = T_a, T_kv = T_v): 24 MB per rank/block (~4x LESS) Break-even: Uly < AG iff (T_q + T_kv) < U * T_kv, equivalently T_q/T_kv < U - 1. With T_v/T_a ~ 122, the wrapper is only profitable when the small modality is on the Q side. Hence v2a only. - `UlyssesCrossAttention` (parallel.py, +107 lines): three-collective wrapper for S_q != S_kv. Q all_to_all_4d, fused K|V 5D all_to_all, output all_to_all_4d. Stacks K|V on a new dim of size 2 so both tensors travel in a single collective. world_size==1 fast path bypasses all three. - `audio_pad_for_ulysses` flag (ParallelConfig, default True): when set, LTXModel pads the audio modality at forward entry so T_a is divisible by ulysses_size, attaches a [B, T_a_padded] bool validity mask to TransformerArgs, and strips the padded tail on exit. Required because T_a is derived from num_frames/fps and is usually NOT divisible (e.g. 121 video frames at 24fps gives T_a=126; 126%4=2). Pad scheme: zero-pad latent, repeat-last for positions / per-token timesteps, repeat-last for the (cos,sin) PE tuple via a seq-dim auto-detect helper. - `LTX2Attention.use_ulysses_cross`: builds the inner backend at (H/U, H_kv/U) heads and wraps with UlyssesCrossAttention. Gated on ulysses_size > 1 (skipped under Attention2D mode which is mutually exclusive with Ulysses per VisualGenMapping). set_ulysses_active toggles between the wrapper and the base plain backend at runtime (used by configure_audio_ulysses and Stage 2 fall-back). - VanillaAttention gains an optional `key_padding_mask` kwarg (expanded to [B,1,1,S_kv] as SDPA attn_mask, non-causal only). The `audio_to_video_attn` call site passes the mask so attention zeros out audio pad slots. `video_to_audio_attn` does not need a mask (video K/V is unpadded; pad audio Q is stripped on output). - `_shard_transformer_args` passes the full-seq audio_padding_mask through unchanged; the mask is identical across Ulysses ranks. - test_ulysses_cross_attention.py: wrapper init, forward shape at world_size in {2, 4}, world_size==1 bit-exact fast path, and parity vs full SDPA on shared seed inputs. - test_ulysses_cross_rope_equivariance.py: verifies the rope-then-a2a == a2a-then-full-rope identity that the wrapper relies on (since rope is applied locally on sharded inputs in LTX2Attention.forward and then a2a'd, the per-position rotation must be preserved by the redistribution). 768x1280, 121 frames, num_inference_steps=40, 3 timed runs after warmup) Wall clock e2e: Baseline (no wrapper): 9.497s v2a wrapper: 9.091s (-0.406s, -4.3%) NCCL kernel time, per rank, per step (1 denoise step capture via TLLM_PROFILE_VISUAL_GEN_START_STOP=5-5 + nsys --capture-range= cudaProfilerApi --cuda-graph-trace=node): Baseline: ~37 ms NCCL (1640 SendRecv + 1653 AllGather summed across 8 ranks; per rank ~205 + ~207 kernels) v2a Uly: ~27 ms NCCL average (528 SendRecv + 196 AllGather per rank; rank-0 outlier 37.7 ms, rest 24-29 ms) Per-rank per-step savings: ~10 ms. 40 steps * 10 ms = 400 ms ~ measured 406 ms e2e. Audio padding cost is negligible at typical T_a (126%4=2 padded to 128; 1.6% padded tokens, masked out in attention). Disabling the wrapper at runtime is supported via `parallel.audio_pad_for_ulysses: false` — `_audio_is_sharded` then only becomes True when T_a%U==0, and the wrapper falls back to the existing all-gather path otherwise. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

…change) Addresses post-review feedback on the v2a Ulysses wrapper. Pure refactoring; no functional changes. 1. PE pad moved into prepare_text_cache (done once per generate() instead of per forward step). 2. _pad_pe takes seq_dim explicitly (was fragile auto-detect by shape match); raises on out-of-range seq_dim. 3. Forward audio pad guard reduced from 4 conjuncts to 2. 4. Removed defensive ``if model_config is not None else True`` fallback. 5. PE pad pre-assign + overwrite block removed from forward (subsumed by NVIDIA#1). 6. configure_audio_ulysses two branches flattened to a single formula. 7. v2a fallback restructured to flat if/elif/else (was nested). Plus a stale comment fix: ``Ulysses wrap (a2v in LTX-2)`` → ``(v2a in LTX-2)``. Diff: ~44 insertions, ~57 deletions in transformer_ltx2.py only. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

…added K/V Adds a single-rank CPU test verifying the math guarantee that audio K/V padding relies on: VanillaAttention.forward with key_padding_mask masking pad columns produces output bit-identical to running SDPA on the unpadded K/V (max|diff|=0 in the realistic O(1) pad-magnitude regime). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

…attn Drops the bespoke UlyssesCrossAttention class. Its unfused path (Q + K + V independent 4D a2a + output a2a) is functionally equivalent to the existing UlyssesAttention._forward_unfused path, which already handles asymmetric S_q != S_kv via the inner backend's seq_len / seq_len_kv kwargs. Since VanillaAttention.support_fused_qkv() is False, UlyssesAttention auto-selects the unfused path, so v2a now flows through the standard wrapper. The dropped K|V 5D fusion (1 collective vs 2) is a perf-only optimization whose isolated contribution was not measured; if a future profile shows it is worthwhile, add a fused_kv flag to UlyssesAttention as a separate PR. Deletes: - tests/.../multi_gpu/test_ulysses_cross_attention.py - tests/.../multi_gpu/test_ulysses_cross_rope_equivariance.py Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

…actor staleness - Drop a copy-paste duplicate of the "Audio padding for Ulysses" comment in LTXModel.forward. - a2v block: drop the "wrapper's three full-video collectives" count, stale after dropping K|V 5D fusion (now 4 collectives), and restore a blank line removed in the original change. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

@yibinl-nvidia

…dio-sharded The v2a forward used ``_audio_is_sharded`` as a proxy for "Ulysses wrapper is active", which is wrong under Attention2D: ``configure_audio_ulysses`` sets ``_audio_is_sharded=True`` whenever ``use_seq_parallel + seq_parallel_size>1``, including the Attention2D case where ``ulysses_size==1`` and the v2a wrapper was never built. The branch then skipped the all-gather and passed seq-sharded K/V to the plain attention backend, which expects full K/V — silently producing wrong outputs for v2a under Attention2D + audio padding. Add ``LTX2Attention.is_ulysses_active()`` as the getter symmetric with ``set_ulysses_active``, and gate the v2a K/V path on it. Returns False whenever no Ulysses wrapper was built or the module is currently routed to the plain attn, so the AG fallback fires in every case where the plain backend will run. Reported by @yibinl-nvidia on NVIDIA#14044. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

@yibinl-nvidia

… audio_attn1 fallback @yibinl-nvidia flagged that audio_attn1 with audio_pad_for_ulysses=True silently drops key_padding_mask under TRTLLM/FA4 backends, letting valid audio Q attend to padded audio K and corrupting the real audio stream. Two surgical fixes, mirroring the existing cross-attn TRTLLM->VANILLA fallback pattern in modules/attention.py: (1) FA4: route key_padding_mask through FlashAttn4Attention.forward as seqused_k = mask.sum(dim=1).int() (FA4 cute interface accepts per-batch valid lengths on padded [B, S, H, D] input - no varlen repack). LTX-2 audio padding is always a True-prefix; assert and lean on that. (2) TRTLLM audio_attn1: downgrade to VANILLA at construction when audio_pad_for_ulysses=True. Audio self-attn is small (T_a ~ 126), so the backend downgrade is negligible, while it keeps cross-attn (already VANILLA via the existing cross-attn fallback) free to use the padding. v2a cross-attn never receives the mask (caller strips pad Q at LTXModel.forward exit), so it is unaffected. Tested: - tests/.../test_fa4_key_padding_mask.py - 3 cases on a B200, all pass. Self-attn parity, cross-attn parity, pad-fill invariance. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

@NVShreyas

…s mode @NVShreyas flagged that this branch's old guard ("Attention2D keeps ulysses_size == 1 — mutual exclusion") is stale: main now enables ring + ulysses, and attn2d + ulysses is coming soon. Under those combined modes seq_size = cp_size * ulysses_size, so the seq axis is sharded by more than ulysses_size and this wrapper's seq-shape assumptions no longer hold. Add an explicit cp_size == 1 gate so v2a only builds the Ulysses dual-attn pair under pure Ulysses, and falls back to plain + AG everywhere else. Updates the now-stale comment accordingly. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

Comment-only pass across the PR's touched files: - Drop dev-process narrative (verbose "why-we-did-this" repeated in code). - Generic attention backends (vanilla, flash_attn4) no longer name LTX-2 in their docstrings. - Tighten field-storage and configure_audio_ulysses sections that were documenting both the YAML knob and the runtime read paths in prose. - Preserve the non-obvious bits: cp_size gate rationale, TRTLLM→VANILLA fallback motivation, is_ulysses_active() vs _audio_is_sharded gate reasoning, seq_dim disambiguation for SPLIT vs INTERLEAVED rope. No behavior change. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

@NVShreyas

…+ extend mask docstring to FA4 Address @NVShreyas review: - Move the inline `from ...parallel import UlyssesAttention` to the module's top-level imports. - Update the `key_padding_mask` docstring on `LTX2Attention.forward` to mention that FA4 also honors the mask (previously the doc only listed VANILLA), and to point to the audio_attn1 TRTLLM->VANILLA construction-time fallback. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

@NVShreyas

…n + key_padding_mask Address @NVShreyas review (LGTM-with-suggestion): "Might be good to have tests validating output with ulysses enabled for audio." Single-card backend tests (test_vanilla_key_padding_mask.py, test_fa4_key_padding_mask.py) cover the math VanillaAttention / FlashAttn4Attention apply when masking padded K/V columns. They do not, however, exercise the UlyssesAttention wrapper's a2a + sharded SDPA + reverse-a2a pipeline with a key_padding_mask in kwargs - the wrapper mutates kwargs (seq_len, seq_len_kv) and the post-a2a K shape must align with the [B, S] mask. This adds a multi-rank parity test that covers that composition. Test pattern mirrors LTX-2's audio_attn1 path under audio_pad_for_ulysses=True: - Audio Q=K=V padded so S_padded is divisible by ulysses_size. - key_padding_mask = True on [0, S_real), False on the pad tail. - UlyssesAttention(VanillaAttention) wrapper; each rank holds [B, S_padded/U, H, D] of the padded sequence. - Reference: plain SDPA on the unpadded [0, S_real) prefix. - Assert: valid Q rows in this rank's shard match the reference within bf16-like tolerance. ws=2 (one sharding boundary), CPU + gloo - runs in pre-merge CI without GPUs. Forces device="cpu" to avoid cuda:rank ordinal errors when fewer GPUs are visible than spawned ranks. Tested locally in the dev container: 1 passed in 44.01s. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

…s=2) End-to-end multi-GPU test that runs ``LTXModel.forward`` on a 2-layer AudioVideo mini-config with ``ulysses_size=2`` and compares to the single-GPU reference (same weights, same inputs). One test logic function covers self-attn (audio_attn1), a2v cross-attn, and the v2a Ulysses wrapper simultaneously. Cases (parametrized over VANILLA + FA4 backends, ws=2): - test_av_ulysses_no_audio_pad: T_a divisible by U, no padding. - test_av_ulysses_audio_pad: T_a NOT divisible by U; exercises audio_pad_for_ulysses + the key_padding_mask thread through audio_attn1 + a2v. head_dim=32 keeps the test runnable in dev images that don't have the LTX-2 fused norm+rope C++ extension built (``fused_dit_split_norm``); the v2a Ulysses + key_padding_mask code paths under VANILLA / FA4 are unchanged by that choice. cross_attention_dim is set to inner_dim (=num_heads * head_dim) because the LTX-2 caption_projection projects caption_channels → inner_dim and feeds straight into ``block.attn2.to_k`` whose input dim is ``cross_attention_dim``. Tolerance is rtol=atol=5e-2: measured BF16 drift through 2 transformer layers + Ulysses collectives is ~1.5% on the failing rank, well within the threshold, and tight enough to catch real numerical regressions. Verified locally: 4/4 PASSED in 73.7s (2 × B200). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

Rebase-conflict-resolution introduced two over-wrapped expressions in transformer_ltx2.py that the project's ruff-format pre-commit hook wants on one line. No behavior change. - LTX2Attention.forward fused-fallback call to _forward_unfused. - BasicAVTransformerBlock.forward v2a AG-fallback gate condition. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

…ysses Per reviewer feedback, remove the ``parallel.audio_pad_for_ulysses`` ParallelConfig knob and derive audio padding behavior internally based on the runtime context. With Ulysses active, audio is always padded to make ``T_a`` divisible by ``ulysses_size`` and a ``[B, T_a_padded]`` validity mask is attached so attention zeros pad slots; with Ulysses inactive nothing changes. The opt-out mode the knob exposed was a silent performance downgrade (disengaged the v2a Ulysses wrapper, fell back to plain attention on non-divisible T_a) and not something users should be deciding. Source changes (all behavior-equivalent to the previous True default): * Drop ``ParallelConfig.audio_pad_for_ulysses``. * ``LTX2Attention._init_audio_modules``: TRTLLM→VANILLA downgrade for ``audio_attn1`` now keys on ``vgm.ulysses_size > 1`` (the condition that actually drives the mask requirement) instead of the dropped flag. * ``LTXModel.__init__``: drop the cached ``self._audio_pad_for_ulysses``. * ``LTXModel.configure_audio_ulysses``: always pad when the sharder is active. Removes the dead "no-pad mode" branch and the CP-without- Ulysses ``ValueError`` (now unreachable because padding always makes audio shardable). * Docstrings + the ``audio_padding_mask`` field comment updated. Tests adjusted: drop the parameter from the three call sites in test_ltx2_ulysses.py / test_ulysses_attention.py / test_fa4_key_padding_mask.py. Verified locally with the multi-GPU e2e LTXModel parity test (VANILLA backend, ws=2): both ``test_av_ulysses_no_audio_pad`` and ``test_av_ulysses_audio_pad`` PASS. The FA4 backend cases failed in this container due to an unrelated env regression already fixed on main (PR NVIDIA#13788, ``nvvm.fmax`` API change in ``nvidia_cutlass_dsl``); CI on a fresh container picks up the fix automatically. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

…class Reviewer suggested using ``enable_ulysses`` (PR NVIDIA#14012) to simplify our PR's hand-rolled v2a Ulysses wrapper. Two surgical changes that together delete the manual wrapper construction + the parallel ``use_ulysses_cross`` ctor flag. 1. Base ``Attention.__init__`` wrap dispatch (modules/attention.py): The combined gate ``(ring_size > 1 or ulysses_size > 1) and qkv_mode != SEPARATE_QKV`` excluded cross-attn from both Ring and Ulysses wrapping. Split it into two single-purpose gates: * Ring stays gated on ``qkv_mode != SEPARATE_QKV`` — Ring shards the seq dim and requires ``S_q == S_kv``. * Ulysses now keys on ``use_ulysses`` only — Ulysses shards the head dim and is value-preserving for ``S_q != S_kv``, so cross-attn is valid. Also fixes a latent bug where ``enable_ulysses=False`` would still build the wrapper while making the inner backend use the unsharded head count. 2. LTX2Attention (models/ltx2/transformer_ltx2.py): * Drop the ``use_ulysses_cross`` ctor flag. Caller passes a single ``use_ulysses=True`` for both self-attn and cross-attn modules. Internally, cross-attn carries an additional ``cp_size == 1`` gate (Ring/Attention2D + cross-attn Ulysses is unvalidated). * Delete the 30-line ``_ulysses_cross_attn = UlyssesAttention(...)`` hand wrap that bypassed the base class. With Step 1 the base class now builds the wrap for cross-attn too. * Collapse the cross-attn dual-attn pair into the existing self-attn dual-attn path (single ``_ulysses_attn`` / ``_plain_attn`` pair, single ``_has_dual_attn`` flag, single branch in ``set_ulysses_active`` / ``is_ulysses_active``). * Rename the locally derived flag from ``wants_ulysses`` to ``enable_ulysses`` so it matches the kwarg passed to the base class. Net: 94→58 lines in transformer_ltx2.py ctor + 30 fewer lines in toggle methods. ``use_ulysses_cross`` call site at ``video_to_audio_attn = LTX2Attention(..., use_ulysses=True)``. Verified: all 4 PR-touched test files PASS (26/26 in 6m57s) — e2e LTXModel parity (VANILLA + FA4, no_pad + audio_pad), multi-rank UlyssesAttention parity, FA4/Vanilla key_padding_mask single-rank. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-05-29T04:38:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-29T04:43:46Z

PR_Github #50968 [ run ] triggered by Bot. Commit: 581760e Link to invocation

tensorrt-cicd · 2026-05-29T08:15:34Z

PR_Github #50968 [ run ] completed with state SUCCESS. Commit: 581760e
/LLM/main/L0_MergeRequest_PR pipeline #40423 completed with status: 'SUCCESS'

CI Report

Link to invocation

NVIDIA#14044 added Ulysses sequence-parallelism for the v2a (video->audio) cross-attention; NVIDIA#13978 (async-Ulysses refactor) silently disabled it, so v2a fell back to all-gathering the full video K/V every block (per-rank comm grows with ulysses_size instead of scaling). Restore the v2a Ulysses path under pure Ulysses (cp_size == 1), preserving NVIDIA#13978's async-Ulysses (self-attn only) and the cp_size > 1 all-gather fallback: - modules/attention.py: head-shard the SEPARATE_QKV cross-attn inner backend when cp_size == 1. - transformer_ltx2.py: build the dual-attn pair for cross-attn under pure Ulysses and pass use_ulysses=True to video_to_audio_attn. Output is numerically unchanged; this is a communication/perf fix. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

github-actions Bot assigned luyiyun1021 May 12, 2026

luyiyun1021 changed the title ~~[None][feat] LTX-2 Ulysses cross-attention for v2a with audio padding~~ [TRTLLM-12653][feat] LTX-2 Ulysses cross-attention for v2a with audio padding May 13, 2026

luyiyun1021 requested a review from yibinl-nvidia May 13, 2026 06:08

luyiyun1021 marked this pull request as ready for review May 13, 2026 06:08

luyiyun1021 requested a review from a team as a code owner May 13, 2026 06:08

coderabbitai Bot reviewed May 13, 2026

View reviewed changes

NVShreyas reviewed May 13, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py Outdated

yibinl-nvidia reviewed May 14, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py

Comment thread tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py Outdated

NVShreyas mentioned this pull request May 15, 2026

[TRTLLM-11410][feat] MoT World Model Support #14012

Merged

1 task

NVShreyas reviewed May 21, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py Outdated

luyiyun1021 force-pushed the feat/ltx2-ulysses-v2a-cross-attn branch from d056233 to 6432850 Compare May 22, 2026 03:03

NVShreyas approved these changes May 22, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py Outdated

Comment thread tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py Outdated

luyiyun1021 force-pushed the feat/ltx2-ulysses-v2a-cross-attn branch from 3709672 to d60dc84 Compare May 23, 2026 12:38

luyiyun1021 requested a review from a team as a code owner May 23, 2026 12:38

luyiyun1021 requested a review from JunyiXu-nv May 23, 2026 12:38

luyiyun1021 force-pushed the feat/ltx2-ulysses-v2a-cross-attn branch from 3844859 to d60dc84 Compare May 24, 2026 11:27

luyiyun1021 force-pushed the feat/ltx2-ulysses-v2a-cross-attn branch from b82436a to 0530756 Compare May 27, 2026 04:14

zhenhuaw-me approved these changes May 28, 2026

View reviewed changes

luyiyun1021 added 14 commits May 28, 2026 19:56

luyiyun1021 force-pushed the feat/ltx2-ulysses-v2a-cross-attn branch from 693348b to 7831a01 Compare May 29, 2026 03:24

luyiyun1021 merged commit b1dfd30 into NVIDIA:main May 29, 2026
7 checks passed

luyiyun1021 mentioned this pull request Jun 12, 2026

[None][fix] LTX-2: re-enable Ulysses for v2a cross-attention #15303

Merged

Uh oh!

Conversation

luyiyun1021 commented May 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Why v2a, not a2v?

Per-rank communication volume (cross-attn K/V path), U = ulysses_size

Performance

NCCL kernel time, per rank, single denoise step (step 5 capture, 4× B200, CFG=1 × Ulysses=4, 768×1280, 121 frames)

End-to-end wall time (8× B200, CFG=2 × Ulysses=4, 768×1280, 121 frames, 40 steps)

Implementation

Backend safety for key_padding_mask

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yibinl-nvidia left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

luyiyun1021 commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

NVShreyas left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

luyiyun1021 commented May 24, 2026

Uh oh!

luyiyun1021 commented May 26, 2026

Uh oh!

luyiyun1021 commented May 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

NVShreyas commented May 26, 2026

Uh oh!

luyiyun1021 commented May 27, 2026

Uh oh!

luyiyun1021 commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

NVShreyas commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

luyiyun1021 commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

luyiyun1021 commented May 12, 2026 •

edited

Loading

Per-rank communication volume (cross-attn K/V path), `U = ulysses_size`

NCCL kernel time, per rank, single denoise step (`step 5` capture, 4× B200, CFG=1 × Ulysses=4, 768×1280, 121 frames)

Backend safety for `key_padding_mask`

coderabbitai Bot commented May 13, 2026 •

edited

Loading

yibinl-nvidia left a comment •

edited

Loading

NVShreyas left a comment •

edited

Loading

luyiyun1021 commented May 26, 2026 •

edited

Loading