From 9479720f81bc931bfc53e601ec619713e87b281c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?= <11222265+fhfuih@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:07:36 +0800 Subject: [PATCH] [Refactor]: clean up diffusion model runner surface (#5218) Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .../diffusion/test_diffusion_model_runner.py | 26 +++-- .../worker/diffusion_model_runner.py | 101 ++++++++---------- 2 files changed, 58 insertions(+), 69 deletions(-) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 5d4ef5d72..e82b8f7f0 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -98,7 +98,7 @@ class _ChunkStepPipeline: return output -def _make_request(skip_cache_refresh: bool = True): +def _make_request(): sampling_params = SimpleNamespace( generator=None, seed=None, @@ -109,7 +109,6 @@ def _make_request(skip_cache_refresh: bool = True): request_id="req-test", prompt="a prompt", sampling_params=sampling_params, - skip_cache_refresh=skip_cache_refresh, kv_sender_info=None, ) @@ -120,7 +119,6 @@ def _make_request_with_params(req_id: str, sampling_params): prompt=f"prompt-{req_id}", prompts=[f"prompt-{req_id}"], sampling_params=sampling_params, - skip_cache_refresh=True, ) @@ -199,7 +197,7 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): runner.pipeline = _ChunkStepPipeline(chunks) runner.od_config.streaming_output = True runner.od_config.step_execution = True - req = _make_request(skip_cache_refresh=True) + req = _make_request() req.request_id = "req" monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) @@ -235,7 +233,7 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch): """Guard cache diagnostics with runtime backend state to avoid stale-config crashes.""" runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit") - req = _make_request(skip_cache_refresh=True) + req = _make_request() cache_summary_calls = [] @@ -259,11 +257,18 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch): class _EnabledCacheBackend: + def __init__(self): + self.refresh_calls = [] + def is_enabled(self): return True - runner = _make_runner(cache_backend=_EnabledCacheBackend(), cache_backend_name="cache_dit") - req = _make_request(skip_cache_refresh=True) + def refresh(self, pipeline, num_inference_steps, verbose=True): + self.refresh_calls.append((pipeline, num_inference_steps, verbose)) + + cache_backend = _EnabledCacheBackend() + runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit") + req = _make_request() cache_summary_calls = [] @@ -281,6 +286,7 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp assert output.output == "ok" assert cache_summary_calls == [(runner.pipeline, True)] + assert cache_backend.refresh_calls == [(runner.pipeline, 4, True)] @pytest.mark.core_model @@ -288,7 +294,7 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch): runner = _make_runner(cache_backend=None, cache_backend_name="none") runner.pipeline = _SingleRequestBatchPipeline() - req = _make_request(skip_cache_refresh=True) + req = _make_request() monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) @@ -304,7 +310,7 @@ def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(mon def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipeline(monkeypatch): runner = _make_runner(cache_backend=None, cache_backend_name="none") runner.pipeline = _SingleRequestDiffusionOutputPipeline() - req = _make_request(skip_cache_refresh=True) + req = _make_request() monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) @@ -515,7 +521,7 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch): """execute_model runs the pipeline forward after the KV receive step.""" runner = _make_runner(cache_backend=None, cache_backend_name=None) runner.kv_transfer_manager.receive_multi_kv_cache_distributed = lambda *a, **k: True - req = _make_request(skip_cache_refresh=True) + req = _make_request() monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 4ac9afe7a..34a191506 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -12,13 +12,13 @@ from __future__ import annotations import copy import time -from collections.abc import Iterable -from contextlib import nullcontext +from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext from typing import Any import torch from torch.profiler import record_function -from vllm.config import LoadConfig +from vllm.config import LoadConfig, VllmConfig from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes @@ -49,6 +49,7 @@ from vllm_omni.diffusion.worker.utils import ( merge_stage_durations, ) from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin @@ -103,7 +104,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): def __init__( self, - vllm_config, + vllm_config: VllmConfig, od_config: OmniDiffusionConfig, device: torch.device, ): @@ -118,10 +119,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.vllm_config = vllm_config self.od_config = od_config self.device = device - self.pipeline = None - self.cache_backend = None - self.offload_backend = None - self.prompt_embed_cache = None + self.pipeline: Any | None = None + self.cache_backend: Any | None = None + self.offload_backend: Any | None = None + self.prompt_embed_cache: Any | None = None + self.input_batch: InputBatch | None = None # Cache for per-request stepwise state. self.state_cache: dict[str, StepRequestState] = {} @@ -140,7 +142,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) @property - def target_device(self) -> torch.device | None: + def _target_device(self) -> torch.device | None: return getattr(self.pipeline, "device", None) def _compile_transformer(self, attr_name: str) -> None: @@ -161,12 +163,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): def load_model( self, - memory_pool_context_fn: callable | None = None, + memory_pool_context_fn: Callable[[str], AbstractContextManager[Any]] | None = None, load_format: str = "default", custom_pipeline_name: str | None = None, ) -> None: """ - Load the diffusion model, apply compilation and offloading. + Launch the diffusion pipeline, applying compilation, offloading, and caching. Args: memory_pool_context_fn: Optional function that returns a context manager @@ -186,9 +188,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): "cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device) ) - def get_memory_context(): + def get_memory_context() -> AbstractContextManager[Any]: if memory_pool_context_fn is not None: - return memory_pool_context_fn(tag="weights") + return memory_pool_context_fn("weights") return nullcontext() # Load model within forward context @@ -217,13 +219,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): logger.warning("streaming_output=True requires step_execution=True; enabling step execution.") self.od_config.step_execution = True - if getattr(self.od_config, "step_execution", False) and not self.supports_step_mode(): + if getattr(self.od_config, "step_execution", False) and not self._supports_step_mode(): raise ValueError( "step_execution=True requires a pipeline implementing " "prepare_encode(), denoise_step(), step_scheduler(), and post_decode(); " f"{self.od_config.model_class_name} does not support that contract." ) - if self.od_config.streaming_output and not self.supports_step_mode(): + if self.od_config.streaming_output and not self._supports_step_mode(): raise ValueError( "streaming_output=True requires step execution support; " f"{self.od_config.model_class_name} does not support that contract." @@ -287,17 +289,19 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): logger.info("Model runner: Initialization complete.") - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights into the pipeline.""" - return self.pipeline.load_weights(weights) - def clear_prompt_embed_cache(self) -> None: - """Evict all cached text-encoder outputs (e.g. between training epochs).""" + """Evict all cached text-encoder outputs (e.g. between training epochs). + + Kept primarily for extension purposes. + """ if self.prompt_embed_cache is not None: self.prompt_embed_cache.clear() def get_prompt_embed_cache_stats(self) -> dict | None: - """Return hit/miss statistics for the prompt-embedding cache, if enabled.""" + """Return hit/miss statistics for the prompt-embedding cache, if enabled. + + Kept primarily for extension purposes. + """ if self.prompt_embed_cache is None: return None return self.prompt_embed_cache.stats() @@ -346,29 +350,32 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if use_prefetch and self._kv_prefetch_enabled: self.kv_transfer_manager.consume_and_distribute_kv_cache( req, - target_device=self.target_device, + target_device=self._target_device, ) else: self.kv_transfer_manager.receive_multi_kv_cache_distributed( req, cfg_kv_collect_func=getattr(od_config, "cfg_kv_collect_func", None), - target_device=self.target_device if use_prefetch else getattr(self.pipeline, "device", None), + target_device=self._target_device if use_prefetch else getattr(self.pipeline, "device", None), ) kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000 logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms) # Kick off the next request's prefetch (+ H2D) to overlap this forward. if use_prefetch and self._kv_prefetch_enabled and kv_prefetch_job is not None: - self.kv_transfer_manager.start_prefetch(kv_prefetch_job, self.target_device) + self.kv_transfer_manager.start_prefetch(kv_prefetch_job, self._target_device) - if req.sampling_params.generator is None and req.sampling_params.seed is not None: - if req.sampling_params.generator_device is not None: - gen_device = req.sampling_params.generator_device + self._initialize_generator(req.sampling_params) + + def _initialize_generator(self, sampling_params: OmniDiffusionSamplingParams) -> None: + if sampling_params.generator is None and sampling_params.seed is not None: + if sampling_params.generator_device is not None: + gen_device = sampling_params.generator_device elif self.device.type == "cpu": gen_device = "cpu" else: gen_device = self.device - req.sampling_params.generator = torch.Generator(device=gen_device).manual_seed(req.sampling_params.seed) + sampling_params.generator = torch.Generator(device=gen_device).manual_seed(sampling_params.seed) def _refresh_cache_for_requests( self, @@ -377,11 +384,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): od_config: OmniDiffusionConfig, ) -> None: first_req = reqs[0] - if ( - getattr(first_req, "skip_cache_refresh", False) - or self.cache_backend is None - or not self.cache_backend.is_enabled() - ): + if self.cache_backend is None or not self.cache_backend.is_enabled(): return # Refresh cache context if needed. Batch admission groups requests by @@ -499,8 +502,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self, state: StepRequestState, output: DiffusionOutput, - *, - is_primary: bool, ) -> None: merge_stage_durations( state, @@ -564,7 +565,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # Step-wise execution # ------------------------------------------------------------------ - def supports_step_mode(self) -> bool: + def _supports_step_mode(self) -> bool: """Return whether current pipeline supports step execution.""" return self.pipeline is not None and supports_step_execution(self.pipeline) @@ -593,7 +594,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.kv_transfer_manager.receive_multi_kv_cache_distributed( state_req, cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), - target_device=self.target_device, + target_device=self._target_device, ) self.state_cache[request_id] = new_state resolved.append(new_state) @@ -615,15 +616,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # process new reqs for state in states: if state.request_id in new_request_ids: - # set generator - if state.sampling.generator is None and state.sampling.seed is not None: - if state.sampling.generator_device is not None: - gen_device = state.sampling.generator_device - elif self.device.type == "cpu": - gen_device = "cpu" - else: - gen_device = self.device - state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed) + self._initialize_generator(state.sampling) clear_pipeline_stage_durations(self.pipeline) # encode self.pipeline.prepare_encode(state) @@ -644,7 +637,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): states: list[StepRequestState], input_batch: InputBatch, interrupted: bool = False, - ): + ) -> None: """Step-after update: clear cached state for completed request.""" gathered_latents = torch.cat([state.latents for state in states], dim=0) if ( @@ -663,19 +656,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if interrupted or state.request_denoise_completed: self.state_cache.pop(state.request_id, None) - def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any: - model_state = getattr(self, "model_state", None) - if model_state is None: - return {} - prepare_attn = getattr(model_state, "prepare_attn", None) - if not callable(prepare_attn): - return {} - return prepare_attn(input_batch) - def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput: """Execute one step for one scheduled request and return runner output.""" assert self.pipeline is not None, "Model not loaded. Call load_model() first." - if not self.supports_step_mode(): + if not self._supports_step_mode(): raise ValueError("Current pipeline does not support step execution.") # Stepwise mode only supports the basic state-driven denoise path for now. # Request-mode extras such as cache backends, editing inputs, and @@ -692,7 +676,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available(): current_omni_platform.reset_peak_memory_stats() input_batch = self._prepare_batch_inputs(states, new_request_ids) - attn_metadata = self._prepare_attn_metadata(input_batch) + attn_metadata = {} with set_forward_context( vllm_config=self.vllm_config, @@ -741,7 +725,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self._attach_stepwise_metrics( req, result, - is_primary=is_primary, ) else: result = None