Zhou TaichangandClaude Fable 5.1
86fcfe95f1
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
1554 lines
60 KiB
Python
1554 lines
60 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Tests for step-level diffusion execution across runner / worker / executor / engine."""
|
|
|
|
import contextlib
|
|
import os
|
|
import queue
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from pytest_mock import MockerFixture
|
|
|
|
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
|
|
from tests.helpers.mark import hardware_test
|
|
from vllm_omni.config.config_factory import StageConfigFactory
|
|
from vllm_omni.diffusion.data import DiffusionOutput
|
|
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
|
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
|
from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
|
|
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
|
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
destroy_distributed_env,
|
|
get_sp_group,
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
|
|
from vllm_omni.diffusion.ipc import (
|
|
pack_diffusion_output_shm,
|
|
unpack_diffusion_output_shm,
|
|
)
|
|
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
|
|
DiffusionPipelineProfilerMixin,
|
|
)
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.sched import StepScheduler
|
|
from vllm_omni.diffusion.sched.interface import (
|
|
CachedRequestData,
|
|
DiffusionSchedulerOutput,
|
|
NewRequestData,
|
|
)
|
|
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
|
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
|
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
|
from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
|
|
from vllm_omni.errors import OmniClientError
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers & fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@contextmanager
|
|
def _noop_forward_context(*args, **kwargs):
|
|
del args, kwargs
|
|
yield
|
|
|
|
|
|
def _update_environment_variables(envs_dict: dict[str, str]) -> None:
|
|
for key, value in envs_dict.items():
|
|
os.environ[key] = value
|
|
|
|
|
|
class _StepPipeline:
|
|
"""Minimal pipeline stub that supports step-wise execution."""
|
|
|
|
supports_step_execution = True
|
|
|
|
def __init__(self):
|
|
self.prepare_calls = 0
|
|
self.denoise_calls = 0
|
|
self.scheduler_calls = 0
|
|
self.decode_calls = 0
|
|
|
|
def prepare_encode(self, state, **kwargs):
|
|
del kwargs
|
|
self.prepare_calls += 1
|
|
state.timesteps = [torch.tensor(10), torch.tensor(5)]
|
|
state.latents = torch.tensor([0.0])
|
|
state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
|
|
return state
|
|
|
|
def denoise_step(self, input_batch, **kwargs):
|
|
self.denoise_calls += 1
|
|
return torch.full_like(input_batch.prompt_embeds, fill_value=0.5)
|
|
|
|
def step_scheduler(self, state, noise_pred, **kwargs):
|
|
del noise_pred, kwargs
|
|
self.scheduler_calls += 1
|
|
state.step_index += 1
|
|
|
|
def post_decode(self, state, **kwargs):
|
|
del kwargs
|
|
self.decode_calls += 1
|
|
return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
|
|
|
|
|
|
class _ChunkedStepPipeline(_StepPipeline):
|
|
"""Streaming stub: two chunks of two denoise steps each, one decode per chunk."""
|
|
|
|
supports_chunk_step_grouping = False
|
|
|
|
def prepare_encode(self, state, **kwargs):
|
|
del kwargs
|
|
self.prepare_calls += 1
|
|
state.timesteps = [torch.tensor(10), torch.tensor(5), torch.tensor(10), torch.tensor(5)]
|
|
state.latents = torch.tensor([0.0])
|
|
state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
|
|
state.chunk_num_steps = 2
|
|
state.total_chunks = 2
|
|
state.chunk_index = 0
|
|
state.step_in_chunk = 0
|
|
return state
|
|
|
|
def step_scheduler(self, state, noise_pred, **kwargs):
|
|
del noise_pred, kwargs
|
|
self.scheduler_calls += 1
|
|
state.step_in_chunk += 1
|
|
state.step_index = state.step_in_chunk
|
|
|
|
def post_decode(self, state, **kwargs):
|
|
del kwargs
|
|
self.decode_calls += 1
|
|
result = DiffusionOutput(output=torch.tensor([float(state.chunk_index)]))
|
|
state.chunk_index += 1
|
|
state.step_in_chunk = 0
|
|
return result
|
|
|
|
|
|
class _BatchDependentChunkPipeline(_ChunkedStepPipeline):
|
|
"""Streaming stub whose noise depends on the batch's latents and timesteps, so a stale batch shows."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.seen: list[tuple[float, float]] = []
|
|
|
|
def prepare_encode(self, state, **kwargs):
|
|
super().prepare_encode(state, **kwargs)
|
|
state.latents = torch.tensor([1.0])
|
|
return state
|
|
|
|
def denoise_step(self, input_batch, **kwargs):
|
|
self.denoise_calls += 1
|
|
latent = float(input_batch.latents.reshape(-1)[0])
|
|
timestep = float(input_batch.timesteps.reshape(-1)[0])
|
|
self.seen.append((latent, timestep))
|
|
return torch.full((input_batch.latents.shape[0], 1), latent + timestep)
|
|
|
|
def step_scheduler(self, state, noise_pred, **kwargs):
|
|
del kwargs
|
|
self.scheduler_calls += 1
|
|
state.latents = state.latents + noise_pred.reshape(-1)[0]
|
|
state.step_in_chunk += 1
|
|
state.step_index = state.step_in_chunk
|
|
|
|
def post_decode(self, state, **kwargs):
|
|
self.final_latent = state.latents.reshape(-1)[0].clone()
|
|
return super().post_decode(state, **kwargs)
|
|
|
|
|
|
class _FailsMidChunkPipeline(_ChunkedStepPipeline):
|
|
"""Streaming stub whose second denoise step of the first chunk raises."""
|
|
|
|
def step_scheduler(self, state, noise_pred, **kwargs):
|
|
super().step_scheduler(state, noise_pred, **kwargs)
|
|
if state.chunk_index == 0 and state.step_in_chunk == 2:
|
|
raise RuntimeError("boom mid chunk")
|
|
|
|
|
|
class _PerRequestErrorStepPipeline(_StepPipeline):
|
|
def prepare_encode(self, state, **kwargs):
|
|
if state.prompt == "fail-prepare":
|
|
raise OmniClientError(
|
|
"invalid request input",
|
|
status_code=422,
|
|
error_type="UnprocessableEntityError",
|
|
)
|
|
return super().prepare_encode(state, **kwargs)
|
|
|
|
def denoise_step(self, input_batch, **kwargs):
|
|
del kwargs
|
|
self.denoise_calls += 1
|
|
return torch.ones_like(input_batch.latents)
|
|
|
|
|
|
class _ProfilingStepPipeline(_StepPipeline):
|
|
enable_diffusion_pipeline_profiler = True
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._stage_durations: dict[str, float] = {}
|
|
|
|
@property
|
|
def stage_durations(self) -> dict[str, float]:
|
|
return dict(self._stage_durations)
|
|
|
|
def clear_profiler_records(self) -> None:
|
|
self._stage_durations.clear()
|
|
|
|
def prepare_encode(self, state, **kwargs):
|
|
result = super().prepare_encode(state, **kwargs)
|
|
self._stage_durations["QwenImagePipeline.text_encoder.forward"] = 1.0
|
|
return result
|
|
|
|
def denoise_step(self, input_batch, **kwargs):
|
|
result = super().denoise_step(input_batch, **kwargs)
|
|
self._stage_durations["QwenImagePipeline.diffuse"] = 2.0
|
|
return result
|
|
|
|
def post_decode(self, state, **kwargs):
|
|
result = super().post_decode(state, **kwargs)
|
|
self._stage_durations["QwenImagePipeline.vae.decode"] = 3.0
|
|
return result
|
|
|
|
|
|
class _AutoDenoiseProfilerPipeline(DiffusionPipelineProfilerMixin):
|
|
_PROFILER_TARGETS: list[str] = []
|
|
|
|
def __init__(self):
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=True,
|
|
)
|
|
|
|
def forward(self):
|
|
return None
|
|
|
|
def denoise_step(self):
|
|
return "ok"
|
|
|
|
|
|
class _InterruptingStepPipeline(_StepPipeline):
|
|
interrupt = True
|
|
|
|
def denoise_step(self, state, **kwargs):
|
|
del state, kwargs
|
|
self.denoise_calls += 1
|
|
return None
|
|
|
|
def step_scheduler(self, state, noise_pred, **kwargs):
|
|
del state, noise_pred, kwargs
|
|
raise AssertionError("step_scheduler should not run after interrupt")
|
|
|
|
def post_decode(self, state, **kwargs):
|
|
del state, kwargs
|
|
raise AssertionError("post_decode should not run after interrupt")
|
|
|
|
|
|
class _FakePeakMemoryPlatform:
|
|
def __init__(self, reserved_mb: list[float]):
|
|
self._reserved_mb = reserved_mb
|
|
self.reset_calls = 0
|
|
|
|
def reset_peak_memory_stats(self):
|
|
self.reset_calls += 1
|
|
|
|
def max_memory_reserved(self):
|
|
index = min(self.reset_calls - 1, len(self._reserved_mb) - 1)
|
|
return int(self._reserved_mb[index] * 1024**2)
|
|
|
|
def max_memory_allocated(self):
|
|
index = min(self.reset_calls - 1, len(self._reserved_mb) - 1)
|
|
return int((self._reserved_mb[index] - 100) * 1024**2)
|
|
|
|
def is_available(self) -> bool:
|
|
return True
|
|
|
|
|
|
class _IdentityNoiseTransformer(torch.nn.Module):
|
|
def forward(self, x: torch.Tensor, **kwargs):
|
|
del kwargs
|
|
return (x,)
|
|
|
|
|
|
class _AdditiveScheduler:
|
|
def step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, return_dict: bool = False):
|
|
del t, return_dict
|
|
return (latents + noise_pred,)
|
|
|
|
|
|
class _DistributedStepPipeline(CFGParallelMixin):
|
|
supports_step_execution = True
|
|
|
|
def __init__(self, mode: str, device: torch.device):
|
|
self.mode = mode
|
|
self.device = device
|
|
self._interrupt = False
|
|
self.scheduler = _AdditiveScheduler()
|
|
self.transformer = _IdentityNoiseTransformer()
|
|
|
|
@property
|
|
def interrupt(self):
|
|
return self._interrupt
|
|
|
|
def prepare_encode(self, state, **kwargs):
|
|
del kwargs
|
|
state.timesteps = [torch.tensor(1.0, device=self.device)]
|
|
state.latents = torch.ones((1, 1), device=self.device)
|
|
state.step_index = 0
|
|
state.scheduler = self.scheduler
|
|
state.do_true_cfg = self.mode == "cfg"
|
|
state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
|
|
return state
|
|
|
|
def denoise_step(self, state, **kwargs):
|
|
del kwargs
|
|
if self.mode == "ulysses":
|
|
sp_group = get_sp_group().ulysses_group
|
|
seq_world_size = torch.distributed.get_world_size(sp_group)
|
|
input_tensor = torch.randn(1, 2, 2 * seq_world_size, 2, device=self.device)
|
|
original = input_tensor.clone()
|
|
intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, False)
|
|
output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, False)
|
|
torch.testing.assert_close(output, original, rtol=1e-5, atol=1e-5)
|
|
return torch.ones_like(state.latents)
|
|
|
|
if self.mode == "ring":
|
|
ring_group = get_sp_group().ring_group
|
|
rank = torch.distributed.get_rank(ring_group)
|
|
world_size = torch.distributed.get_world_size(ring_group)
|
|
comm = RingComm(ring_group)
|
|
input_tensor = torch.full((1, 2, 2), float(rank + 1), device=self.device)
|
|
recv_tensor = comm.send_recv(input_tensor)
|
|
comm.commit()
|
|
comm.wait()
|
|
expected = torch.full_like(recv_tensor, float(((rank - 1) % world_size) + 1))
|
|
torch.testing.assert_close(recv_tensor, expected, rtol=1e-5, atol=1e-5)
|
|
return torch.ones_like(state.latents)
|
|
|
|
positive_kwargs = {"x": state.latents + 1}
|
|
negative_kwargs = {"x": state.latents - 1}
|
|
return self.predict_noise_maybe_with_cfg(
|
|
do_true_cfg=True,
|
|
true_cfg_scale=1.0,
|
|
positive_kwargs=positive_kwargs,
|
|
negative_kwargs=negative_kwargs,
|
|
cfg_normalize=False,
|
|
)
|
|
|
|
def step_scheduler(self, state, noise_pred, **kwargs):
|
|
del kwargs
|
|
if self.mode == "cfg":
|
|
state.latents = self.scheduler_step_maybe_with_cfg(
|
|
noise_pred,
|
|
state.current_timestep,
|
|
state.latents,
|
|
do_true_cfg=True,
|
|
per_request_scheduler=state.scheduler,
|
|
)
|
|
else:
|
|
state.latents = state.latents + noise_pred
|
|
state.step_index += 1
|
|
|
|
def post_decode(self, state, **kwargs):
|
|
del kwargs
|
|
return DiffusionOutput(output=state.latents.detach().cpu())
|
|
|
|
|
|
def _make_step_request(num_inference_steps: int = 2):
|
|
return OmniDiffusionRequest(
|
|
prompt="a prompt",
|
|
request_id="req-1",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
|
|
)
|
|
|
|
|
|
def _assert_aborted_output(output: DiffusionOutput, request_id: str) -> None:
|
|
assert output.output is None
|
|
assert output.error is None
|
|
assert output.aborted is True
|
|
assert output.abort_message == f"Request {request_id} aborted."
|
|
|
|
|
|
def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) -> OmniDiffusionRequest:
|
|
return OmniDiffusionRequest(
|
|
prompt=f"prompt-{req_id}",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
|
|
request_id=req_id,
|
|
)
|
|
|
|
|
|
def _make_vllm_config():
|
|
@contextlib.contextmanager
|
|
def set_priority(*args, **kwargs):
|
|
yield
|
|
|
|
return SimpleNamespace(
|
|
kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)),
|
|
compilation_config=SimpleNamespace(ir_enable_torch_wrap=True),
|
|
)
|
|
|
|
|
|
def _make_runner(
|
|
diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY,
|
|
runner_cls: type[DiffusionModelRunner] = DiffusionModelRunner,
|
|
):
|
|
runner = object.__new__(runner_cls)
|
|
runner.vllm_config = _make_vllm_config()
|
|
runner.od_config = SimpleNamespace(
|
|
cache_backend=None,
|
|
diffusion_kv_mode=diffusion_kv_mode,
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
streaming_output=False,
|
|
)
|
|
runner.device = torch.device("cpu")
|
|
runner.pipeline = _StepPipeline()
|
|
runner.cache_backend = None
|
|
runner.offload_backend = None
|
|
runner.state_cache = {}
|
|
runner.kv_transfer_manager = SimpleNamespace(
|
|
receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None
|
|
)
|
|
return runner
|
|
|
|
|
|
def _make_distributed_runner(mode: str, device: torch.device):
|
|
runner = object.__new__(DiffusionModelRunner)
|
|
runner.vllm_config = _make_vllm_config()
|
|
runner.od_config = SimpleNamespace(
|
|
cache_backend=None,
|
|
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
streaming_output=False,
|
|
)
|
|
runner.device = device
|
|
runner.pipeline = _DistributedStepPipeline(mode=mode, device=device)
|
|
runner.cache_backend = None
|
|
runner.offload_backend = None
|
|
runner.state_cache = {}
|
|
runner.kv_transfer_manager = SimpleNamespace(
|
|
receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None
|
|
)
|
|
return runner
|
|
|
|
|
|
def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None):
|
|
req.request_id = request_id
|
|
return DiffusionSchedulerOutput(
|
|
step_id=step_id,
|
|
scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)],
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
|
|
num_running_reqs=1,
|
|
num_waiting_reqs=0,
|
|
)
|
|
|
|
|
|
def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
|
|
"""Scheduler output for a homogeneous batch (one NewRequestData per req)."""
|
|
new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs]
|
|
return DiffusionSchedulerOutput(
|
|
step_id=step_id,
|
|
scheduled_new_reqs=new_reqs,
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
|
|
num_running_reqs=len(new_reqs),
|
|
num_waiting_reqs=0,
|
|
)
|
|
|
|
|
|
def _make_input_batch_state(request_id: str, latent_value: float) -> StepRequestState:
|
|
state = StepRequestState(
|
|
request_id=request_id,
|
|
sampling=SimpleNamespace(),
|
|
prompt=None,
|
|
)
|
|
state.latents = torch.tensor([[latent_value]])
|
|
state.timesteps = torch.tensor([1.0])
|
|
return state
|
|
|
|
|
|
def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
|
|
return DiffusionSchedulerOutput(
|
|
step_id=step_id,
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
|
|
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
|
|
num_running_reqs=1,
|
|
num_waiting_reqs=0,
|
|
)
|
|
|
|
|
|
def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
|
|
engine = object.__new__(DiffusionEngine)
|
|
engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline", streaming_output=False)
|
|
engine.pre_process_func = None
|
|
engine.post_process_func = None
|
|
engine.scheduler = scheduler
|
|
engine.execute_fn = execute_fn
|
|
engine._rpc_lock = threading.RLock()
|
|
engine._cv = threading.Condition(engine._rpc_lock)
|
|
engine._closed = False
|
|
engine.abort_queue = queue.Queue()
|
|
return engine
|
|
|
|
|
|
def _expected_output_for_mode(mode: str) -> torch.Tensor:
|
|
if mode == "cfg":
|
|
return torch.tensor([[3.0]])
|
|
return torch.tensor([[2.0]])
|
|
|
|
|
|
def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master_port: str):
|
|
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
|
|
current_omni_platform.set_device(device)
|
|
_update_environment_variables(
|
|
{
|
|
"RANK": str(local_rank),
|
|
"LOCAL_RANK": str(local_rank),
|
|
"WORLD_SIZE": str(world_size),
|
|
"MASTER_ADDR": "localhost",
|
|
"MASTER_PORT": master_port,
|
|
}
|
|
)
|
|
model_runner_module.set_forward_context = _noop_forward_context
|
|
|
|
try:
|
|
init_distributed_environment()
|
|
if mode == "ulysses":
|
|
initialize_model_parallel(ulysses_degree=world_size)
|
|
elif mode == "ring":
|
|
initialize_model_parallel(ring_degree=world_size)
|
|
elif mode == "cfg":
|
|
initialize_model_parallel(cfg_parallel_size=world_size)
|
|
else:
|
|
raise ValueError(f"Unsupported distributed test mode: {mode}")
|
|
|
|
runner = _make_distributed_runner(mode, device)
|
|
result = DiffusionModelRunner.execute_stepwise(
|
|
runner,
|
|
_make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
|
|
)
|
|
output = result.get_request_output("req-1")
|
|
|
|
assert output.finished is True
|
|
assert output.result is not None
|
|
torch.testing.assert_close(output.result.output, _expected_output_for_mode(mode), rtol=1e-5, atol=1e-5)
|
|
assert "req-1" not in runner.state_cache
|
|
finally:
|
|
destroy_distributed_env()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runner / Worker
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_input_batch_cached_repack_refreshes_state_references_without_prompt_embeds():
|
|
first_state = _make_input_batch_state("req-1", 1.0)
|
|
batch = InputBatch.make_batch([first_state])
|
|
assert batch.prompt_embeds is None
|
|
|
|
replacement_state = _make_input_batch_state("req-1", 2.0)
|
|
repacked = InputBatch.make_batch([replacement_state], cached_batch=batch)
|
|
|
|
assert repacked is batch
|
|
assert repacked.states[0] is replacement_state
|
|
assert repacked.prompt_embeds is None
|
|
torch.testing.assert_close(repacked.latents, replacement_state.latents)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_composition():
|
|
first_state = _make_input_batch_state("req-1", 1.0)
|
|
first_state.prompt_embeds = torch.ones(1, 2, 3)
|
|
first_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool)
|
|
batch = InputBatch.make_batch([first_state])
|
|
|
|
replacement_state = _make_input_batch_state("req-1", 2.0)
|
|
replacement_state.prompt_embeds = torch.full((1, 2, 3), 2.0)
|
|
replacement_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool)
|
|
repacked = InputBatch.make_batch([replacement_state], cached_batch=batch)
|
|
|
|
assert repacked.states[0] is replacement_state
|
|
torch.testing.assert_close(repacked.latents, replacement_state.latents)
|
|
torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_make_batch_identity_mapping_never_reads_the_device_tensor_back(monkeypatch):
|
|
"""The default index mapping is built on the host; reading it back from the device would sync every step."""
|
|
|
|
def forbidden(self, *args, **kwargs):
|
|
raise AssertionError("make_batch read a tensor back to the host")
|
|
|
|
monkeypatch.setattr(torch.Tensor, "tolist", forbidden)
|
|
monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
|
|
states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)]
|
|
|
|
batch = InputBatch.make_batch(states)
|
|
repacked = InputBatch.make_batch(states, cached_batch=batch)
|
|
|
|
assert repacked is batch
|
|
assert list(batch.idx_mapping_np) == [0, 1]
|
|
assert batch.idx_mapping.dtype is torch.int32 and batch.idx_mapping.shape == (2,)
|
|
assert batch.request_ids == ["req-1", "req-2"]
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_make_batch_explicit_mapping_still_selects_and_orders_states():
|
|
states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)]
|
|
|
|
batch = InputBatch.make_batch(states, idx_mapping=torch.tensor([1, 0]))
|
|
|
|
assert batch.request_ids == ["req-2", "req-1"]
|
|
assert list(batch.idx_mapping_np) == [1, 0]
|
|
torch.testing.assert_close(batch.latents, torch.tensor([[2.0], [1.0]]))
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_step_profiler_reports_denoise_step_as_diffuse():
|
|
pipeline = _AutoDenoiseProfilerPipeline()
|
|
|
|
assert pipeline.denoise_step() == "ok"
|
|
|
|
assert any(key.endswith(".diffuse") for key in pipeline.stage_durations)
|
|
assert not any(key.endswith(".denoise_step") for key in pipeline.stage_durations)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
class TestRunner:
|
|
"""DiffusionModelRunner.execute_stepwise"""
|
|
|
|
@pytest.mark.parametrize(
|
|
("diffusion_kv_mode", "should_cleanup"),
|
|
[
|
|
(DiffusionKVCacheMode.DENSE_LEGACY, False),
|
|
(DiffusionKVCacheMode.PAGED_SCHEDULER, True),
|
|
],
|
|
)
|
|
def test_scheduler_finished_cleanup_only_runs_for_paged_kv(
|
|
self,
|
|
mocker: MockerFixture,
|
|
diffusion_kv_mode: DiffusionKVCacheMode,
|
|
should_cleanup: bool,
|
|
):
|
|
runner = _make_runner(diffusion_kv_mode)
|
|
cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
|
|
expected_output = object()
|
|
mocker.patch.object(runner, "_execute_stepwise_core", return_value=expected_output)
|
|
scheduler_output = DiffusionSchedulerOutput(
|
|
step_id=1,
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
finished_req_ids={"finished-request"},
|
|
num_running_reqs=0,
|
|
num_waiting_reqs=0,
|
|
)
|
|
|
|
output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
|
|
|
assert output is expected_output
|
|
if should_cleanup:
|
|
cleanup.assert_called_once_with(["finished-request"])
|
|
else:
|
|
cleanup.assert_not_called()
|
|
|
|
@pytest.mark.parametrize(
|
|
("diffusion_kv_mode", "should_cleanup"),
|
|
[
|
|
(DiffusionKVCacheMode.DENSE_LEGACY, False),
|
|
(DiffusionKVCacheMode.PAGED_SCHEDULER, True),
|
|
],
|
|
)
|
|
def test_terminal_cleanup_only_runs_for_paged_kv(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
mocker: MockerFixture,
|
|
diffusion_kv_mode: DiffusionKVCacheMode,
|
|
should_cleanup: bool,
|
|
):
|
|
runner = _make_runner(diffusion_kv_mode)
|
|
cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
|
|
install = mocker.patch.object(runner, "install_diffusion_kv_metadata")
|
|
req = _make_step_request()
|
|
scheduler_output = _make_scheduler_output(req, step_id=0)
|
|
if should_cleanup:
|
|
scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata = DiffusionKVMetadata(
|
|
request_id=req.request_id,
|
|
allocation_generation=1,
|
|
sequences=(),
|
|
)
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
|
|
first_output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
|
assert first_output.get_request_output(req.request_id).finished is False
|
|
cleanup.assert_not_called()
|
|
|
|
output = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
|
|
|
|
assert output.get_request_output(req.request_id).finished is True
|
|
if should_cleanup:
|
|
install.assert_called_once_with(scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata)
|
|
cleanup.assert_called_once_with([req.request_id])
|
|
else:
|
|
install.assert_not_called()
|
|
cleanup.assert_not_called()
|
|
|
|
def test_completes_request_and_clears_state(self, monkeypatch):
|
|
runner = _make_runner()
|
|
req = _make_step_request()
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
|
|
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
|
first = result.get_request_output("req-1")
|
|
assert first.request_id == "req-1"
|
|
assert first.step_index == 1
|
|
assert first.finished is False
|
|
assert first.result is None
|
|
assert "req-1" in runner.state_cache
|
|
|
|
result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
|
|
second = result.get_request_output("req-1")
|
|
assert second.request_id == "req-1"
|
|
assert second.step_index == 2
|
|
assert second.finished is True
|
|
assert second.result is not None
|
|
assert second.result.error is None
|
|
assert torch.equal(second.result.output, torch.tensor([2.0]))
|
|
assert "req-1" not in runner.state_cache
|
|
|
|
assert runner.pipeline.prepare_calls == 1
|
|
assert runner.pipeline.denoise_calls == 2
|
|
assert runner.pipeline.scheduler_calls == 2
|
|
assert runner.pipeline.decode_calls == 1
|
|
|
|
@staticmethod
|
|
def _make_grouping_runner(pipeline, *, grouping: bool):
|
|
"""An AR runner without KV (so it inherits the stepwise entry) over a streaming stub pipeline."""
|
|
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
|
|
|
|
runner = _make_runner(runner_cls=ARDiffusionModelRunner)
|
|
runner.od_config.streaming_output = True
|
|
runner.kv_cache = None
|
|
runner.pipeline = pipeline
|
|
runner.pipeline.supports_chunk_step_grouping = grouping
|
|
return runner
|
|
|
|
@staticmethod
|
|
def _drive_to_completion(runner, num_steps=4):
|
|
outputs = [runner.execute_stepwise(_make_scheduler_output(_make_step_request(num_steps)))]
|
|
step = 1
|
|
while not outputs[-1].get_request_output("req-1").finished:
|
|
outputs.append(runner.execute_stepwise(_make_cached_scheduler_output(step_id=step)))
|
|
step += 1
|
|
return [o.get_request_output("req-1") for o in outputs]
|
|
|
|
@pytest.mark.parametrize("grouping", [False, True])
|
|
def test_ar_runner_runs_every_step_of_a_chunk_in_one_call_when_the_pipeline_declares_it(
|
|
self, monkeypatch, grouping
|
|
):
|
|
"""With the capability declared, one AR runner call drives a whole chunk; the outputs are the same either way."""
|
|
runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=grouping)
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
denoise_calls_after_each_call = []
|
|
original = runner.execute_stepwise
|
|
|
|
def counting(scheduler_output):
|
|
output = original(scheduler_output)
|
|
denoise_calls_after_each_call.append(runner.pipeline.denoise_calls)
|
|
return output
|
|
|
|
runner.execute_stepwise = counting
|
|
|
|
per_call = self._drive_to_completion(runner)
|
|
|
|
# Chunk outputs, and the steps they were emitted at, are identical.
|
|
emitted = [(o.step_index, float(o.result.output[0])) for o in per_call if o.result is not None]
|
|
assert emitted == [(2, 0.0), (2, 1.0)]
|
|
assert per_call[-1].finished is True
|
|
assert runner.pipeline.denoise_calls == 4 and runner.pipeline.scheduler_calls == 4
|
|
assert runner.pipeline.decode_calls == 2
|
|
assert runner.pipeline.prepare_calls == 1
|
|
# One call per chunk with the capability, one per step without; a
|
|
# grouped call never crosses into the next chunk.
|
|
assert denoise_calls_after_each_call == ([2, 4] if grouping else [1, 2, 3, 4])
|
|
assert "req-1" not in runner.state_cache
|
|
|
|
def test_ar_runner_continues_a_grouped_chunk_as_a_cached_request(self, monkeypatch):
|
|
"""The first step admits the new request; every further step of the chunk is a cached-request continuation."""
|
|
runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=True)
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
seen = []
|
|
original_core = DiffusionModelRunner._execute_stepwise_core
|
|
|
|
def recording_core(self_, scheduler_output, **kwargs):
|
|
seen.append(scheduler_output)
|
|
return original_core(self_, scheduler_output, **kwargs)
|
|
|
|
monkeypatch.setattr(DiffusionModelRunner, "_execute_stepwise_core", recording_core)
|
|
|
|
first = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4), finished_req_ids={"old"}))
|
|
|
|
assert first.get_request_output("req-1").result is not None
|
|
assert runner.pipeline.prepare_calls == 1
|
|
assert len(seen) == 2
|
|
assert [n.request_id for n in seen[0].scheduled_new_reqs] == ["req-1"]
|
|
continuation = seen[1]
|
|
assert continuation.scheduled_new_reqs == []
|
|
assert continuation.scheduled_cached_reqs.request_ids == ["req-1"]
|
|
assert continuation.finished_req_ids == set()
|
|
assert continuation.kv_prefetch_job is None and continuation.kv_connector_metadata is None
|
|
assert continuation.step_id == seen[0].step_id
|
|
|
|
def test_ar_runner_grouped_steps_see_the_same_inputs_as_per_step_calls(self, monkeypatch):
|
|
"""Every grouped step sees the latents and timestep step_scheduler just advanced, exactly as per-step calls do."""
|
|
finals = {}
|
|
seen = {}
|
|
for grouping in (False, True):
|
|
runner = self._make_grouping_runner(_BatchDependentChunkPipeline(), grouping=grouping)
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
self._drive_to_completion(runner)
|
|
finals[grouping] = float(runner.pipeline.final_latent)
|
|
seen[grouping] = runner.pipeline.seen
|
|
|
|
# 1 -> +(1+10)=12 -> +(12+5)=29 -> +(29+10)=68 -> +(68+5)=141 in both modes.
|
|
assert seen[False] == [(1.0, 10.0), (12.0, 5.0), (29.0, 10.0), (68.0, 5.0)]
|
|
assert seen[True] == seen[False]
|
|
assert finals[True] == finals[False] == 141.0
|
|
|
|
def test_ar_runner_grouped_chunk_stops_at_a_failing_step_and_drops_the_state(self, monkeypatch):
|
|
"""A failure mid-chunk ends the grouped call with that step's error; nothing runs after it."""
|
|
runner = self._make_grouping_runner(_FailsMidChunkPipeline(), grouping=True)
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
|
|
output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4)))
|
|
|
|
failed = output.get_request_output("req-1")
|
|
assert failed.finished is True
|
|
assert failed.result is not None and "boom mid chunk" in failed.result.error
|
|
assert runner.pipeline.denoise_calls == 2 and runner.pipeline.decode_calls == 0
|
|
assert "req-1" not in runner.state_cache
|
|
|
|
def test_ar_runner_fails_a_chunk_that_never_advances(self, monkeypatch):
|
|
pipeline = _ChunkedStepPipeline()
|
|
monkeypatch.setattr(pipeline, "step_scheduler", lambda *args, **kwargs: None)
|
|
runner = self._make_grouping_runner(pipeline, grouping=True)
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4)))
|
|
failed = output.get_request_output("req-1")
|
|
assert failed.finished
|
|
assert "within 2 denoise steps" in failed.result.error
|
|
assert pipeline.denoise_calls == 2
|
|
assert "req-1" not in runner.state_cache
|
|
|
|
def test_base_runner_never_groups_and_the_ar_policy_needs_capability_streaming_and_one_request(self, monkeypatch):
|
|
"""The shared runner is one step per call; the AR runner groups only a lone streaming request that declares it."""
|
|
base = _make_runner()
|
|
base.od_config.streaming_output = True
|
|
base.pipeline = _ChunkedStepPipeline()
|
|
base.pipeline.supports_chunk_step_grouping = True
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
first = DiffusionModelRunner.execute_stepwise(base, _make_scheduler_output(_make_step_request(4)))
|
|
assert first.get_request_output("req-1").result is None and base.pipeline.denoise_calls == 1
|
|
|
|
ar = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=False)
|
|
one = _make_cached_scheduler_output()
|
|
two = _make_batch_scheduler_output([_make_step_request(4), _make_step_request(4)])
|
|
two.scheduled_new_reqs[1].request_id = two.scheduled_new_reqs[1].req.request_id = "req-2"
|
|
assert ar._groups_chunk_steps(one) is False # no capability declared
|
|
ar.pipeline.supports_chunk_step_grouping = True
|
|
assert ar._groups_chunk_steps(one) is True
|
|
assert ar._groups_chunk_steps(two) is False # not for a batch
|
|
ar.od_config.streaming_output = False
|
|
assert ar._groups_chunk_steps(one) is False # not without streaming output
|
|
|
|
def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch):
|
|
runner = _make_runner()
|
|
runner.pipeline = _ProfilingStepPipeline()
|
|
req = _make_step_request()
|
|
reset_calls = []
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
monkeypatch.setattr(
|
|
model_runner_module.current_omni_platform,
|
|
"is_available",
|
|
lambda: True,
|
|
)
|
|
monkeypatch.setattr(
|
|
model_runner_module.current_omni_platform,
|
|
"reset_peak_memory_stats",
|
|
lambda: reset_calls.append(True),
|
|
)
|
|
monkeypatch.setattr(
|
|
model_runner_module.current_omni_platform,
|
|
"max_memory_reserved",
|
|
lambda: 2 * 1024**2,
|
|
)
|
|
monkeypatch.setattr(
|
|
model_runner_module.current_omni_platform,
|
|
"max_memory_allocated",
|
|
lambda: 1024**2,
|
|
)
|
|
|
|
DiffusionModelRunner.execute_stepwise(
|
|
runner,
|
|
_make_scheduler_output(req, step_id=0),
|
|
)
|
|
result = DiffusionModelRunner.execute_stepwise(
|
|
runner,
|
|
_make_cached_scheduler_output(step_id=1),
|
|
)
|
|
|
|
output = result.get_request_output("req-1")
|
|
assert output.finished is True
|
|
assert output.result is not None
|
|
assert output.result.peak_memory_mb == 2
|
|
assert output.result.stage_durations == {
|
|
"QwenImagePipeline.text_encoder.forward": 1.0,
|
|
"QwenImagePipeline.diffuse": 4.0,
|
|
"QwenImagePipeline.vae.decode": 3.0,
|
|
}
|
|
assert reset_calls == [True]
|
|
|
|
def test_carries_peak_memory_across_stepwise_request_lifecycle(self, monkeypatch):
|
|
runner = _make_runner()
|
|
req = _make_step_request()
|
|
fake_platform = _FakePeakMemoryPlatform([1500.0, 1200.0])
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
monkeypatch.setattr(model_runner_module, "current_omni_platform", fake_platform)
|
|
|
|
first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
|
first_output = first.get_request_output("req-1")
|
|
assert first_output.finished is False
|
|
assert first_output.result is None
|
|
|
|
second = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
|
|
second_output = second.get_request_output("req-1")
|
|
assert second_output.finished is True
|
|
assert second_output.result is not None
|
|
assert second_output.result.peak_memory_mb == pytest.approx(1500.0)
|
|
|
|
def test_rejects_multi_request_step_batch(self):
|
|
runner = _make_runner()
|
|
req_1 = _make_step_request()
|
|
req_2 = _make_step_request()
|
|
req_2.request_id = "req-2"
|
|
|
|
scheduler_output = DiffusionSchedulerOutput(
|
|
step_id=0,
|
|
scheduled_new_reqs=[
|
|
NewRequestData(request_id="req-1", req=req_1),
|
|
NewRequestData(request_id="req-2", req=req_2),
|
|
],
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
finished_req_ids=set(),
|
|
num_running_reqs=2,
|
|
num_waiting_reqs=0,
|
|
)
|
|
|
|
result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
|
assert len(result) == 2
|
|
|
|
def test_prepare_error_isolated_from_healthy_request(self, monkeypatch):
|
|
runner = _make_runner()
|
|
runner.pipeline = _PerRequestErrorStepPipeline()
|
|
bad_req = _make_step_request()
|
|
bad_req.request_id = "req-bad"
|
|
bad_req.prompt = "fail-prepare"
|
|
good_req = _make_step_request()
|
|
good_req.request_id = "req-good"
|
|
good_req.prompt = "valid"
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
|
|
result = DiffusionModelRunner.execute_stepwise(
|
|
runner,
|
|
_make_batch_scheduler_output([bad_req, good_req]),
|
|
)
|
|
|
|
bad_output = result.get_request_output("req-bad")
|
|
assert bad_output.finished is True
|
|
assert bad_output.result is not None
|
|
assert bad_output.result.error == "invalid request input"
|
|
assert bad_output.result.error_status_code == 422
|
|
assert bad_output.result.error_type == "UnprocessableEntityError"
|
|
good_output = result.get_request_output("req-good")
|
|
assert good_output.finished is False
|
|
assert good_output.step_index == 1
|
|
assert good_output.result is None
|
|
assert set(runner.state_cache) == {"req-good"}
|
|
assert runner.pipeline.denoise_calls == 1
|
|
assert runner.pipeline.scheduler_calls == 1
|
|
|
|
def test_prepare_error_on_other_rank_is_isolated_via_all_reduce(self, monkeypatch):
|
|
"""P1-B: rank-0-only pipeline work (e.g. H3 reference-video prep)
|
|
can raise on rank 0 while non-zero ranks return cleanly and then hang
|
|
on a downstream ``dist.broadcast``. Beyond the pipeline-side guard,
|
|
the runner all-reduces a per-request failure flag across the DiT
|
|
group so a rank that ``prepare_encode`` returned cleanly on still
|
|
skips the request and does not enter the denoise step batch.
|
|
"""
|
|
runner = _make_runner()
|
|
runner.pipeline = _StepPipeline() # locally always succeeds
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
|
|
# Simulate a peer DiT rank having reported failure. We patch the
|
|
# helper directly (rather than ``torch.distributed`` broadly) because
|
|
# ``execute_stepwise`` also queries ``torch.distributed.get_rank``.
|
|
monkeypatch.setattr(
|
|
model_runner_module,
|
|
"_dit_any_rank_failed",
|
|
lambda local_failed: True,
|
|
)
|
|
|
|
req = _make_step_request(num_inference_steps=1)
|
|
result = DiffusionModelRunner.execute_stepwise(
|
|
runner,
|
|
_make_scheduler_output(req),
|
|
)
|
|
|
|
output = result.get_request_output("req-1")
|
|
assert output.finished is True
|
|
assert output.result is not None
|
|
assert output.result.error is not None
|
|
assert "another DiT rank" in output.result.error
|
|
# A rank that skipped the request must not have driven a denoise step.
|
|
assert runner.pipeline.denoise_calls == 0
|
|
assert runner.state_cache == {}
|
|
|
|
def test_peer_prepare_encode_failure_skips_interaction_session_initialization(self, monkeypatch):
|
|
"""One rank's failure in ``_prepare_batch_inputs -> prepare_encode`` must be handled before all ranks initialize interaction sessions.
|
|
|
|
``maybe_prepare_initial_session`` calls ``synchronized_monotonic_time()``
|
|
(a broadcast). If this rank's ``prepare_encode`` succeeded while a peer
|
|
failed, entering that broadcast while the peer enters the failure
|
|
all-reduce would hang.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from vllm_omni.diffusion.interaction.types import ChunkMediaSpec, InteractionChunkMetadata
|
|
|
|
class _InteractionStepPipeline(_StepPipeline):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.prepare_next_chunk_calls = 0
|
|
|
|
def peek_chunk_media(self, state):
|
|
del state
|
|
return ChunkMediaSpec(num_media_frames=8, fps=16.0, num_latent_frames=8)
|
|
|
|
def apply_interaction_at_chunk_boundary(self, state):
|
|
del state
|
|
|
|
def prepare_next_chunk(self, state):
|
|
del state
|
|
self.prepare_next_chunk_calls += 1
|
|
|
|
runner = _make_runner()
|
|
pipeline = _InteractionStepPipeline()
|
|
runner.pipeline = pipeline
|
|
coordinator = MagicMock()
|
|
coordinator.maybe_prepare_initial_session.return_value = InteractionChunkMetadata(
|
|
started_event_ids=[],
|
|
active_event_ids=[],
|
|
completed_event_ids=[],
|
|
)
|
|
runner._interaction_coordinator = coordinator
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
# Peer failed during prepare_encode; this rank's local encode succeeds.
|
|
monkeypatch.setattr(
|
|
model_runner_module,
|
|
"_dit_any_rank_failed",
|
|
lambda local_failed: True,
|
|
)
|
|
|
|
req = _make_step_request(num_inference_steps=1)
|
|
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req))
|
|
|
|
output = result.get_request_output("req-1")
|
|
assert output.finished is True
|
|
assert output.result is not None
|
|
assert "another DiT rank" in (output.result.error or "")
|
|
assert pipeline.prepare_calls == 1
|
|
coordinator.maybe_prepare_initial_session.assert_not_called()
|
|
assert pipeline.prepare_next_chunk_calls == 0
|
|
assert pipeline.denoise_calls == 0
|
|
assert runner.state_cache == {}
|
|
|
|
def test_receives_kv_payload_before_prepare_encode(self, monkeypatch):
|
|
runner = _make_runner()
|
|
captured: dict[str, object] = {}
|
|
kv_payload = object()
|
|
|
|
class _CapturingStepPipeline(_StepPipeline):
|
|
def prepare_encode(self, state, **kwargs):
|
|
captured["past_key_values"] = getattr(state.sampling, "past_key_values", None)
|
|
return super().prepare_encode(state, **kwargs)
|
|
|
|
class _KVTransferManager:
|
|
def receive_multi_kv_cache_distributed(self, req, cfg_kv_collect_func=None, target_device=None):
|
|
captured["cfg_kv_collect_func"] = cfg_kv_collect_func
|
|
captured["target_device"] = target_device
|
|
req.sampling_params.past_key_values = kv_payload
|
|
|
|
runner.pipeline = _CapturingStepPipeline()
|
|
runner.pipeline.device = torch.device("cpu")
|
|
runner.od_config.cfg_kv_collect_func = "collect-cfg"
|
|
runner.kv_transfer_manager = _KVTransferManager()
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
req = _make_step_request()
|
|
|
|
DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
|
|
|
assert captured["past_key_values"] is kv_payload
|
|
assert captured["cfg_kv_collect_func"] == "collect-cfg"
|
|
assert captured["target_device"] == torch.device("cpu")
|
|
assert getattr(req.sampling_params, "past_key_values", None) is None
|
|
|
|
def test_rejects_missing_cached_state(self):
|
|
runner = _make_runner()
|
|
|
|
with pytest.raises(ValueError, match="Missing cached state"):
|
|
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing"))
|
|
|
|
def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
|
|
runner = _make_runner()
|
|
runner.pipeline = _InterruptingStepPipeline()
|
|
req = _make_step_request()
|
|
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
|
|
|
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
|
output = result.get_request_output("req-1")
|
|
assert output.request_id == "req-1"
|
|
assert output.step_index == 0
|
|
assert output.finished is True
|
|
assert output.result is not None
|
|
assert output.result.error == "stepwise denoise interrupted"
|
|
assert "req-1" not in runner.state_cache
|
|
assert runner.pipeline.prepare_calls == 1
|
|
assert runner.pipeline.denoise_calls == 1
|
|
assert runner.pipeline.scheduler_calls == 0
|
|
assert runner.pipeline.decode_calls == 0
|
|
|
|
def test_load_model_rejects_unsupported_step_execution(self, monkeypatch):
|
|
class _RequestOnlyPipeline:
|
|
pass
|
|
|
|
class _FakeLoader:
|
|
def __init__(self, *args, **kwargs):
|
|
del args, kwargs
|
|
|
|
def load_model(self, **kwargs):
|
|
del kwargs
|
|
return _RequestOnlyPipeline()
|
|
|
|
class _FakeProfiler:
|
|
consumed_memory = 0
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
del exc_type, exc, tb
|
|
return False
|
|
|
|
runner = object.__new__(DiffusionModelRunner)
|
|
runner.vllm_config = _make_vllm_config()
|
|
runner.od_config = SimpleNamespace(
|
|
enable_cpu_offload=False,
|
|
enable_layerwise_offload=False,
|
|
enforce_eager=True,
|
|
cache_backend=None,
|
|
cache_config=None,
|
|
step_execution=True,
|
|
model_class_name="RequestOnlyPipeline",
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
streaming_output=False,
|
|
)
|
|
runner.device = torch.device("cpu")
|
|
runner.pipeline = None
|
|
runner.cache_backend = None
|
|
runner.offload_backend = None
|
|
runner.state_cache = {}
|
|
runner.kv_transfer_manager = SimpleNamespace()
|
|
|
|
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader)
|
|
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler)
|
|
monkeypatch.setattr(
|
|
model_runner_module,
|
|
"enable_offload_backend",
|
|
lambda _config, pipeline, **_kwargs: (pipeline, None),
|
|
)
|
|
monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
|
|
|
|
with pytest.raises(ValueError, match="RequestOnlyPipeline"):
|
|
DiffusionModelRunner.load_model(runner)
|
|
|
|
|
|
class _RecordingLoRAManager:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[object | None, float]] = []
|
|
|
|
def set_active_adapter(self, adapter, scale: float = 1.0) -> None:
|
|
self.calls.append((adapter, scale))
|
|
|
|
|
|
def _make_step_worker(lora_manager=None, *, expected_output=None):
|
|
"""Build a bare DiffusionWorker primed for execute_stepwise tests."""
|
|
worker = object.__new__(DiffusionWorker)
|
|
worker.lora_manager = lora_manager
|
|
worker._step_lora_state = {}
|
|
output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1")
|
|
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
|
|
return worker
|
|
|
|
|
|
@pytest.mark.cpu
|
|
class TestWorker:
|
|
"""DiffusionWorker.execute_stepwise"""
|
|
|
|
def test_delegates_to_model_runner(self):
|
|
expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None)
|
|
worker = _make_step_worker(expected_output=expected)
|
|
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
|
|
|
|
output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
|
|
|
assert output is expected
|
|
|
|
def test_deactivates_lora_when_request_has_no_adapter(self):
|
|
manager = _RecordingLoRAManager()
|
|
worker = _make_step_worker(lora_manager=manager)
|
|
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
|
|
|
|
DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
|
|
|
assert manager.calls == [(None, 1.0)]
|
|
|
|
def test_activates_lora_for_step_requests(self):
|
|
from vllm_omni.lora.request import LoRARequest
|
|
|
|
lora_request = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora")
|
|
request = _make_engine_request("req-1")
|
|
request.sampling_params.lora_request = lora_request
|
|
request.sampling_params.lora_scale = 0.75
|
|
|
|
manager = _RecordingLoRAManager()
|
|
worker = _make_step_worker(lora_manager=manager)
|
|
scheduler_output = _make_scheduler_output(request, request_id="req-1")
|
|
|
|
DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
|
|
|
assert manager.calls == [(lora_request, 0.75)]
|
|
|
|
def test_recovers_lora_for_cached_step_requests(self):
|
|
from vllm_omni.lora.request import LoRARequest
|
|
|
|
lora_request = LoRARequest(lora_name="adapter", lora_int_id=11, lora_path="/tmp/lora")
|
|
request = _make_engine_request("req-1")
|
|
request.sampling_params.lora_request = lora_request
|
|
request.sampling_params.lora_scale = 0.5
|
|
|
|
manager = _RecordingLoRAManager()
|
|
worker = _make_step_worker(lora_manager=manager)
|
|
first = _make_scheduler_output(request, request_id="req-1")
|
|
second = _make_cached_scheduler_output(request_id="req-1", step_id=1)
|
|
|
|
DiffusionWorker.execute_stepwise(worker, first)
|
|
DiffusionWorker.execute_stepwise(worker, second)
|
|
|
|
assert manager.calls == [(lora_request, 0.5), (lora_request, 0.5)]
|
|
|
|
def test_activates_single_lora_for_homogeneous_batch(self):
|
|
"""Multiple requests sharing the same LoRA → exactly one activation,
|
|
and every request id is registered in ``_step_lora_state``."""
|
|
from vllm_omni.lora.request import LoRARequest
|
|
|
|
lora_request = LoRARequest(lora_name="adapter", lora_int_id=9, lora_path="/tmp/lora")
|
|
reqs = []
|
|
for rid in ("req-1", "req-2", "req-3"):
|
|
r = _make_engine_request(rid)
|
|
r.sampling_params.lora_request = lora_request
|
|
r.sampling_params.lora_scale = 0.6
|
|
reqs.append(r)
|
|
|
|
manager = _RecordingLoRAManager()
|
|
worker = _make_step_worker(lora_manager=manager)
|
|
scheduler_output = _make_batch_scheduler_output(reqs)
|
|
|
|
DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
|
|
|
assert manager.calls == [(lora_request, 0.6)]
|
|
assert set(worker._step_lora_state) == {"req-1", "req-2", "req-3"}
|
|
for entry in worker._step_lora_state.values():
|
|
assert entry == (lora_request, 0.6)
|
|
|
|
def test_evicts_step_lora_state_for_finished_requests(self):
|
|
from vllm_omni.lora.request import LoRARequest
|
|
|
|
lora_request = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora")
|
|
finishing = _make_engine_request("req-1")
|
|
finishing.sampling_params.lora_request = lora_request
|
|
next_request = _make_engine_request("req-2")
|
|
next_request.sampling_params.lora_request = lora_request
|
|
|
|
worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
|
|
first = _make_scheduler_output(finishing, request_id="req-1")
|
|
next_batch = _make_scheduler_output(
|
|
next_request,
|
|
request_id="req-2",
|
|
step_id=1,
|
|
finished_req_ids={"req-1"},
|
|
)
|
|
|
|
DiffusionWorker.execute_stepwise(worker, first)
|
|
assert "req-1" in worker._step_lora_state
|
|
|
|
DiffusionWorker.execute_stepwise(worker, next_batch)
|
|
assert "req-1" not in worker._step_lora_state
|
|
assert worker._step_lora_state == {"req-2": (lora_request, 1.0)}
|
|
|
|
|
|
@pytest.mark.cpu
|
|
class TestExecutor:
|
|
"""MultiprocDiffusionExecutor.execute_step"""
|
|
|
|
def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
|
|
executor = object.__new__(MultiprocDiffusionExecutor)
|
|
executor.od_config = SimpleNamespace(streaming_output=False)
|
|
executor._ensure_open = lambda: None
|
|
expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
|
|
executor.collective_rpc = mocker.Mock(return_value=expected)
|
|
|
|
request = _make_engine_request("req-step", num_inference_steps=2)
|
|
scheduler_output = _make_scheduler_output(request, request_id="req-step")
|
|
|
|
output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
|
|
|
|
assert output is expected
|
|
|
|
|
|
@pytest.mark.cpu
|
|
class TestEngine:
|
|
"""Step-execution paths in DiffusionEngine.add_req_and_wait_for_response"""
|
|
|
|
@pytest.mark.parametrize(
|
|
("execute_fn", "expected_error"),
|
|
[
|
|
(
|
|
lambda _: RunnerOutput(
|
|
request_id="req-error",
|
|
step_index=1,
|
|
finished=True,
|
|
result=DiffusionOutput(error="boom"),
|
|
),
|
|
"boom",
|
|
),
|
|
(
|
|
lambda _: (_ for _ in ()).throw(RuntimeError("gpu on fire")),
|
|
"gpu on fire",
|
|
),
|
|
],
|
|
)
|
|
def test_step_engine_returns_error(self, execute_fn, expected_error, mocker: MockerFixture):
|
|
scheduler = StepScheduler()
|
|
scheduler.initialize(SimpleNamespace())
|
|
engine = _make_engine(scheduler, execute_fn=execute_fn)
|
|
|
|
output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2))
|
|
|
|
assert output.output is None
|
|
assert expected_error in output.error
|
|
|
|
def test_step_execution_completes(self, mocker: MockerFixture):
|
|
scheduler = StepScheduler()
|
|
scheduler.initialize(SimpleNamespace())
|
|
engine = _make_engine(scheduler)
|
|
request = _make_engine_request("req-step", num_inference_steps=2)
|
|
|
|
call_count = {"n": 0}
|
|
|
|
def execute_fn(_):
|
|
call_count["n"] += 1
|
|
finished = call_count["n"] == 2
|
|
return RunnerOutput(
|
|
request_id="req-step",
|
|
step_index=call_count["n"],
|
|
finished=finished,
|
|
result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
|
|
)
|
|
|
|
engine.execute_fn = execute_fn
|
|
|
|
output = engine.add_req_and_wait_for_response(request)
|
|
|
|
assert call_count["n"] == 2
|
|
assert output.error is None
|
|
assert torch.equal(output.output, torch.tensor([2.0]))
|
|
|
|
def test_step_abort_stops_rescheduling_after_first_step(self, mocker: MockerFixture):
|
|
scheduler = StepScheduler()
|
|
scheduler.initialize(SimpleNamespace())
|
|
engine = _make_engine(scheduler)
|
|
request = _make_engine_request("req-stop", num_inference_steps=4)
|
|
|
|
step = {"n": 0}
|
|
|
|
def execute_fn(_):
|
|
step["n"] += 1
|
|
engine.abort("req-stop")
|
|
return RunnerOutput(
|
|
request_id="req-stop",
|
|
step_index=1,
|
|
finished=False,
|
|
result=None,
|
|
)
|
|
|
|
engine.execute_fn = execute_fn
|
|
|
|
output = engine.add_req_and_wait_for_response(request)
|
|
|
|
assert step["n"] == 1
|
|
_assert_aborted_output(output, "req-stop")
|
|
|
|
def test_step_abort_after_reschedule_returns_aborted_output(self, mocker: MockerFixture):
|
|
scheduler = StepScheduler()
|
|
scheduler.initialize(SimpleNamespace())
|
|
engine = _make_engine(scheduler)
|
|
request = _make_engine_request("req-mid", num_inference_steps=4)
|
|
|
|
step = {"n": 0}
|
|
|
|
def execute_fn(sched_output):
|
|
step["n"] += 1
|
|
if step["n"] == 2:
|
|
assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
|
|
engine.abort("req-mid")
|
|
return RunnerOutput(
|
|
request_id="req-mid",
|
|
step_index=step["n"],
|
|
finished=False,
|
|
result=None,
|
|
)
|
|
|
|
engine.execute_fn = execute_fn
|
|
|
|
output = engine.add_req_and_wait_for_response(request)
|
|
|
|
assert step["n"] == 2
|
|
_assert_aborted_output(output, "req-mid")
|
|
|
|
def test_finished_step_without_result_returns_error(self, mocker: MockerFixture):
|
|
scheduler = StepScheduler()
|
|
scheduler.initialize(SimpleNamespace())
|
|
engine = _make_engine(
|
|
scheduler,
|
|
execute_fn=lambda _: RunnerOutput(
|
|
request_id="req-missing",
|
|
step_index=1,
|
|
finished=True,
|
|
result=None,
|
|
),
|
|
)
|
|
|
|
output = engine.add_req_and_wait_for_response(_make_engine_request("req-missing", num_inference_steps=1))
|
|
|
|
assert output.output is None
|
|
assert output.error == "Diffusion execution finished without a final output."
|
|
|
|
|
|
@pytest.mark.cpu
|
|
class TestIPC:
|
|
def test_pack_unpack_runner_output_shm(self):
|
|
tensor = torch.zeros(300_000, dtype=torch.float32)
|
|
output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
|
|
|
|
packed = pack_diffusion_output_shm(output)
|
|
assert isinstance(packed.result.output, dict)
|
|
assert packed.result.output["__tensor_shm__"] is True
|
|
|
|
unpacked = unpack_diffusion_output_shm(packed)
|
|
assert isinstance(unpacked.result.output, torch.Tensor)
|
|
torch.testing.assert_close(unpacked.result.output, tensor)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
class TestSupportedPipelines:
|
|
"""Step-execution protocol checks for supported pipelines."""
|
|
|
|
def test_default_stage_config_includes_step_execution(self):
|
|
stage_cfg = StageConfigFactory.create_default_diffusion(
|
|
{
|
|
"step_execution": True,
|
|
}
|
|
)[0]
|
|
|
|
assert stage_cfg["engine_args"]["step_execution"] is True
|
|
|
|
def test_qwen_image_supports_step_execution(self):
|
|
from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
|
|
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline
|
|
|
|
# Avoid loading model weights; protocol membership depends on the class contract.
|
|
pipeline = object.__new__(QwenImagePipeline)
|
|
|
|
assert pipeline.supports_step_execution is True
|
|
assert supports_step_execution(pipeline) is True
|
|
assert isinstance(pipeline, SupportsStepExecution) is True
|
|
|
|
|
|
@hardware_test(
|
|
res={"cuda": "L4"},
|
|
num_cards=2,
|
|
)
|
|
def test_execute_stepwise_with_ulysses_parallel():
|
|
world_size = 2
|
|
if current_omni_platform.get_device_count() < world_size:
|
|
pytest.skip(f"Test requires {world_size} devices")
|
|
|
|
torch.multiprocessing.spawn(
|
|
_distributed_step_worker,
|
|
args=(world_size, "ulysses", "29540"),
|
|
nprocs=world_size,
|
|
)
|
|
|
|
|
|
@hardware_test(
|
|
res={"cuda": "L4"},
|
|
num_cards=2,
|
|
)
|
|
def test_execute_stepwise_with_ring_parallel():
|
|
world_size = 2
|
|
if current_omni_platform.get_device_count() < world_size:
|
|
pytest.skip(f"Test requires {world_size} devices")
|
|
|
|
torch.multiprocessing.spawn(
|
|
_distributed_step_worker,
|
|
args=(world_size, "ring", "29541"),
|
|
nprocs=world_size,
|
|
)
|
|
|
|
|
|
@hardware_test(
|
|
res={"cuda": "L4"},
|
|
num_cards=2,
|
|
)
|
|
def test_execute_stepwise_with_cfg_parallel():
|
|
world_size = 2
|
|
if current_omni_platform.get_device_count() < world_size:
|
|
pytest.skip(f"Test requires {world_size} devices")
|
|
|
|
torch.multiprocessing.spawn(
|
|
_distributed_step_worker,
|
|
args=(world_size, "cfg", "29542"),
|
|
nprocs=world_size,
|
|
)
|