[TRTLLM-14234][perf] LTX-2: replicate audio stream across Ulysses ranks by default

[luyiyun1021]

Summary by CodeRabbit

• Performance and Compatibility
  • Improved distributed audio processing for visual generation workflows.
  • Enhanced video-to-audio attention handling across different sharding configurations.
  • Improved support for FP4-optimized execution in conditional audio-processing scenarios.
  • Added audio padding and synchronization improvements to support varied processing layouts.

Description

Under Ulysses sequence-parallelism the LTX-2 DiT sharded both streams — video (15360 tokens) and audio (126). Audio is ~122x smaller and its self/cross-attention, GEMMs and FFN are latency/occupancy-bound (the audio NVFP4 GEMM is a single 128x128 kernel at ~7% SM / ~9% DRAM SOL, flat across uly1->8), so seq-sharding audio removes ~0 compute while adding per-block collectives (audio-self A2A + a2v all-gather). This PR makes the audio stream replicated on every Ulysses rank by default: audio self/cross-attention and FFN run local (no A2A / all-gather), and only v2a (audio-Q x video-KV — the one audio op that scales, via head-split) keeps Ulysses sharding, implemented as a local per-rank slice of the replicated Q plus an output all-gather around the unchanged async-Ulysses driver, so no attention-backend or C++ change is needed.

The mode is a single internal knob _audio_full_shard (default False = replicate); the resolved runtime state is a 3-value AudioShardMode — NONE (Ulysses off / single rank), CONDITIONAL (replicate, default), FULL (legacy seq-shard). There is no new user-facing config. Setting _audio_full_shard = True restores the exact original sharded path bit-for-bit (same fp4 modulation scales, a2v all-gather, input-shard/output-gather, and v2a async/sync driver), so the change is a default flip over two already-existing, validated code paths.

Changes (all in transformer_ltx2.py): configure_audio_ulysses seeds the mode and pads audio to a multiple of the Ulysses degree; the block forward routes v2a through slice-in -> unchanged driver -> _sp_all_gather in CONDITIONAL and reuses the existing SHARD path in FULL; audio self-attn Ulysses is gated off in replicate; the audio-cross fused-modulation emits bf16 (not fp4) only for the tensor v2a slices (an NVFP4 output carries swizzled per-block scale-factors that cannot be sequence-sliced per rank), and the audio-cross GEMMs re-quantize the tiny audio internally.

Test Coverage

Correctness — 8x B200, seed 42, 768x1280, 121 frames, LPIPS (AlexNet) over decoded mp4 frames:

• replicate vs sharded baseline: mean 0.112, at or below the sharded-vs-sharded run-to-run noise floor 0.129 (two identical-config generations — LTX-2 FA4/autotune is not bit-deterministic); a second cross-check (replicate vs a second sharded baseline) is 0.100. So replicate is numerically equivalent to the sharded path within the model's own non-determinism; the montage below shows the same scene / subject / composition / quality (only fine-detail drift).

Performance — single denoise step, per-layer µs, nsys (--cuda-graph-trace=node), sharded -> replicate:

metric — 8-GPU (cfg2 x uly4)	sharded	replicate	Δ
audio self + v2a NCCL SendRecv	2.90	0.56	−81%
a2v audio all-gather	1.35	0.71	−47%
REAL COMM total	22.86	19.60	−14%
COMPUTE total	87.31	86.49	~flat (within capture noise)
GPU wall / denoise step	115.1	109.0	−5.3%
Ulysses scaling efficiency (1->8 GPU)	65%	69%	+4 pts

4-GPU (cfg2 x uly2) is consistent: REAL COMM 22.37 -> 19.19 (−14%), GPU wall 185.8 -> 177.9 (−4.3%). End-to-end 8x B200 generate (denoise -> VAE -> saved video) completes for both modes: wall 5.56 -> 5.44 s, denoise-only 4.39 -> 4.27 s (single run, indicative — audio is a small fraction of the full pipeline). The gain is modest in relative terms (audio is ~5% of the step) but free and grows with Ulysses degree.

Output Quality (LPIPS)

Replicate vs sharded baseline, 8x B200, seed 42, 768x1280, 121 frames. Header shows mean/min/max LPIPS; each row shows that frame's per-frame LPIPS. Mean 0.112 is at/below the 0.129 sharded-vs-sharded noise floor — same scene and quality, only fine-detail drift.

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.

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58640 [ run ] triggered by Bot. Commit: ac85793 Link to invocation

[coderabbitai]

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: 9fb45a2c-0aa2-4118-9b5f-e215cfbcfd9f

📥 Commits

Reviewing files that changed from the base of the PR and between e2ce7a6 and ac85793.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py

---

📝 Walkthrough

Walkthrough

LTX2 now models audio Ulysses distribution with NONE, CONDITIONAL, and FULL modes. Configuration, positional embeddings, forward-time sharding, v2a attention, and fused-kernel scale handling branch on the derived mode.

Changes

Audio Ulysses shard modes

Layer / File(s)	Summary
Shard mode state and contract tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py	Adds AudioShardMode and replaces _audio_is_sharded with _audio_full_shard plus derived mode properties.
Ulysses configuration and sequence preparation tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py	Pads audio for Ulysses divisibility, configures per-attention activation, updates PE sharding, and gates forward-time audio sharding and gathering on FULL.
Conditional v2a attention execution tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py	Slices replicated audio queries for local v2a execution, supports synchronous and asynchronous attention paths, gathers outputs, and adjusts FP4 scale inputs in conditional mode.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant LTXModel
  participant Ulysses
  participant BasicAVTransformerBlock
  participant VideoToAudioAttention
  LTXModel->>Ulysses: configure and pad audio sequence
  LTXModel->>BasicAVTransformerBlock: pass derived audio shard mode
  BasicAVTransformerBlock->>Ulysses: slice replicated audio Q in CONDITIONAL mode
  Ulysses->>VideoToAudioAttention: execute local v2a attention
  VideoToAudioAttention-->>Ulysses: return local output
  Ulysses->>BasicAVTransformerBlock: all-gather audio output

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 54.55% 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 names the JIRA ticket, type, and the main change to replicate audio by default.
Description check	✅ Passed	The description follows the template with Description, Test Coverage, and PR Checklist sections filled in.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[zhenhuaw-me]

Are we trying to maintain any compatibility? It seems to me this is a "always have benefit optimization", any cases we want to switch to the "legacy mode"?

[luyiyun1021]

I‘ve done a quick test on 15s video shape's audio computation, in this case the audio part starts to scale a little bit(still worse compared to video), but I haven't done a thorough test and detailed analysis on this config since we currently focuses on 5s video. But I think the kind of comparsion between computation saved and comm overhead is a bit tricky while the shape changes. So I leave the full shard path here in case we need to re-enable it.