[Model][Perf] Optimize LTX Ulysses sequence parallelism (#7079)

Signed-off-by: mglyn <1203789601@qq.com>

docs/user_guide/diffusion_features.md

@@ -162,9 +162,9 @@ The following tables show which models support each feature:
 | **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.2-S2V**               |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **LTX-2.5**                  |     ❌     | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) |         ❓         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |    ❓ (FP8)     |        ❌         |
+| **LTX-2**                    |     ❌     |     ✅      | ✅ (Ulysses only) |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      | ✅ (Ulysses only) |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **LTX-2.5**                  |     ❌     | ❓ (one-stage) | ✅* (Ulysses only) | ✅ (Full only) |         ❓         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |    ❓ (FP8)     |        ❌         |
 |  **SANA-Video-2B T2V I2V**   |     ❌     |     ❌      |   ✅ (`--usp`, frame-sharded)   |       ✅       |  ✅ (TP=2 only)   |         ❌         |   ❌   |            ❌             |          ❌          |       ❌       |        ❌        |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |

recipes/LTX/LTX-2.5.md

@@ -110,6 +110,29 @@ vllm serve Lightricks/LTX-2.5-Diffusers \
 
 Use the T2V or I2V request below after the server is ready.
 
+### Multi-GPU low-latency path
+
+For low-latency 1920x1088 distilled two-stage serving, match the Ulysses and
+VAE patch-parallel degrees to the GPU count. For example, on four GPUs:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+vllm serve Lightricks/LTX-2.5-Diffusers \
+  --omni \
+  --model-class-name LTX2DistilledTwoStagePipeline \
+  --usp 4 \
+  --vae-patch-parallel-size 4 \
+  --vae-use-tiling \
+  --host 0.0.0.0 \
+  --port 8000 \
+  --stage-init-timeout 900
+```
+
+Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local
+SP attention head counts must be divisible by `ulysses_degree`. Use each
+phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail.
+Each GPU must still fit the resident weights. Warm up with the production shape.
+
 ## Offline inference
 
 Choose values from the pipeline table. For example, the distilled two-stage

recipes/LTX/LTX-2.md

@@ -101,6 +101,12 @@ sigmas and input latents.
 Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
 and automatically switches to dynamic LoRA when quantization is enabled.
 
+## Sequence Parallelism
+
+Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local
+SP attention head counts must be divisible by `ulysses_degree`. Use each
+phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail.
+
 ## Serving
 
 Start either one-stage checkpoint:
@@ -335,9 +341,8 @@ bundled offline CLI do not currently expose `sigmas`.
   `--max-num-seqs 1`.
 - `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
   modality guidance, and rescale-compatible prediction gathering.
-- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
-  Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
-  divisible request shape.
+- Ulysses masks padded audio latents; video sequences must satisfy the
+  [sequence-parallel divisibility constraints](#sequence-parallelism).
 - Cache-DiT is one-stage only; multi-stage configurations are rejected.
 
 ## Operational Notes

tests/diffusion/models/ltx2/test_ltx25_pipeline.py

@@ -5,6 +5,7 @@
 
 import json
 from dataclasses import replace
+from threading import Lock
 from types import SimpleNamespace
 from typing import Any
 
@@ -500,6 +501,32 @@ def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypa
     assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
 
 
+def test_ltx_runtime_profiles_phase_methods_separately(monkeypatch):
+    pipe = object.__new__(LTX2DistilledTwoStagePipeline)
+    pipe.enable_diffusion_pipeline_profiler = True
+    pipe._profiler_lock = Lock()
+    pipe._stage_durations = {}
+    monkeypatch.setattr(ltx2_runtime.current_omni_platform, "is_available", lambda: False)
+
+    with pipe._profile_phase_method("run_phase", "generate_lowres"):
+        pass
+    with pipe._profile_phase_method("run_phase", "refine"):
+        pass
+    with pipe._profile_phase_method("decode_phase"):
+        pass
+    output = DiffusionOutput()
+    assert pipe._attach_profiled_stage_durations(output) is output
+
+    expected_keys = {
+        "LTX2DistilledTwoStagePipeline.generate_lowres.run_phase",
+        "LTX2DistilledTwoStagePipeline.refine.run_phase",
+        "LTX2DistilledTwoStagePipeline.decode_phase",
+    }
+    assert set(pipe.stage_durations) == expected_keys
+    assert set(output.stage_durations) == expected_keys
+    assert all(duration >= 0.0 for duration in pipe.stage_durations.values())
+
+
 def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -103,7 +103,7 @@ def test_prepare_audio_latents_samples_directly_in_packed_token_space():
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
 
 
-def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
+def test_prepare_audio_latents_pads_short_audio_for_sp():
     pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
 
     latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for LTX guidance and forward-parallel behavior."""
 
@@ -85,14 +85,14 @@ class TestCFGParallelHelpers:
 
         LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size)
 
-    def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch):
+    def test_guidance_parallel_preserves_full_guidance_plan(self, monkeypatch):
         from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
         from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
         pipe = object.__new__(LTX2Pipeline)
-        req = SimpleNamespace(is_dummy_run=lambda: True)
+        req = SimpleNamespace()
         request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
 
         guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -175,7 +175,19 @@ def test_ltx_device_uint8_output_rejects_invalid_contract(video, error):
         _prepare_ltx2_video_output(video, do_normalize=True)
 
 
-def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
+@pytest.mark.parametrize(
+    "parallel_config",
+    [
+        None,
+        SimpleNamespace(ulysses_degree=4),
+        SimpleNamespace(tensor_parallel_size=2),
+        SimpleNamespace(use_hsdp=True, hsdp_shard_size=2),
+        SimpleNamespace(ulysses_degree=2, tensor_parallel_size=2),
+        SimpleNamespace(ulysses_degree=2, use_hsdp=True, hsdp_shard_size=2),
+    ],
+    ids=["single", "ulysses", "tp", "hsdp", "tp-ulysses", "hsdp-ulysses"],
+)
+def test_ltx_supported_parallel_configs_keep_output_profiling(tmp_path, monkeypatch, parallel_config):
     from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
 
     (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
@@ -193,7 +205,13 @@ def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
     monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
     monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup)
 
-    LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+    LTX2Pipeline(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            parallel_config=parallel_config,
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
 
     assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"]
     assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"]
@@ -383,14 +401,42 @@ def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expecte
         )
 
 
-def test_ltx_rejects_advanced_uaa_before_component_initialization():
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+        LTX2T2VDMD2Pipeline,
+        LTX2I2VDMD2Pipeline,
+    ],
+)
+@pytest.mark.parametrize(
+    ("parallel_config", "error"),
+    [
+        (SimpleNamespace(ulysses_mode="advanced_uaa"), "does not support ulysses_mode='advanced_uaa'"),
+        (SimpleNamespace(ring_degree=2), "pure Ulysses sequence parallelism only"),
+        (SimpleNamespace(ulysses_degree=2, ring_degree=2), "pure Ulysses sequence parallelism only"),
+        (SimpleNamespace(allgather_degree=2), "pure Ulysses sequence parallelism only"),
+    ],
+    ids=["advanced-uaa", "ring", "hybrid", "allgather-kv"],
+)
+def test_ltx_rejects_unsupported_sp_before_component_initialization(pipeline_cls, parallel_config, error, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    def unexpected_model_access(*_args, **_kwargs):
+        pytest.fail("Unsupported SP must be rejected before accessing model files or initializing components")
+
+    monkeypatch.setattr(ltx2_runtime, "detect_ltx_model_version", unexpected_model_access)
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", unexpected_model_access)
     od_config = SimpleNamespace(
         model="unused",
-        parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
+        parallel_config=parallel_config,
     )
 
-    with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
-        LTX2Pipeline(od_config=od_config)
+    with pytest.raises(ValueError, match=error):
+        pipeline_cls(od_config=od_config)
 
 
 def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
@@ -1409,39 +1455,19 @@ def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sam
         pipeline.forward(req, **direct_kwargs)
 
 
-def test_ltx_distilled_dummy_run_uses_fixed_recipe_values():
-    from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
-    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-    req = DiffusionRequestBatch(
-        [
-            OmniDiffusionRequest(
-                prompt="dummy run",
-                sampling_params=OmniDiffusionSamplingParams(
-                    height=512,
-                    width=512,
-                    num_frames=LTX2DistilledPipeline.dummy_run_num_frames,
-                    num_inference_steps=1,
-                    guidance_scale=0.0,
-                ),
-                request_id=DUMMY_DIFFUSION_REQUEST_ID,
-            )
-        ]
-    )
-    pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline)
-    seen = {}
-
-    def fake_run_recipe(_req, request_inputs, **_kwargs):
-        seen["request_inputs"] = request_inputs
-        return "dummy-output"
-
-    object.__setattr__(pipeline, "_run_recipe", fake_run_recipe)
-
-    assert pipeline.forward(req) == "dummy-output"
-    assert seen["request_inputs"].num_frames == 1
-    assert seen["request_inputs"].num_inference_steps == 8
-    assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only()
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+        LTX2T2VDMD2Pipeline,
+        LTX2I2VDMD2Pipeline,
+    ],
+)
+def test_ltx_pipelines_disable_startup_dummy_run(pipeline_cls):
+    assert pipeline_cls.dummy_run_num_frames == 0
 
 
 def test_ltx_custom_sigmas_bypass_scheduler_shifting():

tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.models.ltx2.ltx2_sequence_parallel import (
+    LTX2VideoToAudioParallelAttention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _run_video_to_audio_parity(rank: int, world_size: int, master_port: int) -> None:
+    dist.init_process_group(
+        "gloo",
+        init_method=f"tcp://127.0.0.1:{master_port}",
+        rank=rank,
+        world_size=world_size,
+        timeout=timedelta(seconds=60),
+    )
+    try:
+        # Non-contiguous SP subgroups match TP-fastest rank ordering. Different
+        # seeds across groups catch collectives that accidentally use WORLD.
+        group_count = world_size // 2
+        ulysses_group = dist.group.WORLD
+        if group_count > 1:
+            for offset in range(group_count):
+                ranks = list(range(offset, world_size, group_count))
+                group = dist.new_group(ranks)
+                if rank in ranks:
+                    ulysses_group = group
+        ulysses_rank = dist.get_rank(ulysses_group)
+        sp_group = SimpleNamespace(
+            ring_world_size=1,
+            ulysses_group=ulysses_group,
+            ulysses_world_size=2,
+            ulysses_rank=ulysses_rank,
+        )
+        strategy = LTX2VideoToAudioParallelAttention(sp_group)
+
+        for batch_size, masked in [(1, False), (2, True)]:
+            torch.manual_seed(17 + rank % group_count)
+            query = torch.randn(batch_size, 3, 4, 8)
+            global_key = torch.randn(batch_size, 6, 4, 8)
+            global_value = torch.randn(batch_size, 6, 4, 8)
+            key = global_key.chunk(2, dim=1)[ulysses_rank].contiguous()
+            value = global_value.chunk(2, dim=1)[ulysses_rank].contiguous()
+            mask = None
+            if masked:
+                mask = torch.ones(batch_size, 1, 1, 6, dtype=torch.bool)
+                mask[0, ..., -1] = False
+                mask[1, ..., -2:] = False
+            metadata = AttentionMetadata(attn_mask=mask) if masked else None
+
+            local_query, local_key, local_value, local_metadata, ctx = strategy.pre_attention(
+                query, key, value, metadata
+            )
+            assert local_metadata is metadata
+            local_output = F.scaled_dot_product_attention(
+                local_query.transpose(1, 2),
+                local_key.transpose(1, 2),
+                local_value.transpose(1, 2),
+                attn_mask=local_metadata.attn_mask if local_metadata is not None else None,
+            ).transpose(1, 2)
+            actual = strategy.post_attention(local_output, ctx)
+            expected = F.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                global_key.transpose(1, 2),
+                global_value.transpose(1, 2),
+                attn_mask=mask,
+            ).transpose(1, 2)
+
+            torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("world_size", [2, 4], ids=["world", "subgroups"])
+def test_ltx_video_to_audio_sp_matches_replicated_attention(unused_tcp_port, world_size):
+    torch.multiprocessing.spawn(
+        _run_video_to_audio_parity,
+        args=(world_size, unused_tcp_port),
+        nprocs=world_size,
+    )
+
+
+def test_ltx_video_to_audio_sp_rejects_ring_parallelism():
+    sp_group = SimpleNamespace(ring_world_size=2)
+
+    with pytest.raises(NotImplementedError, match="ring_degree must be 1"):
+        LTX2VideoToAudioParallelAttention(sp_group)

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -141,16 +141,54 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
 
 
 @pytest.mark.parametrize("rope_type", ["interleaved", "split"])
-def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
-    root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""]
+def test_ltx_sp_plan_only_shards_video_stream(rope_type):
+    plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)
+    root_plan = plan[""]
 
-    video_timestep = root_plan["timestep"]
-    audio_timestep = root_plan["audio_timestep"]
+    assert set(root_plan) == {"hidden_states", "keyframes_mask", "timestep"}
+    assert root_plan["hidden_states"].split_dim == 1
+    assert root_plan["keyframes_mask"].split_dim == 1
+    assert root_plan["timestep"].split_dim == 1
+    assert root_plan["timestep"].expected_dims == 2
+    assert not root_plan["timestep"].split_output
+    assert set(plan) == {"", "rope", "cross_attn_rope", "proj_out"}
 
-    assert video_timestep.split_dim == audio_timestep.split_dim == 1
-    assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
-    assert not video_timestep.split_output
-    assert not audio_timestep.split_output
+
+def test_ltx_attention_assigns_video_only_sp_modes(monkeypatch):
+    modes = {}
+
+    class FakeAttention(nn.Module):
+        def __init__(self, *, prefix, sequence_parallel_mode="local", **_kwargs):
+            super().__init__()
+            modes[prefix.rsplit(".", 1)[-1]] = sequence_parallel_mode
+
+    class FakeFeedForward(nn.Module):
+        def __init__(self, *_args, **_kwargs):
+            super().__init__()
+
+    monkeypatch.setattr(ltx2_transformer, "LTX2Attention", FakeAttention)
+    monkeypatch.setattr(ltx2_transformer, "LTX2FeedForward", FakeFeedForward)
+
+    LTX2VideoTransformerBlock(
+        dim=8,
+        num_attention_heads=2,
+        attention_head_dim=4,
+        cross_attention_dim=8,
+        audio_dim=8,
+        audio_num_attention_heads=2,
+        audio_attention_head_dim=4,
+        audio_cross_attention_dim=8,
+        prefix="transformer_blocks.0",
+    )
+
+    assert modes == {
+        "attn1": "video_self",
+        "audio_attn1": "local",
+        "attn2": "local",
+        "audio_attn2": "local",
+        "audio_to_video_attn": "local",
+        "video_to_audio_attn": "video_to_audio",
+    }
 
 
 def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):

tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Full-transformer numeric parity for LTX Ulysses sequence parallelism.
+
+This FP32/SDPA check verifies that Ulysses preserves the complete transformer
+mathematical path. It does not guarantee pixel-level E2E parity: production
+BF16 and compiled attention kernels can change rounding order, and diffusion
+can amplify that drift across denoising steps.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.parallel]
+
+_FP32_REL_L2 = 1e-5
+_TINY_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 4,
+    "cross_attention_dim": 8,
+    "vae_scale_factors": (1, 1, 1),
+    "pos_embed_max_pos": 4,
+    "base_height": 2,
+    "base_width": 2,
+    "audio_in_channels": 4,
+    "audio_out_channels": 4,
+    "audio_num_attention_heads": 2,
+    "audio_attention_head_dim": 4,
+    "audio_cross_attention_dim": 8,
+    "audio_scale_factor": 1,
+    "audio_pos_embed_max_pos": 4,
+    "audio_sampling_rate": 16,
+    "audio_hop_length": 4,
+    "num_layers": 1,
+    "caption_channels": 8,
+    "rope_type": "split",
+    "use_prompt_embeddings": False,
+    "use_keyframes_abs_pos_embedding": True,
+}
+
+
+def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float:
+    return ((actual - expected).double().norm() / expected.double().norm()).item()
+
+
+def _omni_config(sp_size: int) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig.from_kwargs(
+        model="test",
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            tensor_parallel_size=1,
+            sequence_parallel_size=sp_size,
+            ulysses_degree=sp_size,
+            ring_degree=1,
+            cfg_parallel_size=1,
+        ),
+        diffusion_attention_backend="TORCH_SDPA",
+    )
+
+
+def _inputs(device: torch.device) -> dict[str, torch.Tensor | int | bool]:
+    generator = torch.Generator().manual_seed(123)
+    return {
+        "hidden_states": torch.randn(1, 8, 4, generator=generator, device="cpu").to(device),
+        "audio_hidden_states": torch.randn(1, 4, 4, generator=generator, device="cpu").to(device),
+        "encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
+        "audio_encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
+        "timestep": torch.linspace(0.1, 0.8, 8, device=device).view(1, 8),
+        "audio_timestep": torch.linspace(0.2, 0.5, 4, device=device).view(1, 4),
+        "keyframes_mask": torch.tensor([[[1.0], [0.0], [1.0], [0.0], [0.0], [1.0], [0.0], [0.0]]], device=device),
+        "sigma": torch.tensor([0.1], device=device),
+        "audio_sigma": torch.tensor([0.2], device=device),
+        "num_frames": 2,
+        "height": 2,
+        "width": 2,
+        "audio_num_frames": 4,
+        "return_dict": False,
+    }
+
+
+def _run_transformer_parity(rank: int, world_size: int, master_port: int) -> None:
+    device = torch.device(f"{current_omni_platform.device_type}:{rank}")
+    current_omni_platform.set_device(device)
+    os.environ.update(
+        {
+            "RANK": str(rank),
+            "LOCAL_RANK": str(rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+
+    try:
+        init_distributed_environment()
+        initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
+
+        sp_config = _omni_config(world_size)
+        torch.manual_seed(42)
+        with set_forward_context(omni_diffusion_config=sp_config), set_current_diffusion_config(sp_config):
+            model = LTX2VideoTransformer3DModel(**_TINY_CONFIG).to(device).eval()
+        with torch.no_grad():
+            for _, parameter in sorted(model.named_parameters()):
+                torch.nn.init.normal_(parameter, mean=0.0, std=0.02)
+
+        inputs = _inputs(device)
+        reference_config = _omni_config(1)
+        with torch.no_grad(), set_forward_context(omni_diffusion_config=reference_config):
+            expected_video, expected_audio = model(**inputs)
+
+        apply_sequence_parallel(
+            model,
+            SequenceParallelConfig(ulysses_degree=world_size, ring_degree=1),
+            model._sp_plan,
+        )
+        with torch.no_grad(), set_forward_context(omni_diffusion_config=sp_config):
+            get_forward_context().sp_plan_hooks_applied = True
+            actual_video, actual_audio = model(**inputs)
+
+        v2a_attention = model.transformer_blocks[0].video_to_audio_attn.attn
+        assert v2a_attention._video_to_audio_strategy is not None
+        assert torch.isfinite(expected_video).all()
+        assert torch.isfinite(expected_audio).all()
+        for modality, actual, expected in (
+            ("video", actual_video, expected_video),
+            ("audio", actual_audio, expected_audio),
+        ):
+            drift = _relative_l2(actual, expected)
+            assert drift <= _FP32_REL_L2, (
+                f"LTX SP{world_size} {modality} output drifted rel_l2={drift:.3e} from SP1 (bound {_FP32_REL_L2:.1e})"
+            )
+    finally:
+        destroy_distributed_env()
+
+
+@pytest.mark.full_model
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ltx_ulysses_transformer_matches_sp1(unused_tcp_port) -> None:
+    """Exercise real plan hooks, RoPE/timestep sharding, and attention routing."""
+    world_size = 2
+    torch.multiprocessing.spawn(
+        _run_transformer_parity,
+        args=(world_size, unused_tcp_port),
+        nprocs=world_size,
+    )

vllm_omni/diffusion/models/ltx2/ltx2_denoise.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Shared denoise execution primitives for LTX pipelines."""
 
@@ -578,12 +578,6 @@ class LTXPhaseExecutor:
             video_audio_step_adapter=video_audio_step_adapter,
         )
         video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents)
