Merge remote-tracking branch 'origin/main' into qwen-image-2.1
# Conflicts: # tests/entrypoints/openai_api/test_image_server.py # vllm_omni/entrypoints/openai/api_server.py
This commit is contained in:
Isotr0py committed
commit
3e274a0024
220 files changed
+22342
-1079
No files matched your search
@@ -42,6 +42,14 @@ from vllm_omni.diffusion.data import (
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_inherited_attention_backend(monkeypatch):
|
||||
# Weekly CPU exports DIFFUSION_ATTENTION_BACKEND; OmniDiffusionConfig
|
||||
# treats that as an explicit default so auto-backend cases break.
|
||||
# Tests that setenv the same variable still work: this fixture runs first.
|
||||
monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
|
||||
|
||||
|
||||
class TestAttentionSpec:
|
||||
def test_construct_no_skip_softmax(self):
|
||||
spec = AttentionSpec(backend="FLASH_ATTN")
|
||||
@@ -416,13 +424,6 @@ class TestBuildAttentionConfig:
|
||||
class TestOmniDiffusionConfigAttentionParsing:
|
||||
"""Test OmniDiffusionConfig attention shorthand and structured config."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_diffusion_attention_backend_env(self, monkeypatch):
|
||||
# OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via
|
||||
# build_attention_config(); clear it so these tests assert config defaults,
|
||||
# not whatever the process inherited from CI / sibling tests.
|
||||
monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
|
||||
|
||||
def test_diffusion_attention_backend_sets_default(self):
|
||||
config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
|
||||
assert isinstance(config.diffusion_attention_config, AttentionConfig)
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""
|
||||
Two-rank ``pause_generation(mode="keep")``: every rank runs the pause barrier
|
||||
before the ACK, and a barrier failure on a non-zero rank fails the pause
|
||||
without aborting or dropping the running batch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections import defaultdict
|
||||
from contextlib import ExitStack
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.e2e.features.custom_pipeline.test_diffusion_pause_keep import (
|
||||
CUSTOM_PIPELINE_CLASS,
|
||||
MODEL,
|
||||
WORKER_EXTENSION_CLASS,
|
||||
_generate,
|
||||
_pause_round,
|
||||
_wait_for,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
|
||||
|
||||
def _all_true(results) -> bool:
|
||||
flat = []
|
||||
stack = list(results)
|
||||
while stack:
|
||||
item = stack.pop()
|
||||
if isinstance(item, list):
|
||||
stack.extend(item)
|
||||
else:
|
||||
flat.append(item)
|
||||
return bool(flat) and all(item is True for item in flat)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=2)
|
||||
@pytest.mark.asyncio
|
||||
async def test_keep_pause_two_ranks_ack_and_rank_failure():
|
||||
with ExitStack() as after:
|
||||
engine = AsyncOmni(
|
||||
model=MODEL,
|
||||
custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
|
||||
worker_extension_cls=WORKER_EXTENSION_CLASS,
|
||||
enforce_eager=True,
|
||||
enable_sleep_mode=True,
|
||||
num_gpus=2,
|
||||
parallel_config={"ulysses_degree": 2, "ulysses_mode": "advanced_uaa"},
|
||||
)
|
||||
after.callback(engine.shutdown)
|
||||
events: dict[str, dict[str, float]] = defaultdict(dict)
|
||||
assert _all_true(await engine.collective_rpc(method="start_counting_synchronize_device"))
|
||||
|
||||
pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(3)]
|
||||
# The result is the all-rank AND, so True means both ranks ran the barrier.
|
||||
assert _all_true(await engine.collective_rpc(method="synchronize_device_seen"))
|
||||
|
||||
assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, True)))
|
||||
a = asyncio.create_task(_generate(engine, "a-fail", events))
|
||||
await _wait_for(lambda: "started" in events["a-fail"])
|
||||
with pytest.raises(RuntimeError, match="pause_scheduler failed"):
|
||||
await engine.pause_generation(mode="keep", clear_cache=False)
|
||||
assert await engine.is_paused()
|
||||
a_output = await asyncio.wait_for(a, 60.0)
|
||||
assert a_output.images, "a failed barrier must not abort the running batch"
|
||||
|
||||
assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, False)))
|
||||
await engine.resume_generation()
|
||||
assert not await engine.is_paused()
|
||||
recovered = await asyncio.wait_for(_generate(engine, "after-recovery", events), 60.0)
|
||||
assert recovered.images
|
||||
|
||||
assert not engine.request_states
|
||||
print(f"[sp2] rounds=3 pause_latency_s min/max={min(pause_latencies):.3f}/{max(pause_latencies):.3f}")
|
||||
@@ -493,6 +493,84 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
|
||||
assert last.shape == (1, 1, 2, 3)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
class _PlainDecoder(torch.nn.Module):
|
||||
"""A decoder with nothing to patch, so the install only wraps ``forward``."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.calls: list[dict] = []
|
||||
|
||||
def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False):
|
||||
self.calls.append({"feat_cache": feat_cache, "feat_idx": feat_idx, "first_chunk": first_chunk})
|
||||
return x
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("world_size, dst", [(4, -1), (4, 4), (4, 5), (1, 1), (4, 1.5)])
|
||||
def test_install_rejects_invalid_destination_before_mutating_decoder(monkeypatch, world_size, dst):
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, world_size))
|
||||
vae = SimpleNamespace(decoder=torch.nn.Identity())
|
||||
original_forward = vae.decoder.forward
|
||||
|
||||
def unexpected_patch(*args, **kwargs):
|
||||
pytest.fail("invalid destination must fail before patching modules")
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "_patch_decoder_modules", unexpected_patch)
|
||||
with pytest.raises(ValueError, match="dst must be None or an integer"):
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), dst=dst)
|
||||
assert vae.decoder.forward == original_forward
|
||||
assert not getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False)
|
||||
|
||||
|
||||
def _install_on_plain_decoder(monkeypatch: pytest.MonkeyPatch, *, dst):
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2))
|
||||
monkeypatch.setattr(
|
||||
wan_spatial_shard,
|
||||
"split_for_parallel_decode",
|
||||
lambda x, *, upsample_count, split_dim, group: (x, x.shape[wan_spatial_shard._spatial_dim(split_dim)]),
|
||||
)
|
||||
gathers: list[dict] = []
|
||||
|
||||
def gather(x, *, expected_extent, split_dim, group, dst=None):
|
||||
gathers.append({"expected_extent": expected_extent, "split_dim": split_dim, "dst": dst})
|
||||
return x
|
||||
|
||||
monkeypatch.setattr(wan_spatial_shard, "gather_and_trim_extent", gather)
|
||||
vae = SimpleNamespace(decoder=_PlainDecoder())
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=dst)
|
||||
return vae, gathers
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("dst", [0, 1, None])
|
||||
def test_install_passes_the_assembling_rank_through_to_the_gather(monkeypatch: pytest.MonkeyPatch, dst):
|
||||
vae, gathers = _install_on_plain_decoder(monkeypatch, dst=dst)
|
||||
cache = [None]
|
||||
|
||||
out = vae.decoder(torch.zeros(1, 3, 1, 4, 6), feat_cache=cache, feat_idx=[0], first_chunk=True)
|
||||
|
||||
assert out.shape == (1, 3, 1, 4, 6)
|
||||
assert gathers == [{"expected_extent": 6, "split_dim": "width", "dst": dst}]
|
||||
# The streaming arguments reach the wrapped decoder untouched.
|
||||
assert vae.decoder.calls == [{"feat_cache": cache, "feat_idx": [0], "first_chunk": True}]
|
||||
assert vae._vllm_omni_wan_spatial_shard_dst == dst
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_install_refuses_to_change_the_assembling_rank(monkeypatch: pytest.MonkeyPatch):
|
||||
vae, _ = _install_on_plain_decoder(monkeypatch, dst=None)
|
||||
|
||||
# Same settings again: a no-op, as before.
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=None)
|
||||
with pytest.raises(ValueError, match="already patched to assemble on None"):
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=0)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pytest.MonkeyPatch):
|
||||
@@ -626,6 +704,105 @@ def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
_STREAMING_SHARD_CHUNKS = 3
|
||||
|
||||
|
||||
def _streaming_shard_worker(rank: int, return_dict, master_port: str) -> None:
|
||||
"""Sharded streaming decode with every rank keeping the frame, against an unsharded streaming reference.
|
||||
|
||||
Two sessions are interleaved chunk by chunk, so the per-session temporal cache
|
||||
and the halo exchange are exercised across successive calls, then one session
|
||||
is reset and must reproduce its opening chunk. Every rank checks every chunk.
|
||||
"""
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
destroy_model_parallel,
|
||||
get_sp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
|
||||
|
||||
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
||||
os.environ["MASTER_PORT"] = master_port
|
||||
device = current_omni_platform.get_torch_device(rank)
|
||||
current_omni_platform.set_device(device)
|
||||
dtype = torch.float32
|
||||
backend = current_omni_platform.dist_backend
|
||||
init_distributed_environment(world_size=_SPATIAL_SHARD_WORLD_SIZE, rank=rank, local_rank=rank, backend=backend)
|
||||
initialize_model_parallel(
|
||||
sequence_parallel_size=_SPATIAL_SHARD_WORLD_SIZE, ulysses_degree=_SPATIAL_SHARD_WORLD_SIZE, backend=backend
|
||||
)
|
||||
# VLLM_OMNI_TEST_WAN_VAE may point at a local diffusers checkpoint (with or
|
||||
# without a ``vae`` subfolder) for machines without hub access.
|
||||
model = os.environ.get("VLLM_OMNI_TEST_WAN_VAE", _SPATIAL_SHARD_MODEL)
|
||||
local_vae_only = os.path.isdir(model) and not os.path.isdir(os.path.join(model, "vae"))
|
||||
load = {"torch_dtype": dtype} if local_vae_only else {"subfolder": _SPATIAL_SHARD_SUBFOLDER, "torch_dtype": dtype}
|
||||
try:
|
||||
reference_vae = DistributedAutoencoderKLWan.from_pretrained(model, **load).to(device=device, dtype=dtype).eval()
|
||||
sharded_vae = DistributedAutoencoderKLWan.from_pretrained(model, **load).to(device=device, dtype=dtype).eval()
|
||||
wan_spatial_shard.install_wan_spatial_shard_decode(
|
||||
sharded_vae, get_sp_group().device_group, split_dim="width", dst=None
|
||||
)
|
||||
reference = WanStreamingDecoder(reference_vae, bytes_per_pixel_fp32=1.0)
|
||||
sharded = WanStreamingDecoder(sharded_vae, bytes_per_pixel_fp32=1.0)
|
||||
|
||||
generator = torch.Generator(device=device).manual_seed(0)
|
||||
# Width 104 latents = 832 px: not divisible by two ranks' halo-padded shards without padding.
|
||||
shape = (1, reference_vae.config.z_dim, 1, _SPATIAL_SHARD_LATENT_HEIGHT, _SPATIAL_SHARD_LATENT_WIDTH)
|
||||
chunks = {
|
||||
name: [
|
||||
torch.randn(shape, generator=generator, device=device, dtype=dtype)
|
||||
for _ in range(_STREAMING_SHARD_CHUNKS)
|
||||
]
|
||||
for name in ("a", "b")
|
||||
}
|
||||
max_diff = 0.0
|
||||
with torch.inference_mode():
|
||||
ref_states = {name: reference.new_decode_state(name) for name in chunks}
|
||||
shard_states = {name: sharded.new_decode_state(name) for name in chunks}
|
||||
for index in range(_STREAMING_SHARD_CHUNKS):
|
||||
for name in ("a", "b"):
|
||||
expected = reference.decode_chunk(chunks[name][index], ref_states[name]).float()
|
||||
got = sharded.decode_chunk(chunks[name][index], shard_states[name]).float()
|
||||
if got.shape != expected.shape:
|
||||
raise AssertionError(
|
||||
f"rank {rank} {name}[{index}]: {tuple(got.shape)} != {tuple(expected.shape)}"
|
||||
)
|
||||
max_diff = max(max_diff, (got - expected).abs().max().item())
|
||||
# Reset session a: a fresh state must reproduce its opening chunk, on this rank too.
|
||||
shard_states["a"].release()
|
||||
ref_states["a"].release()
|
||||
shard_states["a"] = sharded.new_decode_state("a")
|
||||
ref_states["a"] = reference.new_decode_state("a")
|
||||
expected = reference.decode_chunk(chunks["a"][0], ref_states["a"]).float()
|
||||
got = sharded.decode_chunk(chunks["a"][0], shard_states["a"]).float()
|
||||
max_diff = max(max_diff, (got - expected).abs().max().item())
|
||||
return_dict[f"max_abs_diff_rank{rank}"] = max_diff
|
||||
return_dict[f"shape_rank{rank}"] = tuple(got.shape)
|
||||
finally:
|
||||
destroy_model_parallel()
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parallel
|
||||
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
|
||||
def test_streaming_shard_decode_matches_unsharded_streaming_on_every_rank():
|
||||
manager = mp.get_context("spawn").Manager()
|
||||
return_dict = manager.dict()
|
||||
mp.spawn(_streaming_shard_worker, args=(return_dict, "29510"), nprocs=_SPATIAL_SHARD_WORLD_SIZE, join=True)
|
||||
|
||||
for rank in range(_SPATIAL_SHARD_WORLD_SIZE):
|
||||
assert f"max_abs_diff_rank{rank}" in return_dict, f"rank {rank} did not report"
|
||||
diff = return_dict[f"max_abs_diff_rank{rank}"]
|
||||
shape = return_dict[f"shape_rank{rank}"]
|
||||
print(f"streaming shard rank {rank}: max_abs_diff={diff:.6e} shape={shape}")
|
||||
assert shape[-1] == _SPATIAL_SHARD_LATENT_WIDTH * 8 and shape[-2] == _SPATIAL_SHARD_LATENT_HEIGHT * 8
|
||||
assert diff <= _SPATIAL_SHARD_TOLERANCE, f"rank {rank}: {diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
|
||||
|
||||
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parallel
|
||||
|
||||
@@ -3,21 +3,21 @@
|
||||
|
||||
import contextlib
|
||||
|
||||
import huggingface_hub
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mock import patch_hf_snapshot_download
|
||||
from vllm_omni.diffusion.model_loader import hub_prefetch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_prefetch_subfolders_propagates_revision(monkeypatch):
|
||||
def test_prefetch_subfolders_propagates_revision(monkeypatch, tmp_path):
|
||||
calls = []
|
||||
|
||||
def fake_snapshot_download(self, **kwargs):
|
||||
def fake_snapshot_download(**kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", fake_snapshot_download)
|
||||
patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
|
||||
monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
|
||||
|
||||
hub_prefetch.prefetch_subfolders(
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Focused CPU tests for BAGEL step execution and packed batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -18,12 +18,14 @@ from vllm_omni.diffusion.models.bagel.bagel_transformer import Bagel, NaiveCache
|
||||
from vllm_omni.diffusion.models.bagel.pipeline_bagel import (
|
||||
BagelGenParams,
|
||||
BagelPipeline,
|
||||
_bagel_canvas_requested,
|
||||
get_bagel_pre_process_func,
|
||||
)
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.sched import StepScheduler
|
||||
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -118,7 +120,7 @@ def test_bagel_step_preprocessor_separates_incompatible_cfg_settings(first_extra
|
||||
assert first.batch_compatibility_key != second.batch_compatibility_key
|
||||
|
||||
|
||||
def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
|
||||
def _pre_process_with_stub_checkpoint(tmp_path):
|
||||
(tmp_path / "config.json").write_text(
|
||||
'{"vae_config":{"downsample":8},"latent_patch_size":2,"max_latent_size":32}',
|
||||
encoding="utf-8",
|
||||
@@ -131,9 +133,22 @@ def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
|
||||
'{"weight_map":{"latent_pos_embed.pos_embed":"ema.safetensors"}}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
pre_process = get_bagel_pre_process_func(
|
||||
types.SimpleNamespace(model=str(tmp_path), revision=None, step_execution=True)
|
||||
return get_bagel_pre_process_func(types.SimpleNamespace(model=str(tmp_path), revision=None, step_execution=True))
|
||||
|
||||
|
||||
def _edit_request(source_size, sampling_params, request_id="req"):
|
||||
return OmniDiffusionRequest(
|
||||
prompt={
|
||||
"prompt": "edit this image",
|
||||
"multi_modal_data": {"image": [Image.new("RGB", source_size)]},
|
||||
},
|
||||
sampling_params=sampling_params,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
|
||||
def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
requests = [
|
||||
OmniDiffusionRequest(
|
||||
prompt={
|
||||
@@ -163,6 +178,172 @@ def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
|
||||
assert schedule.num_waiting_reqs == 1
|
||||
|
||||
|
||||
def test_bagel_pre_process_keeps_explicitly_requested_canvas(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request((512, 1024), OmniDiffusionSamplingParams(num_inference_steps=2, height=1024, width=1024))
|
||||
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
|
||||
assert _bagel_canvas_requested(request.sampling_params)
|
||||
|
||||
|
||||
def test_bagel_pre_process_aligns_source_derived_canvas_for_size_auto(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request(
|
||||
(500, 700),
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=2,
|
||||
height=700,
|
||||
width=500,
|
||||
height_not_provided=True,
|
||||
width_not_provided=True,
|
||||
),
|
||||
)
|
||||
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (704, 496)
|
||||
assert not _bagel_canvas_requested(request.sampling_params)
|
||||
|
||||
|
||||
def test_bagel_pre_process_derives_canvas_when_caller_asked_for_none(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request((800, 400), OmniDiffusionSamplingParams(num_inference_steps=2))
|
||||
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (400, 800)
|
||||
assert (request.sampling_params.height_not_provided, request.sampling_params.width_not_provided) == (True, True)
|
||||
assert not _bagel_canvas_requested(request.sampling_params)
|
||||
|
||||
|
||||
def test_bagel_pre_process_marks_kv_metadata_canvas_as_requested(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request(
|
||||
(500, 700),
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=2,
|
||||
height=700,
|
||||
width=500,
|
||||
height_not_provided=True,
|
||||
width_not_provided=True,
|
||||
kv_metadata={"image_shape": (1024, 1024)},
|
||||
),
|
||||
)
|
||||
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
|
||||
assert (request.sampling_params.height_not_provided, request.sampling_params.width_not_provided) == (False, False)
|
||||
assert _bagel_canvas_requested(request.sampling_params)
|
||||
|
||||
|
||||
def test_bagel_pre_process_prefers_explicit_canvas_over_kv_metadata(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request(
|
||||
(1280, 720),
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=2,
|
||||
height=1024,
|
||||
width=1024,
|
||||
kv_metadata={"image_shape": (576, 1024)},
|
||||
),
|
||||
)
|
||||
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
|
||||
assert _bagel_canvas_requested(request.sampling_params)
|
||||
|
||||
|
||||
def test_bagel_pre_process_floors_unaligned_explicit_canvas_with_warning(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request(
|
||||
(1280, 720),
|
||||
OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
|
||||
)
|
||||
|
||||
with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.logger.warning") as warning:
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (688, 992)
|
||||
assert _bagel_canvas_requested(request.sampling_params)
|
||||
warning.assert_called_once()
|
||||
assert "width 1000->992, height 700->688" in warning.call_args.args[-1]
|
||||
|
||||
|
||||
def test_bagel_pre_process_keeps_aligned_explicit_canvas_silent(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request(
|
||||
(1280, 720),
|
||||
OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=1024),
|
||||
)
|
||||
|
||||
with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.logger.warning") as warning:
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (512, 1024)
|
||||
warning.assert_not_called()
|
||||
|
||||
|
||||
def test_bagel_pre_process_rejects_explicit_canvas_over_checkpoint_limit(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = _edit_request(
|
||||
(1280, 720),
|
||||
OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=2048),
|
||||
)
|
||||
|
||||
with pytest.raises(OmniClientError, match="width=2048 exceeds the BAGEL checkpoint limit of 1024") as excinfo:
|
||||
pre_process(request)
|
||||
|
||||
assert excinfo.value.status_code == 400
|
||||
|
||||
|
||||
def test_bagel_pre_process_aligns_text2img_canvas_too(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = OmniDiffusionRequest(
|
||||
prompt="a cat",
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
|
||||
request_id="req",
|
||||
)
|
||||
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (688, 992)
|
||||
|
||||
|
||||
def test_bagel_pre_process_skips_canvas_alignment_for_text_output(tmp_path):
|
||||
pre_process = _pre_process_with_stub_checkpoint(tmp_path)
|
||||
request = OmniDiffusionRequest(
|
||||
prompt={
|
||||
"prompt": "describe this image",
|
||||
"modalities": ["text"],
|
||||
"multi_modal_data": {"image": [Image.new("RGB", (1280, 720))]},
|
||||
},
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
|
||||
request_id="req",
|
||||
)
|
||||
|
||||
pre_process(request)
|
||||
|
||||
assert (request.sampling_params.height, request.sampling_params.width) == (700, 1000)
|
||||
assert request.use_step_execution is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("sampling_kwargs", "requested"),
|
||||
[
|
||||
({"height": 1024, "width": 1024}, True),
|
||||
({"height": 1024, "width": 512, "height_not_provided": True, "width_not_provided": True}, False),
|
||||
({}, False),
|
||||
],
|
||||
ids=["explicit_size", "size_auto", "unset"],
|
||||
)
|
||||
def test_bagel_canvas_requested_separates_asked_for_from_derived(sampling_kwargs, requested):
|
||||
assert _bagel_canvas_requested(OmniDiffusionSamplingParams(**sampling_kwargs)) is requested
|
||||
|
||||
|
||||
def _cache(length: int, value: float, num_layers: int = 1) -> NaiveCache:
|
||||
cache = NaiveCache(num_layers)
|
||||
if length:
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Multi-image input regression for HunyuanImage3 IT2I prompt construction.
|
||||
|
||||
The official HunyuanImage-3.0-Instruct supports up to 3 reference images
|
||||
@@ -41,7 +41,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
"""Recording fake tokenizer mirroring the one in test_prompt_utils.
|
||||
"""Recording fake tokenizer mirroring the one in test_hunyuan_image3_prompt_utils.
|
||||
|
||||
Special token ids: `<|startoftext|>`=1, `<img>`=2, `<think>`=3,
|
||||
`<recaption>`=4. encode() returns one id per character starting at
|
||||
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
@@ -134,6 +134,8 @@ def _install_vllm_stubs() -> None:
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.quant_config = kwargs.pop("quant_config", None)
|
||||
self.prefix = kwargs.pop("prefix", "")
|
||||
del kwargs
|
||||
self.return_bias = return_bias
|
||||
self.weight = nn.Parameter(torch.empty(output_size, input_size))
|
||||
@@ -164,6 +166,8 @@ def _install_vllm_stubs() -> None:
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.quant_config = kwargs.pop("quant_config", None)
|
||||
self.prefix = kwargs.pop("prefix", "")
|
||||
del kwargs
|
||||
self.num_heads = total_num_heads
|
||||
self.num_kv_heads = total_num_kv_heads or total_num_heads
|
||||
|
||||
@@ -35,6 +35,8 @@ def _tiny_model(
|
||||
num_layers: int = 2,
|
||||
num_frames_per_block: int = 1,
|
||||
sliding_window_num_frames: int = 3,
|
||||
quant_config=None,
|
||||
prefix: str = "",
|
||||
attention_head_dim: int = 2,
|
||||
rope_max_seq_len: int = 16,
|
||||
):
|
||||
@@ -55,6 +57,8 @@ def _tiny_model(
|
||||
num_frames_per_block=num_frames_per_block,
|
||||
sliding_window_num_frames=sliding_window_num_frames,
|
||||
local_attn_size=-1,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
|
||||
@@ -485,11 +489,52 @@ def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
|
||||
@pytest.mark.parametrize("prefix", ["", "transformer"])
|
||||
@pytest.mark.parametrize("quantized", [False, True])
|
||||
def test_constructor_propagates_quantization_to_parallel_linears(prefix: str, quantized: bool) -> None:
|
||||
module = attention_tests._load_module()
|
||||
quant_config = object() if quantized else None
|
||||
model = _tiny_model(module, num_layers=1, quant_config=quant_config, prefix=prefix)
|
||||
block = model.blocks[0]
|
||||
quantized_linears = [
|
||||
model.c2ws_hidden_states_layer1,
|
||||
model.c2ws_hidden_states_layer2,
|
||||
block.self_attn.qkv,
|
||||
block.self_attn.o,
|
||||
block.cross_attn.q,
|
||||
block.cross_attn.k,
|
||||
block.cross_attn.v,
|
||||
block.cross_attn.o,
|
||||
block.ffn[0],
|
||||
block.ffn[2],
|
||||
block.cam_injector_layer1,
|
||||
block.cam_injector_layer2,
|
||||
]
|
||||
|
||||
with pytest.raises(RuntimeError, match="quant_config.*not supported"):
|
||||
module.CausalLingBotWorldTransformer3DModel(quant_config=object())
|
||||
assert len(quantized_linears) == 12
|
||||
assert all(layer.quant_config is quant_config for layer in quantized_linears)
|
||||
expected_names = {
|
||||
"c2ws_hidden_states_layer1",
|
||||
"c2ws_hidden_states_layer2",
|
||||
*(
|
||||
f"blocks.0.{name}"
|
||||
for name in (
|
||||
"self_attn.qkv",
|
||||
"self_attn.o",
|
||||
"cross_attn.q",
|
||||
"cross_attn.k",
|
||||
"cross_attn.v",
|
||||
"cross_attn.o",
|
||||
"ffn.0",
|
||||
"ffn.2",
|
||||
"cam_injector_layer1",
|
||||
"cam_injector_layer2",
|
||||
)
|
||||
),
|
||||
}
|
||||
assert {layer.prefix for layer in quantized_linears} == {
|
||||
f"{prefix}.{name}" if prefix else name for name in expected_names
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
|
||||
@@ -737,6 +737,99 @@ def test_component_discovery_uses_official_checkpoint_contract() -> None:
|
||||
assert module._loader_state.prefetch_calls == [("checkpoint", ("tokenizer", "text_encoder", "vae"), False)]
|
||||
|
||||
|
||||
def _ulysses_config(size: int):
|
||||
parallel_config = _od_config().parallel_config
|
||||
parallel_config.sequence_parallel_size = size
|
||||
parallel_config.ulysses_degree = size
|
||||
return _od_config(parallel_config=parallel_config)
|
||||
|
||||
|
||||
def _record_shard_install(monkeypatch, module, *, world_size: int):
|
||||
"""Stand in for the Ulysses group and the decoder patch; return what the install was asked for."""
|
||||
installs: list[dict] = []
|
||||
group = object()
|
||||
|
||||
def install(vae, group_arg, split_dim, *, dst):
|
||||
installs.append({"vae": vae, "group": group_arg, "split_dim": split_dim, "dst": dst})
|
||||
|
||||
monkeypatch.setattr(module, "install_wan_spatial_shard_decode", install)
|
||||
monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", lambda self: (group, world_size))
|
||||
return installs, group
|
||||
|
||||
|
||||
def test_a_multi_rank_deployment_shards_the_decoder_across_the_ulysses_ranks(monkeypatch) -> None:
|
||||
"""Multi-rank deployments shard decode by default."""
|
||||
module = _load_pipeline_module()
|
||||
installs, group = _record_shard_install(monkeypatch, module, world_size=2)
|
||||
|
||||
pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
|
||||
|
||||
# Along the width, and assembled on every rank: each rank's post_decode consumes the frame.
|
||||
assert installs == [{"vae": pipeline.vae, "group": group, "split_dim": "width", "dst": None}]
|
||||
assert pipeline._vae_shard_split_dim == "width"
|
||||
|
||||
|
||||
def test_a_single_rank_deployment_leaves_the_decoder_alone(monkeypatch) -> None:
|
||||
module = _load_pipeline_module()
|
||||
installs, _ = _record_shard_install(monkeypatch, module, world_size=1)
|
||||
|
||||
pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(1))
|
||||
|
||||
assert installs == [] and pipeline._vae_shard_split_dim is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enabled", [True, False])
|
||||
@pytest.mark.parametrize("width, local_width", [(832, 208), (840, 216)])
|
||||
def test_streaming_decode_reservation_uses_padded_shards(monkeypatch, enabled, width, local_width):
|
||||
module = _load_pipeline_module()
|
||||
installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
|
||||
config = _ulysses_config(4)
|
||||
config.model_config.update(lingbot_vae_spatial_sharding=enabled, ar_diffusion_height=480, ar_diffusion_width=width)
|
||||
pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
|
||||
from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
|
||||
|
||||
# Exercise the real byte estimator without constructing a checkpoint VAE.
|
||||
decoder = object.__new__(WanStreamingDecoder)
|
||||
decoder._bytes_per_pixel_fp32 = 16.0
|
||||
monkeypatch.setattr(pipeline, "_streaming_decoder", lambda: decoder)
|
||||
expected_width = local_width if enabled else width
|
||||
assert pipeline._streaming_decode_bytes_per_session() == 16 * 480 * expected_width
|
||||
assert len(installs) == int(enabled)
|
||||
assert pipeline._vae_shard_world_size == (4 if enabled else 1)
|
||||
assert config.parallel_config.sequence_parallel_size == config.parallel_config.ulysses_degree == 4
|
||||
|
||||
|
||||
def test_disabling_vae_sharding_does_not_require_a_shard_group(monkeypatch):
|
||||
module = _load_pipeline_module()
|
||||
config = _ulysses_config(4)
|
||||
config.model_config["lingbot_vae_spatial_sharding"] = False
|
||||
|
||||
def unexpected_group(self):
|
||||
pytest.fail("disabled VAE sharding must not access its process group")
|
||||
|
||||
monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", unexpected_group)
|
||||
pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
|
||||
assert pipeline._vae_shard_split_dim is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["false", 0, None])
|
||||
def test_vae_sharding_switch_requires_a_boolean(value):
|
||||
module = _load_pipeline_module()
|
||||
config = _ulysses_config(4)
|
||||
config.model_config["lingbot_vae_spatial_sharding"] = value
|
||||
with pytest.raises(ValueError, match="lingbot_vae_spatial_sharding must be a boolean"):
|
||||
module.LingBotWorldCausalDMDPipeline(od_config=config)
|
||||
|
||||
|
||||
def test_vae_shard_refuses_a_group_of_the_wrong_size(monkeypatch) -> None:
|
||||
module = _load_pipeline_module()
|
||||
installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
|
||||
|
||||
with pytest.raises(RuntimeError, match="sequence_parallel_size=2 but the Ulysses group has 4 ranks"):
|
||||
module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
|
||||
assert installs == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "feature"),
|
||||
[
|
||||
@@ -758,8 +851,9 @@ def test_unsupported_parallel_modes_fail_before_component_loading(field: str, va
|
||||
assert module._loader_state.prefetch_calls == []
|
||||
|
||||
|
||||
def test_pure_ulysses_parallel_config_is_supported() -> None:
|
||||
def test_pure_ulysses_parallel_config_is_supported(monkeypatch) -> None:
|
||||
module = _load_pipeline_module()
|
||||
_record_shard_install(monkeypatch, module, world_size=2)
|
||||
parallel_config = _od_config().parallel_config
|
||||
parallel_config.sequence_parallel_size = 2
|
||||
parallel_config.ulysses_degree = 2
|
||||
@@ -795,13 +889,14 @@ def test_unsupported_sp_config_fails_before_component_loading(overrides):
|
||||
assert module._loader_state.prefetch_calls == []
|
||||
|
||||
|
||||
def test_unsupported_quantization_fails_before_component_loading() -> None:
|
||||
def test_quantization_reaches_transformer_factory() -> None:
|
||||
module = _load_pipeline_module()
|
||||
quant_config = object()
|
||||
|
||||
with pytest.raises(NotImplementedError, match="quantization"):
|
||||
module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=object()))
|
||||
pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=quant_config))
|
||||
|
||||
assert module._loader_state.prefetch_calls == []
|
||||
assert pipeline.transformer is not None
|
||||
assert _FakeTransformerFactory.last_call[1:] == (quant_config, "transformer")
|
||||
|
||||
|
||||
def test_official_scheduler_config_matches_fixed_dmd_contract() -> None:
|
||||
@@ -1358,10 +1453,34 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N
|
||||
assert expected_camera.shape == (1, 384, 3, 2, 2)
|
||||
|
||||
|
||||
def test_fixed_dmd_transition_and_cache_commit_trace() -> None:
|
||||
@pytest.mark.parametrize("value", ["true", "false", 1, 0, None])
|
||||
def test_reuse_last_step_kv_rejects_non_boolean_config(value: object) -> None:
|
||||
config = _od_config()
|
||||
config.model_config["lingbot_reuse_last_step_kv"] = value
|
||||
with pytest.raises(ValueError, match="lingbot_reuse_last_step_kv must be a bool"):
|
||||
_pipeline(_load_pipeline_module(), od_config=config)
|
||||
|
||||
|
||||
def test_reuse_last_step_kv_is_fixed_per_instance(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("VLLM_OMNI_LINGBOT_REUSE_LAST_STEP_KV", "1")
|
||||
module = _load_pipeline_module()
|
||||
config = _od_config()
|
||||
config.model_config["lingbot_reuse_last_step_kv"] = True
|
||||
reuse = _pipeline(module, od_config=config)
|
||||
default = _pipeline(module)
|
||||
config.model_config["lingbot_reuse_last_step_kv"] = False
|
||||
assert reuse._dmd_blocks.reuse_last_step_kv is True
|
||||
assert default._dmd_blocks.reuse_last_step_kv is False
|
||||
assert reuse._dmd_blocks is reuse._dmd_blocks
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reuse_last_step_kv", [False, True])
|
||||
def test_fixed_dmd_transition_and_cache_commit_trace(reuse_last_step_kv: bool) -> None:
|
||||
module = _load_pipeline_module()
|
||||
transformer = _RecordingTransformer()
|
||||
pipeline = _pipeline(module, transformer=transformer)
|
||||
config = _od_config()
|
||||
config.model_config["lingbot_reuse_last_step_kv"] = reuse_last_step_kv
|
||||
pipeline = _pipeline(module, transformer=transformer, od_config=config)
|
||||
|
||||
result = pipeline(_request())
|
||||
|
||||
@@ -1383,14 +1502,18 @@ def test_fixed_dmd_transition_and_cache_commit_trace() -> None:
|
||||
current = x0
|
||||
torch.testing.assert_close(result.output, current)
|
||||
|
||||
expected_timesteps = [timestep for timestep, _ in warped_schedule]
|
||||
if not reuse_last_step_kv:
|
||||
expected_timesteps.append(0.0)
|
||||
torch.testing.assert_close(
|
||||
torch.cat([call["timestep"] for call in transformer.calls]),
|
||||
torch.tensor([*(timestep for timestep, _ in warped_schedule), 0.0]),
|
||||
torch.cat([call["timestep"] for call in transformer.calls]), torch.tensor(expected_timesteps)
|
||||
)
|
||||
assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True]
|
||||
assert [call["start_frame"] for call in transformer.calls] == [0, 0, 0, 0, 0]
|
||||
commit_flags = [False] * (len(expected_timesteps) - 1) + [True]
|
||||
assert [call["update_cache"] for call in transformer.calls] == commit_flags
|
||||
assert [call["start_frame"] for call in transformer.calls] == [0] * len(expected_timesteps)
|
||||
assert len({call["cache_id"] for call in transformer.calls}) == 1
|
||||
torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], result.output)
|
||||
committed_latent = result.output + warped_schedule[-1][1] if reuse_last_step_kv else result.output
|
||||
torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], committed_latent)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -2389,10 +2512,13 @@ def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None:
|
||||
pipeline(_request(sampling=sampling))
|
||||
|
||||
|
||||
def test_stepwise_progress_metadata_and_commit_trace() -> None:
|
||||
@pytest.mark.parametrize("reuse_last_step_kv", [False, True])
|
||||
def test_stepwise_progress_metadata_and_commit_trace(reuse_last_step_kv: bool) -> None:
|
||||
module = _load_pipeline_module()
|
||||
transformer = _RecordingTransformer()
|
||||
pipeline = _pipeline(module, transformer=transformer)
|
||||
config = _od_config()
|
||||
config.model_config["lingbot_reuse_last_step_kv"] = reuse_last_step_kv
|
||||
pipeline = _pipeline(module, transformer=transformer, od_config=config)
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
state = _stepwise_state(num_frames=21)
|
||||
@@ -2406,8 +2532,9 @@ def test_stepwise_progress_metadata_and_commit_trace() -> None:
|
||||
assert [output.chunk_index for output in outputs] == [0, 1]
|
||||
assert all(output.total_chunks == 2 for output in outputs)
|
||||
assert outputs[-1].finished is True
|
||||
assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2
|
||||
assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
|
||||
calls_per_chunk = 4 if reuse_last_step_kv else 5
|
||||
assert [call["update_cache"] for call in transformer.calls] == ([False] * (calls_per_chunk - 1) + [True]) * 2
|
||||
assert [call["start_frame"] for call in transformer.calls] == [0] * calls_per_chunk + [3] * calls_per_chunk
|
||||
assert fake.commits == ["main", "main"]
|
||||
for chunk_index, output in enumerate(outputs):
|
||||
metadata = output.output["metadata"]["ar_diffusion"]
|
||||
|
||||
@@ -171,7 +171,7 @@ def test_huggingface_url_resolves_pinned_snapshot(tmp_path, monkeypatch):
|
||||
else pytest.fail("unexpected snapshot request")
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.transformers_utils.repo_utils.hf_api",
|
||||
"vllm_omni.transformers_utils.repo_utils.hf_api",
|
||||
lambda: api,
|
||||
)
|
||||
assert _resolve_checkpoint_root(
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
|
||||
import json
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from multiprocessing.reduction import ForkingPickler
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, TypeVar
|
||||
@@ -2837,41 +2837,63 @@ def test_dit_encoder_selection_keeps_vae_resident():
|
||||
component.offload_to_cpu.assert_not_called()
|
||||
|
||||
|
||||
def test_keyframe_encode_pins_and_restores_cudnn_settings():
|
||||
@pytest.mark.parametrize(
|
||||
("capability", "allow_tf32"),
|
||||
[((8, 6), True), ((9, 0), False), ((10, 0), True), ((10, 3), True), (None, True)],
|
||||
)
|
||||
@pytest.mark.parametrize("initial_tf32", [False, True])
|
||||
@pytest.mark.parametrize("fail_encode", [False, True])
|
||||
def test_keyframe_encode_pins_and_restores_cudnn_settings(
|
||||
monkeypatch, capability, allow_tf32, initial_tf32, fail_encode
|
||||
):
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3.vae import (
|
||||
_minimax_h3_keyframe_encode_context,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
def get_capability(device_id):
|
||||
assert device_id == 1
|
||||
return DeviceCapability(*capability) if capability is not None else None
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
||||
monkeypatch.setattr(current_omni_platform, "get_device_capability", get_capability)
|
||||
|
||||
cudnn = torch.backends.cudnn
|
||||
original = (
|
||||
cudnn.enabled,
|
||||
cudnn.benchmark,
|
||||
cudnn.deterministic,
|
||||
cudnn.allow_tf32,
|
||||
)
|
||||
try:
|
||||
cudnn.enabled = False
|
||||
cudnn.benchmark = True
|
||||
cudnn.deterministic = False
|
||||
cudnn.allow_tf32 = False
|
||||
|
||||
with _minimax_h3_keyframe_encode_context(torch.device("cuda")):
|
||||
with cudnn.flags(enabled=False, benchmark=True, deterministic=False, allow_tf32=initial_tf32):
|
||||
error_context = pytest.raises(RuntimeError, match="encode failed") if fail_encode else nullcontext()
|
||||
with error_context, _minimax_h3_keyframe_encode_context(torch.device("cuda:1")):
|
||||
assert cudnn.enabled
|
||||
assert not cudnn.benchmark
|
||||
assert cudnn.deterministic
|
||||
assert cudnn.allow_tf32
|
||||
assert cudnn.allow_tf32 is allow_tf32
|
||||
if fail_encode:
|
||||
raise RuntimeError("encode failed")
|
||||
|
||||
assert not cudnn.enabled
|
||||
assert cudnn.benchmark
|
||||
assert not cudnn.deterministic
|
||||
assert not cudnn.allow_tf32
|
||||
finally:
|
||||
(
|
||||
cudnn.enabled,
|
||||
cudnn.benchmark,
|
||||
cudnn.deterministic,
|
||||
cudnn.allow_tf32,
|
||||
) = original
|
||||
assert cudnn.allow_tf32 is initial_tf32
|
||||
|
||||
|
||||
def test_keyframe_encode_cpu_keeps_cudnn_settings(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3.vae import (
|
||||
_minimax_h3_keyframe_encode_context,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
def unexpected_capability_query(device_id):
|
||||
pytest.fail("CPU keyframe encode must not query CUDA capability")
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "get_device_capability", unexpected_capability_query)
|
||||
cudnn = torch.backends.cudnn
|
||||
with cudnn.flags(enabled=False, benchmark=True, deterministic=False, allow_tf32=False):
|
||||
with _minimax_h3_keyframe_encode_context(torch.device("cpu")):
|
||||
assert not cudnn.enabled
|
||||
assert cudnn.benchmark
|
||||
assert not cudnn.deterministic
|
||||
assert not cudnn.allow_tf32
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fail_encode", [False, True])
|
||||
|
||||
@@ -10,7 +10,6 @@ from vllm_omni.diffusion.layers.indexed_modulation import (
|
||||
_MAX_1D_GRID_SIZE,
|
||||
_iter_row_chunks,
|
||||
_launch_row_chunks,
|
||||
_use_hopper_bf16_affine_semantics,
|
||||
indexed_gate,
|
||||
indexed_gate_rms_norm_scale_shift,
|
||||
indexed_scale_shift_,
|
||||
@@ -194,52 +193,28 @@ def test_fused_modulation_preserves_bf16_residual_boundary() -> None:
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.skipif(
|
||||
torch.cuda.is_available() and torch.cuda.get_device_capability()[0] != 9,
|
||||
reason="SM90-specific precision path",
|
||||
)
|
||||
@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA required")
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
|
||||
def test_hopper_fused_modulation_matches_pytorch_reference() -> None:
|
||||
torch.manual_seed(42)
|
||||
rows, hidden_size, conditions = 128, 3072, 4
|
||||
residual = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
gate = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
branch = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
shift = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
scale = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
|
||||
indices = torch.arange(rows, device="cuda") % conditions
|
||||
@pytest.mark.parametrize("fused_gate", [False, True])
|
||||
def test_fused_modulation_keeps_fp32_affine_intermediates(fused_gate: bool) -> None:
|
||||
device = current_omni_platform.get_torch_device()
|
||||
rows, hidden_size = 2, 5376
|
||||
x = torch.ones(rows, hidden_size, device=device, dtype=torch.bfloat16)
|
||||
weight = torch.full((hidden_size,), 1.125, device=device, dtype=x.dtype)
|
||||
shift = torch.zeros(1, hidden_size, device=device, dtype=x.dtype)
|
||||
scale = torch.full_like(shift, 0.09375)
|
||||
indices = torch.zeros(rows, device=device, dtype=torch.int64)
|
||||
eps = 1e-6
|
||||
if fused_gate:
|
||||
residual, actual = indexed_gate_rms_norm_scale_shift(
|
||||
x, torch.zeros_like(shift), x, weight, shift, scale, indices, eps
|
||||
)
|
||||
assert torch.equal(residual, x)
|
||||
else:
|
||||
actual = rms_norm_indexed_scale_shift(x, weight, shift, scale, indices, eps)
|
||||
|
||||
reference_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
|
||||
normalized = reference_residual.float()
|
||||
variance = normalized.pow(2).mean(-1, keepdim=True)
|
||||
normalized = (weight.float() * (normalized * torch.rsqrt(variance + eps))).to(torch.bfloat16)
|
||||
reference_modulated = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(
|
||||
torch.bfloat16
|
||||
)
|
||||
|
||||
actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift(
|
||||
residual,
|
||||
gate,
|
||||
branch,
|
||||
weight,
|
||||
shift,
|
||||
scale,
|
||||
indices,
|
||||
eps,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual_residual, reference_residual, atol=5e-2, rtol=5e-2)
|
||||
torch.testing.assert_close(actual_modulated, reference_modulated, atol=5e-2, rtol=5e-2)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
def test_hopper_precision_dispatch_is_architecture_scoped(monkeypatch) -> None:
|
||||
tensor = torch.empty(0, device="cuda")
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0))
|
||||
assert _use_hopper_bf16_affine_semantics(tensor)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
|
||||
assert not _use_hopper_bf16_affine_semantics(tensor)
|
||||
# Without the RMS epsilon, 1.125 * 1.09375 = 1.23046875 is exactly
|
||||
# halfway between two BF16 values. FP32 normalization keeps the epsilon
|
||||
# and rounds down to 1.2265625. Premature BF16 rounding drops the epsilon
|
||||
# and instead rounds the tie up to 1.234375, changing every output row.
|
||||
assert torch.equal(actual, torch.full_like(x, 1.2265625))
|
||||
@@ -16,6 +16,7 @@ from tests.diffusion.quantization.test_quantization_quality import (
|
||||
_maybe_save_output,
|
||||
)
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.mock import patch_hf_snapshot_download
|
||||
|
||||
_MINIMAX_H3_REPO = "MiniMaxAI/MiniMax-H3"
|
||||
_MINIMAX_H3_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc"
|
||||
@@ -258,13 +259,13 @@ def test_resolve_fl2va_model_ref(tmp_path, monkeypatch):
|
||||
fl2va_root.mkdir()
|
||||
(fl2va_root / "model_index.json").write_text("{}", encoding="utf-8")
|
||||
|
||||
def fake_snapshot_download(self, *, repo_id, revision, allow_patterns):
|
||||
def fake_snapshot_download(*, repo_id, revision, allow_patterns):
|
||||
assert repo_id == _MINIMAX_H3_REPO
|
||||
assert revision == _MINIMAX_H3_REVISION
|
||||
assert allow_patterns == ["FL2VA/**"]
|
||||
return str(tmp_path)
|
||||
|
||||
monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", fake_snapshot_download)
|
||||
patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
|
||||
assert _resolve_fl2va_model_ref() == str(fl2va_root)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import MethodType, SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.nextstep_1_1 import pipeline_nextstep_1_1 as pipeline_module
|
||||
from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import NextStepModel
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("height,width", [(32, 32), (32, 48), (48, 32)])
|
||||
@pytest.mark.parametrize("batch", [1, 2])
|
||||
def test_forward_preserves_latent_grid(monkeypatch, height, width, batch):
|
||||
pipeline = object.__new__(pipeline_module.NextStep11Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline._NextStep11Pipeline__device = torch.device("cpu")
|
||||
pipeline.config = SimpleNamespace(use_gen_pos_embed=False)
|
||||
pipeline.down_factor = 16
|
||||
pipeline.scaling_factor = 2.0
|
||||
pipeline.shift_factor = 0.25
|
||||
grid = SimpleNamespace(config=SimpleNamespace(latent_channels=2, latent_patch_size=2))
|
||||
expected = torch.arange(batch * 2 * (height // 8) * (width // 8), dtype=torch.float32)
|
||||
expected = expected.reshape(batch, 2, height // 8, width // 8)
|
||||
tokens = NextStepModel.patchify(grid, expected)
|
||||
pipeline.model = Mock(
|
||||
return_value=SimpleNamespace(last_hidden_state=torch.zeros(batch, 1, 4), past_key_values=None)
|
||||
)
|
||||
pipeline.model.forward_model.return_value = pipeline.model.return_value
|
||||
pipeline.model.unpatchify = MethodType(NextStepModel.unpatchify, grid)
|
||||
pipeline.tokenizer = Mock(
|
||||
bos_token=None,
|
||||
return_value=SimpleNamespace(
|
||||
input_ids=torch.ones(batch, 2, dtype=torch.long), attention_mask=torch.ones(batch, 2, dtype=torch.long)
|
||||
),
|
||||
)
|
||||
pipeline._check_input = Mock(return_value=(["teapot"], None))
|
||||
pipeline._build_captions = Mock(return_value=(["teapot"] * batch, None, 1, 1.0))
|
||||
pipeline._add_prefix_ids = lambda _hw, ids, mask: (ids, mask)
|
||||
pipeline.decoding = Mock(return_value=tokens)
|
||||
pipeline.vae = SimpleNamespace(dtype=torch.float32, decode=lambda latent: SimpleNamespace(sample=latent))
|
||||
monkeypatch.setattr(pipeline_module, "StaticCache", lambda **_kwargs: None)
|
||||
params = OmniDiffusionSamplingParams(
|
||||
height=height, width=width, guidance_scale=1.0, num_outputs_per_prompt=batch, seed=None
|
||||
)
|
||||
output = pipeline.forward(SimpleNamespace(prompts=["teapot"], sampling_params=params))
|
||||
torch.testing.assert_close(output.output, expected / 2 + 0.25, rtol=0, atol=0)
|
||||
@@ -497,3 +497,99 @@ class TestWorkerProcShutdown:
|
||||
result_mq.shutdown.assert_called_once_with()
|
||||
assert proc.mq is None
|
||||
assert proc.result_mq is None
|
||||
|
||||
|
||||
class TestSynchronizeDeviceBarrier:
|
||||
"""The pause barrier must not report a quiet device while outputs still pack."""
|
||||
|
||||
@staticmethod
|
||||
def _rpc(timeout):
|
||||
return {
|
||||
"method": "synchronize_device",
|
||||
"args": (),
|
||||
"kwargs": {"timeout": timeout},
|
||||
"output_rank": 0,
|
||||
"exec_all_ranks": False,
|
||||
"collect_rank_status": False,
|
||||
}
|
||||
|
||||
def test_rejects_undrained_output(self, mocker):
|
||||
proc = _make_worker_proc(step_execution=False)
|
||||
proc.worker = MagicMock()
|
||||
mocker.patch.object(proc, "drain_async_outputs", return_value=False)
|
||||
|
||||
with pytest.raises(TimeoutError, match="did not drain"):
|
||||
proc._execute_rpc(self._rpc(0.01))
|
||||
|
||||
proc.worker.execute_method.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize("timeout", [0.5, None])
|
||||
def test_drains_with_the_call_timeout_then_synchronizes(self, mocker, timeout):
|
||||
proc = _make_worker_proc(step_execution=False)
|
||||
proc.worker = MagicMock()
|
||||
call_order: list = []
|
||||
mocker.patch.object(
|
||||
proc,
|
||||
"drain_async_outputs",
|
||||
side_effect=lambda **kw: call_order.append(("drain", kw)) or True,
|
||||
)
|
||||
proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append(("sync", a, kw))
|
||||
|
||||
_, should_reply = proc._execute_rpc(self._rpc(timeout))
|
||||
|
||||
assert call_order == [
|
||||
("drain", {"timeout": timeout}),
|
||||
("sync", ("synchronize_device",), {"timeout": timeout}),
|
||||
]
|
||||
assert should_reply is True
|
||||
|
||||
def test_completes_only_after_output_ready_is_published(self, mocker):
|
||||
proc = _make_worker_proc(step_execution=False)
|
||||
proc.worker = MagicMock()
|
||||
|
||||
mock_platform = mocker.MagicMock()
|
||||
mock_platform.record_device_event.return_value = MagicMock()
|
||||
mocker.patch(
|
||||
"vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
|
||||
mock_platform,
|
||||
)
|
||||
mock_torch = mocker.MagicMock()
|
||||
mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
|
||||
mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
|
||||
release = threading.Event()
|
||||
mocker.patch(
|
||||
"vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
|
||||
side_effect=lambda *a, **kw: release.wait(timeout=5.0),
|
||||
)
|
||||
|
||||
proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
|
||||
loop_thread = threading.Thread(target=proc._async_output_loop, daemon=True)
|
||||
loop_thread.start()
|
||||
|
||||
barrier_done = threading.Event()
|
||||
|
||||
def run_barrier():
|
||||
proc._execute_rpc(self._rpc(5.0))
|
||||
barrier_done.set()
|
||||
|
||||
barrier_thread = threading.Thread(target=run_barrier, daemon=True)
|
||||
barrier_thread.start()
|
||||
try:
|
||||
assert not barrier_done.wait(timeout=0.2)
|
||||
proc.worker.execute_method.assert_not_called()
|
||||
|
||||
release.set()
|
||||
assert barrier_done.wait(timeout=5.0)
|
||||
|
||||
kinds = [
|
||||
call.args[0].kind
|
||||
for call in proc.result_mq.enqueue.call_args_list
|
||||
if isinstance(call.args[0], AsyncDiffusionOutput)
|
||||
]
|
||||
assert kinds == [AsyncOutputKind.COMPUTE_DONE, AsyncOutputKind.OUTPUT_READY]
|
||||
proc.worker.execute_method.assert_called_once_with("synchronize_device", timeout=5.0)
|
||||
finally:
|
||||
release.set()
|
||||
proc._async_output_queue.put(None)
|
||||
loop_thread.join(timeout=2.0)
|
||||
barrier_thread.join(timeout=2.0)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Regression tests for the busy-loop-based RPC routing in DiffusionEngine.
|
||||
|
||||
These tests guard the invariants established when ``collective_rpc`` was
|
||||
@@ -60,6 +60,12 @@ class _ConcurrencyTrackingExecutor:
|
||||
self.is_failed = False
|
||||
self._closed = False
|
||||
self.od_config = SimpleNamespace()
|
||||
# Pause tests: a call whose gate is registered (by request id, or by
|
||||
# method name for ``synchronize_device``) blocks inside the executor
|
||||
# until the gate is set; ``version`` tags each call so a test can tell
|
||||
# which "weights" a request ran on.
|
||||
self.gates: dict[str, threading.Event] = {}
|
||||
self.version = "v1"
|
||||
|
||||
def collective_rpc(
|
||||
self,
|
||||
@@ -81,16 +87,25 @@ class _ConcurrencyTrackingExecutor:
|
||||
"kwargs": kwargs,
|
||||
"unique_reply_rank": unique_reply_rank,
|
||||
"thread": threading.get_ident(),
|
||||
"version": self.version,
|
||||
}
|
||||
)
|
||||
try:
|
||||
if self.rpc_delay:
|
||||
time.sleep(self.rpc_delay)
|
||||
if method == "synchronize_device":
|
||||
gate = self.gates.get(method)
|
||||
if gate is not None:
|
||||
gate.wait(timeout=5.0)
|
||||
return [None]
|
||||
# Distinguish per-request execution from raw RPC by method name.
|
||||
if method in {"execute_model", "execute_stepwise", "generate"}:
|
||||
# args[0] is the request-like object for execute_model.
|
||||
req = args[0] if args else None
|
||||
tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
|
||||
gate = self.gates.get(tag)
|
||||
if gate is not None:
|
||||
gate.wait(timeout=5.0)
|
||||
return DiffusionOutput(error=f"result_for_{tag}", finished=True)
|
||||
tag = args[0] if args else method
|
||||
return DiffusionOutput(error=f"rpc_result_for_{tag}")
|
||||
@@ -133,21 +148,27 @@ def _make_request(tag: str):
|
||||
|
||||
|
||||
def _make_engine_with_loop(
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
loop: asyncio.AbstractEventLoop | None,
|
||||
rpc_delay: float = 0.0,
|
||||
*,
|
||||
scheduler: RequestScheduler | None = None,
|
||||
start_loop: bool = True,
|
||||
):
|
||||
"""Construct a ``DiffusionEngine`` skeleton with a real busy loop.
|
||||
|
||||
The engine is wired with a fake executor that asserts no concurrent
|
||||
calls and a real ``RequestScheduler``.
|
||||
calls and a real ``RequestScheduler``. ``start_loop=False`` leaves the
|
||||
engine in its bootstrap state (no busy-loop thread).
|
||||
"""
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine._closed = False
|
||||
engine.od_config = SimpleNamespace(streaming_output=False)
|
||||
engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
|
||||
|
||||
sched = RequestScheduler()
|
||||
sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
|
||||
sched = scheduler
|
||||
if sched is None:
|
||||
sched = RequestScheduler()
|
||||
sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
|
||||
engine.scheduler = sched
|
||||
engine.step_execution = False
|
||||
engine.supports_request_batch = False
|
||||
@@ -160,9 +181,14 @@ def _make_engine_with_loop(
|
||||
engine._closed = False
|
||||
engine.abort_queue = queue.Queue()
|
||||
engine._rpc_queue = queue.Queue()
|
||||
engine._scheduling_paused = False
|
||||
|
||||
engine.main_loop = loop
|
||||
engine.stop_event = threading.Event()
|
||||
if not start_loop:
|
||||
engine._loop_started = False
|
||||
engine.worker_thread = None
|
||||
return engine
|
||||
engine._loop_started = True
|
||||
engine.worker_thread = threading.Thread(target=engine._busy_loop, daemon=True)
|
||||
engine.worker_thread.start()
|
||||
@@ -170,6 +196,10 @@ def _make_engine_with_loop(
|
||||
|
||||
|
||||
def _stop_engine(engine: DiffusionEngine) -> None:
|
||||
if engine.worker_thread is None:
|
||||
return
|
||||
for gate in engine.executor.gates.values():
|
||||
gate.set()
|
||||
with engine._cv:
|
||||
engine.stop_event.set()
|
||||
engine._cv.notify_all()
|
||||
@@ -565,3 +595,319 @@ async def test_rpc_and_request_results_do_not_swap():
|
||||
assert req_res.error == "result_for_req1"
|
||||
# Belt and braces: never overlapped on the executor.
|
||||
assert engine.executor.max_concurrent == 1
|
||||
|
||||
|
||||
# ───────────────────────── scheduler pause (keep) ─────────────────────────
|
||||
|
||||
|
||||
class _GatedAdmissionScheduler(RequestScheduler):
|
||||
"""Batching wait that only ends when the test opens admission."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.admission_open = threading.Event()
|
||||
|
||||
def get_admission_wait_decision(self, *, now: float, dp_concurrent: bool = False):
|
||||
decision = super().get_admission_wait_decision(now=now, dp_concurrent=dp_concurrent)
|
||||
return decision.__class__(should_wait=True, deadline=None, max_batch=self.max_num_running_reqs)
|
||||
|
||||
def should_end_admission_wait(self, decision, *, now: float, stable_since: float) -> bool:
|
||||
return self.admission_open.is_set()
|
||||
|
||||
|
||||
def _pause(engine: DiffusionEngine, timeout: float | None = None):
|
||||
return engine.async_collective_rpc("pause_scheduler", timeout=timeout, kwargs={"mode": "keep"})
|
||||
|
||||
|
||||
def _resume(engine: DiffusionEngine):
|
||||
return engine.async_collective_rpc("resume_scheduler")
|
||||
|
||||
|
||||
def _methods(engine: DiffusionEngine) -> list[str]:
|
||||
return [call["method"] for call in engine.executor.calls]
|
||||
|
||||
|
||||
def _executed(engine: DiffusionEngine) -> list[tuple[str, str]]:
|
||||
return [
|
||||
(call["args"][0].request_id, call["version"])
|
||||
for call in engine.executor.calls
|
||||
if call["method"] == "execute_model"
|
||||
]
|
||||
|
||||
|
||||
def _submit(engine: DiffusionEngine, tag: str) -> asyncio.Task:
|
||||
return asyncio.create_task(_consume_final_output(engine.async_add_req_and_stream_response(_make_request(tag))))
|
||||
|
||||
|
||||
async def _wait_until(predicate, timeout: float = 3.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while not predicate():
|
||||
if time.monotonic() > deadline:
|
||||
raise AssertionError("condition not reached in time")
|
||||
await asyncio.sleep(0.005)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keep_pause_retains_waiting_request_and_serves_control_rpcs():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
try:
|
||||
assert await _pause(engine) is None
|
||||
queued = _submit(engine, "queued")
|
||||
result = await engine.async_collective_rpc("list_loras", unique_reply_rank=0)
|
||||
assert result.error == "rpc_result_for_list_loras"
|
||||
await asyncio.sleep(0.05)
|
||||
assert engine.scheduler.num_waiting_requests() == 1
|
||||
assert _executed(engine) == []
|
||||
assert not queued.done()
|
||||
|
||||
await _resume(engine)
|
||||
assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_queued"
|
||||
assert _methods(engine) == ["synchronize_device", "list_loras", "execute_model"]
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keep_pause_acks_after_running_batch_and_freezes_queue():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
gate = engine.executor.gates["A"] = threading.Event()
|
||||
try:
|
||||
a = _submit(engine, "A")
|
||||
await _wait_until(lambda: _executed(engine) == [("A", "v1")])
|
||||
b = _submit(engine, "B")
|
||||
await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
|
||||
|
||||
pause = asyncio.create_task(_pause(engine))
|
||||
await asyncio.sleep(0.1)
|
||||
assert not pause.done()
|
||||
|
||||
gate.set()
|
||||
await asyncio.wait_for(pause, 3.0)
|
||||
assert _methods(engine) == ["execute_model", "synchronize_device"]
|
||||
assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
|
||||
assert engine.scheduler.num_waiting_requests() == 1
|
||||
assert not b.done()
|
||||
|
||||
engine.executor.version = "v2"
|
||||
await _resume(engine)
|
||||
assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
|
||||
assert _executed(engine) == [("A", "v1"), ("B", "v2")]
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keep_pause_ack_does_not_release_later_rpcs_before_the_batch_is_emitted():
|
||||
"""A pause ACK must not open the queue to RPCs that touch device memory
|
||||
(sleep) before the batch that just ran has been delivered.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
barrier_gate = engine.executor.gates["synchronize_device"] = threading.Event()
|
||||
methods_at_emit: list[list[str]] = []
|
||||
emit_outputs = engine._emit_outputs
|
||||
|
||||
def _recording_emit(*args, **kwargs):
|
||||
methods_at_emit.append(_methods(engine))
|
||||
return emit_outputs(*args, **kwargs)
|
||||
|
||||
engine._emit_outputs = _recording_emit
|
||||
try:
|
||||
gate = engine.executor.gates["A"] = threading.Event()
|
||||
a = _submit(engine, "A")
|
||||
await _wait_until(lambda: _executed(engine) == [("A", "v1")])
|
||||
pause = asyncio.create_task(_pause(engine))
|
||||
await _wait_until(lambda: engine._scheduling_paused)
|
||||
gate.set()
|
||||
# The barrier holds the busy loop inside the post-execute RPC drain,
|
||||
# so the sleep is queued before the pause ACK is published.
|
||||
await _wait_until(lambda: "synchronize_device" in _methods(engine))
|
||||
sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
|
||||
await _wait_until(lambda: not engine._rpc_queue.empty())
|
||||
barrier_gate.set()
|
||||
|
||||
await asyncio.wait_for(pause, 3.0)
|
||||
assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
|
||||
assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
assert _methods(engine) == ["execute_model", "synchronize_device", "handle_sleep_task"]
|
||||
assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timed_out_pause_still_ends_the_drain_before_the_batch_is_emitted():
|
||||
"""A pause whose caller already gave up is still a drain barrier: its
|
||||
cancelled task must not let a later RPC run before the batch is emitted.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
methods_at_emit: list[list[str]] = []
|
||||
emit_outputs = engine._emit_outputs
|
||||
|
||||
def _recording_emit(*args, **kwargs):
|
||||
methods_at_emit.append(_methods(engine))
|
||||
return emit_outputs(*args, **kwargs)
|
||||
|
||||
engine._emit_outputs = _recording_emit
|
||||
try:
|
||||
gate = engine.executor.gates["A"] = threading.Event()
|
||||
a = _submit(engine, "A")
|
||||
await _wait_until(lambda: _executed(engine) == [("A", "v1")])
|
||||
|
||||
# The caller times out while A is still executing: the pause future is
|
||||
# cancelled, but its task stays queued ahead of the sleep.
|
||||
with pytest.raises(TimeoutError):
|
||||
await _pause(engine, timeout=0.1)
|
||||
sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
|
||||
await _wait_until(lambda: engine._rpc_queue.qsize() == 2)
|
||||
gate.set()
|
||||
|
||||
assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
|
||||
assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
assert _methods(engine) == ["execute_model", "handle_sleep_task"]
|
||||
assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_rejects_unsupported_modes_without_closing_the_gate():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
try:
|
||||
for kwargs in ({"mode": "abort"}, {"mode": "wait"}, {}):
|
||||
with pytest.raises(ValueError):
|
||||
await engine.async_collective_rpc("pause_scheduler", kwargs=kwargs)
|
||||
engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
|
||||
with pytest.raises(NotImplementedError):
|
||||
await _pause(engine)
|
||||
engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
|
||||
|
||||
assert (await asyncio.wait_for(_submit(engine, "still-runs"), 3.0)).error == "result_for_still-runs"
|
||||
assert _methods(engine) == ["execute_model"]
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keep_pause_during_batching_wait_holds_the_batch():
|
||||
loop = asyncio.get_running_loop()
|
||||
sched = _GatedAdmissionScheduler()
|
||||
sched.initialize(SimpleNamespace(max_num_seqs=4, request_batch_max_wait_ms=1000.0))
|
||||
engine = _make_engine_with_loop(loop, scheduler=sched)
|
||||
try:
|
||||
waiting = _submit(engine, "W")
|
||||
await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
await asyncio.wait_for(_pause(engine), 3.0)
|
||||
assert _methods(engine) == ["synchronize_device"]
|
||||
assert engine.scheduler.num_waiting_requests() == 1
|
||||
|
||||
await _resume(engine)
|
||||
await asyncio.sleep(0.05)
|
||||
assert _executed(engine) == []
|
||||
sched.admission_open.set()
|
||||
assert (await asyncio.wait_for(waiting, 3.0)).error == "result_for_W"
|
||||
finally:
|
||||
sched.admission_open.set()
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_timeout_leaves_gate_closed_until_explicit_resume():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
gate = engine.executor.gates["A"] = threading.Event()
|
||||
try:
|
||||
a = _submit(engine, "A")
|
||||
await _wait_until(lambda: _executed(engine) == [("A", "v1")])
|
||||
b = _submit(engine, "B")
|
||||
await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
await _pause(engine, timeout=0.1)
|
||||
gate.set()
|
||||
assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
|
||||
await asyncio.sleep(0.1)
|
||||
assert _executed(engine) == [("A", "v1")]
|
||||
assert not b.done()
|
||||
|
||||
await _resume(engine)
|
||||
assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
|
||||
assert "synchronize_device" not in _methods(engine)
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_pause_and_resume_are_idempotent():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
try:
|
||||
await _resume(engine)
|
||||
await _pause(engine)
|
||||
await _pause(engine)
|
||||
queued = _submit(engine, "Q")
|
||||
await asyncio.sleep(0.05)
|
||||
assert _executed(engine) == []
|
||||
await _resume(engine)
|
||||
await _resume(engine)
|
||||
assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_Q"
|
||||
assert _methods(engine) == ["synchronize_device", "synchronize_device", "execute_model"]
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_abort_of_waiting_request_while_paused_delivers_terminal():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
try:
|
||||
await _pause(engine)
|
||||
b = _submit(engine, "B")
|
||||
await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
|
||||
engine.abort("B")
|
||||
|
||||
result = await asyncio.wait_for(b, 3.0)
|
||||
assert result.aborted is True
|
||||
assert engine.scheduler.get_request_state("B") is None
|
||||
assert _executed(engine) == []
|
||||
|
||||
await _resume(engine)
|
||||
assert (await asyncio.wait_for(_submit(engine, "C"), 3.0)).error == "result_for_C"
|
||||
assert _executed(engine) == [("C", "v1")]
|
||||
assert engine.worker_thread.is_alive()
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
def test_pause_and_resume_before_loop_starts_use_engine_local_path():
|
||||
engine = _make_engine_with_loop(None, start_loop=False)
|
||||
|
||||
assert engine.collective_rpc("pause_scheduler", kwargs={"mode": "keep"}) is None
|
||||
assert _methods(engine) == ["synchronize_device"]
|
||||
|
||||
assert engine.collective_rpc("resume_scheduler") is None
|
||||
assert engine.add_req_and_wait_for_response(_make_request("X")).error == "result_for_X"
|
||||
assert _executed(engine) == [("X", "v1")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_requests_without_pause_never_trigger_the_barrier():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
try:
|
||||
first = await asyncio.wait_for(_submit(engine, "r1"), 3.0)
|
||||
await engine.async_collective_rpc("ping", args=("p",), unique_reply_rank=0)
|
||||
second = await asyncio.wait_for(_submit(engine, "r2"), 3.0)
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
assert (first.error, second.error) == ("result_for_r1", "result_for_r2")
|
||||
assert _methods(engine) == ["execute_model", "ping", "execute_model"]
|
||||
@@ -2092,3 +2092,19 @@ class TestStepScheduler:
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
self.scheduler.add_request(request)
|
||||
|
||||
|
||||
class TestPendingFinishedRequestIds:
|
||||
def test_reports_finished_ids_that_still_hold_state_without_clearing_them(self):
|
||||
sched = RequestScheduler()
|
||||
sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
|
||||
sched.add_request(_make_request("a"))
|
||||
sched.add_request(_make_request("b"))
|
||||
|
||||
sched.finish_requests("a", DiffusionRequestStatus.FINISHED_ABORTED)
|
||||
assert sched.pending_finished_request_ids() == {"a"}
|
||||
|
||||
sched.pop_request_state("a")
|
||||
assert sched.pending_finished_request_ids() == set()
|
||||
# The next wave still ships the id to the worker for its own cleanup.
|
||||
assert sched.schedule().finished_req_ids == {"a"}
|
||||
@@ -0,0 +1,53 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for the device step of the diffusion pause barrier."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_worker(runner: object) -> DiffusionWorker:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.model_runner = runner
|
||||
worker.rank = 0
|
||||
return worker
|
||||
|
||||
|
||||
def test_synchronize_device_waits_for_prefetch_before_device_sync(mocker):
|
||||
platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||||
order: list = []
|
||||
manager = MagicMock()
|
||||
manager.wait_prefetch.side_effect = lambda timeout: order.append(("prefetch", timeout)) or True
|
||||
platform.synchronize.side_effect = lambda: order.append("sync")
|
||||
worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
|
||||
|
||||
worker.synchronize_device(timeout=2.5)
|
||||
|
||||
assert order == [("prefetch", 2.5), "sync"]
|
||||
|
||||
|
||||
def test_synchronize_device_fails_when_prefetch_is_still_outstanding(mocker):
|
||||
platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||||
manager = MagicMock()
|
||||
manager.wait_prefetch.return_value = False
|
||||
worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
|
||||
|
||||
with pytest.raises(TimeoutError, match="prefetch"):
|
||||
worker.synchronize_device(timeout=0.01)
|
||||
|
||||
platform.synchronize.assert_not_called()
|
||||
|
||||
|
||||
def test_synchronize_device_without_transfer_manager_only_syncs(mocker):
|
||||
platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||||
worker = _make_worker(SimpleNamespace(kv_transfer_manager=None))
|
||||
|
||||
worker.synchronize_device(timeout=None)
|
||||
|
||||
platform.synchronize.assert_called_once_with()
|
||||
Reference in new issue
Block a user