anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[diffusion][feature] Add component-selective offload policies (#5929)

Browse Source 
Signed-off-by: lishunyang12 <lishunyang12@163.com>
This commit is contained in:
SYLAR authored and GitHub committed 2026-09-06 05:40:09 +08:00
1 parent 44d3ae100a
commit ea3cf99894
74 files changed
+5121 -1584
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_diffusion_engine.py
+56
View File
@@ -600,6 +600,34 @@ class TestDiffusionCompileConfig:
),
({"enable_cpu_offload": True}, "CPU offload"),
({"enable_layerwise_offload": True}, "layerwise offload"),
(
{
"diffusion_offload_config": {
"mode": "module",
"components": ["dit"],
}
},
"CPU offload",
),
(
{
"diffusion_offload_config": {
"mode": "layer",
"components": ["dit"],
}
},
"layerwise offload",
),
(
{
"diffusion_offload_config": {
"mode": "layer",
"components": ["dit"],
"layer_options": {"dit": {"weight_transfer": "allgather"}},
}
},
"distributed layerwise offload",
),
],
)
def test_full_compile_rejects_incompatible_features(self, kwargs, feature) -> None:
@@ -610,6 +638,34 @@ class TestDiffusionCompileConfig:
**kwargs,
)
def test_raw_parallel_mapping_is_normalized_before_allgather_cache_validation(self) -> None:
with pytest.raises(ValueError, match="rank-local cache decisions"):
OmniDiffusionConfig(
model="test",
num_gpus=2,
parallel_config={"data_parallel_size": 2},
cache_backend="tea_cache",
diffusion_offload_config={
"mode": "layer",
"components": ["dit"],
"layer_options": {"dit": {"weight_transfer": "allgather"}},
},
)
def test_auto_dp_is_resolved_before_encoder_prompt_cache_validation(self) -> None:
with pytest.raises(ValueError, match="rank-local cache hits"):
OmniDiffusionConfig(
model="test",
num_gpus=2,
parallel_config={"data_parallel_size": None},
enable_prompt_embed_cache=True,
diffusion_offload_config={
"mode": "layer",
"components": ["text_encoder"],
"layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
},
)
class TestRequestBatchAdmission:
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]