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