[None][fix] LTX-2 audio PE pad: use token-axis seq_dim=1 for token-major rope by luyiyun1021 · Pull Request #14818 · NVIDIA/TensorRT-LLM

luyiyun1021 · 2026-06-01T09:07:22Z

Summary

Fix LTX-2 audio PE padding regression: prepare_text_cache was padding the head axis (dim 2) instead of the token axis (dim 1) for 4D positional embeddings, leaving cos.shape[1] short of the padded latent. The fused split-norm-rope kernel then trips its cos_emb.size(0) == num_tokens check at runtime with cos=N vs num_tokens=N+1 (e.g. cos=31, num_tokens=32 at WAN-style warmup; cos=7, num_tokens=8 at the unit-test scale).
Update the audio-padding unit test to default to rope_type=SPLIT so it actually exercises the production token-major rope layout. Previously the test fixture used INTERLEAVED (3D PE), which made the ndim == 4 -> seq_dim=2 heuristic fall through to the correct seq_dim=1 by accident and missed the regression entirely.

Root cause

a173175ac2 [None][perf] ltx2: switch SPLIT-rope cos/sin to token-major [B, T, H, D] changed the SPLIT-rope layout from [B, H, T, D] to token-major [B, T, H, D] and removed all swapaxes(1, 2) plumbing; at the same time it correctly set pe_seq_dim = 1 in the (then-existing) code that touched PE shape.

b1dfd3094c [TRTLLM-12653][feat] LTX-2 Ulysses cross-attention for v2a with audio padding (#14044) later introduced the _pad_pe helper plus its call site in prepare_text_cache to extend audio PE under Ulysses. The new call site picked seq_dim via:

pe_seq_dim = 2 if (a_pe is not None and a_pe[0].ndim == 4) else 1

That heuristic is the pre-token-major convention ([B, H, T, D]). For the post-token-major SPLIT layout the actual seq axis is dim 1, so the pad grows the head axis instead and downstream _make_pe_local's cos[:, s:e] still shards the un-padded token axis — exactly the off-by-_audio_pad we observe at runtime.

The existing test_av_ulysses_audio_pad did not catch this because:

LTXModel defaults rope_type=LTXRopeType.INTERLEAVED, and the test fixture left that default.
INTERLEAVED produces 3D PE, so the ndim == 4 branch is never taken and the call lands on seq_dim=1 by accident.

Fix

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

# cos/sin are token-major: SPLIT rope -> [B, T, H, D]; INTERLEAVED -> [B, T, D].
# The token (seq) axis is dim 1 in both, matching `_make_pe_local`'s
# `cos[:, s:e]` shard.
if self._audio_pad > 0:
    a_pe = self._pad_pe(a_pe, self._audio_pad, seq_dim=1)
    a_cross_pe = self._pad_pe(a_cross_pe, self._audio_pad, seq_dim=1)

This matches _make_pe_local's slicing on dim 1 and is uniform across SPLIT/INTERLEAVED.

tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py:

Thread rope_type_name through _logic_ltx2_av_ulysses_vs_single_gpu and default to "SPLIT" so the audio-padding parity test exercises the production layout that motivated the original _pad_pe plumbing.
Without the fix, test_av_ulysses_audio_pad[VANILLA] and [FA4] both fail with the expected cos_emb.size(0) (7) must equal num_tokens (8) mismatch (verified locally on dev image).
With the fix, both pass.

Test plan

pytest tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py -v (4 cases pass — no-pad / pad × VANILLA / FA4) on a 2-GPU dev image.
Reproduced the bug on tip of main with the new SPLIT-default test (fails with the production-side cos_emb=N, num_tokens=N+1 message).
Confirmed both updated tests pass after the fix.

PR Checklist

Please check this if you reviewed the Coding Guidelines
PR title format follows the convention [JIRA/NVBUG/None][type] description
Test coverage updated; the regression test now defaults to the production rope layout
No documentation update needed; the fix is internal to LTX-2 audio padding

Summary by CodeRabbit

Bug Fixes
- Fixed padding inconsistency in audio text positional embeddings for LTX-2 visual generation.
Tests
- Enhanced test coverage for RoPE layout configurations in LTX-2 testing.

luyiyun1021 · 2026-06-01T09:08:22Z

/bot run --disable-fail-fast

coderabbitai · 2026-06-01T09:11:31Z

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: c7771a92-cf7a-460d-b1ce-4bf00a52330c

📥 Commits

Reviewing files that changed from the base of the PR and between 2e6f602 and 54502ef.

📒 Files selected for processing (2)

tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py
tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py

📝 Walkthrough

Walkthrough

This PR fixes audio text rotary positional embedding padding in LTXModel's prepare_text_cache() method. When audio padding is active, the code now unconditionally uses seq_dim=1 for embedding padding, matching the token-major layout expected by attention logic. Test infrastructure is extended to support configurable RoPE layout selection and documents the padding behavior.

Changes

Audio PE padding fix and test improvements

Layer / File(s)	Summary
Audio PE padding fix in prepare_text_cache() `tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py`	When audio padding is active, `a_pe` and `a_cross_pe` are now padded using `seq_dim=1` unconditionally, replacing prior heuristic-based `seq_dim` selection. Adds clarifying comments about prior shape/layout mismatches.
Test helper RoPE configurability and validation `tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py`	Helper function signature extended with `rope_type_name` parameter (default `"SPLIT"`); both reference and Ulysses model constructions now pass the selected `rope_type`. Test docstrings expanded to document token-axis padding behavior and prior regressions.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the primary fix: changing seq_dim to 1 for token-axis padding in LTX-2 audio PE with token-major rope layouts.
Description check	✅ Passed	The PR description is comprehensive and complete, covering root cause, fix, test coverage, and addressing all checklist items from the template.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.}

