# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Diffusion Model Runner for vLLM-Omni. Handles model loading, compilation, caching, and execution of diffusion model forward passes. This follows the AR pattern where the Runner handles all model-related operations. """ from __future__ import annotations import copy import gc import time from collections.abc import Callable, Sequence from contextlib import AbstractContextManager, nullcontext from typing import TYPE_CHECKING, Any, cast import torch from torch.profiler import record_function from vllm.config import LoadConfig, VllmConfig from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, cache_summary from vllm_omni.diffusion.cache.prompt_embed_cache import ( install_prompt_embed_cache, resolve_prompt_embed_cache_config, ) from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.cancellation import check_request_cancellation, request_cancellation_scope from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionMetadata, DiffusionPagedAttentionRow, ) from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_rank from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator from vllm_omni.diffusion.interaction.types import InteractionPayload from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import ( SupportsInteractionApply, adopt_request_scoped_cache_dit, is_request_scoped_cache_dit_enabled, supports_interaction_apply, supports_step_execution, ) from vllm_omni.diffusion.offloader import enable_offload_backend from vllm_omni.diffusion.offloader.config import ( TEXT_ENCODER_COMPONENT, OffloadStrategy, offload_enabled, resolve_offload, resolve_offload_strategy, ) from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( CachedRequestData, DiffusionSchedulerOutput, KVPrefetchJob, NewRequestData, validate_new_request_data_identity, ) from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.diffusion.worker.stage_payload import DiffusionStagePayloadMixin from vllm_omni.diffusion.worker.utils import ( BatchRunnerOutput, RunnerOutput, StepRequestState, attach_stage_durations, clear_pipeline_stage_durations, consume_pipeline_stage_durations, 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 if TYPE_CHECKING: from vllm_omni.inputs.data import OmniInteractionPrompt logger = init_logger(__name__) def _dit_any_rank_failed(local_failed: bool) -> bool: """All-reduce a per-request failure flag across the DiT process group. Every DiT rank must reach this point in lockstep; the caller wraps ``pipeline.prepare_encode`` so both the success and the exception paths call the helper. In single-rank execution this collapses to the local flag with no collectives issued. """ if not torch.distributed.is_initialized(): return local_failed try: from vllm_omni.diffusion.distributed import parallel_state get_dit_group = getattr(parallel_state, "get_dit_group", None) group = get_dit_group() if get_dit_group is not None else None except (AssertionError, ImportError): group = None if group is None: return local_failed signal = torch.tensor(1 if local_failed else 0, dtype=torch.int32) if current_omni_platform.is_available(): signal = signal.to(device=current_omni_platform.device_type) torch.distributed.all_reduce(signal, op=torch.distributed.ReduceOp.MAX, group=group) return bool(signal.item()) def _normalize_pipeline_outputs( outputs: object, *, expected_count: int, allow_single_output: bool, pipeline_name: str, ) -> list[DiffusionOutput]: if isinstance(outputs, DiffusionOutput): if allow_single_output and expected_count == 1: return [outputs] raise RuntimeError( f"{pipeline_name}.forward returned a single DiffusionOutput; " "request-batch forward must return list[DiffusionOutput]." ) if not isinstance(outputs, list): raise RuntimeError( f"{pipeline_name}.forward returned {type(outputs).__name__}; " "expected DiffusionOutput or list[DiffusionOutput]." ) if len(outputs) != expected_count: raise RuntimeError( f"{pipeline_name}.forward returned {len(outputs)} outputs for {expected_count} requests; " "expected exactly one DiffusionOutput per request." ) bad_index = next((idx for idx, output in enumerate(outputs) if not isinstance(output, DiffusionOutput)), None) if bad_index is not None: raise RuntimeError( f"{pipeline_name}.forward returned list item {bad_index} with type " f"{type(outputs[bad_index]).__name__}; expected DiffusionOutput." ) return outputs class DiffusionModelRunner(DiffusionStagePayloadMixin): """ Model runner that handles model loading and execution for diffusion models. This class follows the AR pattern where the Runner handles all model-related operations including loading, compilation, offloading, caching, and execution. The Worker only handles infrastructure (device, distributed env). """ def __init__( self, vllm_config: VllmConfig, od_config: OmniDiffusionConfig, device: torch.device, ): """ Initialize the diffusion model runner. Args: vllm_config: vLLM configuration. od_config: OmniDiffusion configuration. device: The device to run on. """ self.vllm_config = vllm_config self.od_config = od_config self.device = device 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 self._interaction_coordinator: InteractionCoordinator | None = None self.model_memory_usage = 0 self.diffusion_kv_backend = DiffusionKVModelRunnerBackend( vllm_config=vllm_config, od_config=od_config, device=device, ) # Compatibility view for callers that inspect the installed config; # physical ownership remains in ``diffusion_kv_backend``. self.kv_cache_config: KVCacheConfig | None = None # Cache for per-request stepwise state. self.state_cache: dict[str, StepRequestState] = {} # Initialize KV cache manager for connector management. payload_transfer_manager = OmniKVTransferManager.from_od_config(od_config) self.kv_transfer_manager = ( payload_transfer_manager if getattr(od_config, "kv_transfer_config", None) is None else None ) self.init_omni_connectors(od_config, payload_transfer_manager, synchronous=True) self._kv_connector = None from vllm_omni.diffusion.diffusion_kv.kv_connector import KVReceiveProgress, native_prefetch_enabled self._kv_receive_progress = KVReceiveProgress() if native_prefetch_enabled(od_config) else None # Prefetch covers TP / SP / CFG-Parallel / HSDP. Disabled when a CFG # companion KV collector is set (that KV is not backgrounded). has_cfg_companion_kv = getattr(od_config, "cfg_kv_collect_func", None) is not None self._kv_prefetch_enabled = ( self.kv_transfer_manager is not None and bool(self.kv_transfer_manager.config.enable_kv_async_prefetch) and not has_cfg_companion_kv and self.kv_transfer_manager.config.need_recv_cache ) @property def _target_device(self) -> torch.device | None: return getattr(self.pipeline, "device", None) def _validate_diffusion_kv_metadata( self, *, request_id: str, metadata: DiffusionKVMetadata | None, ) -> None: cache_mode = getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) if cache_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is None: raise ValueError(f"paged_scheduler request {request_id!r} requires Diffusion KV metadata") if cache_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is not None: raise ValueError(f"{cache_mode.value} request {request_id!r} must not carry Diffusion KV metadata") if metadata is not None and metadata.request_id != request_id: raise ValueError( f"Diffusion KV metadata request mismatch: expected={request_id!r}, got={metadata.request_id!r}" ) def _compile_transformer(self, attr_name: str) -> None: """Compile a transformer attribute on the pipeline with torch.compile.""" model = getattr(self.pipeline, attr_name, None) if model is None: return compile_granularity = self.od_config.diffusion_compile_granularity compile_dynamic = self.od_config.diffusion_compile_dynamic try: if compile_granularity == "full": model.compile(dynamic=compile_dynamic) compiled_model = model else: compiled_model = regionally_compile(model, dynamic=compile_dynamic) setattr(self.pipeline, attr_name, compiled_model) except Exception as e: logger.warning( "Model runner: %s torch.compile setup for %s failed before activation: %s. " "Continuing with the uncompiled model; lazy compilation errors can still " "surface on the first request.", compile_granularity, attr_name, e, ) return logger.info( "Model runner: %s configured for lazy %s torch.compile with dynamic=%s; " "compilation errors may surface on the first request.", attr_name, compile_granularity, compile_dynamic, ) def load_model( self, memory_pool_context_fn: Callable[[str], AbstractContextManager[Any]] | None = None, load_format: str = "default", custom_pipeline_name: str | None = None, ) -> None: """ Launch the diffusion pipeline, applying compilation, offloading, and caching. Args: memory_pool_context_fn: Optional function that returns a context manager for memory pool allocation (used for sleep mode). load_format: Format for loading model weights. Supported formats: - "default" (default): Automatically detect and use the default format based on configuration - "custom_pipeline": Init model from a custom pipeline class specified by `custom_pipeline_name` - "dummy": Skip actual weight loading, useful for testing and custom pipelines that don't require default weights. custom_pipeline_name: Optional custom pipeline class name to use. """ if load_format == "dummy": return # Resolve environment overrides before model loading. Rank-local # prompt-cache hits cannot safely skip text-encoder collectives. enable_pec, pec_size = resolve_prompt_embed_cache_config( enable=getattr(self.od_config, "enable_prompt_embed_cache", False), max_size=getattr(self.od_config, "prompt_embed_cache_size", 32), ) resolved_offload = resolve_offload(self.od_config) dp_size = int(getattr(getattr(self.od_config, "parallel_config", None), "data_parallel_size", 1)) if ( enable_pec and dp_size > 1 and resolved_offload.offloads(TEXT_ENCODER_COMPONENT) and resolved_offload.uses_allgather(TEXT_ENCODER_COMPONENT) ): raise ValueError( "Prompt embedding cache cannot be combined with text_encoder " "AllGather across data-parallel ranks; disable the cache or use " "rank-local text_encoder transfer." ) current_omni_platform.init_diffusion_model_runner_runtime( vllm_config=self.vllm_config, od_config=self.od_config, device=self.device, ) load_device = "cpu" if offload_enabled(self.od_config) else str(self.device) def get_memory_context() -> AbstractContextManager[Any]: if memory_pool_context_fn is not None: return memory_pool_context_fn("weights") return nullcontext() # Load model within forward context load_config = LoadConfig() model_loader = DiffusersPipelineLoader(load_config, od_config=self.od_config) time_before_load = time.perf_counter() with get_memory_context(): with DeviceMemoryProfiler() as m: self.pipeline = model_loader.load_model( load_device=load_device, load_format=load_format, custom_pipeline_name=custom_pipeline_name, device=self.device, ) time_after_load = time.perf_counter() logger.info( "Model loading took %.4f GiB and %.6f seconds", m.consumed_memory / GiB_bytes, time_after_load - time_before_load, ) self.model_memory_usage = int(m.consumed_memory) logger.info("Model runner: Model loaded successfully.") if self.od_config.streaming_output and not getattr(self.od_config, "step_execution", False): 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(): 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(): raise ValueError( "streaming_output=True requires step execution support; " f"{self.od_config.model_class_name} does not support that contract." ) # The offloader owns loader-plan handoff and startup recovery. The # runner only receives the pipeline/backend pair that is ready to use. self.pipeline, self.offload_backend = enable_offload_backend( self.od_config, self.pipeline, device=self.device, ) # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: if current_omni_platform.supports_torch_inductor(): if hasattr(self.pipeline, "setup_compile"): try: self.pipeline.setup_compile() except Exception as exc: logger.warning( "Model runner: setup_compile() failed (%s); running without compile.", exc, ) else: transformer_attrs = getattr(self.pipeline, "_dit_modules", None) if not transformer_attrs: transformer_attrs = ("transformer", "transformer_2") for attr_name in transformer_attrs: self._compile_transformer(attr_name) else: logger.warning( "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", current_omni_platform.get_torch_device(), ) # Setup cache backend self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config) if self.cache_backend is not None: if self.od_config.model_class_name in _NO_CACHE_ACCELERATION: logger.warning( "Cache backend '%s' is not supported for %s; disabling cache acceleration.", self.od_config.cache_backend, self.od_config.model_class_name, ) self.cache_backend = None self.od_config.cache_backend = None else: # Install configured cache capability once at startup. A model # may explicitly adopt the enabled Cache-DiT backend and then # own all later request-boundary enable/disable transitions. self.cache_backend.enable(self.pipeline) if isinstance(self.cache_backend, CacheDiTBackend) and adopt_request_scoped_cache_dit( self.pipeline, self.cache_backend, ): logger.info( "Pipeline %s owns request-scoped Cache-DiT transitions.", type(self.pipeline).__name__, ) self.cache_backend = None # Install prompt-embedding cache (transparent wrapper around # ``pipeline.encode_prompt``). Enabled via config or env var; a no-op # when the pipeline does not expose ``encode_prompt``. if enable_pec: self.prompt_embed_cache = install_prompt_embed_cache( self.pipeline, max_size=pec_size, enabled=True, model_tag=self.od_config.model_class_name, ) self._interaction_coordinator = InteractionCoordinator.build(self.pipeline, self.od_config) if hasattr(self.pipeline, "_interaction_coordinator"): self.pipeline._interaction_coordinator = self._interaction_coordinator logger.info("Model runner: Initialization complete.") def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: """Collect native specs from cache-enabled loaded attention modules.""" if self.od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER: return {} if self.pipeline is None: raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs") cache_layers: dict[str, tuple[Attention, KVCacheSpec]] = {} for module_path, module in self.pipeline.named_modules(): if not isinstance(module, Attention): continue spec = module.get_kv_cache_spec(self.vllm_config) if spec is not None: layer_name = module.prefix if not isinstance(layer_name, str) or not layer_name: raise RuntimeError( "Paged Diffusion Attention must expose a non-empty canonical prefix; " f"module_path={module_path!r}" ) if layer_name in cache_layers: raise RuntimeError(f"Duplicate canonical paged Diffusion Attention prefix {layer_name!r}") cache_layers[layer_name] = (module, spec) if not cache_layers: raise RuntimeError( "paged_scheduler Diffusion KV found no cache-enabled Attention modules " f"in {type(self.pipeline).__name__}" ) return self.diffusion_kv_backend.register_kv_cache_layers(cache_layers) def set_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: """Physically initialize the Engine-generated rank-local config.""" self.diffusion_kv_backend.initialize_kv_cache(kv_cache_config) self.kv_cache_config = self.diffusion_kv_backend.kv_cache_config if getattr(self.od_config, "kv_transfer_config", None) is not None: from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector self._kv_connector = ActiveKVConnector(self.vllm_config, self.diffusion_kv_backend.kv_caches_by_layer) def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput): from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load assert self.od_config.kv_transfer_config is not None timeout = self.od_config.kv_transfer_config.kv_connector_extra_config.get("transfer_timeout", 60.0) if self._kv_receive_progress is not None: return self._kv_receive_progress.prepare(self._kv_connector, scheduler_output, timeout) return wait_for_kv_load(self._kv_connector, scheduler_output, timeout) def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool: return self.diffusion_kv_backend.install_diffusion_kv_metadata(metadata) def get_diffusion_kv_row( self, request_id: str, sequence_id: int | None, context_id: str | None = None, ) -> int: return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id) def remove_diffusion_kv_requests(self, request_ids: Sequence[str | tuple[str, int]]) -> int: return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids) def refresh_diffusion_kv_block_table_layout(self) -> None: self.diffusion_kv_backend.refresh_block_table_layout() def _build_paged_attention_metadata( self, metadata: list[DiffusionKVMetadata], ) -> DiffusionPagedAttentionMetadata: """Translate Scheduler sequence state into ordered request-level rows.""" cfg_size = int(getattr(self.od_config.parallel_config, "cfg_parallel_size", 1) or 1) cfg_rank = get_classifier_free_guidance_rank() if cfg_size > 1 else None prefill_rows: list[DiffusionPagedAttentionRow] = [] denoise_rows: list[DiffusionPagedAttentionRow] = [] for request_metadata in metadata: sequences = request_metadata.sequences if cfg_rank is not None and len(sequences) > 1: if len(sequences) != cfg_size: raise ValueError( "Paged CFG parallel execution requires one Scheduler sequence per CFG rank: " f"request={request_metadata.request_id!r}, cfg_size={cfg_size}, rows={len(sequences)}" ) sequences = (sequences[cfg_rank],) for sequence in sequences: active_seq_len = sequence.prefix_len + sequence.target_len if active_seq_len > sequence.seq_len: raise ValueError( "Paged denoise span exceeds its Scheduler allocation: " f"request={request_metadata.request_id!r}, sequence={sequence.sequence_id}, " f"active={active_seq_len}, allocated={sequence.seq_len}" ) # Imported AR KV and local hits have separate owners. kv_start_pos = ( sequence.num_computed_tokens if getattr(self.od_config, "kv_transfer_config", None) is not None else sequence.cached_prefix_len ) prefill_rows.append( DiffusionPagedAttentionRow( request_id=request_metadata.request_id, sequence_id=sequence.sequence_id, kv_start_pos=kv_start_pos, query_len=sequence.seq_len - kv_start_pos, seq_len=sequence.seq_len, ) ) denoise_rows.append( DiffusionPagedAttentionRow( request_id=request_metadata.request_id, sequence_id=sequence.sequence_id, query_len=sequence.target_len, seq_len=active_seq_len, kv_start_pos=sequence.prefix_len, ) ) return DiffusionPagedAttentionMetadata(tuple(prefill_rows), tuple(denoise_rows)) def clear_prompt_embed_cache(self) -> None: """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 release_captured_graphs(self) -> None: """Drop every CUDA graph held for this model, wherever it is kept. Sleep level 2 discards the memory a capture was recorded against, so a graph that outlives it replays over freed storage. The runner owns compilation and execution resources, so it owns the release: a pipeline that captures graphs of its own implements ``release_captured_graphs`` and is collected here, instead of every caller having to know which pipelines have one. """ runners = getattr(self, "graph_runners", None) if runners is not None: runners.clear() release = getattr(self.pipeline, "release_captured_graphs", None) if callable(release): release() def get_prompt_embed_cache_stats(self) -> dict | None: """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() def _sample_peak_memory_mb(self) -> float: """Return peak GPU memory for the current forward pass in MB. Must be called immediately after the measured forward/step work, with reset_peak_memory_stats() called just before it, so the measurement reflects the current execution slice and not the global historical maximum. Uses max_memory_reserved (CUDA memory pool high-water mark) rather than max_memory_allocated so that allocator fragmentation is also visible. See: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html """ peak_reserved_bytes = current_omni_platform.max_memory_reserved() peak_allocated_bytes = current_omni_platform.max_memory_allocated() peak_memory_mb = peak_reserved_bytes / (1024**2) peak_reserved_gb = peak_reserved_bytes / (1024**3) peak_allocated_gb = peak_allocated_bytes / (1024**3) pool_overhead_gb = peak_reserved_gb - peak_allocated_gb logger.debug( "Peak GPU memory (this request): %.2f GB reserved, %.2f GB allocated, %.2f GB pool overhead (%.1f%%)", peak_reserved_gb, peak_allocated_gb, pool_overhead_gb, pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb > 0 else 0.0, ) return peak_memory_mb def _prepare_request_for_forward( self, req: OmniDiffusionRequest, *, od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, use_prefetch: bool = False, ) -> None: # Fetch upstream conditioning before anything else: the pipeline reads # it out of the prompt during the forward below. self._maybe_recv_stage_payload(req) if self.kv_transfer_manager is None: self._initialize_generator(req.sampling_params) return # Receive AR KV. Single-request execution can use the prefetch path: # consume prior-forward payload, sync-fallback on miss; request-batch # execution keeps the synchronous per-request receive path. kv_recv_t0 = time.perf_counter() if use_prefetch and self._kv_prefetch_enabled: self.kv_transfer_manager.consume_and_distribute_kv_cache( req, 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), ) kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000 req.kv_recv_ms = kv_recv_ms 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._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 sampling_params.generator = torch.Generator(device=gen_device).manual_seed(sampling_params.seed) def _refresh_cache_for_requests( self, reqs: list[OmniDiffusionRequest], *, od_config: OmniDiffusionConfig, ) -> None: first_req = reqs[0] if self.cache_backend is None or not self.cache_backend.is_enabled(): return # Refresh cache context if needed. Batch admission groups requests by # RequestBatchSamplingParamsKey, so the first request's num_inference_steps applies # to the whole runner batch. num_inference_steps = first_req.sampling_params.num_inference_steps if num_inference_steps is None and first_req.sampling_params.timesteps is not None: num_inference_steps = len(first_req.sampling_params.timesteps) if num_inference_steps is None and first_req.sampling_params.sigmas is not None: num_inference_steps = len(first_req.sampling_params.sigmas) if num_inference_steps is None: num_inference_steps = getattr(self.pipeline, "default_num_inference_steps", None) if num_inference_steps is None and od_config.cache_backend in ( "tea_cache", "sea_cache", "step_cache", ): # When num_inference_steps is None, some pipelines defer to their # own defaults. These backends use refresh to reset request state; # runtime step metadata is either unused or resolved in the # pipeline. Use the pipeline default when available to keep refresh # behavior aligned with single-request execution. num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0 if num_inference_steps is not None: self.cache_backend.refresh(self.pipeline, num_inference_steps) else: logger.warning( "Failed to refresh the diffusion transformer cache; backend %s " "currently requires num_inference_steps to be passed explicitly", od_config.cache_backend, ) def _prepare_output_for_transport( self, output: DiffusionOutput, sampling_params: OmniDiffusionSamplingParams, ) -> DiffusionOutput: if output.media is not None: if output.output is not None: raise ValueError("DiffusionOutput cannot contain both media and legacy output") output.media = prepare_diffusion_media_for_transport( output.media, od_config=self.od_config, sampling_params=sampling_params, ) return output def _runner_output_from_outputs( self, reqs: list[OmniDiffusionRequest], outputs: list[DiffusionOutput], ) -> BatchRunnerOutput: for i in range(len(reqs)): # Carry the runner-measured KV-recv timing onto the output so the # engine's step_streaming can surface it as diffusion_kv_load_s. outputs[i].kv_recv_ms = reqs[i].kv_recv_ms return BatchRunnerOutput.from_list( [ RunnerOutput( request_id=reqs[i].request_id, step_index=None, finished=True, result=outputs[i], ) for i in range(len(reqs)) ] ) def _execute_request_list( self, reqs: list[OmniDiffusionRequest], *, od_config: OmniDiffusionConfig, allow_single_output: bool, require_request_batch_support: bool, kv_prefetch_job: KVPrefetchJob | None = None, record_name: str, record_output_peak_memory: bool = True, in_diffusion_kv_memory_profile: bool = False, diffusion_kv_metadata: list[DiffusionKVMetadata] | None = None, ) -> BatchRunnerOutput: assert self.pipeline is not None, "Model not loaded. Call load_model() first." if not reqs: return BatchRunnerOutput.from_list([]) for req in reqs: if req.prompt is None: raise ValueError("Cannot execute model with empty prompt") if require_request_batch_support and not getattr(self.pipeline, "supports_request_batch", False): raise RuntimeError(f"{type(self.pipeline).__name__} does not support request-batch forward.") # Use no_grad() for HSDP compatibility, inference_mode() otherwise for # better perf. HSDP2's fully_shard pre-forward hooks need tensor version # counters, which inference tensors do not track. use_hsdp = od_config.parallel_config.use_hsdp use_distributed_offload = resolve_offload_strategy(self.od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE grad_context = torch.no_grad() if (use_hsdp or use_distributed_offload) else torch.inference_mode() with grad_context: for req in reqs: self._prepare_request_for_forward( req, od_config=od_config, kv_prefetch_job=kv_prefetch_job, use_prefetch=allow_single_output, ) self._refresh_cache_for_requests(reqs, od_config=od_config) batch = DiffusionRequestBatch(requests=reqs) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 if is_primary and record_output_peak_memory: current_omni_platform.reset_peak_memory_stats() paged_kv_runtime = None paged_kv_context: AbstractContextManager[Any] = nullcontext() if diffusion_kv_metadata is not None: if len(diffusion_kv_metadata) != len(reqs): raise ValueError( "Diffusion KV metadata count must match the request batch: " f"metadata={len(diffusion_kv_metadata)}, requests={len(reqs)}" ) native_kv_transfer = getattr(self.od_config, "kv_transfer_config", None) is not None if native_kv_transfer: for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True): # This field is consumed by Hunyuan only for native # AR->DiT transfer. Local prefix hits must still run # their VAE/ViT conditioning path. req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences) paged_metadata = self._build_paged_attention_metadata(diffusion_kv_metadata) paged_kv_cached_prefix_len = 0 if not native_kv_transfer: cached_prefix_lens = {row.kv_start_pos for row in paged_metadata.prefill_rows} if len(cached_prefix_lens) != 1: raise ValueError( "One paged request-level forward requires a uniform cached prefix boundary; " f"got {sorted(cached_prefix_lens)}" ) paged_kv_cached_prefix_len = next(iter(cached_prefix_lens)) paged_kv_runtime, paged_kv_context = self.diffusion_kv_backend.activate_paged_attention_metadata( paged_metadata ) if is_primary: # Trace the boundary actually passed to the model, not a # speculative lookup. Useful for warm-cache regressions. for request_metadata in diffusion_kv_metadata: for sequence in request_metadata.sequences: logger.debug( "Diffusion prefix prefill: request_id=%s sequence_id=%d " "cached_prefix_len=%d prefix_len=%d query_len=%d", request_metadata.request_id, sequence.sequence_id, sequence.cached_prefix_len, sequence.prefix_len, sequence.seq_len - sequence.cached_prefix_len, ) else: paged_kv_cached_prefix_len = 0 with ( set_forward_context( vllm_config=self.vllm_config, omni_diffusion_config=od_config, paged_kv_runtime=paged_kv_runtime, paged_kv_cached_prefix_len=paged_kv_cached_prefix_len, in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, ), paged_kv_context, request_cancellation_scope( [getattr(req, "cancellation_signal", None) for req in reqs], enabled=getattr(self.pipeline, "supports_request_cancellation", False) is True, ), ): with record_function(record_name): try: check_request_cancellation() raw_outputs = self.pipeline.forward(batch) outputs = _normalize_pipeline_outputs( raw_outputs, expected_count=len(reqs), allow_single_output=allow_single_output, pipeline_name=type(self.pipeline).__name__, ) except DiffusionRequestAbortedError as exc: # The checkpoint aborts only a fully cancelled wave; # a mixed batch must keep running for its live peers. logger.info( "Stopped cancelled diffusion request(s) %s at a model execution boundary", [req.request_id for req in reqs], ) outputs = [DiffusionOutput(aborted=True, abort_message=str(exc)) for _ in reqs] with record_function("prepare_output_for_transport"): outputs = [ self._prepare_output_for_transport(output, req.sampling_params) for req, output in zip(reqs, outputs, strict=True) ] if is_primary and outputs and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() for output in outputs: output.peak_memory_mb = max(output.peak_memory_mb, batch_peak_memory_mb) # Log prompt-embed cache activity; hits/misses accumulate across requests. prompt_embed_cache = getattr(self, "prompt_embed_cache", None) if is_primary and prompt_embed_cache is not None: logger.debug("prompt-embed cache: %s", prompt_embed_cache.stats()) runner_cache_dit_enabled = self.cache_backend is not None and self.cache_backend.is_enabled() if ( od_config.cache_backend == "cache_dit" and od_config.enable_cache_dit_summary and (runner_cache_dit_enabled or is_request_scoped_cache_dit_enabled(self.pipeline)) ): cache_summary(self.pipeline, details=True) self._maybe_send_stage_payload(reqs, outputs) return self._runner_output_from_outputs(reqs, outputs) def _attach_stepwise_metadata( self, state: StepRequestState, output: DiffusionOutput, ) -> None: merge_stage_durations( state, consume_pipeline_stage_durations(self.pipeline), ) attach_stage_durations(state, output) # In streaming output mode with interaction, acknowledge which events are handled in this chunk. meta = state.interaction_chunk_metadata state.interaction_chunk_metadata = None if meta is not None: output.started_event_ids = list(meta.started_event_ids) output.active_event_ids = list(meta.active_event_ids) output.completed_event_ids = list(meta.completed_event_ids) def execute_model( self, req: OmniDiffusionRequest, kv_prefetch_job: KVPrefetchJob | None = None, diffusion_kv_metadata: DiffusionKVMetadata | None = None, ) -> DiffusionOutput: """ Execute a forward pass for the given requests. Args: req: A diffusion request containing a list of prompts to process. Returns: DiffusionOutput with generated results. Note: We use torch.no_grad() for HSDP because HSDP2's fully_shard requires access to tensor version counters in pre_forward hooks, which inference tensors do not track. For non-HSDP inference, we use torch.inference_mode() for better performance. """ self._validate_diffusion_kv_metadata( request_id=req.request_id, metadata=diffusion_kv_metadata, ) installed_request = False if diffusion_kv_metadata is not None: self.install_diffusion_kv_metadata(diffusion_kv_metadata) installed_request = True try: runner_output = self._execute_request_list( [req], od_config=self.od_config, allow_single_output=True, require_request_batch_support=False, kv_prefetch_job=kv_prefetch_job, record_name="pipeline_forward", diffusion_kv_metadata=[diffusion_kv_metadata] if diffusion_kv_metadata is not None else None, ) output = runner_output.runner_outputs[0].result assert output is not None return output finally: if installed_request: self.remove_diffusion_kv_requests([req.request_id]) def profile_run(self, requests: list[OmniDiffusionRequest]) -> None: """Run the maximum per-rank request batch for memory profiling. This deliberately bypasses Scheduler admission and Diffusion KV metadata validation because cache capacity has not been sized yet. It otherwise uses the normal execution-mode path so model inputs, collective communication, backend workspaces, denoising, and decode allocations contribute to the observed peak. Step execution profiles one fused ``InputBatch`` instead of sequential single-request forwards. """ if not requests: raise ValueError("Diffusion memory profiling requires at least one request.") runner_output: BatchRunnerOutput | None = None request_ids = [request.request_id for request in requests] try: if getattr(self.od_config, "step_execution", False): scheduler_output = DiffusionSchedulerOutput( step_id=0, scheduled_new_reqs=[ NewRequestData(request_id=request.request_id, req=request) for request in requests ], scheduled_cached_reqs=CachedRequestData.make_empty(), finished_req_ids=set(), num_running_reqs=len(requests), num_waiting_reqs=0, ) runner_output = self._execute_stepwise( scheduler_output, validate_kv_metadata=False, record_output_peak_memory=False, in_diffusion_kv_memory_profile=True, ) else: runner_output = self._execute_request_list( requests, od_config=self.od_config, allow_single_output=len(requests) == 1, require_request_batch_support=len(requests) > 1, record_name="pipeline_memory_profile", # The enclosing native memory_profiling context owns the # peak counters. Resetting them here would discard request # preparation allocations and understate the budget. record_output_peak_memory=False, in_diffusion_kv_memory_profile=True, ) current_omni_platform.synchronize() finally: for request_id in request_ids: self.state_cache.pop(request_id, None) self.input_batch = None del runner_output gc.collect() def execute_model_batch( self, scheduler_output: DiffusionSchedulerOutput, od_config: OmniDiffusionConfig, ) -> BatchRunnerOutput: """Execute scheduled request-mode requests through the batch forward path. Builds a ``DiffusionRequestBatch`` from scheduled new requests, runs per-request setup, and calls ``pipeline.forward(batch)``. The pipeline must declare ``supports_request_batch = True``. """ for new_req in scheduler_output.scheduled_new_reqs: validate_new_request_data_identity(new_req) self._validate_diffusion_kv_metadata( request_id=new_req.req.request_id, metadata=new_req.diffusion_kv_metadata, ) installed_request_ids: list[str] = [] try: for new_req in scheduler_output.scheduled_new_reqs: if new_req.diffusion_kv_metadata is not None: self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata) installed_request_ids.append(new_req.request_id) reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs] return self._execute_request_list( reqs, od_config=od_config, allow_single_output=False, require_request_batch_support=True, record_name="pipeline_forward_batch", diffusion_kv_metadata=[ new_req.diffusion_kv_metadata for new_req in scheduler_output.scheduled_new_reqs if new_req.diffusion_kv_metadata is not None ] or None, ) finally: if installed_request_ids: self.remove_diffusion_kv_requests(installed_request_ids) # ------------------------------------------------------------------ # Step-wise execution # ------------------------------------------------------------------ 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) def _cleanup_finished_step_requests(self, scheduler_output: DiffusionSchedulerOutput) -> None: """Retire state and paged-KV rows released by the scheduler wave.""" finished_req_ids = scheduler_output.finished_req_ids for request_id in finished_req_ids: self.state_cache.pop(request_id, None) if ( getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) is DiffusionKVCacheMode.PAGED_SCHEDULER and finished_req_ids ): self.remove_diffusion_kv_requests(list(finished_req_ids)) def _update_states(self, scheduler_output: DiffusionSchedulerOutput) -> tuple[list[StepRequestState], list[str]]: """Resolve cached state and create state for newly admitted requests.""" resolved: list[StepRequestState] = [] new_request_ids: list[str] = [] try: # process new requests for sched_new_req in scheduler_output.scheduled_new_reqs: request_id = sched_new_req.request_id new_request_ids.append(request_id) if request_id in self.state_cache: raise ValueError(f"Received duplicate new-request payload for cached request {request_id}.") self._maybe_recv_stage_payload(sched_new_req.req) new_state = StepRequestState( request_id=request_id, sampling=copy.deepcopy(sched_new_req.req.sampling_params), prompt=sched_new_req.req.prompt, kv_sender_info=sched_new_req.req.kv_sender_info, prepared_layout=getattr(sched_new_req.req, "prepared_layout", None), external_req_id=getattr(sched_new_req.req, "external_req_id", None), ) if ( sched_new_req.diffusion_kv_metadata is not None and getattr(self.od_config, "kv_transfer_config", None) is not None ): new_state.extra["kv_computed_tokens"] = tuple( seq.num_computed_tokens for seq in sched_new_req.diffusion_kv_metadata.sequences ) state_req = copy.copy(sched_new_req.req) state_req.sampling_params = new_state.sampling if self.kv_transfer_manager is not None: 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, ) self.state_cache[request_id] = new_state resolved.append(new_state) # process cached requests for request_id in scheduler_output.scheduled_cached_reqs.request_ids: state = self.state_cache.get(request_id) if state is None: raise ValueError(f"Missing cached state for request {request_id}.") resolved.append(state) except Exception: for request_id in new_request_ids: self.state_cache.pop(request_id, None) raise return resolved, new_request_ids def _prepare_batch_inputs( self, states: list[StepRequestState], new_request_ids: list[str], ) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]: # process new reqs pipeline = self.pipeline assert pipeline is not None, "Model not loaded. Call load_model() first." prepared_states: list[StepRequestState] = [] error_outputs: list[RunnerOutput] = [] for state in states: if state.request_id in new_request_ids: # Everything that requires rank-synchronization must be called # inside a try, record the exception and handle with `_dit_any_rank_failed`. # Reason (example): An exception in ``_initialize_generator`` or # ``clear_pipeline_stage_durations`` on one rank would skip the # all-reduce here while every peer proceeds into it, and the # peers then hang on the NCCL collective until timeout. def _abort_prep_failure(per_req_exc: BaseException | None) -> None: self.state_cache.pop(state.request_id, None) if per_req_exc is None: per_req_exc = RuntimeError( f"Stepwise preparation failed on another DiT rank for {state.request_id}" ) logger.error( "Stepwise request preparation failed for %s: %s", state.request_id, per_req_exc, exc_info=isinstance(per_req_exc, Exception), ) error_outputs.append( RunnerOutput( request_id=state.request_id, step_index=state.step_index, finished=True, result=DiffusionOutput.from_exception(per_req_exc), ) ) per_req_exc: BaseException | None = None try: self._initialize_generator(state.sampling) clear_pipeline_stage_durations(pipeline) pipeline.prepare_encode(state) except Exception as exc: per_req_exc = exc # Pipelines that do rank-0-only work (e.g. MiniMax H3 # reference-video prep) must broadcast per-request failures # internally so downstream collectives stay in step; even so, # cross-check that every DiT rank agrees so a rank-local error # (or a future pipeline that omits the guard) does not leave # the process group half-way through a new request. if _dit_any_rank_failed(per_req_exc is not None): _abort_prep_failure(per_req_exc) continue # If the pipeline supports interaction, the interaction session initialization also needs to call # synchronized_monotonic_time(). Wrap in another try-block to not block on prepare_encode failures. try: if supports_interaction_apply(pipeline) and state.chunk_index == 0: pipe = cast(SupportsInteractionApply, pipeline) assert self._interaction_coordinator is not None, "Model not loaded. Call load_model() first." state.interaction_chunk_metadata = self._interaction_coordinator.maybe_prepare_initial_session( state, pipe ) pipe.prepare_next_chunk(state) merge_stage_durations( state, consume_pipeline_stage_durations(pipeline), ) except Exception as exc: per_req_exc = exc if _dit_any_rank_failed(per_req_exc is not None): _abort_prep_failure(per_req_exc) continue prepared_states.append(state) if not prepared_states: return prepared_states, None, error_outputs input_batch = InputBatch.make_batch( prepared_states, cached_batch=getattr(self, "input_batch", None), ) self.input_batch = input_batch return prepared_states, input_batch, error_outputs def _update_states_after( self, 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 ( input_batch.latents.size() == gathered_latents.size() and input_batch.latents.dtype == gathered_latents.dtype and input_batch.latents.device == gathered_latents.device ): input_batch.latents.copy_(gathered_latents) else: input_batch.latents = gathered_latents.clone() self.input_batch = input_batch scatter_latents(states, input_batch) for state in states: if interrupted or state.request_denoise_completed: self.state_cache.pop(state.request_id, None) def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput: """Execute one step for one scheduled request and return runner output.""" return self._execute_stepwise( scheduler_output, validate_kv_metadata=True, record_output_peak_memory=True, ) def _execute_non_step_requests(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput: """Run a complete legacy forward for requests excluded from step mode.""" if scheduler_output.scheduled_cached_reqs.request_ids: raise ValueError("A non-step fallback batch cannot contain cached stepwise requests.") runner_outputs: list[RunnerOutput] = [] for index, new_req in enumerate(scheduler_output.scheduled_new_reqs): try: result = self.execute_model( new_req.req, kv_prefetch_job=getattr(scheduler_output, "kv_prefetch_job", None) if index == 0 else None, diffusion_kv_metadata=new_req.diffusion_kv_metadata, ) except Exception as exc: logger.error( "Non-step fallback execution failed for %s", new_req.request_id, exc_info=True, ) result = DiffusionOutput(error=str(exc)) step_index = getattr(new_req.req.sampling_params, "step_index", None) runner_outputs.append( RunnerOutput( request_id=new_req.request_id, step_index=0 if step_index is None else step_index, finished=True, result=result, ) ) return BatchRunnerOutput.from_list(runner_outputs) def _execute_stepwise( self, scheduler_output: DiffusionSchedulerOutput, *, validate_kv_metadata: bool, record_output_peak_memory: bool, in_diffusion_kv_memory_profile: bool = False, ) -> BatchRunnerOutput: """Execute one step with explicit validation and profiling policy.""" assert self.pipeline is not None, "Model not loaded. Call load_model() first." # A scheduler wave can release a previous stepwise request while it # admits a full-forward fallback request. Do this before dispatch so # the fallback does not bypass normal state/KV retirement. self._cleanup_finished_step_requests(scheduler_output) for new_req in scheduler_output.scheduled_new_reqs: validate_new_request_data_identity(new_req) non_step_requests = [ new_req for new_req in scheduler_output.scheduled_new_reqs if not getattr(new_req.req, "use_step_execution", True) ] if non_step_requests: scheduled_request_count = len(scheduler_output.scheduled_new_reqs) + len( scheduler_output.scheduled_cached_reqs.request_ids ) if len(non_step_requests) != scheduled_request_count: raise ValueError("Cannot mix stepwise and non-step fallback requests in one scheduler batch.") # This wave has no stepwise requests, so the previous batch cannot # be reused and must not keep its step tensors alive. self.input_batch = None return self._execute_non_step_requests(scheduler_output) for new_req in scheduler_output.scheduled_new_reqs: if validate_kv_metadata: self._validate_diffusion_kv_metadata( request_id=new_req.req.request_id, metadata=new_req.diffusion_kv_metadata, ) 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 # similar features are not supported here yet. if self.od_config.cache_backend not in (None, "none"): raise ValueError("Step mode does not support cache_backend yet.") # Scheduler metadata is installed only for newly admitted requests; # cached requests continue to use the row installed on their first # step. Finished-state retirement is shared with full-forward # fallback dispatch above. installed_request_ids: list[str] = [] try: for new_req in scheduler_output.scheduled_new_reqs: if new_req.diffusion_kv_metadata is not None: self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata) installed_request_ids.append(new_req.request_id) return self._execute_stepwise_core( scheduler_output, record_output_peak_memory=record_output_peak_memory, in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, ) except Exception: if installed_request_ids: self.remove_diffusion_kv_requests(installed_request_ids) raise def _execute_stepwise_core( self, scheduler_output: DiffusionSchedulerOutput, *, record_output_peak_memory: bool, in_diffusion_kv_memory_profile: bool = False, ) -> BatchRunnerOutput: """Run the denoise step after metadata admission and row installation.""" use_hsdp = self.od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: pipeline = self.pipeline assert pipeline is not None, "Model not loaded. Call load_model() first." had_active_states = bool(self.state_cache) states, new_request_ids = self._update_states(scheduler_output) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 if ( record_output_peak_memory and new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available() ): current_omni_platform.reset_peak_memory_stats() states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids) if input_batch is None: return BatchRunnerOutput.from_list(runner_output_list) attn_metadata: dict[str, Any] = {} kv_backend = getattr(self, "diffusion_kv_backend", None) paged_kv_runtime = kv_backend if getattr(kv_backend, "paged_attention_adapter", None) is not None else None with set_forward_context( vllm_config=self.vllm_config, omni_diffusion_config=self.od_config, attn_metadata=attn_metadata, paged_kv_runtime=paged_kv_runtime, in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, ): clear_pipeline_stage_durations(pipeline) noise_pred = pipeline.denoise_step(input_batch, states=states) denoise_stage_durations = consume_pipeline_stage_durations(pipeline) for state in states: merge_stage_durations( state, denoise_stage_durations, ) pipeline_interrupted = getattr(pipeline, "interrupt", False) if noise_pred is None and pipeline_interrupted: for state in states: runner_output_list.append( RunnerOutput( request_id=state.request_id, step_index=state.step_index, finished=True, result=DiffusionOutput(error="stepwise denoise interrupted"), ) ) else: offset = 0 for req in states: if req.latents is None: raise RuntimeError(f"Stepwise request {req.request_id} has no latent state.") row_num = req.latents.shape[0] try: pipeline.step_scheduler( req, noise_pred[offset : offset + row_num] if noise_pred is not None else None ) if self.od_config.streaming_output: should_decode = req.chunk_denoise_completed or req.request_denoise_completed else: should_decode = req.denoise_completed if should_decode: clear_pipeline_stage_durations(pipeline) result = pipeline.post_decode(req) if result is not None: result = self._prepare_output_for_transport(result, req.sampling) self._attach_stepwise_metadata( req, result, ) # After consuming this chunk's interaction metadata, apply pending interactions and # prepare the next chunk (prepare_next_chunk may be a no-op---depending on pipeline) if supports_interaction_apply(pipeline) and not req.request_denoise_completed: pipe = cast(SupportsInteractionApply, pipeline) pipe.apply_interaction_at_chunk_boundary(req) pipe.prepare_next_chunk(req) else: result = None # finished should be computed after post_decode() advanced chunk_index finished = ( req.request_denoise_completed if self.od_config.streaming_output else req.denoise_completed ) if finished and result is not None: self._maybe_send_stage_payload([req], [result]) runner_output_list.append( RunnerOutput( request_id=req.request_id, step_index=req.step_index, finished=finished, result=result, ) ) offset = offset + row_num except Exception as per_req_exc: offset = offset + row_num self.state_cache.pop(req.request_id, None) logger.error( "Stepwise per-request error for %s: %s", req.request_id, per_req_exc, exc_info=True, ) runner_output_list.append( RunnerOutput( request_id=req.request_id, step_index=req.step_index, finished=True, result=DiffusionOutput.from_exception(per_req_exc), ) ) if noise_pred is not None and offset != noise_pred.shape[0]: raise ValueError( f"Stepwise noise_pred consumed {offset} rows, " f"but batched noise_pred has {noise_pred.shape[0]} rows." ) if is_primary and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() states_by_id = {state.request_id: state for state in states} for state in states: state.peak_memory_mb = max(state.peak_memory_mb, batch_peak_memory_mb) for runner_output in runner_output_list: if runner_output.result is None: continue matched_state = states_by_id.get(runner_output.request_id) if matched_state is None: continue runner_output.result.peak_memory_mb = max( runner_output.result.peak_memory_mb, matched_state.peak_memory_mb, ) terminal_request_ids = [ runner_output.request_id for runner_output in runner_output_list if runner_output.finished ] self._update_states_after(states, input_batch, pipeline_interrupted) if ( getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) is DiffusionKVCacheMode.PAGED_SCHEDULER and terminal_request_ids ): self.remove_diffusion_kv_requests(terminal_request_ids) return BatchRunnerOutput.from_list(runner_output_list) def submit_interaction( self, request_id: str, interaction: OmniInteractionPrompt, ) -> None: """Route a midway interaction through the pipeline interaction coordinator.""" assert self.pipeline is not None and self._interaction_coordinator is not None, ( "Model not loaded. Call load_model() first." ) if not self.od_config.streaming_output: raise ValueError("submit_interaction requires streaming_output=True") if not self._supports_step_mode(): raise ValueError("submit_interaction requires step execution support") state = self.state_cache.get(request_id) if state is None: raise ValueError(f"No active request state for interaction: {request_id!r}") event = interaction["event"] parts: list[tuple[str, InteractionPayload]] = [] prompt = event.get("prompt") if prompt is not None: parts.append(("prompt", {"prompt": prompt})) multi_modal_data = event.get("multi_modal_data") if multi_modal_data: parts.extend( (str(modality), cast(InteractionPayload, payload)) for modality, payload in multi_modal_data.items() ) self._interaction_coordinator.enqueue_parts( state, parts=parts, event_id=interaction["event_id"], received_at=time.monotonic(), transition_chunks=interaction.get("transition_chunks"), )