Merge pull request #16 from tzhouam/feat/diffusion_gpu_worker_and_model_runner

[Worker]Refactor GPU diffusion model runner and worker

vllm_omni/worker/gpu_diffusion_model_runner.py

@@ -1,133 +1,362 @@
-"""Model runner that wraps a diffusers text-to-image pipeline."""
-
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Optional, Union, List
+import numpy as np
+import gc
+import logging
 
-@dataclass
-class DiffusionRunnerOutput:
-    """Structured output returned by the diffusion runner."""
+import torch
 
-    prompt: str
-    images: list
-    request_id: str = "diffusion_request"
-    finished: bool = True
-    output_type: str = "image"
+from vllm.v1.worker.gpu_model_runner import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    IntermediateTensors,
+    get_pp_group,
+    set_forward_context,
+)
+from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
+from vllm.multimodal.inputs import MultiModalKwargs
+
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 
-class DiffusionModelRunner:
-    """Lightweight runner that loads and executes diffusers pipelines."""
+logger = logging.getLogger(__name__)
 
-    def __init__(
+
+class GPUDiffusionModelRunner(OmniGPUModelRunner):
+    """Diffusion model runner for vLLM-omni (non-autoregressive).
+
+    - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
+    - Does not compute logits or perform token sampling.
+    - Executes diffusion process and returns tensors via `pooler_output`.
+    """
+
+    @torch.inference_mode()
+    def execute_model(
         self,
-        model_path: str,
-        *,
-        pipeline_name: Optional[str] = None,
-        device: Optional[str] = None,
-        dtype: Optional[str] = None,
-    ) -> None:
-        self.model_path = model_path
-        self.pipeline_name = pipeline_name or "auto"
-        self.device, self.torch_dtype = self._resolve_device_and_dtype(device, dtype)
-        self._pipeline = self._load_pipeline()
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
+        self._update_states(scheduler_output)
 
-    def _resolve_device_and_dtype(
-        self,
-        device: Optional[str],
-        dtype: Optional[str],
-    ):
-        import torch
+        if not scheduler_output.total_num_scheduled_tokens:
 
-        resolved_device = device
-        if resolved_device is None:
-            if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-                resolved_device = "mps"
-            elif torch.cuda.is_available():
-                resolved_device = "cuda"
-            else:
-                resolved_device = "cpu"
+            return EMPTY_MODEL_RUNNER_OUTPUT
+           
+        # Prepare decoder inputs and attention metadata (for batch/order mapping)
+        (attn_metadata, attention_cuda_graphs, logits_indices,
+         spec_decode_metadata, num_scheduled_tokens_np,
+         spec_decode_common_attn_metadata) = self._prepare_inputs(
+             scheduler_output)
 
-        if dtype is not None:
-            resolved_dtype = getattr(torch, dtype, None)
+        # Input token count for this iteration (not used by diffusion, but
+        # retained to keep DP padding/ordering consistent)
+        num_input_tokens = scheduler_output.total_num_scheduled_tokens
+        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
+        num_input_tokens += num_pad
+
+        # Multimodal conditioning (e.g., text/audio/video encoders)
+        if self.is_multimodal_model:
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
         else:
-            if resolved_device in {"cuda", "mps"}:
-                resolved_dtype = torch.float16
-            else:
-                resolved_dtype = torch.float32
+            mm_embeds = []
 
-        return resolved_device, resolved_dtype
-
-    def _load_pipeline(self):
-        import diffusers
-
-        if self.pipeline_name != "auto":
-            PipelineCls = getattr(diffusers, self.pipeline_name)
+        # Build inputs to mirror AR runner: input_ids/positions/embeds
+        if self.is_multimodal_model and get_pp_group().is_first_rank:
+            inputs_embeds_scheduled = self.model.get_input_embeddings(
+                input_ids=self.input_ids[:scheduler_output.total_num_scheduled_tokens],
+                multimodal_embeddings=mm_embeds or None,
+            )
+            self.inputs_embeds[:scheduler_output.total_num_scheduled_tokens].copy_(
+                inputs_embeds_scheduled)
+            input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds[:num_input_tokens]
+            model_mm_kwargs = self._extract_mm_kwargs(scheduler_output)
         else:
-            PipelineCls = diffusers.AutoPipelineForText2Image
+            input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = None
+            model_mm_kwargs = {}
 
-        pipeline = PipelineCls.from_pretrained(
-            self.model_path,
-            torch_dtype=self.torch_dtype,
+        # Positions (mrope or standard)
+        if self.uses_mrope:
+            positions = self.mrope_positions[:, :num_input_tokens]
+        else:
+            positions = self.positions[:num_input_tokens]
+
+        # Intermediate tensors sync for PP (if any)
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True)
+
+        # Set forward context mainly for resource management and kv connector
+        skip_cuda_graphs = True  # diffusion path does not rely on cuda graphs here
+        with set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                skip_cuda_graphs=skip_cuda_graphs,
+        ), self.maybe_get_kv_connector_output(
+                scheduler_output) as kv_connector_output:
+
+            if not get_pp_group().is_last_rank:
+                # For non-last PP stages, pass through intermediate tensors.
+                assert intermediate_tensors is not None
+                intermediate_tensors.kv_connector_output = kv_connector_output
+                return intermediate_tensors
+
+            outputs = self._run_diffusion(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                multimodal_kwargs=model_mm_kwargs,
+                logits_indices=logits_indices,
+            )
+        _, multimodal_outputs = (
+            self.extract_multimodal_outputs(outputs))
+
+        # Ensure one tensor per request, map to CPU for output struct
+        pooler_output: List[Optional[torch.Tensor]] = []
+        if isinstance(multimodal_outputs, torch.Tensor):
+            # If model returned a single stacked tensor, split by requests
+            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
+            for i in range(self.input_batch.num_reqs):
+                pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
+        elif isinstance(multimodal_outputs, list):
+            for out in multimodal_outputs:
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+        elif isinstance(multimodal_outputs, dict):
+            for out in multimodal_outputs.values():
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+        else:
+            raise RuntimeError("Unsupported diffusion output type")
+
+        self.eplb_step()
+
+        return OmniModelRunnerOutput(
+            req_ids=self.input_batch.req_ids,
+            req_id_to_index=self.input_batch.req_id_to_index,
+            sampled_token_ids=[],
+            spec_token_ids=None,
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=pooler_output,
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits={},
         )
 