-        ring_degree = getattr(pipeline.od_config.parallel_config, "ring_degree", 1) or 1
-        if padded_audio_num_frames > original_audio_num_frames and ring_degree > 1:
-            raise ValueError(
-                "LTX audio padding requires an attention mask, which Ring sequence parallelism does not support. "
-                "Use Ulysses-only SP or choose a request whose audio latent length is divisible by the SP size."
-            )
         denoise_ctx = LTXDenoiseContext(
             latents=latents,
             audio_latents=audio_latents,

vllm_omni/diffusion/models/ltx2/ltx2_request.py

@@ -497,7 +497,6 @@ class LTXRequestMixin:
         image_crf = _normalize_image_crf(first_image_crf if first_image_crf is not None else image_crf)
 
         sampling = sampling_params_list[0]
-        is_dummy_run = req.is_dummy_run()
         prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
         negative_prompt = _resolve_negative_prompts(
             req.prompts,
@@ -511,8 +510,6 @@ class LTXRequestMixin:
         num_inference_steps = (
             sampling.num_inference_steps or num_inference_steps or self.pipeline_recipe.num_inference_steps
         )
-        if is_dummy_run and self.pipeline_recipe.fixed_num_inference_steps:
-            num_inference_steps = self.pipeline_recipe.num_inference_steps
         num_inference_steps = max(int(num_inference_steps), 2)
 
         num_videos_per_prompt = (
@@ -528,8 +525,6 @@ class LTXRequestMixin:
             for item in sampling_params_list
         ]
         guidance = guidance_specs[0]
-        if is_dummy_run and not self.pipeline_recipe.phases[0].allow_guidance_override:
-            guidance = self.pipeline_recipe.request_guidance
         if any(item != guidance for item in guidance_specs[1:]):
             raise ValueError("Batched LTX requests must use identical video/audio guidance parameters.")
 

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -5,14 +5,16 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
-from contextlib import nullcontext
+import time
+from collections.abc import Iterable, Iterator
+from contextlib import contextmanager, nullcontext
 from dataclasses import replace
 from typing import Any, ClassVar
 
 import torch
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
+from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -60,6 +62,8 @@ from .ltx2_request import (
     validate_pipeline_request,
 )
 
+logger = init_logger(__name__)
+
 
 def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.Tensor:
     """Run the BWE vocoder in FP32, matching the official LTX pipeline."""
@@ -195,7 +199,7 @@ class LTXRuntime(
     connector_batches_cfg = False
     distributed_video_decode = True
     support_image_input = False
-    dummy_run_num_frames = 1
+    dummy_run_num_frames: ClassVar[int] = 0
     preserve_sp_padded_audio_duration = False
     reports_stage_durations = False
 
@@ -207,6 +211,14 @@ class LTXRuntime(
                 f"{self.__class__.__name__} does not support ulysses_mode='advanced_uaa'. "
                 "Use the default ulysses_mode='strict' for LTX sequence parallelism."
             )
+        ring_degree = getattr(parallel_config, "ring_degree", 1)
+        allgather_degree = getattr(parallel_config, "allgather_degree", 1)
+        if ring_degree != 1 or allgather_degree != 1:
+            raise ValueError(
+                f"{self.__class__.__name__} supports pure Ulysses sequence parallelism only. "
+                "Set ring_degree=1 and allgather_degree=1; "
+                f"got {ring_degree=} and {allgather_degree=}."
+            )
         self.model_version = detect_ltx_model_version(od_config.model, revision=getattr(od_config, "revision", None))
         self.use_diffusion_decoder = _ltx2_use_diffusion_decoder(od_config, self.model_version)
         self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version)
@@ -354,27 +366,29 @@ class LTXRuntime(
                 )
             )
             self._enter_phase(phase_recipe)
-            phase_inputs = self._build_phase_inputs(
-                request_inputs,
-                phase_recipe,
-                phase_results[-1] if phase_results else None,
-            )
+            with self._profile_phase_method("_build_phase_inputs", phase_recipe.name):
+                phase_inputs = self._build_phase_inputs(
+                    request_inputs,
+                    phase_recipe,
+                    phase_results[-1] if phase_results else None,
+                )
             if phase_sigmas is not None:
                 phase_inputs = replace(phase_inputs, num_inference_steps=len(phase_sigmas) - 1)
             noise_scale = phase_recipe.noise_scale
             if override_sigmas is not None and phase_recipe.input_transform == "spatial_upsample":
                 noise_scale = float(override_sigmas[0])
-            phase_result = self.run_phase(
-                req,
-                phase_inputs,
-                noise_scale=noise_scale,
-                sigmas=phase_sigmas,
-                timesteps=None,
-                attention_kwargs=None,
-                phase_recipe=phase_recipe,
-                image=image,
-                prompt_context=prompt_context,
-            )
+            with self._profile_phase_method("run_phase", phase_recipe.name):
+                phase_result = self.run_phase(
+                    req,
+                    phase_inputs,
+                    noise_scale=noise_scale,
+                    sigmas=phase_sigmas,
+                    timesteps=None,
+                    attention_kwargs=None,
+                    phase_recipe=phase_recipe,
+                    image=image,
+                    prompt_context=prompt_context,
+                )
             phase_results.append(phase_result)
             prompt_context = phase_result.forward_context.prompt_context
 
@@ -384,7 +398,45 @@ class LTXRuntime(
             video=phase_results[self.pipeline_recipe.video_output_phase].video,
             audio=phase_results[self.pipeline_recipe.audio_output_phase].audio,
         )
-        return self.decode_phase(output_phase)
+        with self._profile_phase_method("decode_phase"):
+            output = self.decode_phase(output_phase)
+        return self._attach_profiled_stage_durations(output)
+
+    @contextmanager
+    def _profile_phase_method(self, method_name: str, phase_name: str | None = None) -> Iterator[None]:
+        if not getattr(self, "enable_diffusion_pipeline_profiler", False):
+            yield
+            return
+
+        metric_parts = [self.__class__.__name__]
+        if phase_name is not None:
+            metric_parts.append(phase_name)
+        metric_parts.append(method_name)
+        metric_name = ".".join(metric_parts)
+
+        if current_omni_platform.is_available():
+            current_omni_platform.synchronize()
+        start_time = time.perf_counter()
+        try:
+            yield
+        finally:
+            if current_omni_platform.is_available():
+                current_omni_platform.synchronize()
+            duration = time.perf_counter() - start_time
+            logger.info("[DiffusionPipelineProfiler] %s took %.6fs", metric_name, duration)
+            with self._profiler_lock:
+                self._stage_durations[metric_name] = self._stage_durations.get(metric_name, 0.0) + duration
+
+    def _attach_profiled_stage_durations(
+        self, output: DiffusionOutput | list[DiffusionOutput]
+    ) -> DiffusionOutput | list[DiffusionOutput]:
+        if not getattr(self, "enable_diffusion_pipeline_profiler", False):
+            return output
+        stage_durations = self.stage_durations
+        outputs = output if isinstance(output, list) else [output]
+        for item in outputs:
+            item.stage_durations = dict(stage_durations)
+        return output
 
     def _enter_phase(self, phase: LTXPhaseRecipe) -> None:
         self._active_phase_name = phase.name

vllm_omni/diffusion/models/ltx2/ltx2_sequence_parallel.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.parallel.base import (
+    ParallelAttentionContext,
+)
+from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
+from vllm_omni.diffusion.distributed.group_coordinator import (
+    SequenceParallelGroupCoordinator,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _VideoToAudioContext(ParallelAttentionContext):
+    ulysses_pg: dist.ProcessGroup
+    use_sync: bool
+
+
+class LTX2VideoToAudioParallelAttention:
+    """LTX video-to-audio SP with replicated audio queries.
+
+    Video K/V arrive sequence-sharded, while the much shorter audio query is
+    replicated. Redistribute only K/V across sequence/head dimensions, slice
+    audio query heads locally, then gather output heads. This uses three
+    collectives instead of standard Ulysses' four and keeps audio replicated.
+    """
+
+    def __init__(
+        self,
+        sp_group: SequenceParallelGroupCoordinator,
+        *,
+        use_sync: bool = False,
+    ) -> None:
+        if sp_group.ring_world_size != 1:
+            raise NotImplementedError(
+                "LTX video-to-audio sequence parallelism currently supports pure Ulysses only (ring_degree must be 1)."
+            )
+        self._sp_group = sp_group
+        self._ulysses_pg = sp_group.ulysses_group
+        self._use_sync = use_sync
+
+    @property
+    def enabled(self) -> bool:
+        return True
+
+    @property
+    def name(self) -> str:
+        return "ltx2_video_to_audio"
+
+    def pre_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ):
+        if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
+            raise ValueError(
+                "LTX video-to-audio SP expects Q/K/V in (B, S, H, D) layout, "
+                f"got {query.shape=}, {key.shape=}, {value.shape=}."
+            )
+        if key.shape != value.shape:
+            raise ValueError(
+                "LTX video-to-audio SP requires matching K/V shapes, "
+                f"got key={tuple(key.shape)} and value={tuple(value.shape)}."
+            )
+
+        world_size = self._sp_group.ulysses_world_size
+        rank = self._sp_group.ulysses_rank
+        query_heads = int(query.shape[2])
+        if query_heads % world_size != 0:
+            raise ValueError(
+                "LTX video-to-audio SP requires query heads divisible by "
+                f"ulysses_degree, got heads={query_heads}, ulysses_degree={world_size}."
+            )
+        if int(key.shape[2]) % world_size != 0:
+            raise ValueError(
+                "LTX video-to-audio SP requires K/V heads divisible by "
+                f"ulysses_degree, got heads={key.shape[2]}, ulysses_degree={world_size}."
+            )
+
+        heads_per_rank = query_heads // world_size
+        head_start = rank * heads_per_rank
+        query = query[:, :, head_start : head_start + heads_per_rank, :].contiguous()
+        key = SeqAllToAll4D.apply(self._ulysses_pg, key, 2, 1, self._use_sync)
+        value = SeqAllToAll4D.apply(self._ulysses_pg, value, 2, 1, self._use_sync)
+
+        ctx = _VideoToAudioContext(
+            name=self.name,
+            ulysses_pg=self._ulysses_pg,
+            use_sync=self._use_sync,
+        )
+        return query, key, value, attn_metadata, ctx
+
+    def post_attention(
+        self,
+        attn_output: torch.Tensor,
+        ctx: ParallelAttentionContext | None,
+    ) -> torch.Tensor:
+        if not isinstance(ctx, _VideoToAudioContext):
+            raise TypeError(f"Unexpected LTX video-to-audio context: {type(ctx)!r}.")
+
+        # Gather along heads by moving H to the leading dimension required by
+        # all_gather_into_tensor, then restore the backend's (B, S, H, D) layout.
+        local_heads_first = attn_output.permute(2, 0, 1, 3).contiguous()
+        world_size = dist.get_world_size(ctx.ulysses_pg)
+        gathered = torch.empty(
+            (local_heads_first.shape[0] * world_size, *local_heads_first.shape[1:]),
+            dtype=attn_output.dtype,
+            device=attn_output.device,
+        )
+        dist.all_gather_into_tensor(gathered, local_heads_first, group=ctx.ulysses_pg)
+        if ctx.use_sync:
+            from vllm_omni.platforms import current_omni_platform
+
+            current_omni_platform.synchronize()
+        return gathered.permute(1, 2, 0, 3).contiguous()

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2025 The Lightricks team and The HuggingFace Team.
 # All rights reserved.
 #
@@ -47,9 +50,12 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
+from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 
+from .ltx2_sequence_parallel import LTX2VideoToAudioParallelAttention
+
 logger = init_logger(__name__)
 
 _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters
@@ -359,6 +365,34 @@ def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor:
 class _LTX2ParallelAttention(Attention):
     """Preserve LTX SP collectives while applying model-specific padding masks."""
 
+    _SP_MODES = frozenset({"local", "video_self", "video_to_audio"})
+
+    def __init__(self, *args, ltx_sp_mode: str = "local", **kwargs):
+        if ltx_sp_mode not in self._SP_MODES:
+            raise ValueError(
+                f"Unsupported LTX sequence-parallel mode {ltx_sp_mode!r}; expected one of {sorted(self._SP_MODES)}."
+            )
+        super().__init__(*args, **kwargs)
+        self.ltx_sp_mode = ltx_sp_mode
+        self._video_to_audio_strategy = None
+
+    def _get_active_parallel_strategy(self):
+        base_strategy = super()._get_active_parallel_strategy()
+        if self.ltx_sp_mode == "local" or not base_strategy.enabled:
+            return self._no_parallel_strategy
+        if self.ltx_sp_mode == "video_self":
+            return base_strategy
+        if base_strategy.name != "ulysses":
+            raise NotImplementedError(
+                f"LTX video-to-audio sequence parallelism requires pure Ulysses; got strategy={base_strategy.name!r}."
+            )
+        if self._video_to_audio_strategy is None:
+            self._video_to_audio_strategy = LTX2VideoToAudioParallelAttention(
+                get_sp_group(),
+                use_sync=self.use_sync,
+            )
+        return self._video_to_audio_strategy
+
     def _run_local_attention(self, query, key, value, attn_metadata):
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         if attention_mask is not None:
@@ -597,6 +631,7 @@ class LTX2Attention(torch.nn.Module):
         quant_config: "QuantizationConfig | None" = None,
         prefix: str = "",
         disable_kv_quant: bool = False,
+        sequence_parallel_mode: str = "local",
     ):
         super().__init__()
         # LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both
