fywc
d01eaca39a
Signed-off-by: hanzheli <hanzheli@kuaishou.com> Signed-off-by: fywc <hanzheli@kuaishou.com>
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Model specific tests for CacheDiT enablement.
|
|
"""
|
|
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
|
|
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
SEPARATE_CFG_ENABLERS = [
|
|
cd_backend.enable_cache_for_dreamid_omni,
|
|
cd_backend.enable_cache_for_ltx2,
|
|
cd_backend.enable_cache_for_helios,
|
|
cd_backend.enable_cache_for_wan22,
|
|
cd_backend.enable_cache_for_longcat_image,
|
|
]
|
|
|
|
SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
|
|
|
|
|
|
@pytest.mark.parametrize("enabler", SEPARATE_CFG_ENABLERS)
|
|
@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
|
|
@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
|
|
def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
|
|
"""Ensure that custom enablers for models with separate CFG pass
|
|
the param through to cache_dit correctly.
|
|
|
|
Regression test for: https://github.com/vllm-project/vllm-omni/pull/2860
|
|
"""
|
|
mock_pipeline = Mock()
|
|
enabler(mock_pipeline, SAMPLE_CACHE_CONFIG)
|
|
|
|
mock_cache_dit.enable_cache.assert_called_once()
|
|
adapter_kwargs = mock_block_adapter.call_args.kwargs
|
|
assert adapter_kwargs["has_separate_cfg"] is True
|