-        try:
-            pipeline = pipeline.to(self.device)
-        except Exception:
-            pipeline = pipeline.to("cpu")
-
-        self._apply_optimisations(pipeline)
-        return pipeline
-
-    @staticmethod
-    def _apply_optimisations(pipeline) -> None:
-        # Enable attention slicing / VAE tiling where available to reduce memory.
-        try:
-            if hasattr(pipeline, "enable_attention_slicing"):
-                pipeline.enable_attention_slicing()
-        except Exception:
-            pass
-        try:
-            vae = getattr(pipeline, "vae", None)
-            if vae and hasattr(vae, "enable_tiling"):
-                vae.enable_tiling()
-        except Exception:
-            pass
-
-    def generate(
+    def _run_diffusion(
         self,
-        prompt: str,
         *,
-        height: int = 512,
-        width: int = 512,
-        num_inference_steps: int = 30,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[str] = None,
-        seed: Optional[int] = None,
-        image: Optional[Any] = None,
-    ) -> DiffusionRunnerOutput:
-        import torch
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors],
+        inputs_embeds: Optional[torch.Tensor],
+        multimodal_kwargs: dict,
+        logits_indices: torch.Tensor,
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        """Runs the diffusion process and returns per-request tensors.
 
-        generator = None
-        if seed is not None:
-            try:
-                generator = torch.Generator(device=self.device).manual_seed(int(seed))
-            except Exception:
-                # Some backends (e.g., CPU without full support) may not accept the device.
-                generator = torch.Generator().manual_seed(int(seed))
-
-        output = self._pipeline(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=int(num_inference_steps),
-            guidance_scale=float(guidance_scale),
-            height=int(height),
-            width=int(width),
-            generator=generator,
-            image=image,
+        Tries model interfaces in the following order for maximal compatibility:
+        1) model.sample(condition=..., **kwargs)
+        2) model.forward(condition=..., **kwargs)
+        3) model.diffuse(condition=..., **kwargs)
+        """
+        # Keep inputs identical to AR runner
+        kwargs = dict(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
+            sampling_metadata=self.input_batch.sampling_metadata,
+            logits_index=logits_indices,
+            sampler=self.sampler,
         )
 