@@ -721,6 +756,7 @@ class LTX2Attention(torch.nn.Module):
             causal=False,
             prefix=prefix,
             disable_kv_quant=disable_kv_quant,
+            ltx_sp_mode=sequence_parallel_mode,
         )
 
         # LTX-2.3: per-head gated attention
@@ -876,6 +912,7 @@ class LTX2VideoTransformerBlock(nn.Module):
             apply_gated_attention=video_gated_attn,
             quant_config=quant_config,
             prefix=f"{prefix}.attn1" if prefix else "attn1",
+            sequence_parallel_mode="video_self",
         )
 
         self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -965,6 +1002,7 @@ class LTX2VideoTransformerBlock(nn.Module):
             apply_gated_attention=audio_gated_attn,
             quant_config=quant_config,
             prefix=f"{prefix}.video_to_audio_attn" if prefix else "video_to_audio_attn",
+            sequence_parallel_mode="video_to_audio",
         )
 
         # 4. Feedforward layers
@@ -1568,38 +1606,25 @@ class LTX2VideoTransformer3DModel(nn.Module):
 
         return {
             "": {
-                # Shard video/audio latents across sequence
+                # Keep only the long video stream sequence-sharded. Audio and
+                # prompt streams remain replicated on every SP rank.
                 "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
-                "audio_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
                 "keyframes_mask": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
-                # Shard prompt embeds across sequence
-                "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
-                "audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
-                # Keep per-token video/audio timestep embeddings aligned with
-                # their corresponding latent sequence shards. Official LTX
-                # multi-guidance supplies both tensors as (B, seq_len).
+                # Keep per-token video timestep embeddings aligned with the
+                # video latent shard. LTX guidance supplies a separate,
+                # replicated audio_timestep under SP.
                 "timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False),
-                "audio_timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False),
             },
             "rope": {
                 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
                 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
             },
-            "audio_rope": {
-                0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
-                1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
-            },
             "cross_attn_rope": {
                 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
                 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
             },
-            "cross_attn_audio_rope": {
-                0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
-                1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
-            },
-            # Gather outputs before returning
+            # Video output is sharded; audio output is already replicated.
             "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-            "audio_proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
         }
 
     def __init__(