5f95115e70
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>
90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
# 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
|