-        return DiffusionRunnerOutput(
-            prompt=prompt,
-            images=getattr(output, "images", []) or [],
-        )
+        # For Qwen 2.5 Omni's current implementation, we only support the forward method
+        if hasattr(self.model, "forward"):
+            return self.model.forward(**kwargs)
+        #TODO: add the diffuse method for other models
+
+        raise RuntimeError(
+            "The loaded model does not expose diffusion interfaces 'sample', "
+            "'forward', or 'diffuse'. Please implement one of them or adapt the runner.")
+
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        capture_attn_cudagraph: bool = False,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Padding for DP
+        num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+        num_tokens += num_pad
+
+        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+        # for dummy run with LoRA so that the num_reqs collectively
+        # has num_tokens in total.
+        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        max_num_reqs = self.scheduler_config.max_num_seqs
+        num_reqs = min(num_tokens, max_num_reqs)
+        min_tokens_per_req = num_tokens // num_reqs
+        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+        assert sum(num_scheduled_tokens_list) == num_tokens
+        assert len(num_scheduled_tokens_list) == num_reqs
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
+                                        dtype=np.int32)
+
+        attn_metadata: Optional[dict[str, dict]] = None
+        if capture_attn_cudagraph:
+            attn_metadata = {}
+
+            # Make sure max_model_len is used at the graph capture time.
+            self.seq_lens_np[:num_reqs] = self.max_model_len
+            self.seq_lens_np[num_reqs:] = 0
+            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
+                                           non_blocking=True)
+
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                    self.kv_cache_config.kv_cache_groups):
+                common_attn_metadata = CommonAttentionMetadata(
+                    query_start_loc=self.query_start_loc[:num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
+                    seq_lens=self.seq_lens[:num_reqs],
+                    seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.
+                    num_computed_tokens_cpu_tensor[:num_reqs],
+                    num_reqs=num_reqs,
+                    num_actual_tokens=num_tokens,
+                    max_query_len=num_tokens,
+                    block_table_tensor=self.input_batch.block_table[
+                        kv_cache_group_id].get_device_tensor()[:num_reqs],
+                    slot_mapping=self.input_batch.block_table[
+                        kv_cache_group_id].slot_mapping[:num_tokens],
+                    causal=True)
+
+                for attn_group in self.attn_groups[kv_cache_group_id]:
+                    attn_metadata_i = attn_group.metadata_builder \
+                        .build_for_cudagraph_capture(common_attn_metadata)
+                    for layer_name in kv_cache_group_spec.layer_names:
+                        attn_metadata[layer_name] = attn_metadata_i
+
+        with self.maybe_dummy_run_with_lora(self.lora_config,
+                                            num_scheduled_tokens):
+            if self.is_multimodal_model:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds[:num_tokens]
+                model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
+            else:
+                input_ids = self.input_ids[:num_tokens]
+                inputs_embeds = None
+                model_mm_kwargs = {}
+
+            if self.uses_mrope:
+                positions = self.mrope_positions[:, :num_tokens]
+            else:
+                positions = self.positions[:num_tokens]
+
+            if get_pp_group().is_first_rank:
+                intermediate_tensors = None
+            else:
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = (
+                        self.model.make_empty_intermediate_tensors(
+                            batch_size=self.max_num_tokens,
+                            dtype=self.model_config.dtype,
+                            device=self.device))
+
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                    num_tokens, None, False)
+
+            # Diffusion path: avoid CUDA graphs; we only use context for resource wiring
+            with self.maybe_randomize_inputs(input_ids), set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens,
+                    num_tokens_across_dp=num_tokens_across_dp):
+                outputs = self.model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **MultiModalKwargs.as_kwargs(
+                        model_mm_kwargs,
+                        device=self.device,
+                    ),
+                    sampler=None,
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, _ = outputs
+            else:
+                hidden_states = outputs
+
+            # Extract multimodal outputs if present; we ignore them here because
+            # dummy run returns tensors only. The actual diffusion runner returns
+            # multimodal outputs via pooler_output in execute_model.
+            text_hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+
+            if not skip_eplb:
+                self.eplb_step(is_dummy=True, is_profile=is_profile)
+
+        logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        return text_hidden_states, None
+    
+    @torch.inference_mode()
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        logger.warning("Dummy sampler run is not implemented for diffusion model")
+        return None
+
+    def profile_run(self) -> None:
+        # Profile with multimodal encoder & encoder cache, similar to base but
+        # without any logits/sampler warming.
+        if self.is_multimodal_model:
+            mm_budget = self.mm_budget
+            assert mm_budget is not None
+
+            # TODO: handle encoder-decoder models once supported.
+            if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
+                (
+                    dummy_modality,
+                    max_tokens,
+                ) = mm_budget.get_modality_with_max_tokens()
+                (
+                    max_mm_items_per_prompt,
+                    max_mm_items_per_batch,
+                ) = mm_budget.get_max_items(dummy_modality, max_tokens)
+
+                batched_dummy_mm_inputs = self._get_mm_dummy_batch(
+                    dummy_modality,
+                    max_mm_items_per_batch,
+                )
+
+                dummy_encoder_outputs = self.model.get_multimodal_embeddings(
+                    **batched_dummy_mm_inputs)
+
+                sanity_check_mm_encoder_outputs(
+                    dummy_encoder_outputs,
+                    expected_num_items=max_mm_items_per_batch,
+                )
+
+                self.encoder_cache["tmp"] = dict(
+                    enumerate(dummy_encoder_outputs))
+
+        hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
+        if get_pp_group().is_last_rank:
+            pass  # No sampler/pooler warmup for diffusion
+        self._sync_device()
+        del hidden_states
+        self.encoder_cache.clear()
+        gc.collect()

