From 3204a0b2ade0f05424b7f11e3bcc03cb3542e0c6 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Mon, 7 Sep 2026 18:48:26 +0800 Subject: [PATCH] [Model][Perf] Optimize LTX Ulysses sequence parallelism (#7079) Signed-off-by: mglyn <1203789601@qq.com> --- docs/user_guide/diffusion_features.md | 6 +- recipes/LTX/LTX-2.5.md | 23 +++ recipes/LTX/LTX-2.md | 11 +- .../models/ltx2/test_ltx25_pipeline.py | 27 +++ .../models/ltx2/test_ltx2_latents.py | 2 +- .../models/ltx2/test_ltx2_parallel.py | 6 +- .../models/ltx2/test_ltx2_pipeline.py | 104 +++++++---- .../ltx2/test_ltx2_sequence_parallel.py | 100 +++++++++++ .../models/ltx2/test_ltx2_transformer.py | 54 +++++- .../ltx2/test_ltx2_transformer_ulysses.py | 167 ++++++++++++++++++ .../diffusion/models/ltx2/ltx2_denoise.py | 8 +- .../diffusion/models/ltx2/ltx2_request.py | 5 - .../diffusion/models/ltx2/ltx2_runtime.py | 92 +++++++--- .../models/ltx2/ltx2_sequence_parallel.py | 125 +++++++++++++ .../diffusion/models/ltx2/ltx2_transformer.py | 63 +++++-- 15 files changed, 685 insertions(+), 108 deletions(-) create mode 100644 tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py create mode 100644 tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py create mode 100644 vllm_omni/diffusion/models/ltx2/ltx2_sequence_parallel.py diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 00c245489..a8fb6d972 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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) | ✅ | ❌ | diff --git a/recipes/LTX/LTX-2.5.md b/recipes/LTX/LTX-2.5.md index f0be6948b..538eed1ce 100644 --- a/recipes/LTX/LTX-2.5.md +++ b/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 diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index ec4e8fb3b..5a6a049ee 100644 --- a/recipes/LTX/LTX-2.md +++ b/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 diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py index 950906190..ae90ea1a3 100644 --- a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py +++ b/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 diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index fbac0bfe5..724498f7c 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/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( diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 0573dfb37..272cdc7d8 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/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) diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index c2fe7cbfa..3e630d8dd 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/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(): diff --git a/tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py new file mode 100644 index 000000000..395b53ba2 --- /dev/null +++ b/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) diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index 937d487e6..f12fbc466 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/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): diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py b/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py new file mode 100644 index 000000000..9f09b72a7 --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index 7ebd94143..f35ce25c5 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/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, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index 22230e486..f42c43709 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/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.") diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 0ffa09dcb..2c51b17d9 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_sequence_parallel.py b/vllm_omni/diffusion/models/ltx2/ltx2_sequence_parallel.py new file mode 100644 index 000000000..874df90e3 --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 0b0e7fc08..1d95fc3df 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/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__(