[Model] Add LTX-2.5 support (#6070)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/models/ltx2/test_ltx25_pipeline.py

@@ -0,0 +1,881 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Focused LTX-2.5 pipeline correctness tests."""
+
+import json
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime
+from vllm_omni.diffusion.models.ltx2.ltx2_components import (
+    LTX25_DISTILLED_COMPONENT_PROFILE,
+    LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
+    LTX25_FULL_COMPONENT_PROFILE,
+    LTX25_TWO_STAGE_COMPONENT_PROFILE,
+    detect_ltx_model_version,
+    resolve_ltx_checkpoint_kind,
+    resolve_ltx_component_profile,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
+    LTXPhaseResult,
+    _first_frame_keyframes_mask,
+    _official_ltx_sigmas,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
+    LTX_GUIDANCE_EXECUTOR,
+    LTXGuidanceSpec,
+    LTXModalityGuidance,
+    velocity_from_x0,
+    x0_from_velocity,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
+from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+    LTX2_DISTILLED_TWO_STAGE_RECIPE,
+    LTX2_ONE_STAGE_RECIPE,
+    LTX23_ONE_STAGE_RECIPE,
+    LTX25_DISTILLED_ONE_STAGE_RECIPE,
+    LTX25_DISTILLED_TWO_STAGE_RECIPE,
+    LTX25_FULL_RECIPE,
+    LTX25_TWO_STAGE_RECIPE,
+    LTX_DISTILLED_SIGMAS,
+    LTX_POSITIVE_ONLY_RECIPE,
+    resolve_ltx_pipeline_recipe,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
+from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2DistilledOneStagePipeline,
+    LTX2Pipeline,
+)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+    LTX2DistilledTwoStagePipeline,
+    LTX2TwoStagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"])
+def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version):
+    pipe = object.__new__(LTX2Pipeline)
+    pipe.model_version = model_version
+    pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+    ts = torch.tensor([750.0, 500.0])
+
+    kwargs = LTXRuntime._denoise_timestep_kwargs(
+        pipe,
+        ts,
+        SimpleNamespace(),
+        SimpleNamespace(),
+        video_token_count=4,
+        audio_token_count=3,
+    )
+
+    torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4))
+    torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3))
+    assert kwargs["sigma"] is ts
+    assert kwargs["audio_sigma"] is ts
+
+
+def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch):
+    state = LTXAVState(
+        video=torch.zeros(1, 2, 2),
+        audio=torch.zeros(1, 4, 2),
+    )
+    forward_ctx = SimpleNamespace(
+        original_audio_num_frames=3,
+        sampler="euler_ancestral",
+    )
+    denoise_ctx = SimpleNamespace(latents=None, audio_latents=None)
+    pipe = SimpleNamespace(
+        _predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio))
+    )
+
+    def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args):
+        assert actual_forward_ctx.sampler == "euler_ancestral"
+        updated_audio = torch.ones_like(state.audio)
+        updated_audio[:, 3:] = 99.0
+        return torch.ones_like(state.video), updated_audio
+
+    monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step)
+
+    actual = LTXRuntime._denoise_step(
+        pipe,
+        0,
+        torch.tensor(1.0),
+        state,
+        forward_ctx,
+        denoise_ctx,
+    )
+
+    torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3]))
+    torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:]))
+
+
+def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding():
+    reference = torch.zeros(2, 12, 8)
+
+    mask = _first_frame_keyframes_mask(reference, latent_num_frames=3)
+
+    assert mask.shape == (2, 12, 1)
+    torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1))
+    torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1))
+
+
+def _make_ltx_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    pipe.vae_spatial_compression_ratio = 32
+    pipe.vae_temporal_compression_ratio = 8
+    return pipe
+
+
+def _resolve_request_inputs_for_test(
+    pipe,
+    req,
+    *,
+    guidance_scale=4.0,
+    negative_prompt=None,
+):
+    return pipe._resolve_request_inputs(
+        req,
+        prompt=None,
+        negative_prompt=negative_prompt,
+        height=None,
+        width=None,
+        num_frames=None,
+        frame_rate=None,
+        num_inference_steps=None,
+        guidance_scale=guidance_scale,
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.0,
+        decode_noise_scale=None,
+        output_type="np",
+        max_sequence_length=None,
+    )
+
+
+def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
+    pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
+    od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
+
+    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object())
+
+    with pytest.raises(ImportError) as exc_info:
+        ltx2_components.initialize_pipeline_components(pipe, od_config)
+
+    assert str(exc_info.value) == (
+        "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15."
+    )
+
+
+def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
+    revision = "pinned-revision"
+    calls = {"components": []}
+    profile = replace(
+        LTX25_DISTILLED_COMPONENT_PROFILE,
+        text_encoder_cls=object,
+        video_vae_cls=object,
+        vocoder_cls=object,
+        vocoder_fallback_cls=None,
+        scheduler_use_dynamic_shifting=False,
+    )
+    pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
+    od_config = SimpleNamespace(
+        model="org/converted-ltx25",
+        revision=revision,
+        dtype=torch.bfloat16,
+        quantization_config=None,
+    )
+
+    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None)
+    monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None)
+
+    def fake_prefetch(model, subfolders, **kwargs):
+        calls["prefetch"] = (model, tuple(subfolders), kwargs)
+
+    def fake_tokenizer(*args, **kwargs):
+        calls["tokenizer"] = (args, kwargs)
+        return SimpleNamespace(model_max_length=1024)
+
+    def fake_component(component_cls, model, subfolder, **kwargs):
+        calls["components"].append((component_cls, model, subfolder, kwargs))
+        if subfolder == "vae":
+            return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8)
+        if subfolder == "audio_vae":
+            return SimpleNamespace(
+                mel_compression_ratio=4,
+                temporal_compression_ratio=4,
+                config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160),
+            )
+        return object()
+
+    def fake_transformer_config(model, subfolder, local_files_only, *, revision):
+        calls["transformer_config"] = (model, subfolder, local_files_only, revision)
+        return {"component": "transformer"}
+
+    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
+    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
+    monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
+    monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
+    monkeypatch.setattr(
+        ltx2_components,
+        "create_transformer_from_config",
+        lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)),
+    )
+
+    def fake_scheduler(*args, **kwargs):
+        calls["scheduler"] = (args, kwargs)
+        return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None})
+
+    monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler)
+
+    ltx2_components.initialize_pipeline_components(pipeline, od_config)
+
+    assert pipeline.weights_sources[0].revision == revision
+    assert calls["prefetch"][2]["revision"] == revision
+    assert calls["tokenizer"][1]["revision"] == revision
+    assert {call[2] for call in calls["components"]} == {
+        "text_encoder",
+        "connectors",
+        "vae",
+        "audio_vae",
+        "vocoder",
+        "latent_upsampler",
+    }
+    assert all(call[3]["revision"] == revision for call in calls["components"])
+    assert calls["transformer_config"] == (
+        od_config.model,
+        profile.transformer_subfolder,
+        False,
+        revision,
+    )
+    assert calls["scheduler"][1]["revision"] == revision
+
+
+def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
+    (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
+    (tmp_path / "transformer").mkdir()
+    (tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False}))
+
+    assert detect_ltx_model_version(str(tmp_path)) == "2.3"
+
+
+def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+
+    assert pipe.model_version == "2.5"
+    assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE
+    assert pipe.pipeline_recipe is LTX25_FULL_RECIPE
+    assert pipe.preserve_sp_padded_audio_duration
+    assert pipe.reports_stage_durations
+
+
+def test_ltx25_full_recipe_matches_official_sft_defaults():
+    (phase,) = LTX25_FULL_RECIPE.phases
+
+    assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960)
+    assert LTX25_FULL_RECIPE.num_frames == 121
+    assert LTX25_FULL_RECIPE.frame_rate == 24.0
+    assert LTX25_FULL_RECIPE.num_inference_steps == 30
+    assert phase.guidance.video == LTXModalityGuidance(
+        cfg_scale=3.0,
+        stg_scale=1.0,
+        modality_scale=3.0,
+        rescale_scale=0.7,
+        stg_blocks=(28,),
+    )
+    assert phase.guidance.audio == LTXModalityGuidance(
+        cfg_scale=7.0,
+        stg_scale=1.0,
+        modality_scale=3.0,
+        rescale_scale=0.7,
+        stg_blocks=(28,),
+    )
+    assert phase.sigmas is None
+    assert phase.noise_scale == 1.0
+    assert phase.use_official_sigma_schedule
+    assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting
+    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1
+
+
+def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor():
+    scheduler = SimpleNamespace(
+        config={
+            "base_image_seq_len": 1024,
+            "max_image_seq_len": 4096,
+            "base_shift": 0.95,
+            "max_shift": 2.05,
+            "shift_terminal": 0.1,
+        }
+    )
+
+    sigmas = _official_ltx_sigmas(
+        scheduler,
+        steps=30,
+        device=torch.device("cpu"),
+    )
+
+    torch.testing.assert_close(
+        sigmas[[0, 1, -2, -1]],
+        torch.tensor([1.0, 0.9949570, 0.1, 0.0]),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+
+def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = LTX2DistilledTwoStagePipeline(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
+
+    assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE
+    assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+
+def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"}))
+    for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")):
+        path = tmp_path / subfolder
+        path.mkdir()
+        (path / "config.json").write_text(json.dumps({"marker": marker}))
+
+    config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline")
+    config.enrich_config()
+
+    assert config.tf_model_config.get("marker") == "full"
+
+
+def test_ltx25_four_public_pipeline_semantics_are_disjoint():
+    assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE
+    assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE
+    assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE
+    assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+    assert LTX2Pipeline.pipeline_kind == "one_stage"
+    assert LTX2TwoStagePipeline.pipeline_kind == "two_stage"
+    assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage"
+    assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
+
+
+@pytest.mark.parametrize(
+    ("pipeline_kind", "expected"),
+    [
+        ("one_stage", "regular"),
+        ("two_stage", "regular"),
+        ("distilled_one_stage", "distilled"),
+        ("distilled_two_stage", "distilled"),
+        ("dmd2", None),
+    ],
+)
+def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected):
+    assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected
+
+
+def test_ltx_checkpoint_kind_rejects_unknown_pipeline():
+    with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"):
+        resolve_ltx_checkpoint_kind("unknown")
+
+
+def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution():
+    stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases
+
+    assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920)
+    assert stage1.spatial_downscale == 2
+    assert stage1.sigmas == LTX_DISTILLED_SIGMAS
+    assert stage1.sampler == "euler_ancestral"
+    assert stage2.sampler == "euler"
+    assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
+    assert stage2.input_transform == "spatial_upsample"
+
+    assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == ""
+    assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt
+
+
+def test_ltx25_distilled_one_stage_uses_native_low_resolution():
+    (phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases
+
+    assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960)
+    assert phase.spatial_downscale == 1
+    assert phase.sigmas == LTX_DISTILLED_SIGMAS
+    assert phase.sampler == "euler_ancestral"
+    assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit
+
+
+def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora():
+    stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases
+
+    assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920)
+    assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30
+    assert stage1.spatial_downscale == 2
+    assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance
+    assert stage2.adapter_slot == "ltx_distilled"
+    assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas
+    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None
+
+
+@pytest.mark.parametrize(
+    "recipe",
+    [
+        LTX2_ONE_STAGE_RECIPE,
+        LTX23_ONE_STAGE_RECIPE,
+        LTX_POSITIVE_ONLY_RECIPE,
+    ],
+)
+def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe):
+    assert recipe.supports_cache_dit
+
+
+@pytest.mark.parametrize(
+    "recipe",
+    [
+        LTX25_FULL_RECIPE,
+        LTX25_DISTILLED_ONE_STAGE_RECIPE,
+        LTX2_DISTILLED_TWO_STAGE_RECIPE,
+        LTX25_DISTILLED_TWO_STAGE_RECIPE,
+    ],
+)
+def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe):
+    assert not recipe.supports_cache_dit
+
+
+def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip():
+    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+    sigma = torch.tensor(0.5)
+
+    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+        sample,
+        {"cond": velocity},
+        sigma,
+        LTXModalityGuidance(),
+    )
+    expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma)
+
+    assert actual.dtype == torch.float32
+    torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+
+
+def test_ltx_velocity_from_x0_materializes_official_scalar_sigma():
+    item_calls = []
+
+    class TrackingTensor(torch.Tensor):
+        @staticmethod
+        def __new__(cls):
+            return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False)
+
+        def item(self):
+            item_calls.append(True)
+            return super().item()
+
+    sample = torch.tensor([[1.0]], dtype=torch.bfloat16)
+    x0 = torch.tensor([[0.5]], dtype=torch.bfloat16)
+
+    velocity_from_x0(sample, x0, TrackingTensor())
+
+    assert item_calls == [True]
+
+
+def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity():
+    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+
+    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+        sample,
+        {"cond": velocity},
+        torch.tensor(0.5),
+        LTXModalityGuidance(),
+        preserve_positive_velocity=True,
+    )
+
+    assert actual is velocity
+
+
+@pytest.mark.parametrize("model_version", ["2.5", "2.3"])
+def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version):
+    expected_dtype = torch.bfloat16
+    captured = {}
+
+    def prepare_latents(**kwargs):
+        captured["video_dtype"] = kwargs["dtype"]
+        return torch.empty(1, 1, 1, dtype=kwargs["dtype"])
+
+    def prepare_audio_latents(*_args, **kwargs):
+        captured["audio_dtype"] = kwargs["dtype"]
+        return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1
+
+    pipeline = SimpleNamespace(
+        model_version=model_version,
+        transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)),
+        prepare_latents=prepare_latents,
+        prepare_audio_latents=prepare_audio_latents,
+        audio_sampling_rate=48_000,
+        audio_hop_length=160,
+        audio_vae_temporal_compression_ratio=4,
+        audio_vae_mel_compression_ratio=4,
+        audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)),
+        _resolve_audio_latent_length=lambda length, _latents: length,
+    )
+    request = SimpleNamespace(
+        num_videos_per_prompt=1,
+        height=64,
+        width=64,
+        num_frames=1,
+        frame_rate=24.0,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        audio_latents_normalized=False,
+    )
+    prompt = SimpleNamespace(
+        batch_size=1,
+        positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+        positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+    )
+
+    LTXRuntime._prepare_video_latents_stage(
+        pipeline,
+        request,
+        prompt,
+        device=torch.device("cpu"),
+        noise_scale=1.0,
+    )
+    LTXRuntime._prepare_audio_latents_stage(
+        pipeline,
+        request,
+        prompt,
+        device=torch.device("cpu"),
+        noise_scale=1.0,
+    )
+
+    assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype}
+
+
+def test_ltx_seeded_latents_match_official_packed_rng_layout():
+    pipeline = SimpleNamespace(
+        model_version="2.5",
+        vae_spatial_compression_ratio=2,
+        vae_temporal_compression_ratio=2,
+        transformer_spatial_patch_size=1,
+        transformer_temporal_patch_size=1,
+        audio_vae_mel_compression_ratio=2,
+        od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)),
+    )
+    device = torch.device("cpu")
+    actual_generator = torch.Generator(device=device).manual_seed(123)
+    expected_generator = torch.Generator(device=device).manual_seed(123)
+
+    actual_video = ltx2_latents.prepare_video_latents(
+        pipeline,
+        batch_size=1,
+        num_channels_latents=2,
+        height=4,
+        width=6,
+        num_frames=3,
+        dtype=torch.float32,
+        device=device,
+        generator=actual_generator,
+    )
+    actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents(
+        pipeline,
+        batch_size=1,
+        num_channels_latents=2,
+        audio_latent_length=3,
+        num_mel_bins=4,
+        dtype=torch.float32,
+        device=device,
+        generator=actual_generator,
+    )
+
+    expected_video = torch.randn((1, 12, 2), generator=expected_generator)
+    expected_audio = torch.randn((1, 3, 4), generator=expected_generator)
+
+    torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0)
+    assert (original_length, padded_length) == (3, 3)
+
+
+def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
+    request_inputs = LTXRequestInputs(
+        prompt="prompt",
+        negative_prompt="",
+        height=64,
+        width=64,
+        num_frames=1,
+        frame_rate=24.0,
+        num_inference_steps=8,
+        guidance=LTXGuidanceSpec.positive_only(),
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.25,
+        decode_noise_scale=0.5,
+        output_type="np",
+        max_sequence_length=16,
+    )
+    prompt_context = object()
+    phase_calls = []
+
+    def resolve_request_inputs(req, **kwargs):
+        return request_inputs
+
+    def run_phase(req, inputs, *, prompt_context=None, **kwargs):
+        phase_recipe = kwargs["phase_recipe"]
+        assert kwargs["image"] is source_image
+        phase_calls.append((phase_recipe, inputs, prompt_context))
+        if len(phase_calls) == 1:
+            assert prompt_context is None
+            assert inputs.decode_timestep == 0.25
+            assert inputs.decode_noise_scale == 0.5
+            context = prompt_context_sentinel
+            assert phase_recipe.name == "generate_lowres"
+            assert (inputs.height, inputs.width) == (32, 32)
+            assert inputs.num_inference_steps == 2
+            assert kwargs["sigmas"] == stage_1_sigmas
+            assert kwargs["noise_scale"] == 1.0
+            video = torch.ones(1, 128, 1, 1, 1)
+            audio = torch.full((1, 8, 1, 2), 2.0)
+            audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0)
+        else:
+            assert prompt_context is prompt_context_sentinel
+            assert phase_recipe.name == "refine"
+            assert (inputs.height, inputs.width) == (64, 64)
+            assert inputs.latents.shape == (1, 128, 1, 2, 2)
+            torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0))
+            assert inputs.audio_latents_normalized
+            assert inputs.guidance_scale == 1.0
+            assert inputs.num_inference_steps == 2
+            assert kwargs["sigmas"] == stage_2_sigmas
+            assert kwargs["noise_scale"] == stage_2_sigmas[0]
+            assert inputs.decode_timestep == 0.0
+            assert inputs.decode_noise_scale is None
+            context = prompt_context
+            video = torch.full((1, 128, 1, 2, 2), 3.0)
+            audio = torch.full((1, 8, 1, 2), 4.0)
+            audio_for_next_phase = None
+        return LTXPhaseResult(
+            forward_context=SimpleNamespace(prompt_context=context),
+            video=video,
+            audio=audio,
+            audio_for_next_phase=audio_for_next_phase,
+        )
+
+    def decode_phase(phase):
+        return DiffusionOutput(output=(phase.video, phase.audio))
+
+    class FakeUpsampler(torch.nn.Module):
+        dtype = torch.float32
+
+        def forward(self, latents):
+            return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
+
+    prompt_context_sentinel = prompt_context
+    stage_1_sigmas = [1.0, 0.5, 0.0]
+    stage_2_sigmas = [0.8, 0.2, 0.0]
+    pipeline = object.__new__(LTX2TwoStagePipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE
+    pipeline.device = torch.device("cpu")
+    pipeline.vae_spatial_compression_ratio = 32
+    pipeline.vae_temporal_compression_ratio = 8
+    pipeline.latent_upsampler = FakeUpsampler()
+    object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
+    object.__setattr__(pipeline, "run_phase", run_phase)
+    object.__setattr__(pipeline, "decode_phase", decode_phase)
+
+    source_image = object()
+    output = pipeline.forward(
+        SimpleNamespace(sampling_params_list=[], prompts=[]),
+        image=source_image,
+        stage_1_sigmas=stage_1_sigmas,
+        stage_2_sigmas=stage_2_sigmas,
+    )
+
+    assert len(phase_calls) == 2
+    assert phase_calls[1][2] is prompt_context_sentinel
+    torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
+    torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
+
+
+class TestLTXRequestParsing:
+    def test_request_resolves_custom_sigmas_from_extra_args(self):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        sigmas = [1.0, 0.5, 0.0]
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}),
+                    request_id="ltx-custom-sigmas",
+                )
+            ]
+        )
+
+        assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas
+
+    def test_request_resolves_per_stage_sigmas_from_extra_args(self):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        stage_1_sigmas = [1.0, 0.4, 0.0]
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}),
+                    request_id="ltx-custom-phase-sigmas",
+                )
+            ]
+        )
+
+        resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas(
+            req, [1.0, 0.0], [0.8, 0.0]
+        )
+
+        assert resolved == (stage_1_sigmas, [0.8, 0.0])
+
+    @pytest.mark.parametrize("value", [-1, 52, 1.5, True])
+    def test_request_rejects_invalid_image_crf(self, value):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}),
+                    request_id="ltx-invalid-image-crf",
+                )
+            ]
+        )
+
+        with pytest.raises(ValueError, match="image_crf"):
+            _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req)
+
+
+class TestLTXForwardStages:
+    def test_gemma_tokenization_matches_official_bos_and_left_padding(self):
+        pipe = _make_ltx_request_pipe(LTX2Pipeline)
+
+        class FakeTokenizer:
+            bos_token_id = 2
+            eos_token_id = 1
+            eos_token = "<eos>"
+            pad_token_id = 0
+            pad_token = "<pad>"
+            padding_side = "right"
+
+            def __init__(self):
+                self.calls = []
+
+            def __call__(self, text, **kwargs):
+                self.calls.append((text, kwargs))
+                ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id]
+                return SimpleNamespace(input_ids=ids)
+
+            def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask):
+                assert padding == "max_length"
+                assert return_tensors == "pt"
+                assert return_attention_mask is True
+                rows = []
+                masks = []
+                for ids in encoded["input_ids"]:
+                    pad = max_length - len(ids)
+                    rows.append([self.pad_token_id] * pad + ids)
+                    masks.append([0] * pad + [1] * len(ids))
+                return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks))
+
+        class FakeTextEncoder(torch.nn.Module):
+            dtype = torch.float32
+
+            def __init__(self):
+                super().__init__()
+                self.seen_input_ids = None
+                self.seen_attention_mask = None
+
+            def forward(self, *, input_ids, attention_mask, output_hidden_states):
+                assert output_hidden_states is True
+                self.seen_input_ids = input_ids.clone()
+                self.seen_attention_mask = attention_mask.clone()
+                return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),))
+
+        tokenizer = FakeTokenizer()
+        text_encoder = FakeTextEncoder()
+        object.__setattr__(pipe, "tokenizer", tokenizer)
+        object.__setattr__(pipe, "text_encoder", text_encoder)
+
+        pipe._get_gemma_prompt_embeds(
+            [" long ", "short"],
+            max_sequence_length=4,
+            device=torch.device("cpu"),
+            dtype=torch.float32,
+        )
+
+        torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]]))
+        torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]))
+        assert tokenizer.padding_side == "left"
+        assert [call[0] for call in tokenizer.calls] == ["long", "short"]

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -3,6 +3,7 @@
 
 """Unit tests for LTX image-to-video input and conditioning behavior."""
 
+import sys
 from types import SimpleNamespace
 
 import pytest
@@ -23,7 +24,7 @@ def _make_ltx23_request_pipe(cls):
 
 
 class TestLTXImageToVideoForwardStages:
-    def test_i2v_pil_preprocessing_matches_official_center_crop(self):
+    def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
         from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
 
         pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
@@ -35,6 +36,136 @@ class TestLTXImageToVideoForwardStages:
 
         torch.testing.assert_close(actual, expected)
 
+    def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+
+        pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
+        pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
+        image = Image.fromarray(pixels.numpy())
+        monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array)
+
+        actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18)
+        expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
+
+        torch.testing.assert_close(actual, expected)
+
+    def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        captured_crfs = []
+        captured_prepare_kwargs = []
+
+        def fake_preprocess(image, *, height, width, crf, device, dtype):
+            del image, device
+            captured_crfs.append(crf)
+            return torch.zeros(1, 3, height, width, dtype=dtype)
+
+        def fake_prepare_latents(**kwargs):
+            captured_prepare_kwargs.append(kwargs)
+            return kwargs["image"], None
+
+        monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess)
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+        request_inputs = SimpleNamespace(
+            height=4,
+            width=4,
+            num_frames=3,
+            num_videos_per_prompt=1,
+            generator=None,
+            latents=None,
+        )
+        prompt_context = SimpleNamespace(
+            batch_size=1,
+            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+        )
+
+        pipe._prepare_video_latents_stage(
+            request_inputs,
+            prompt_context,
+            device=torch.device("cpu"),
+            noise_scale=0.0,
+            image=Image.new("RGB", (8, 4)),
+        )
+
+        assert captured_crfs == [18]
+        assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16
+        assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16
+
+    def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+
+        with pytest.raises(ValueError, match="requires PIL images"):
+            pipe._prepare_video_latents_stage(
+                SimpleNamespace(),
+                SimpleNamespace(),
+                device=torch.device("cpu"),
+                noise_scale=0.0,
+                image=torch.zeros(3, 4, 4),
+            )
+
+    def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        captured_prepare_kwargs = []
+
+        def fake_prepare_latents(**kwargs):
+            captured_prepare_kwargs.append(kwargs)
+            return kwargs["image"], None
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+        request_inputs = SimpleNamespace(
+            height=4,
+            width=4,
+            num_frames=3,
+            num_videos_per_prompt=1,
+            generator=None,
+            latents=None,
+            image_crf=0,
+        )
+        prompt_context = SimpleNamespace(
+            batch_size=1,
+            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+        )
+
+        pipe._prepare_video_latents_stage(
+            request_inputs,
+            prompt_context,
+            device=torch.device("cpu"),
+            noise_scale=0.0,
+            image=torch.zeros(3, 4, 4),
+        )
+
+        assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4)
+
+    def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+        image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy()
+        monkeypatch.setitem(sys.modules, "av", None)
+
+        assert _apply_image_conditioning_crf(image, 18) is image
+
+    def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+        monkeypatch.setitem(sys.modules, "av", None)
+
+        with pytest.raises(ImportError, match="PyAV with a libx264 encoder"):
+            _apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18)
+
     def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
         from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -98,7 +229,7 @@ class TestLTXImageToVideoForwardStages:
         with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
             pipe._resolve_request_image(mixed_req, None, request_inputs)
 
-    def test_denoise_timestep_kwargs_masks_video_only(self):
+    def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         pipe = object.__new__(LTX2Pipeline)
@@ -189,13 +320,18 @@ class TestLTXImageToVideoConditioning:
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
         torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
 
-    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+    @pytest.mark.parametrize(
+        ("model_version", "sampled_shape"),
+        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+    )
+    def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
         import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
         import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
+        pipe.model_version = model_version
         pipe.vae_spatial_compression_ratio = 1
         pipe.vae_temporal_compression_ratio = 1
         pipe.transformer_spatial_patch_size = 1
@@ -237,14 +373,20 @@ class TestLTXImageToVideoConditioning:
 
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
         torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
-        assert sampled_shapes == [(1, 3, 2)]
+        assert sampled_shapes == [sampled_shape]
 
-    def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
+    @pytest.mark.parametrize(
+        ("model_version", "sampled_shape"),
+        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+    )
+    def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
         import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
+        pipe.model_version = model_version
         pipe.vae_spatial_compression_ratio = 1
         pipe.vae_temporal_compression_ratio = 1
         pipe.transformer_spatial_patch_size = 1
@@ -266,6 +408,7 @@ class TestLTXImageToVideoConditioning:
             return torch.ones(shape, device=device, dtype=dtype)
 
         monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
 
         out, conditioning_mask = pipe.prepare_latents(
             image=torch.zeros(1, 3, 1, 1),
@@ -279,10 +422,54 @@ class TestLTXImageToVideoConditioning:
             device=torch.device("cpu"),
         )
 
-        assert sampled_shapes == [(1, 3, 2)]
+        assert sampled_shapes == [sampled_shape]
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
         torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
 
+    def test_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        object.__setattr__(
+            pipe,
+            "_encode_i2v_image_latents",
+            lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16),
+        )
+        captured = {}
+
+        def fake_create_noised_state(latents, noise_scale, generator=None):
+            del generator
+            captured["noise_scale"] = noise_scale
+            return latents
+
+        monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state)
+        pipe.prepare_latents(
+            image=torch.zeros(1, 3, 1, 1),
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=0.909375,
+            dtype=torch.bfloat16,
+            device=torch.device("cpu"),
+            latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16),
+        )
+
+        assert captured["noise_scale"] == 0.909375
+
     def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
@@ -385,6 +572,99 @@ class TestLTXImageToVideoConditioning:
         torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
         assert pipeline.step_indices == [0, 1]
 
+    def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
+
+        device = torch.device("cpu")
+        sigmas = torch.tensor([1.0, 0.5, 0.0])
+        pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas))
+        mask = torch.tensor([[1.0, 0.0]])
+        adapter = LTXVideoAudioStepAdapter(
+            pipeline,
+            SimpleNamespace(sigmas=sigmas.clone()),
+            latent_num_frames=1,
+            latent_height=1,
+            latent_width=1,
+            image_conditioned=True,
+            sampler="euler_ancestral",
+            generator=torch.Generator(device=device).manual_seed(42),
+            conditioning_mask=mask,
+        )
+        video = torch.tensor([[[3.0], [2.0]]])
+        audio = torch.tensor([[[4.0]]])
+        video_velocity = torch.tensor([[[7.0], [0.25]]])
+        audio_velocity = torch.tensor([[[0.5]]])
+
+        reference_generator = torch.Generator(device=device).manual_seed(10042)
+        video_noise = torch.randn(video.shape, generator=reference_generator)
+        audio_noise = torch.randn(audio.shape, generator=reference_generator)
+
+        def official_step(sample, velocity, noise):
+            sigma = sigmas[0]
+            sigma_next = sigmas[1]
+            denoised = sample.float() - velocity.float() * sigma
+            sigma_down = sigma_next * (sigma_next / sigma)
+            ratio = sigma_down / sigma
+            deterministic = ratio * sample.float() + (1.0 - ratio) * denoised
+            alpha_next = 1.0 - sigma_next
+            alpha_down = 1.0 - sigma_down
+            coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+            return (alpha_next / alpha_down) * deterministic + noise * coeff
+
+        expected_video = official_step(video, video_velocity, video_noise)
+        expected_video[:, :1] = video[:, :1]
+        expected_audio = official_step(audio, audio_velocity, audio_noise)
+
+        ((actual_video, actual_audio),) = adapter.step(
+            (video_velocity, audio_velocity),
+            (sigmas[0], sigmas[0]),
+            (video, audio),
+        )
+        torch.testing.assert_close(actual_video, expected_video)
+        torch.testing.assert_close(actual_audio, expected_audio)
+
+        terminal_video_velocity = torch.full_like(actual_video, 0.125)
+        terminal_audio_velocity = torch.full_like(actual_audio, 0.25)
+        ((terminal_video, terminal_audio),) = adapter.step(
+            (terminal_video_velocity, terminal_audio_velocity),
+            (sigmas[1], sigmas[1]),
+            (actual_video, actual_audio),
+        )
+        expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1]
+        expected_terminal_video[:, :1] = actual_video[:, :1]
+        expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1]
+        torch.testing.assert_close(terminal_video, expected_terminal_video)
+        torch.testing.assert_close(terminal_audio, expected_terminal_audio)
+
+    def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity
+
+        sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16)
+        velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16)
+        sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32)
+        generator = torch.Generator(device="cpu").manual_seed(10042)
+        reference_generator = torch.Generator(device="cpu").manual_seed(10042)
+
+        sigma, sigma_next = sigmas[:2]
+        denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float()
+        sigma_down = sigma_next * (sigma_next / sigma)
+        sigma_down_ratio = sigma_down / sigma
+        expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
+        alpha_next = 1.0 - sigma_next
+        alpha_down = 1.0 - sigma_down
+        coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+        noise = torch.randn(
+            sample.shape,
+            generator=reference_generator,
+            dtype=sample.dtype,
+            device=sample.device,
+        )
+        expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype)
+
+        actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator)
+
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
     def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
         from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -202,3 +202,68 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length():
             device=torch.device("cpu"),
             latents=latents,
         )
+
+
+def test_create_noised_state_matches_official_fp32_lerp():
+    latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128)
+    expected_generator = torch.Generator().manual_seed(42)
+    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+    expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to(
+        latents.dtype
+    )
+
+    actual = latent_ops.create_noised_state(
+        latents,
+        torch.tensor(0.15, dtype=torch.bfloat16),
+        torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_create_conditioned_noised_state_matches_official_two_lerps():
+    latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8)
+    clean_latents = latents.clone()
+    clean_latents[:, 0] = 3.0
+    denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]])
+    expected_generator = torch.Generator().manual_seed(42)
+    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+    noised = torch.lerp(latents.float(), noise.float(), 0.909375)
+    expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype)
+
+    actual = latent_ops.create_conditioned_noised_state(
+        latents,
+        clean_latents,
+        denoise_mask,
+        0.909375,
+        torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_prepare_supplied_video_latents_uses_official_token_major_layout():
+    pipeline = _make_pipeline(LTX2Pipeline)
+    supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16)
+    actual_video = pipeline.prepare_latents(
+        batch_size=1,
+        num_channels_latents=4,
+        height=64,
+        width=64,
+        num_frames=9,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=supplied,
+    )
+
+    torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0)
+    assert actual_video.stride()[1:] == (1, actual_video.shape[1])
+
+
+def test_clear_audio_padding_keeps_padding_outside_sampler_state():
+    updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]])
+
+    actual = latent_ops.clear_audio_padding(updated, 2)
+
+    torch.testing.assert_close(actual[:, :2], updated[:, :2])
+    torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:]))

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -28,6 +28,15 @@ def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
 class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""
 
+    def test_single_guidance_pass_preserves_official_token_layout(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch
+
+        tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2)
+        repeated = _repeat_batch(tokens, 1)
+
+        assert repeated is tokens
+        assert repeated.stride() == tokens.stride()
+
     def test_combine_cfg_noise_matches_x0_space_formula(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -20,6 +20,7 @@ from vllm_omni.diffusion.models.ltx2 import (
 from vllm_omni.diffusion.models.ltx2.ltx2_components import (
     LTX2_COMPONENT_PROFILE,
     LTX2_DISTILLED_COMPONENT_PROFILE,
+    LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
     LTX2_TWO_STAGE_COMPONENT_PROFILE,
     LTX23_COMPONENT_PROFILE,
     LTX23_DISTILLED_COMPONENT_PROFILE,
@@ -47,6 +48,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
 from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+    LTX2_DISTILLED_ONE_STAGE_RECIPE,
     LTX2_DISTILLED_TWO_STAGE_RECIPE,
     LTX2_ONE_STAGE_RECIPE,
     LTX2_TWO_STAGE_RECIPE,
@@ -64,12 +66,14 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2DistilledOneStagePipeline,
     LTX2I2VDMD2Pipeline,
     LTX2Pipeline,
     LTX2T2VDMD2Pipeline,
 )
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
     LTX2DistilledPipeline,
+    LTX2DistilledTwoStagePipeline,
     LTX2TwoStagePipeline,
 )
 
@@ -125,15 +129,21 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
 
     assert issubclass(LTX2Pipeline, LTXRuntime)
     assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
+    assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime)
+    assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime)
     assert issubclass(LTX2DistilledPipeline, LTXRuntime)
+    assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline
     assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
     assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
     assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
     assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
+    assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+    assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE
     assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
     for pipeline_cls, profile in (
         (LTX2Pipeline, LTX2_COMPONENT_PROFILE),
         (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
+        (LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE),
         (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
     ):
         assert pipeline_cls._dit_modules == list(profile.dit_modules)
@@ -254,7 +264,7 @@ def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_pa
 
     monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
 
-    with pytest.raises(ValueError, match="only one-stage LTX recipes"):
+    with pytest.raises(ValueError, match="not qualified for this LTX recipe"):
         LTX2TwoStagePipeline(
             od_config=SimpleNamespace(
                 model=str(tmp_path),
@@ -709,15 +719,48 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch):
 
     assert len(installed) == 2
     assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed)
-    assert all(processor.has_learned_registers for processor in installed)
     assert all(kernel["role"] == "ltx2.connector" for kernel in kernels)
     assert all(kernel["role_category"] == "self" for kernel in kernels)
     assert all(kernel["skip_sequence_parallel"] for kernel in kernels)
     assert all(kernel["disable_kv_quant"] for kernel in kernels)
+    assert all(processor.has_learned_registers for processor in installed)
+    assert not any(processor.preserve_learned_register_mask for processor in installed)
 
 
-@pytest.mark.parametrize("has_learned_registers", [False, True])
-def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers):
+def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch):
+    installed = []
+
+    class OmniAttention:
+        def __init__(self, **_kwargs):
+            pass
+
+    class Attention:
+        heads = 32
+        head_dim = 128
+        inner_kv_dim = 4096
+
+        def set_processor(self, processor):
+            installed.append(processor)
+
+    monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention)
+    connectors = SimpleNamespace(
+        video_connector=SimpleNamespace(
+            learnable_registers=object(),
+            transformer_blocks=[SimpleNamespace(attn1=Attention())],
+        ),
+        audio_connector=SimpleNamespace(
+            learnable_registers=object(),
+            transformer_blocks=[SimpleNamespace(attn1=Attention())],
+        ),
+    )
+
+    _install_connector_attention(connectors, preserve_learned_register_mask=True)
+
+    assert all(processor.has_learned_registers for processor in installed)
+    assert all(processor.preserve_learned_register_mask for processor in installed)
+
+
+def test_ltx_connector_attention_dispatches_through_omni_kernel():
     calls = []
 
     class Backend:
@@ -732,12 +775,16 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
             calls.append((query, key, value, metadata))
             return query
 
+    class UpcastingNorm(torch.nn.Module):
+        def forward(self, tensor):
+            return tensor.float()
+
     attention = SimpleNamespace(
         heads=2,
         head_dim=4,
         inner_kv_dim=8,
-        norm_q=torch.nn.Identity(),
-        norm_k=torch.nn.Identity(),
+        norm_q=UpcastingNorm(),
+        norm_k=UpcastingNorm(),
         to_q=torch.nn.Identity(),
         to_k=torch.nn.Identity(),
         to_v=torch.nn.Identity(),
@@ -746,24 +793,29 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
         rope_type="split",
         omni_attention=OmniAttention(),
     )
-    hidden_states = torch.randn(2, 3, 8)
+    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
     additive_mask = torch.zeros(2, 1, 3, 3)
+    rotary = (
+        torch.randn(2, 2, 3, 2, dtype=torch.float32),
+        torch.randn(2, 2, 3, 2, dtype=torch.float32),
+    )
 
-    output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)(
+    output = ltx2_components._LTXConnectorAttnProcessor()(
         attention,
         hidden_states,
         attention_mask=additive_mask,
+        query_rotary_emb=rotary,
     )
 
     query, key, value, metadata = calls[0]
     assert query.shape == key.shape == value.shape == (2, 3, 2, 4)
-    if has_learned_registers:
-        assert metadata is None
-    else:
-        assert metadata.attn_mask.shape == (2, 3)
-        assert metadata.attn_mask.dtype == torch.bool
-        assert metadata.attn_mask.all()
-    torch.testing.assert_close(output, hidden_states)
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert metadata.attn_mask.shape == (2, 3)
+    assert metadata.attn_mask.dtype == torch.bool
+    assert metadata.attn_mask.all()
+    expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary)
+    expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4)
+    torch.testing.assert_close(output, expected, rtol=0, atol=0)
 
 
 def test_ltx2_distilled_pipeline_shares_recipe_runtime():
@@ -1707,6 +1759,7 @@ class TestLTXForwardStages:
         assert seen["request_inputs"].output_type == "latent"
         assert seen["kwargs"] == {
             "request_sigmas": request_sigmas,
+            "request_phase_sigmas": None,
             "image": None,
         }
 
@@ -1824,6 +1877,16 @@ class TestRegistryIntegration:
             "pipeline_ltx2_two_stage",
             "LTX2TwoStagePipeline",
         )
+        assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == (
+            "ltx2",
+            "pipeline_ltx2",
+            "LTX2DistilledOneStagePipeline",
+        )
+        assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == (
+            "ltx2",
+            "pipeline_ltx2_two_stage",
+            "LTX2DistilledTwoStagePipeline",
+        )
         assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
             "ltx2",
             "pipeline_ltx2_two_stage",
@@ -1848,6 +1911,8 @@ class TestRegistryIntegration:
         expected = [
             "LTX2Pipeline",
             "LTX2TwoStagePipeline",
+            "LTX2DistilledOneStagePipeline",
+            "LTX2DistilledTwoStagePipeline",
             "LTX2DistilledPipeline",
             "LTX2T2VDMD2Pipeline",
             "LTX2I2VDMD2Pipeline",
@@ -1861,6 +1926,8 @@ class TestRegistryIntegration:
         [
             "LTX2Pipeline",
             "LTX2TwoStagePipeline",
+            "LTX2DistilledOneStagePipeline",
+            "LTX2DistilledTwoStagePipeline",
             "LTX2DistilledPipeline",
             "LTX2T2VDMD2Pipeline",
             "LTX2I2VDMD2Pipeline",

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -7,18 +7,54 @@ import pytest
 import torch
 from torch import nn
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
+from vllm_omni.diffusion.models.ltx2 import ltx2_transformer
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
-    LTX2AudioVideoAttnProcessor,
+    LTX2FeedForward,
     LTX2VideoTransformer3DModel,
-    _LTX2ParallelAttention,
+    LTX2VideoTransformerBlock,
     _make_rms_norm,
+    apply_interleaved_rotary_emb,
+    apply_keyframes_absolute_embedding,
+    apply_split_rotary_emb,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
+def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic():
+    torch.manual_seed(0)
+    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    cos = torch.randn(2, 3, 8, dtype=x.dtype)
+    sin = torch.randn(2, 3, 8, dtype=x.dtype)
+    x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1)
+    rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+    expected = x * cos + rotated * sin
+
+    actual = apply_interleaved_rotary_emb(x, (cos, sin))
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_ltx_split_rope_matches_official_model_dtype_arithmetic():
+    torch.manual_seed(1)
+    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    cos = torch.randn(2, 3, 4, dtype=x.dtype)
+    sin = torch.randn(2, 3, 4, dtype=x.dtype)
+    split_x = x.reshape(2, 3, 2, 4)
+    first_x, second_x = split_x.chunk(2, dim=-2)
+    cos_u = cos.unsqueeze(-2)
+    sin_u = sin.unsqueeze(-2)
+    expected = split_x * cos_u
+    expected[..., :1, :].addcmul_(-sin_u, second_x)
+    expected[..., 1:, :].addcmul_(sin_u, first_x)
+    expected = expected.reshape_as(x)
+
+    actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8)
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
 def test_ltx_rms_norm_no_affine_matches_official_torch_module():
     norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
 
@@ -36,11 +72,53 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
     assert "weight" in norm.state_dict()
 
 
-def test_ltx_transformer_has_separate_cfg_cache_dit_config():
+def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms():
+    calls = []
+
+    class UpcastingNorm(nn.Module):
+        def forward(self, tensor):
+            return tensor.float()
+
+    class RecordingAttention:
+        def __call__(self, query, key, value, metadata):
+            calls.append((query, key, value, metadata))
+            return query
+
+    attention = SimpleNamespace(
+        heads=2,
+        head_dim=4,
+        to_qkv=None,
+        to_q=nn.Identity(),
+        to_k=nn.Identity(),
+        to_v=nn.Identity(),
+        norm_q=UpcastingNorm(),
+        norm_k=UpcastingNorm(),
+        to_gate_logits=None,
+        attn=RecordingAttention(),
+        to_out=(nn.Identity(), nn.Identity()),
+        rope_type="split",
+    )
+    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    rotary = (
+        torch.ones(2, 2, 3, 2, dtype=torch.float32),
+        torch.zeros(2, 2, 3, 2, dtype=torch.float32),
+    )
+
+    processor = ltx2_transformer.LTX2AudioVideoAttnProcessor()
+    processor._slice_rope_for_tp = lambda rope, _: rope
+    output = processor(attention, hidden_states, query_rotary_emb=rotary)
+
+    query, key, value, metadata = calls[0]
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert metadata is None
+    torch.testing.assert_close(output, hidden_states)
+
+
+def test_ltx_transformer_has_fused_cfg_cache_dit_config():
     adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
 
     assert isinstance(adapter_config, CacheDiTAdapterConfig)
-    assert adapter_config.has_separate_cfg
+    assert not adapter_config.has_separate_cfg
 
 
 def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
@@ -76,8 +154,9 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
 
 
 def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
-    monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
-    processor = LTX2AudioVideoAttnProcessor()
+    processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor
+    monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True))
+    processor = processor_cls()
     hidden_states = torch.zeros(1, 3, 4)
     encoder_hidden_states = torch.zeros(1, 4, 4)
     key_padding_mask = torch.tensor([[True, True, True, False]])
@@ -113,7 +192,7 @@ def _fake_ltx_parallel_attention():
 
         return forward
 
-    layer = object.__new__(_LTX2ParallelAttention)
+    layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
     nn.Module.__init__(layer)
     layer.attn_backend = Backend
     layer.attention = SimpleNamespace(forward=run("native"))
@@ -123,33 +202,281 @@ def _fake_ltx_parallel_attention():
 
 
 @pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
-def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
+def test_ltx_masked_attention_uses_sdpa(query_length, key_length):
     layer, calls = _fake_ltx_parallel_attention()
     query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
     key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
     key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
-    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
+    metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
 
     layer._run_local_attention(query, key, value, metadata)
 
     assert calls == ["sdpa"]
 
 
-def test_ltx_masked_self_attention_uses_sdpa():
-    layer, calls = _fake_ltx_parallel_attention()
-    query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
-    key_padding_mask = torch.tensor([[True, True, True, False]])
-    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
-
-    layer._run_local_attention(query, key, value, metadata)
-
-    assert calls == ["sdpa"]
-
-
-def test_ltx_unmasked_cross_attention_keeps_native_backend():
+def test_ltx_unmasked_attention_keeps_native_backend():
     layer, calls = _fake_ltx_parallel_attention()
     query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
 
     layer._run_local_attention(query, key, value, None)
 
     assert calls == ["native"]
+
+
+def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch):
+    class FakeParallelLinear(nn.Linear):
+        def __init__(self, dim_in, dim_out, *, bias=True, **kwargs):
+            super().__init__(dim_in, dim_out, bias=bias)
+
+    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear)
+
+    without_bias = LTX2FeedForward(8, inner_dim=16, bias=False)
+    with_bias = LTX2FeedForward(8, inner_dim=16, bias=True)
+
+    assert without_bias.net[0].proj.bias is None
+    assert without_bias.net[2].bias is None
+    assert with_bias.net[0].proj.bias is not None
+    assert with_bias.net[2].bias is not None
+
+
+def _build_tiny_ltx_transformer(monkeypatch, **overrides):
+    class FakeBlock(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+            self.init_kwargs = kwargs
+
+        def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs):
+            return hidden_states, audio_hidden_states
+
+    monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock)
+    kwargs = {
+        "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": 8,
+        "base_width": 8,
+        "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,
+        "use_prompt_embeddings": False,
+        "cross_attn_mod": True,
+        "audio_cross_attn_mod": True,
+    }
+    kwargs.update(overrides)
+    return LTX2VideoTransformer3DModel(**kwargs)
+
+
+def test_ltx25_keyframe_absolute_embedding_matches_official():
+    hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8)
+    keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]])
+    embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8)
+
+    actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding)
+    expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states
+
+
+def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch):
+    legacy_model = _build_tiny_ltx_transformer(monkeypatch)
+
+    assert legacy_model.config.ff_bias is True
+    assert legacy_model.config.audio_ff_bias is True
+    assert legacy_model.config.use_prompt_adaln_single is True
+    assert legacy_model.config.use_keyframes_abs_pos_embedding is False
+    assert hasattr(legacy_model, "prompt_adaln")
+    assert hasattr(legacy_model, "audio_prompt_adaln")
+    assert not hasattr(legacy_model, "keyframes_abs_pos_embedding")
+    assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True
+    assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+    ltx25_model = _build_tiny_ltx_transformer(
+        monkeypatch,
+        ff_bias=False,
+        audio_ff_bias=True,
+        use_prompt_adaln_single=False,
+        use_keyframes_abs_pos_embedding=True,
+    )
+
+    assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8)
+    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8))
+    assert not hasattr(ltx25_model, "prompt_adaln")
+    assert not hasattr(ltx25_model, "audio_prompt_adaln")
+    assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False
+    assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1)
+    expected = torch.arange(8, dtype=torch.float32).view(1, 8)
+    loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)])
+
+    assert loaded == {"keyframes_abs_pos_embedding"}
+    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected)
+
+
+def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch):
+    class CaptureAda(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.calls = []
+
+        def forward(self, timestep, *, batch_size, hidden_dtype):
+            self.calls.append(timestep.detach().clone())
+            embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype)
+            return embedding, embedding
+
+    class CaptureRope(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.output_dtypes = []
+
+        def forward(self, _coords, *, device, out_dtype=None):
+            self.output_dtypes.append(out_dtype)
+            empty = torch.empty(0, device=device, dtype=out_dtype)
+            return empty, empty
+
+    model = _build_tiny_ltx_transformer(monkeypatch)
+    rope = CaptureRope()
+    model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope
+    model.prompt_modulation = False
+    video_scale_shift = CaptureAda()
+    video_gate = CaptureAda()
+    audio_scale_shift = CaptureAda()
+    audio_gate = CaptureAda()
+    model.av_cross_attn_video_scale_shift = video_scale_shift
+    model.av_cross_attn_video_a2v_gate = video_gate
+    model.av_cross_attn_audio_scale_shift = audio_scale_shift
+    model.av_cross_attn_audio_v2a_gate = audio_gate
+
+    model(
+        hidden_states=torch.zeros(1, 2, 4),
+        audio_hidden_states=torch.zeros(1, 1, 4),
+        encoder_hidden_states=torch.zeros(1, 1, 8),
+        audio_encoder_hidden_states=torch.zeros(1, 1, 8),
+        timestep=torch.tensor([[0.0, 2.0]]),
+        audio_timestep=torch.tensor([[5.0]]),
+        sigma=torch.tensor([2.0]),
+        audio_sigma=torch.tensor([5.0]),
+        video_coords=torch.zeros(1, 3, 2, 2),
+        audio_coords=torch.zeros(1, 1, 1, 2),
+    )
+
+    torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0]))
+    torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0]))
+    torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0]))
+    torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0]))
+    assert len(rope.output_dtypes) == 4
+    assert all(dtype is torch.float32 for dtype in rope.output_dtypes)
+
+
+def test_ltx25_static_prompt_table_is_used_without_prompt_adaln():
+    class CaptureAttention(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.encoder_hidden_states = None
+
+        def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+            if encoder_hidden_states is not None:
+                self.encoder_hidden_states = encoder_hidden_states.detach().clone()
+            return torch.zeros_like(hidden_states)
+
+    class ZeroOutput(nn.Module):
+        def forward(self, hidden_states):
+            return torch.zeros_like(hidden_states)
+
+    block = object.__new__(LTX2VideoTransformerBlock)
+    nn.Module.__init__(block)
+    block.video_cross_attn_adaln = True
+    block.audio_cross_attn_adaln = True
+    block.cross_attn_adaln = True
+    block.scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+    block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+    block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+    block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]]))
+    block.norm1 = nn.Identity()
+    block.audio_norm1 = nn.Identity()
+    block.norm2 = nn.Identity()
+    block.audio_norm2 = nn.Identity()
+    block.norm3 = nn.Identity()
+    block.audio_norm3 = nn.Identity()
+    block.attn1 = CaptureAttention()
+    block.audio_attn1 = CaptureAttention()
+    block.attn2 = CaptureAttention()
+    block.audio_attn2 = CaptureAttention()
+    block.ff = ZeroOutput()
+    block.audio_ff = ZeroOutput()
+
+    hidden_states = torch.ones(1, 1, 2)
+    audio_hidden_states = torch.ones(1, 1, 2)
+    encoder_hidden_states = torch.ones(1, 1, 2)
+    audio_encoder_hidden_states = torch.ones(1, 1, 2)
+    block(
+        hidden_states,
+        audio_hidden_states,
+        encoder_hidden_states,
+        audio_encoder_hidden_states,
+        temb=torch.zeros(1, 1, 18),
+        temb_audio=torch.zeros(1, 1, 18),
+        temb_ca_scale_shift=torch.empty(0),
+        temb_ca_audio_scale_shift=torch.empty(0),
+        temb_ca_gate=torch.empty(0),
+        temb_ca_audio_gate=torch.empty(0),
+        temb_prompt=None,
+        temb_prompt_audio=None,
+        use_a2v_cross_attention=False,
+        use_v2a_cross_attention=False,
+    )
+
+    torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]]))
+    torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]]))
+
+
+def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch):
+    column_calls = []
+
+    class FakeColumnParallelLinear(nn.Linear):
+        def __init__(self, input_size, output_size, *, bias=True, **kwargs):
+            column_calls.append((output_size, kwargs))
+            super().__init__(input_size, output_size, bias=bias)
+
+    class FakeRowParallelLinear(nn.Linear):
+        def __init__(self, input_size, output_size, *, bias=True, **_kwargs):
+            super().__init__(input_size, output_size, bias=bias)
+
+    class FakeParallelAttention(nn.Module):
+        def __init__(self, **_kwargs):
+            super().__init__()
+
+    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention)
+    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
+
+    attention = ltx2_transformer.LTX2Attention(
+        query_dim=8,
+        heads=4,
+        kv_heads=4,
+        dim_head=2,
+        qk_norm="rms_norm",
+        apply_gated_attention=True,
+        pack_qkv=False,
+    )
+
+    gate_output_size, gate_kwargs = column_calls[-1]
+    assert gate_output_size == 4
+    assert gate_kwargs["gather_output"] is False
+    assert gate_kwargs["return_bias"] is False
+    assert attention.query_num_heads == 2