vllm_omni/worker/gpu_diffusion_worker.py

@@ -1,56 +1,67 @@
-"""Worker abstraction for executing diffusion model runners."""
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+import torch
+import os
+import gc
+from vllm.utils import GiB_bytes, MemorySnapshot
+from vllm.platforms import current_platform
+from vllm.worker.worker import _check_if_gpu_supports_dtype
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.model_executor import set_random_seed
+from vllm.v1.utils import report_usage_stats
 
-from __future__ import annotations
-
-from typing import Any, Optional
-
-from vllm.v1.worker.worker_base import WorkerBase
-
-from .gpu_diffusion_model_runner import DiffusionModelRunner, DiffusionRunnerOutput
+from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
 
 
-class DiffusionGPUWorker(WorkerBase):
-    """Minimal worker wrapping DiffusionModelRunner.
+class GPUDiffusionWorker(GPUWorker):
+    def init_device(self):
+        if self.device_config.device.type == "cuda":
+            # torch.distributed.all_reduce does not free the input tensor until
+            # the synchronization point. This causes the memory usage to grow
+            # as the number of all_reduce calls increases. This env var disables
+            # this behavior.
+            # Related issue:
+            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
+            os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
 
-    The worker interface mirrors the shape expected by vLLM executors: it owns
-    a model runner instance and exposes a `generate` method that can be called
-    by an executor implementation.
-    """
+            # This env var set by Ray causes exceptions with graph building.
+            os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+            self.device = torch.device(f"cuda:{self.local_rank}")
+            current_platform.set_device(self.device)
 
-    def __init__(
-        self,
-        model_path: str,
-        *,
-        pipeline_name: Optional[str] = None,
-        device: Optional[str] = None,
-        dtype: Optional[str] = None,
-    ) -> None:
-        self.model_runner = DiffusionModelRunner(
-            model_path,
-            pipeline_name=pipeline_name,
-            device=device,
-            dtype=dtype,
-        )
+            _check_if_gpu_supports_dtype(self.model_config.dtype)
+            gc.collect()
+            torch.cuda.empty_cache()
 
-    def generate(
-        self,
-        prompt: str,
-        *,
-        height: int = 512,
-        width: int = 512,
-        num_inference_steps: int = 30,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[str] = None,
-        seed: Optional[int] = None,
-        image: Optional[Any] = None,
-    ) -> DiffusionRunnerOutput:
-        return self.model_runner.generate(
-            prompt=prompt,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            negative_prompt=negative_prompt,
-            seed=seed,
-            image=image,
-        )
+            # take current memory snapshot
+            self.init_snapshot = MemorySnapshot()
+            self.requested_memory = (self.init_snapshot.total_memory *
+                                     self.cache_config.gpu_memory_utilization)
+            if self.init_snapshot.free_memory < self.requested_memory:
+                GiB = lambda b: round(b / GiB_bytes, 2)
+                raise ValueError(
+                    f"Free memory on device "
+                    f"({GiB(self.init_snapshot.free_memory)}/"
+                    f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
+                    f"is less than desired GPU memory utilization "
+                    f"({self.cache_config.gpu_memory_utilization}, "
+                    f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
+                    f"utilization or reduce GPU memory used by other processes."
+                )
+        else:
+            raise RuntimeError(
+                f"Not support device type: {self.device_config.device}")
+        # Initialize the distributed environment.
+        init_worker_distributed_environment(self.vllm_config, self.rank,
+                                            self.distributed_init_method,
+                                            self.local_rank,
+                                            current_platform.dist_backend)
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+        # Construct the model runner
+        self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
+            self.vllm_config, self.device)
+
+        if self.rank == 0:
+            # If usage stat is enabled, collect relevant info.
+            report_usage_stats(self.vllm_config)