[None][fix] LTX-2: re-enable Ulysses for v2a cross-attention by luyiyun1021 · Pull Request #15303 · NVIDIA/TensorRT-LLM

luyiyun1021 · 2026-06-12T08:36:29Z

Description

#14044 added Ulysses sequence-parallelism for the v2a (video→audio) cross-attention; #13978 (async-Ulysses refactor) silently disabled it — cross-attn was excluded from the Ulysses path at both the base-class use_ulysses gate and the LTX-2 dual-attn gate, and use_ulysses=True was dropped from the v2a construction. As a result v2a fell back to all-gathering the full video K/V every block, so its per-rank communication grows with ulysses_size instead of scaling.

This restores the v2a Ulysses path under pure Ulysses (cp_size == 1), while preserving #13978's async-Ulysses (self-attn only) and keeping the combined ring/attn2d + Ulysses (cp_size > 1) all-gather fallback unchanged:

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 (v2a uses the Ulysses 4-collective a2a instead of the full-video all-gather).

Test Coverage

Existing multi-GPU parity tests (multi_gpu/test_ltx2_ulysses.py, multi_gpu/test_ltx2_async_ulysses.py) cover correctness/parity for the v2a + async paths.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

Release Notes

Improvements
- Enhanced attention mechanism configuration in visual generation models for improved performance.
- Optimized cross-attention setup for audio-visual task processing.

coderabbitai · 2026-06-12T08:40:26Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0a4c9aff-45f7-4360-8a80-fab73d431462

📥 Commits

Reviewing files that changed from the base of the PR and between 44550bc and c0fb1c2.

📒 Files selected for processing (2)

tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py
tensorrt_llm/_torch/visual_gen/modules/attention.py

📝 Walkthrough

Walkthrough

This PR updates Ulysses sequence parallelism wiring across visual generation attention modules. The base Attention module extends its Ulysses gating logic to account for context parallelism size, while LTX2Attention broadens its dual-attention pair creation for cross-attention and enables Ulysses for audio-visual cross-attention paths.

Changes

Ulysses Attention Configuration

Layer / File(s)	Summary
Base Attention module Ulysses gating with cp_size `tensorrt_llm/_torch/visual_gen/modules/attention.py`	`Attention.__init__` extends the `use_ulysses` condition to incorporate `cp_size` from `visual_gen_mapping`, enabling the Ulysses path when `cp_size == 1` alongside existing `qkv_mode` and `async_ulysses` checks.
LTX2 Attention dual-pair and AV cross-attention Ulysses setup `tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py`	`LTX2Attention.__init__` broadens the condition for creating runtime-togglable Ulysses/plain dual-attention pairs to include cross-attention when `cp_size == 1`. `BasicAVTransformerBlock` now configures `video_to_audio_attn` with `use_ulysses=True` to enable Ulysses for the audio-visual cross-attention path.

Possibly related PRs

NVIDIA/TensorRT-LLM#15141: Modifies the same visual_gen Ulysses wiring paths in attention initialization by updating how Attention.__init__ gates use_ulysses and how LTX2Attention applies it.

Suggested reviewers

brb-nv
yuxianq
chang-l

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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: re-enabling Ulysses for v2a cross-attention in the LTX-2 model, following the proper format with [None][fix].
Description check	✅ Passed	The description comprehensively explains the issue (Ulysses disabled by `#13978` regression), the solution (restore it for pure Ulysses), specific code changes, impact, and test coverage, though PR Checklist items are not explicitly marked.
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.

_{✏️ 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.}

luyiyun1021 · 2026-06-12T08:42:11Z

/bot run

tensorrt-cicd · 2026-06-12T08:48:18Z

PR_Github #53857 [ run ] triggered by Bot. Commit: c0fb1c2 Link to invocation

NVShreyas

LGTM after requested change

tensorrt-cicd · 2026-06-12T18:25:37Z

PR_Github #53857 [ run ] completed with state SUCCESS. Commit: c0fb1c2
/LLM/main/L0_MergeRequest_PR pipeline #42964 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

…es + log) Address NVShreyas's review on PR NVIDIA#15303: - wrap_parallel_attention now takes a use_ulysses argument and gates the UlyssesAttention wrap on it, so the outer wrap stays consistent with the inner backend's head-count decision (avoids a full-head backend being wrapped in Ulysses). The single use_ulysses flag now drives both the inner head count and the wrap. - Add a debug log when Ulysses is disabled despite ulysses_size>1 and sequence-parallel on (a SEPARATE_QKV cross-attn that is neither async nor pure-Ulysses, cp_size>1), so the all-gather fallback is observable. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-06-15T05:43:00Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-15T05:48:18Z

PR_Github #54227 [ run ] triggered by Bot. Commit: 42cc7f1 Link to invocation

…es + log) Address NVShreyas's review on PR NVIDIA#15303: - wrap_parallel_attention now takes a use_ulysses argument and gates the UlyssesAttention wrap on it, so the outer wrap stays consistent with the inner backend's head-count decision (avoids a full-head backend being wrapped in Ulysses). The single use_ulysses flag now drives both the inner head count and the wrap. - Add a debug log when Ulysses is disabled despite ulysses_size>1 and sequence-parallel on (a SEPARATE_QKV cross-attn that is neither async nor pure-Ulysses, cp_size>1), so the all-gather fallback is observable. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-06-15T10:22:39Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-15T10:29:01Z

