WeiQing ChenandHongsheng Liu
1753a878e0
Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from vllm_omni.entrypoints import omni as omni_module
|
|
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
|
|
|
|
|
def test_default_stage_config_includes_cache_backend(monkeypatch):
|
|
"""Ensure cache_backend/cache_config are preserved in default diffusion stage."""
|
|
monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
|
|
monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
|
|
monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
|
|
monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
|
|
|
|
omni = AsyncOmni(
|
|
model="dummy-model",
|
|
cache_backend="cache_dit",
|
|
cache_config='{"Fn_compute_blocks": 2}',
|
|
vae_use_slicing=True,
|
|
ulysses_degree=2,
|
|
)
|
|
|
|
stage_cfg = omni.stage_configs[0]
|
|
engine_args = stage_cfg.engine_args
|
|
|
|
assert engine_args.get("cache_backend") == "cache_dit"
|
|
cache_config = engine_args.get("cache_config")
|
|
assert cache_config["Fn_compute_blocks"] == 2
|
|
assert engine_args.get("vae_use_slicing") is True
|
|
parallel_config = engine_args.get("parallel_config")
|
|
if hasattr(parallel_config, "get"):
|
|
ulysses_degree = parallel_config.get("ulysses_degree")
|
|
else:
|
|
ulysses_degree = getattr(parallel_config, "ulysses_degree", None)
|
|
assert ulysses_degree == 2
|
|
|
|
|
|
def test_default_cache_config_used_when_missing(monkeypatch):
|
|
"""Ensure default cache_config is applied when cache_backend is set."""
|
|
monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
|
|
monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
|
|
monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
|
|
monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
|
|
|
|
omni = AsyncOmni(
|
|
model="dummy-model",
|
|
cache_backend="cache_dit",
|
|
)
|
|
|
|
engine_args = omni.stage_configs[0].engine_args
|
|
cache_config = engine_args.get("cache_config")
|
|
assert cache_config is not None
|
|
assert cache_config["Fn_compute_blocks"] == 1
|
|
|
|
|
|
def test_default_stage_devices_from_sequence_parallel(monkeypatch):
|
|
"""Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
|
|
monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
|
|
monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
|
|
monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
|
|
monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
|
|
|
|
omni = AsyncOmni(
|
|
model="dummy-model",
|
|
ulysses_degree=2,
|
|
ring_degree=2,
|
|
)
|
|
|
|
stage_cfg = omni.stage_configs[0]
|
|
runtime = stage_cfg.runtime
|
|
if hasattr(runtime, "get"):
|
|
devices = runtime.get("devices")
|
|
else:
|
|
devices = getattr(runtime, "devices", None)
|
|
assert devices == "0,1,2,3"
|