Kandinsky6 (#8537)

Signed-off-by: leffff <levnovitskiy@gmail.com>
Signed-off-by: Artem Borisov <divotionchannel@gmail.com>
Co-authored-by: leffff <levnovitskiy@gmail.com>
Co-authored-by: Artem Borisov <divotionchannel@gmail.com>

tests/diffusion/models/kandinsky6/test_acceleration.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU checks that Kandinsky 6 acceleration hooks are actually wired."""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.kandinsky6.cache_accel import K6_PRO_MAG_RATIOS, K6StepCache
+from vllm_omni.diffusion.models.kandinsky6.kandinsky6_transformer import Kandinsky6Transformer3DModel
+from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import _shard_loaded_weight
+from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_shard_loaded_weight_narrows_column_parallel_output():
+    class _Owner:
+        tp_rank = 1
+
+        def weight_loader(self, param, loaded):
+            del param, loaded
+
+    param = nn.Parameter(torch.empty(2, 4))
+    param.output_dim = 0
+    param.weight_loader = _Owner().weight_loader
+    loaded = torch.arange(16, dtype=torch.float32).reshape(4, 4)
+
+    shard = _shard_loaded_weight(param, loaded)
+
+    assert tuple(shard.shape) == (2, 4)
+    torch.testing.assert_close(shard, loaded[2:4])
+
+
+def test_row_bias_is_not_sharded_when_shapes_match():
+    class _Owner:
+        tp_rank = 1
+
+        def weight_loader(self, param, loaded):
+            del param, loaded
+
+    param = nn.Parameter(torch.empty(4))
+    param.output_dim = 0
+    param.weight_loader = _Owner().weight_loader
+    loaded = torch.arange(4, dtype=torch.float32)
+
+    shard = _shard_loaded_weight(param, loaded)
+
+    torch.testing.assert_close(shard, loaded)
+
+
+def test_hsdp_predicate_matches_visual_blocks_only_when_indexed():
+    assert Kandinsky6Transformer3DModel._is_transformer_block("visual_transformer_blocks.3", nn.Identity())
+    assert Kandinsky6Transformer3DModel._is_transformer_block("video_text_transformer_blocks.0", nn.Identity())
+    assert not Kandinsky6Transformer3DModel._is_transformer_block("visual_embeddings", nn.Identity())
+    assert Kandinsky6Transformer3DModel._hsdp_shard_conditions
+
+
+def test_sp_plan_shards_visual_tokens_and_gathers_before_the_head():
+    plan = Kandinsky6Transformer3DModel._sp_plan
+    assert "_sp_visual_shard" in plan
+    assert "_sp_visual_rope" in plan
+    assert plan["_sp_visual_gather"].gather_dim == 1
+    assert Kandinsky6Transformer3DModel._magcache_block_attrs == ("visual_transformer_blocks",)
+
+
+def test_k6_is_not_on_the_no_cache_list_and_enablers_are_registered():
+    import vllm_omni.diffusion.cache.cachedit  # noqa: F401  registers enablers
+    from vllm_omni.diffusion.cache.cachedit.backend import CUSTOM_DIT_ENABLERS
+    from vllm_omni.diffusion.cache.teacache.backend import CUSTOM_TEACACHE_ENABLERS
+
+    assert "Kandinsky6TI2VAPipeline" not in _NO_CACHE_ACCELERATION
+    assert "Kandinsky6TI2VAPipeline" in CUSTOM_TEACACHE_ENABLERS
+    assert "Kandinsky6TI2VAPipeline" in CUSTOM_DIT_ENABLERS
+    assert len(K6_PRO_MAG_RATIOS) == 100
+
+
+def test_step_cache_skips_only_after_a_real_step():
+    cache = K6StepCache(
+        kind="mag",
+        threshold=0.24,
+        max_skip_steps=3,
+        retention_ratio=0.0,
+        mag_ratios=(1.0, 1.0),
+    )
+    assert cache.should_skip(0) is False
+    cache.store(torch.zeros(1), torch.zeros(1))
+    assert cache.should_skip(1) is True

tests/diffusion/models/kandinsky6/test_pipeline_kandinsky6.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import numpy as np
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_kandinsky6_pipeline_import_and_registry():
+    from vllm_omni.diffusion.models.kandinsky6 import (
+        AutoencoderKLHunyuanVideo,
+        Kandinsky6AudioVAE,
+        Kandinsky6TI2VAPipeline,
+        Kandinsky6Transformer3DModel,
+        KandinskyFlowMatchScheduler,
+        get_kandinsky6_post_process_func,
+        get_kandinsky6_pre_process_func,
+    )
+    from vllm_omni.diffusion.registry import (
+        _DIFFUSION_MODELS,
+        _DIFFUSION_POST_PROCESS_FUNCS,
+        _DIFFUSION_PRE_PROCESS_FUNCS,
+        _NO_CACHE_ACCELERATION,
+    )
+
+    assert Kandinsky6TI2VAPipeline is not None
+    assert Kandinsky6Transformer3DModel is not None
+    assert AutoencoderKLHunyuanVideo is not None
+    assert Kandinsky6AudioVAE is not None
+    assert KandinskyFlowMatchScheduler is not None
+    assert get_kandinsky6_post_process_func is not None
+    assert get_kandinsky6_pre_process_func is not None
+
+    assert _DIFFUSION_MODELS["Kandinsky6TI2VAPipeline"] == (
+        "kandinsky6",
+        "pipeline_kandinsky6",
+        "Kandinsky6TI2VAPipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_post_process_func"
+    assert _DIFFUSION_PRE_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_pre_process_func"
+    assert "Kandinsky6TI2VAPipeline" not in _NO_CACHE_ACCELERATION
+
+
+def test_kandinsky6_component_discovery_declarations():
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline
+
+    assert Kandinsky6TI2VAPipeline._dit_modules == ["transformer"]
+    assert Kandinsky6TI2VAPipeline._encoder_modules == ["text_encoder", "text_encoder_2"]
+    assert Kandinsky6TI2VAPipeline._vae_modules == ["vae", "audio_vae"]
+    assert Kandinsky6TI2VAPipeline.supports_step_execution is True
+    assert Kandinsky6TI2VAPipeline.support_audio_output is True
+    assert Kandinsky6TI2VAPipeline.support_image_input is True
+
+
+def test_kandinsky6_pipeline_satisfies_capability_protocols():
+    """isinstance checks against the @runtime_checkable protocols the
+    framework actually uses to detect these capabilities (io_support.py's
+    supports_audio_output, etc.) — stronger than just checking the class
+    attribute exists."""
+    from vllm_omni.diffusion.models.interface import (
+        SupportAudioOutput,
+        SupportImageInput,
+        SupportsComponentDiscovery,
+    )
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline
+
+    assert isinstance(Kandinsky6TI2VAPipeline, SupportAudioOutput)
+    assert isinstance(Kandinsky6TI2VAPipeline, SupportImageInput)
+    assert isinstance(Kandinsky6TI2VAPipeline, SupportsComponentDiscovery)
+
+
+class _FakeAudioVAE:
+    downsample_factor = 1024
+
+
+def test_kandinsky6_post_process_func_packages_video_and_audio():
+    """Pure-function test: no model/weights needed. Confirms the post-process
+    payload shape matches what io_support.py -> output_formatter.py ->
+    media_utils.py's PyAV muxing expects (the same flat {"video", "audio",
+    "audio_sample_rate", "fps"} shape MiniMax H3's own post-process function
+    produces), and that int16 PCM from ``postprocess_audio`` is rescaled to
+    the float32 [-1, 1] waveform the muxer consumes."""
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 4, 8, 8, 3), dtype=np.uint8)
+    audio = np.array([0, 32767, -32767], dtype=np.int16)
+
+    result = post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np")
+
+    assert result["video"] is video
+    assert result["audio"].dtype == np.float32
+    np.testing.assert_allclose(result["audio"], [0.0, 1.0, -1.0])
+    assert result["audio_sample_rate"] == 44100
+    assert result["fps"] == 24.0
+
+
+def test_kandinsky6_post_process_sample_rate_survives_output_formatter():
+    """The formatter only lifts ``audio_sample_rate`` into metadata from the
+    flat payload form; an envelope would silently lose it and the MP4 would
+    be muxed at the 24 kHz default."""
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+    from vllm_omni.diffusion.output_formatter import normalize_diffusion_postprocess_output
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
+    audio = np.zeros((10,), dtype=np.int16)
+
+    normalized = normalize_diffusion_postprocess_output(
+        post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np")
+    )
+
+    assert normalized.primary_key == "video"
+    assert normalized.metadata["audio"]["sample_rate"] == 44100
+    assert normalized.metadata["video"]["fps"] == 24.0
+    assert "audio_sample_rate" not in normalized.outputs
+
+
+def test_kandinsky6_post_process_func_omits_audio_when_none():
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
+
+    result = post_process({"video": video, "audio": None, "audio_sample_rate": None}, output_type="np")
+
+    assert result["video"] is video
+    assert "audio" not in result
+    assert "audio_sample_rate" not in result
+
+
+def test_kandinsky6_post_process_func_unwraps_batched_audio_list():
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
+
+    post_process = get_kandinsky6_post_process_func(od_config=None)
+    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
+    audio_item = np.full((50,), 16384, dtype=np.int16)
+
+    result = post_process({"video": video, "audio": [audio_item], "audio_sample_rate": 44100}, output_type="np")
+
+    assert result["audio"].shape == (50,)
+    np.testing.assert_allclose(result["audio"], 16384 / 32767, rtol=1e-6)
+
+
+class _StubAudioVAE:
+    """Audio VAE stand-in: identity "decode" that returns the latent as a waveform."""
+
+    scaling_factor = 0.5
+    mean_value = 0.0
+    device = "cpu"
+
+    def wrapped_decode(self, latents):
+        # (1, audio_dim, A) -> (A,) waveform: take the first channel.
+        return latents[0, 0]
+
+
+def test_postprocess_audio_normalizes_by_default_and_clips_on_request():
+    """The default mode must match the k6_video production pipeline, which
+    peak-normalizes each decoded waveform to full scale; ``clip`` keeps the
+    raw amplitude (saturated to [-1, 1]) for callers that want it."""
+    import torch
+
+    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import LatentBundle, postprocess_audio
+
+    # Latent (A=4, audio_dim=1); the stub scales by 1/scaling_factor -> [0.2, -0.4, 0.6, 3.0].
+    audio = torch.tensor([[0.1], [-0.2], [0.3], [1.5]], dtype=torch.float32)
+    bundle = LatentBundle(
+        video=None,
+        audio=audio,
+        video_cu_seqlens=None,
+        audio_cu_seqlens=torch.tensor([0, 4], dtype=torch.int32),
+    )
+    vae = _StubAudioVAE()
+
+    normalized = postprocess_audio(bundle, vae)
+    assert normalized is not None and len(normalized) == 1
+    assert normalized[0].dtype == np.int16
+    np.testing.assert_allclose(
+        normalized[0].astype(np.float32) / 32767,
+        np.array([0.2, -0.4, 0.6, 3.0], dtype=np.float32) / 3.0,
+        atol=1e-4,
+    )
+
+    clipped = postprocess_audio(bundle, vae, normalization_mode="clip")
+    np.testing.assert_allclose(
+        clipped[0].astype(np.float32) / 32767,
+        np.array([0.2, -0.4, 0.6, 1.0], dtype=np.float32),
+        atol=1e-4,
+    )
+
+    assert postprocess_audio(LatentBundle(None, None, None, None), vae) is None
+    with pytest.raises(ValueError, match="normalization_mode"):
+        postprocess_audio(bundle, vae, normalization_mode="loud")
+
+
+def test_kandinsky6_pre_process_func_is_identity_for_now():
+    """v1 scope: the pre-process hook is registered (matching the framework's
+    convention for image-conditioned models like SanaImageToVideoPipeline)
+    but does not yet validate/transform the request — documented follow-up,
+    not a silent gap."""
+    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_pre_process_func
+
+    pre_process = get_kandinsky6_pre_process_func(od_config=None)
+    sentinel = object()
+
+    assert pre_process(sentinel) is sentinel
+
+
+def test_hub_component_keys_map_onto_pipeline_modules():
+    """Hub folder keys from Kandinsky-6.0-Pro-5s-Diffusers land on the modules."""
+    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import (
+        _WEIGHT_SUBFOLDERS,
+        _adapt_k6_weight_name,
+    )
+
+    assert _WEIGHT_SUBFOLDERS == (
+        ("transformer", "transformer."),
+        ("vae", "vae."),
+        ("text_encoder", "text_encoder."),
+        ("text_encoder_2", "text_encoder_2."),
+        ("audio_vae", "audio_vae."),
+    )
+    assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.videoT.self_attention.to_query.weight") == (
+        "transformer.visual_transformer_blocks.0.video_dec_block.self_attention.to_query.weight"
+    )
+    assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.audioT.feed_forward.in_layer.weight") == (
+        "transformer.visual_transformer_blocks.0.audio_dec_block.feed_forward.in_layer.weight"
+    )
+    assert _adapt_k6_weight_name("vae.decoder.conv_in.conv.weight") == "vae.decoder.conv_in.conv.weight"
+    assert _adapt_k6_weight_name("text_encoder_2.encoder.layers.0.mlp.fc1.weight") == (
+        "text_encoder_2.encoder.layers.0.mlp.fc1.weight"
+    )
+    assert (
+        _adapt_k6_weight_name("text_encoder.model.layers.0.input_layernorm.weight")
+        == "text_encoder.model.language_model.layers.0.input_layernorm.weight"
+    )
+    assert _adapt_k6_weight_name("text_encoder.visual.blocks.0.norm1.weight") == (
+        "text_encoder.model.visual.blocks.0.norm1.weight"
+    )
+    assert (
+        _adapt_k6_weight_name("audio_vae.vae.decoder.conv_in.weight")
+        == "audio_vae.native.tod.vae.decoder.conv_in.weight"
+    )
+    assert _adapt_k6_weight_name("audio_vae.vocoder.conv_pre.weight") == "audio_vae.native.tod.vocoder.conv_pre.weight"
+    assert _adapt_k6_weight_name("audio_vae.mel_converter.hann_window") == "audio_vae.native.mel_converter.hann_window"

tests/diffusion/models/kandinsky6/test_transformer_kandinsky6.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed(monkeypatch):
+    """The native transformer uses vLLM parallel linear layers, which require a
+    tensor-parallel group; initialize a single-process group for CPU tests
+    (mirrors tests/diffusion/models/sana_video/test_transformer_sana_video.py)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda *_args, **_kwargs: default_unquantized_gemm,
+    )
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29502")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+# model_dim must be an exact multiple of head_dim = sum(axes_dims); this is
+# an architectural constraint of the native model (DiffusionTransformer3D),
+# not specific to this port.
+_TINY_T2V_CONFIG = {
+    "in_visual_dim": 4,
+    "out_visual_dim": 4,
+    "in_text_dim": 8,
+    "in_text_dim2": 6,
+    "time_dim": 16,
+    "patch_size": (1, 2, 2),
+    "model_dim": 24,
+    "ff_dim": 32,
+    "num_text_blocks": 1,
+    "num_visual_blocks": 2,
+    "axes_dims": (4, 4, 4),
+    "visual_cond": False,
+    "is_multimodal": False,
+    "attention_engine": "sdpa",
+}
+
+# time_dim_a is intentionally omitted (defaults to time_dim): the fused
+# block's cross-modal modulation (va_mod/av_mod) consumes the *other*
+# modality's time embedding by default (fix_modulation=False), so time_dim
+# and time_dim_a must match unless fix_modulation=True.
+_TINY_T2VA_CONFIG = {
+    **_TINY_T2V_CONFIG,
+    "is_multimodal": True,
+    "in_audio_dim": 6,
+    "model_dim_a": 12,
+    "ff_dim_a": 16,
+    "axes_dims_a": (2, 2, 2),
+}
+
+
+def _sdpa_config():
+    from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+
+    return SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")),
+        parallel_config=SimpleNamespace(ring_degree=1, allgather_degree=1),
+    )
+
+
+def _build_transformer(config):
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6Transformer3DModel
+
+    with set_current_diffusion_config(_sdpa_config()):
+        model = Kandinsky6Transformer3DModel(**config)
+    model.eval()
+    return model
+
+
+def test_t2v_only_forward_shape():
+    model = _build_transformer(_TINY_T2V_CONFIG)
+    time_frames, height, width = 2, 4, 4
+    x_video = torch.randn(time_frames, height, width, 4)
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+    visual_rope = model.visual_rope_embeddings(
+        shape=(time_frames, height // 2, width // 2),
+        pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+        scale_factor=(1.0, 1.0, 1.0),
+    )
+    text_rope = model.text_rope_embeddings(torch.arange(5))
+
+    with torch.no_grad():
+        out = model(
+            x_video=x_video,
+            x_audio=None,
+            text_embed=text_embed,
+            pooled_text_embed=pooled,
+            time=time,
+            visual_rope=visual_rope,
+            audio_rope=None,
+            text_rope=text_rope,
+        )
+
+    assert out.shape == (time_frames, height, width, 4)
+
+
+def test_visual_cond_channel_concat_path():
+    """visual_cond=True doubles+1 the visual input channel count (cond +
+    mask channels appended by _build_video_input, stripped again by
+    OutLayer) — the I2V conditioning path."""
+    config = {**_TINY_T2V_CONFIG, "visual_cond": True}
+    model = _build_transformer(config)
+    time_frames, height, width = 2, 4, 4
+    x_video = torch.randn(time_frames, height, width, 2 * 4 + 1)
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+    visual_rope = model.visual_rope_embeddings(
+        shape=(time_frames, height // 2, width // 2),
+        pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+        scale_factor=(1.0, 1.0, 1.0),
+    )
+    text_rope = model.text_rope_embeddings(torch.arange(5))
+
+    with torch.no_grad():
+        out = model(
+            x_video=x_video,
+            x_audio=None,
+            text_embed=text_embed,
+            pooled_text_embed=pooled,
+            time=time,
+            visual_rope=visual_rope,
+            audio_rope=None,
+            text_rope=text_rope,
+        )
+
+    assert out.shape == (time_frames, height, width, 4)
+
+
+def test_t2va_joint_forward_shapes():
+    """Exercises the fused cross-modal decoder blocks (Kandinsky6FusedTransformerDecoderBlock)."""
+    model = _build_transformer(_TINY_T2VA_CONFIG)
+    time_frames, height, width = 2, 4, 4
+    audio_len = 7
+    x_video = torch.randn(time_frames, height, width, 4)
+    x_audio = torch.randn(audio_len, 6)
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+    visual_rope = model.visual_rope_embeddings(
+        shape=(time_frames, height // 2, width // 2),
+        pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+        scale_factor=(1.0, 1.0, 1.0),
+    )
+    audio_rope = model.audio_rope_embeddings(torch.arange(audio_len))
+    video_text_rope = model.video_text_rope_embeddings(torch.arange(5))
+    audio_text_rope = model.audio_text_rope_embeddings(torch.arange(5))
+
+    with torch.no_grad():
+        video_out, audio_out = model(
+            x_video=x_video,
+            x_audio=x_audio,
+            text_embed=[text_embed, text_embed],
+            pooled_text_embed=[pooled, pooled],
+            time=[time, time],
+            visual_rope=visual_rope,
+            audio_rope=audio_rope,
+            text_rope=[video_text_rope, audio_text_rope],
+        )
+
+    assert video_out.shape == (time_frames, height, width, 4)
+    assert audio_out.shape == (audio_len, 6)
+
+
+@pytest.mark.parametrize("modality", ["video", "audio"])
+def test_single_modality_through_multimodal_model(modality):
+    """A multimodal (is_multimodal=True) model must still support a pure
+    video-only or audio-only forward call (denoise_loop's partial-sampling
+    path uses this)."""
+    model = _build_transformer(_TINY_T2VA_CONFIG)
+    time_frames, height, width = 2, 4, 4
+    audio_len = 7
+    text_embed = torch.randn(5, 8)
+    pooled = torch.randn(1, 6)
+    time = torch.tensor([500.0])
+
+    if modality == "video":
+        x_video = torch.randn(time_frames, height, width, 4)
+        visual_rope = model.visual_rope_embeddings(
+            shape=(time_frames, height // 2, width // 2),
+            pos=[torch.arange(time_frames), torch.arange(height // 2), torch.arange(width // 2)],
+            scale_factor=(1.0, 1.0, 1.0),
+        )
+        text_rope = model.video_text_rope_embeddings(torch.arange(5))
+        with torch.no_grad():
+            out = model(
+                x_video=x_video,
+                x_audio=None,
+                text_embed=text_embed,
+                pooled_text_embed=pooled,
+                time=time,
+                visual_rope=visual_rope,
+                audio_rope=None,
+                text_rope=text_rope,
+            )
+        assert out.shape == (time_frames, height, width, 4)
+    else:
+        x_audio = torch.randn(audio_len, 6)
+        audio_rope = model.audio_rope_embeddings(torch.arange(audio_len))
+        text_rope = model.audio_text_rope_embeddings(torch.arange(5))
+        with torch.no_grad():
+            out = model(
+                x_video=None,
+                x_audio=x_audio,
+                text_embed=text_embed,
+                pooled_text_embed=pooled,
+                time=time,
+                visual_rope=None,
+                audio_rope=audio_rope,
+                text_rope=text_rope,
+            )
+        assert out.shape == (audio_len, 6)
+
+
+@pytest.mark.parametrize("batched", [False, True])
+def test_self_attention_attends_over_tokens_for_batched_and_unbatched_inputs(batched):
+    """Self-attention must contract over the *token* axis whether its input is
+    unbatched text ``(S, C)`` or batched=1 visual/audio ``(1, N, C)``.
+
+    Regression: an unconditional ``unsqueeze(0)`` on the already-batched
+    visual path handed SDPA a 5-D ``(1, 1, N, H, D)`` tensor, which it
+    accepts silently but then treats the head axis as the sequence — every
+    generated video decoded to pure noise."""
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.kandinsky6.kandinsky6_transformer import Kandinsky6Attention, apply_rotary
+
+    torch.manual_seed(0)
+    channels, head_dim, tokens = 24, 12, 7
+    with set_current_diffusion_config(_sdpa_config()):
+        attn = Kandinsky6Attention(channels, head_dim, engine="sdpa").eval()
+    with torch.no_grad():
+        # vLLM parallel linears allocate with torch.empty (no init).
+        for param in attn.parameters():
+            param.normal_(std=0.2)
+    x = torch.randn(tokens, channels)
+    rope = torch.randn(tokens, 1, head_dim // 2, 2, 2)
+    hidden = x.unsqueeze(0) if batched else x
+
+    with torch.no_grad():
+        out = attn(hidden, rotary_emb=rope)
+
+        heads = channels // head_dim
+        q = attn.query_norm(attn.to_query(x).reshape(tokens, heads, head_dim))
+        k = attn.key_norm(attn.to_key(x).reshape(tokens, heads, head_dim))
+        v = attn.to_value(x).reshape(tokens, heads, head_dim)
+        q = apply_rotary(q, rope).type_as(q)
+        k = apply_rotary(k, rope).type_as(k)
+        ref = torch.nn.functional.scaled_dot_product_attention(
+            q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1)
+        ).transpose(0, 1)
+        expected = attn.out_layer(ref.reshape(tokens, channels))
+
+    assert out.shape == (hidden.shape[:-1] + (channels,))
+    torch.testing.assert_close(out.reshape(tokens, channels), expected, rtol=1e-4, atol=1e-5)
+
+
+def test_modulation_projections_are_zero_initialized():
+    """Modulation layers are zero-initialized at construction (standard
+    AdaLN-zero init) — confirms the zero-init in Kandinsky6Modulation.__init__
+    reaches the tensor-parallel out_layer correctly (it's a
+    ColumnParallelLinear, not plain nn.Linear)."""
+    model = _build_transformer(_TINY_T2V_CONFIG)
+    for block in model.text_transformer_blocks:
+        assert torch.all(block.text_modulation.out_layer.weight == 0)
+    for block in model.visual_transformer_blocks:
+        assert torch.all(block.visual_modulation.out_layer.weight == 0)
+
+
+def test_dit_constructs_from_a_diffusers_config_dict():
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6Transformer3DModel
+
+    with set_current_diffusion_config(_sdpa_config()):
+        model = Kandinsky6Transformer3DModel.from_diffusers_config(
+            {**_TINY_T2VA_CONFIG, "patch_size": [1, 2, 2], "axes_dims": [4, 4, 4]}
+        )
+    assert tuple(model.patch_size) == (1, 2, 2)
+    assert model._repeated_blocks == [
+        "Kandinsky6FusedTransformerDecoderBlock",
+        "Kandinsky6TransformerEncoderBlock",
+        "Kandinsky6TransformerDecoderBlock",
+    ]
+    visual = model.visual_transformer_blocks[0].video_dec_block.self_attention.attn
+    audio = model.visual_transformer_blocks[0].audio_dec_block.self_attention.attn
+    text = model.video_text_transformer_blocks[0].self_attention.attn
+    cross = model.visual_transformer_blocks[0].va_cross_attention.attn
+    assert visual.role == "kandinsky6.visual_self"
+    assert visual.skip_sequence_parallel is False
+    assert audio.skip_sequence_parallel is True
+    assert text.role == "kandinsky6.text_self"
+    assert text.skip_sequence_parallel is True
+    assert cross.role == "kandinsky6.video_audio_cross"
+    assert cross.skip_sequence_parallel is True
+
+
+def test_cfg_combine_matches_apply_cfg_without_normalization():
+    from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import apply_cfg
+
+    positive = torch.tensor([1.0, 3.0])
+    negative = torch.tensor([0.25, -1.0])
+    combined = CFGParallelMixin.combine_cfg_noise(object(), positive, negative, 5.0, False)
+    torch.testing.assert_close(combined, apply_cfg(positive, negative, 5.0))

tests/diffusion/models/kandinsky6/test_vae_patch_parallel.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Kandinsky 6 VAE patch parallel keeps the single-GPU tile grid."""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from vllm_omni.diffusion.models.kandinsky6.modeling_kandinsky6_vae import AutoencoderKLHunyuanVideo
+
+
+def _stub_vae() -> AutoencoderKLHunyuanVideo:
+    vae = AutoencoderKLHunyuanVideo.__new__(AutoencoderKLHunyuanVideo)
+    vae.spatial_compression_ratio = 8
+    vae.temporal_compression_ratio = 4
+    vae.use_tiling = True
+    vae.use_framewise_decoding = True
+    vae.tile_sample_min_height = 32
+    vae.tile_sample_min_width = 32
+    vae.tile_sample_min_num_frames = 16
+    vae.tile_sample_stride_height = 24
+    vae.tile_sample_stride_width = 24
+    vae.tile_sample_stride_num_frames = 12
+    vae.tile_size = None
+    vae.distributed_executor = None
+    vae.post_quant_conv = lambda tensor: tensor
+    vae._seen_widths = []
+
+    def decoder(tensor: torch.Tensor) -> torch.Tensor:
+        vae._seen_widths.append(int(tensor.shape[-1]))
+        return tensor + 1
+
+    vae.decoder = decoder
+    return vae
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_distributed_decode_tiles_match_local_tiled_decode():
+    vae = _stub_vae()
+    latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12)
+
+    local = vae.tiled_decode(latent).sample
+    tasks, spec = vae._decode_tile_split(latent)
+    assert len(tasks) > 1
+    assert max(int(task.tensor.shape[-1]) for task in tasks) < latent.shape[-1]
+
+    decoded = {task.grid_coord: vae._decode_tile_exec(task) for task in tasks}
+    merged = vae._decode_tile_merge(decoded, spec)
+
+    torch.testing.assert_close(merged, local)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_decode_still_applies_optimal_tiling_when_patch_parallel_is_on():
+    vae = _stub_vae()
+    seen = {}
+
+    def tiling(shape, device=None):
+        del device
+        seen["shape"] = tuple(shape)
+        return (1, 17, 32, 32), (8, 24, 24)
+
+    from diffusers.models.autoencoders.vae import DecoderOutput
+
+    vae.get_dec_optimal_tiling = tiling
+    vae._decode = lambda z, return_dict=True: DecoderOutput(sample=z)
+    vae.distributed_executor = SimpleNamespace(parallel_size=2)
+    vae.is_distributed_enabled = lambda: True
+
+    out = vae.decode(torch.zeros(1, 16, 8, 8, 12))
+
+    assert seen["shape"] == (1, 16, 8, 8, 12)
+    assert tuple(out.sample.shape) == (1, 16, 8, 8, 12)
+
+
+def _patch_parallel_worker(rank: int, port: str, shapes: dict) -> None:
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = port
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = "2"
+    dist.init_process_group("gloo", rank=rank, world_size=2)
+    try:
+        import vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor as executor_mod
+
+        world = SimpleNamespace(device_group=dist.group.WORLD)
+        executor_mod.get_world_group = lambda: world
+
+        vae = _stub_vae()
+        latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12)
+        reference = vae.tiled_decode(latent).sample.clone()
+        vae._seen_widths.clear()
+        vae.set_parallel_size(2, mode="tile")
+        decoded = vae.tiled_decode(latent).sample
+        torch.testing.assert_close(decoded, reference)
+        if rank == 0:
+            shapes["widths"] = list(vae._seen_widths)
+            shapes["full_width"] = int(latent.shape[-1])
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_patch_parallel_decode_splits_tiles_across_ranks():
+    manager = mp.Manager()
+    shapes = manager.dict()
+    port = str(29500 + (os.getpid() % 1000))
+    mp.spawn(_patch_parallel_worker, args=(port, shapes), nprocs=2, join=True)
+
+    assert shapes["widths"]
+    assert max(shapes["widths"]) < shapes["full_width"]