PR_Github #54284 [ run ] triggered by Bot. Commit: a6166db Link to invocation

tensorrt-cicd · 2026-06-15T10:33:12Z

PR_Github #54227 [ run ] completed with state ABORTED. Commit: 42cc7f1

Link to invocation

tensorrt-cicd · 2026-06-15T18:42:31Z

PR_Github #54284 [ run ] completed with state SUCCESS. Commit: a6166db
/LLM/main/L0_MergeRequest_PR pipeline #43355 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

luyiyun1021 · 2026-06-16T02:54:31Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T02:59:55Z

PR_Github #54436 [ run ] triggered by Bot. Commit: a6166db Link to invocation

…es + log) Address NVShreyas's review on PR NVIDIA#15303: - wrap_parallel_attention now takes a use_ulysses argument and gates the UlyssesAttention wrap on it, so the outer wrap stays consistent with the inner backend's head-count decision (avoids a full-head backend being wrapped in Ulysses). The single use_ulysses flag now drives both the inner head count and the wrap. - Add a debug log when Ulysses is disabled despite ulysses_size>1 and sequence-parallel on (a SEPARATE_QKV cross-attn that is neither async nor pure-Ulysses, cp_size>1), so the all-gather fallback is observable. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-06-16T05:15:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T05:20:41Z

PR_Github #54471 [ run ] triggered by Bot. Commit: 9d701e3 Link to invocation

tensorrt-cicd · 2026-06-16T05:24:44Z

PR_Github #54436 [ run ] completed with state ABORTED. Commit: a6166db

Link to invocation

tensorrt-cicd · 2026-06-16T13:06:29Z

PR_Github #54471 [ run ] completed with state SUCCESS. Commit: 9d701e3
/LLM/main/L0_MergeRequest_PR pipeline #43533 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

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>

…es + log) Address NVShreyas's review on PR NVIDIA#15303: - wrap_parallel_attention now takes a use_ulysses argument and gates the UlyssesAttention wrap on it, so the outer wrap stays consistent with the inner backend's head-count decision (avoids a full-head backend being wrapped in Ulysses). The single use_ulysses flag now drives both the inner head count and the wrap. - Add a debug log when Ulysses is disabled despite ulysses_size>1 and sequence-parallel on (a SEPARATE_QKV cross-attn that is neither async nor pure-Ulysses, cp_size>1), so the all-gather fallback is observable. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-06-24T07:31:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-24T07:38:49Z

PR_Github #55441 [ run ] triggered by Bot. Commit: 811b044 Link to invocation

tensorrt-cicd · 2026-06-24T15:46:55Z

PR_Github #55441 [ run ] completed with state SUCCESS. Commit: 811b044
/LLM/main/L0_MergeRequest_PR pipeline #44377 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

luyiyun1021 · 2026-06-24T16:06:37Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-24T16:13:03Z

PR_Github #55534 [ run ] triggered by Bot. Commit: 811b044 Link to invocation

tensorrt-cicd · 2026-06-24T20:22:01Z

PR_Github #55534 [ run ] completed with state SUCCESS. Commit: 811b044
/LLM/main/L0_MergeRequest_PR pipeline #44457 completed with status: 'SUCCESS'

CI Report

Link to invocation

…15303) Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 requested review from a team as code owners June 12, 2026 08:36

luyiyun1021 requested a review from QiJune June 12, 2026 08:36

github-actions Bot assigned luyiyun1021 Jun 12, 2026

NVShreyas reviewed Jun 12, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/modules/attention.py

NVShreyas approved these changes Jun 12, 2026

View reviewed changes

litaotju approved these changes Jun 15, 2026

View reviewed changes

luyiyun1021 force-pushed the fix/ltx2-v2a-ulysses-regression branch from 42cc7f1 to a6166db Compare June 15, 2026 10:22

luyiyun1021 force-pushed the fix/ltx2-v2a-ulysses-regression branch from a6166db to 9d701e3 Compare June 16, 2026 05:15

luyiyun1021 added 2 commits June 24, 2026 00:21

luyiyun1021 force-pushed the fix/ltx2-v2a-ulysses-regression branch from 9d701e3 to 811b044 Compare June 24, 2026 07:25

luyiyun1021 enabled auto-merge (squash) June 24, 2026 16:06

luyiyun1021 merged commit bd6db6c into NVIDIA:main Jun 24, 2026
8 checks passed

BrianLi23 pushed a commit to BrianLi23/TensorRT-LLM that referenced this pull request Jul 9, 2026

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

25b0152

…15303) Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 mentioned this pull request Jul 13, 2026

[None][feat] LTX-2: run v2a cross-attention under distributed Attention2D #16301

Draft

Uh oh!

Conversation

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

Uh oh!

Description

Test Coverage

GitHub Bot Help

Summary by CodeRabbit

Release Notes

Uh oh!

coderabbitai Bot commented Jun 12, 2026

Walkthrough

Changes

Possibly related PRs

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

luyiyun1021 commented Jun 12, 2026

Uh oh!

tensorrt-cicd commented Jun 12, 2026

Uh oh!

Uh oh!

NVShreyas left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jun 12, 2026

Uh oh!

luyiyun1021 commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

luyiyun1021 commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

luyiyun1021 commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

luyiyun1021 commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

luyiyun1021 commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

luyiyun1021 commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

luyiyun1021 commented Jun 12, 2026 •

edited by coderabbitai Bot

Loading