tensorrt-cicd · 2026-06-01T09:13:54Z

PR_Github #51346 [ run ] triggered by Bot. Commit: 54502ef Link to invocation

…jor rope `prepare_text_cache` extended the audio PE with a `seq_dim` derived via `ndim == 4 -> seq_dim=2`, which predates the SPLIT-rope switch to the token-major `[B, T, H, D]` layout. With token-major rope the token axis is dim 1; the heuristic instead pads the head axis and `_make_pe_local` still shards the unpadded token axis, leaving `cos.shape[1]` short of the padded latent. The fused split-norm-rope kernel then rejects the cos/num_tokens mismatch at runtime (e.g. `cos=31, num_tokens=32` for WAN-style warmup; `cos=7, num_tokens=8` at the unit-test scale). Pin the seq_dim to 1 to match `_make_pe_local`'s `cos[:, s:e]` shard; this is uniform across SPLIT (4D) and INTERLEAVED (3D) layouts. Default the audio-padding unit test to `rope_type=SPLIT` so the regression that motivated the original `_pad_pe` plumbing is actually exercised on the production layout. The previous INTERLEAVED default produced 3D PE, never hit the buggy 4D branch, and silently passed. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 · 2026-06-01T09:32:03Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-01T09:37:37Z

PR_Github #51349 [ run ] triggered by Bot. Commit: 27589d8 Link to invocation

tensorrt-cicd · 2026-06-01T09:42:36Z

PR_Github #51346 [ run ] completed with state ABORTED. Commit: 54502ef

Link to invocation

tensorrt-cicd · 2026-06-01T14:34:30Z

PR_Github #51349 [ run ] completed with state SUCCESS. Commit: 27589d8
/LLM/main/L0_MergeRequest_PR pipeline #40760 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-01T16:17:12Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-01T16:22:55Z

PR_Github #51379 [ run ] triggered by Bot. Commit: 27589d8 Link to invocation

tensorrt-cicd · 2026-06-01T21:29:39Z

PR_Github #51379 [ run ] completed with state FAILURE. Commit: 27589d8
/LLM/main/L0_MergeRequest_PR pipeline #40790 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

Link to invocation

luyiyun1021 · 2026-06-02T04:15:43Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-02T04:21:12Z

PR_Github #51499 [ run ] triggered by Bot. Commit: 27589d8 Link to invocation

luyiyun1021 · 2026-06-02T10:55:34Z

/bot run --disable-fail-fast

luyiyun1021 · 2026-06-02T11:01:35Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-02T11:50:50Z

PR_Github #51499 [ run ] completed with state SUCCESS. Commit: 27589d8
/LLM/main/L0_MergeRequest_PR pipeline #40904 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-02T13:50:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-02T13:58:41Z

PR_Github #51600 [ run ] triggered by Bot. Commit: 27589d8 Link to invocation

tensorrt-cicd · 2026-06-02T14:08:23Z

PR_Github #51603 [ run ] triggered by Bot. Commit: 27589d8 Link to invocation

tensorrt-cicd · 2026-06-02T14:13:00Z

PR_Github #51600 [ run ] completed with state ABORTED. Commit: 27589d8

Link to invocation

tensorrt-cicd · 2026-06-02T14:13:16Z

PR_Github #51613 [ run ] triggered by Bot. Commit: 27589d8 Link to invocation

tensorrt-cicd · 2026-06-02T14:16:56Z

PR_Github #51603 [ run ] completed with state ABORTED. Commit: 27589d8

Link to invocation

tensorrt-cicd · 2026-06-02T18:24:13Z

PR_Github #51613 [ run ] completed with state SUCCESS. Commit: 27589d8
/LLM/main/L0_MergeRequest_PR pipeline #40997 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-03T02:11:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-03T02:16:44Z

PR_Github #51729 [ run ] triggered by Bot. Commit: 27589d8 Link to invocation

tensorrt-cicd · 2026-06-03T03:35:06Z

PR_Github #51729 [ run ] completed with state SUCCESS. Commit: 27589d8
/LLM/main/L0_MergeRequest_PR pipeline #41104 completed with status: 'SUCCESS'

CI Report

Link to invocation

luyiyun1021 requested a review from a team as a code owner June 1, 2026 09:07

github-actions Bot assigned luyiyun1021 Jun 1, 2026

luyiyun1021 requested a review from zhenhuaw-me June 1, 2026 09:08

luyiyun1021 force-pushed the fix-ltx2-audio-pe-pad-seq-dim branch from 54502ef to 8609028 Compare June 1, 2026 09:22

luyiyun1021 force-pushed the fix-ltx2-audio-pe-pad-seq-dim branch from 8609028 to 27589d8 Compare June 1, 2026 09:30

zhenhuaw-me approved these changes Jun 1, 2026

View reviewed changes

luyiyun1021 enabled auto-merge (squash) June 1, 2026 16:17

luyiyun1021 merged commit 328ef0b into NVIDIA:main Jun 3, 2026
8 checks passed

Uh oh!

Conversation

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

Uh oh!

Summary

Root cause

Fix

Test plan

PR Checklist

Summary by CodeRabbit

Uh oh!

luyiyun1021 commented Jun 1, 2026

Uh oh!

coderabbitai Bot commented Jun 1, 2026

Walkthrough

Changes

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

luyiyun1021 commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

luyiyun1021 commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

luyiyun1021 commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

luyiyun1021 commented Jun 2, 2026

Uh oh!

luyiyun1021 commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

luyiyun1021 commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

luyiyun1021 commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

luyiyun1021 commented Jun 1, 2026 •

edited by coderabbitai Bot

Loading