comment

Signed-off-by: Roger Wang <hey@rogerw.io>

.buildkite/pipeline.yml

@@ -1,6 +1,63 @@
 steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge_us_east_1"
+
   - label: "Simple Unit Test"
     commands:
       - ".buildkite/scripts/simple_test.sh"
     agents:
       queue: "cpu_queue_premerge"
+  - label: "Diffusion Model Test"
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/single_stage/test_diffusion_model.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+  - label: "Large Model Test (H100)"
+    depends_on: image-build
+    agents:
+      queue: mithril-h100-pool
+    commands:
+      - pytest -s -v tests/single_stage/test_diffusion_model.py # update this to your actual test file
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

.gitignore

@@ -156,9 +156,6 @@ cython_debug/
 # cursor
 .cursor/
 
-# docker
-docker/
-
 # scripts
 /scripts/
 

docker/Dockerfile.ci

@@ -0,0 +1,14 @@
+ARG VLLM_BASE_IMAGE=vllm/vllm-openai
+ARG VLLM_BASE_TAG=v0.11.0
+FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
+ARG APP_DIR=/workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Install vllm-omni into the same uv-managed Python environment used by the base image.
+RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENTRYPOINT []

docs/community/volunteers.md

@@ -0,0 +1,12 @@
+# Volunteers for Bugfix and CI
+
+We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
+
+If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
+
+| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
+|----------|----------|----------|----------|----------|----------|
+| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
+| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
+
+We kindly welcome more contributors to fix bugs and contribute new features!

docs/design/vllm_omni_design.md

@@ -16,7 +16,7 @@ vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressiv
 
 API Server --> OmniLLM/AsyncOmni (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
  --> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
- --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessoer (New one for final multimodal output)
+ --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessor (New one for final multimodal output)
 
 ## Core Components (Ordered by Data Flow)
 
@@ -137,7 +137,7 @@ class OmniLLM(LLM):
             prompt_str, engine_request, tokenization_kwargs = self._process_stage_inputs(stage_config, **stage_args)
 
             # Add inputs to Engine
-            stage_engine.add_request(requesy_id, prompt_str, tokenization_kwargs)
+            stage_engine.add_request(request_id, prompt_str, tokenization_kwargs)
             # Run Engine
             stage_output = stage_engine.step()
 
@@ -157,7 +157,7 @@ class OmniLLM(LLM):
             return self._process_ar_inputs(**stage_args)
         elif stage_config.engine_type == "DiT":
             return self._process_dit_inputs(**stage_args)
-        else
+        else:
             raise NotImplementedError
 
     def _process_dit_inputs(**stage_args)-> Any:

docs/getting_started/installation/README.md

@@ -1,6 +1,6 @@
 # Installation
 
-vLLM supports the following hardware platforms:
+vLLM-Omni supports the following hardware platforms:
 
 - [GPU](gpu.md)
     - [NVIDIA CUDA](gpu.md)

docs/getting_started/installation/gpu.md

@@ -37,4 +37,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 ## Set up using Docker
 
 ### Pre-built images
-To be released... Stay tuned!
+
+=== "NVIDIA CUDA"
+
+    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"

docs/getting_started/installation/gpu/cuda.inc.md

@@ -73,3 +73,24 @@ uv pip install --no-build-isolation --editable .
 </details>
 
 # --8<-- [end:build-wheel-from-source]
+
+
+# --8<-- [start:pre-built-images]
+
+vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
+
+Here's an example deployment command that has been verified on 2 x H100's:
+```bash
+docker run --runtime nvidia --gpus 2 \
+    -v ~/.cache/huggingface:/root/.cache/huggingface \
+    --env "HF_TOKEN=$HF_TOKEN" \
+    -p 8091:8091 \
+    --ipc=host \
+    vllm/vllm-omni:v0.11.0rc1 \
+    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+```
+
+!!! tip
+    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+
+# --8<-- [end:pre-built-images]

docs/mkdocs/stylesheets/extra.css

@@ -25,8 +25,11 @@ a:not(:has(svg)):not(.md-icon):not(.autorefs-external) {
 
 a[href*="localhost"]::after,
 a[href*="127.0.0.1"]::after,
-a[href*="org.readthedocs.build"]::after,
-a[href*="docs.vllm.ai"]::after {
+
+/* Hide external link icons for all links */
+a[href^="//"]::after,
+a[href^="http://"]::after,
+a[href^="https://"]::after {
     display: none !important;
 }
 

docs/models/supported_models.md

@@ -23,7 +23,7 @@ th {
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 
 
-## List of Supported Models for Ascend NPU
+## List of Supported Models for NPU
 
 <style>
 th {

docs/usage/faq.md

@@ -2,7 +2,7 @@
 
 > Q: How many chips do I need to infer a model in vLLM-Omni?
 
-A: Now, we support natively disaggregated deployment for different model stages within a model. There is a restriction that one chip can only has one AutoRegressive model stage. This is because the unified KV cache management of vLLM. Stages of other types can coexist with in a chip. The restriction will be resolved in later version.
+A: Now, we support natively disaggregated deployment for different model stages within a model. There is a restriction that one chip can only have one AutoRegressive model stage. This is because the unified KV cache management of vLLM. Stages of other types can coexist within a chip. The restriction will be resolved in later version.
 
 > Q: When trying to run examples, I encounter error about backend of librosa or soundfile. How to solve it?
 
@@ -11,3 +11,15 @@ A: If you encounter error about backend of librosa, try to install ffmpeg with c
 sudo apt update
 sudo apt install ffmpeg
 ```
+
+> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
+
+A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
+
+> Q: Does vLLM-Omni support AWQ or any other quantization?
+
+A: vLLM-Omni partitions model into several stages. For AR stages, it will reuse main logic of LLMEngine in vLLM. So current quantization supported in vLLM should be also supported in vLLM-Omni for them. But systematic verification is ongoing. For quantization for DiffusionEngine, we are working on it. Please stay tuned and welcome contribution!
+
+> Q: Does vLLM-Omni support multimodal streaming input and output?
+
+A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -8,6 +8,7 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
 Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
 ## Run examples (Qwen2.5-Omni)
+
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -26,6 +27,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
+
 ### Single Prompt
 Get into the example folder
 ```bash
@@ -36,6 +38,34 @@ Then run the command below.
 bash run_single_prompt.sh
 ```
 
+#### Using Local Media Files
+The `end2end.py` script supports local media files (audio, video, image) via CLI arguments:
+
+```bash
+# Use single local media files
+python end2end.py --query-type use_image --image-path /path/to/image.jpg
+python end2end.py --query-type use_video --video-path /path/to/video.mp4
+python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+
+# Combine multiple local media files
+python end2end.py --query-type mixed_modalities \
+    --video-path /path/to/video.mp4 \
+    --image-path /path/to/image.jpg \
+    --audio-path /path/to/audio.wav
+
+# Use audio from video file
+python end2end.py --query-type use_audio_in_video --video-path /path/to/video.mp4
+
+```
+
+If media file paths are not provided, the script will use default assets. Supported query types:
+- `use_image`: Image input only
+- `use_video`: Video input only
+- `use_audio`: Audio input only
+- `mixed_modalities`: Audio + image + video
+- `use_audio_in_video`: Extract audio from video
+- `text`: Text-only query
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -8,6 +8,7 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
 Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
 ## Run examples (Qwen3-Omni)
+
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
 Get into the example folder
@@ -32,6 +33,26 @@ If you have not enough memory, you can set thinker with tensor parallel. Just ru
 bash run_single_prompt_tp.sh
 ```
 
+#### Using Local Media Files
+The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
+
+```bash
+# Use local video file
+python end2end.py --query-type use_video --video-path /path/to/video.mp4
+
+# Use local image file
+python end2end.py --query-type use_image --image-path /path/to/image.jpg
+
+# Use local audio file
+python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+```
+
+If media file paths are not provided, the script will use default assets. Supported query types:
+- `use_video`: Video input
+- `use_image`: Image input
+- `use_audio`: Audio input
+- `text`: Text-only query
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -0,0 +1,36 @@
+# Stage config for running Qwen3-Omni-MoE-Thinking (text-only output)
+# This config is for models like Qwen3-Omni-30B-A3B-Thinking that only have the
+# thinker component and do not support audio output.
+#
+# Single stage: Thinker (multimodal understanding + text generation)
+
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05

tests/single_stage/test_diffusion_model.py

@@ -0,0 +1,28 @@
+import pytest
+import torch
+
+from vllm_omni import Omni
+
+models = ["Tongyi-MAI/Z-Image-Turbo"]
+
+
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_model(model_name: str):
+    m = Omni(model=model_name)
+    # high resolution may cause OOM on L4
+    height = 256
+    width = 256
+    images = m.generate(
+        "a photo of a cat sitting on a laptop keyboard",
+        height=height,
+        width=width,
+        num_inference_steps=2,
+        guidance_scale=0.0,
+        generator=torch.Generator("cuda").manual_seed(42),
+        num_outputs_per_prompt=2,
+    )
+    assert len(images) == 2
+    # check image size
+    assert images[0].width == width
+    assert images[0].height == height
+    images[0].save("z_image_output.png")

vllm_omni/core/sched/__init__.py

@@ -2,12 +2,12 @@
 Scheduling components for vLLM-Omni.
 """
 
-from .diffusion_scheduler import DiffusionScheduler
+from .omni_ar_scheduler import OmniARScheduler
+from .omni_generation_scheduler import OmniGenerationScheduler
 from .output import OmniNewRequestData
-from .scheduler import OmniScheduler
 
 __all__ = [
-    "OmniScheduler",
-    "DiffusionScheduler",
+    "OmniARScheduler",
+    "OmniGenerationScheduler",
     "OmniNewRequestData",
 ]

vllm_omni/core/sched/diffusion_scheduler.py

@@ -1,315 +0,0 @@
-import time
-from collections import defaultdict
-from typing import Optional
-
-from vllm.v1.core.kv_cache_manager import KVCacheBlocks
-from vllm.v1.core.sched.request_queue import create_request_queue
-from vllm.v1.core.sched.scheduler import Request, RequestStatus, SchedulerOutput, SpecDecodingStats
-from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
-
-from vllm_omni.core.sched.output import OmniNewRequestData
-from vllm_omni.core.sched.scheduler import OmniScheduler
-from vllm_omni.outputs import OmniModelRunnerOutput
-
-
-class DiffusionScheduler(OmniScheduler):
-    """Scheduler for diffusion-based non-autoregressive stages.
-
-    Implements a fast-path scheduling strategy for diffusion models where
-    all input tokens are fed at once. Falls back to standard vLLM scheduling
-    if the token budget cannot accommodate all tokens in a single iteration.
-    """
-
-    def schedule(self) -> SchedulerOutput:
-        """Schedule requests for diffusion model execution.
-
-        Uses a fast-path approach that feeds all input tokens at once
-        (or 1 placeholder token if input is empty). Falls back to
-        standard scheduling if token budget is insufficient.
-
-        Returns:
-            SchedulerOutput containing scheduled requests and token allocations
-        """
-
-        # Select requests with zero prompt and using pooling
-        # (diffusion results returned via pooler_output)
-        token_budget = self.max_num_scheduled_tokens
-        capacity = self.max_num_running_reqs - len(self.running)
-        scheduled_timestamp = time.monotonic()
-
-        scheduled_new_reqs: list[Request] = []
-        scheduled_resumed_reqs: list[Request] = []
-        scheduled_running_reqs: list[Request] = []
-
-        req_to_new_blocks: dict[str, KVCacheBlocks] = {}
-        num_scheduled_tokens: dict[str, int] = {}
-        scheduled_spec_decode_tokens: dict[str, list[int]] = {}
-        scheduled_encoder_inputs: dict[str, list[int]] = {}
-
-        # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
-        skipped_waiting_requests = create_request_queue(self.policy)
-
-        # Fast path selection and scheduling (treat all as diffusion requests,
-        # independent of pooling_params)
-        while self.waiting and token_budget > 0 and capacity > 0:
-            request = self.waiting.peek_request()
-            # Uniformly treat as diffusion. A feature flag can be added later
-            # via config or request tag.
-
-            # Allocate all input tokens for the request in one shot
-            # (allocate 1 placeholder if zero)
-            required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
-            if required_tokens > token_budget:
-                # Insufficient budget to process all inputs at once;
-                # stop fast path attempt
-                break
-            num_new_tokens = required_tokens
-            new_blocks = self.kv_cache_manager.allocate_slots(
-                request,
-                num_new_tokens,
-                num_lookahead_tokens=self.num_lookahead_tokens,
-            )
-            if new_blocks is None:
-                # Allocation failed (e.g., VRAM pressure); stop fast path and
-                # fall back to default scheduling
-                # Put the current request back to the head of the waiting queue
-                # Note: the original queue order is preserved
-                break
-
-            # Officially schedule this request
-            request = self.waiting.pop_request()
-            self.running.append(request)
-            request.status = RequestStatus.RUNNING
-            if self.log_stats:
-                request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
-
-            req_to_new_blocks[request.request_id] = new_blocks
-            num_scheduled_tokens[request.request_id] = num_new_tokens
-            token_budget -= num_new_tokens
-            capacity -= 1
-            scheduled_new_reqs.append(request)
-
-        # Return skipped waiting requests
-        if skipped_waiting_requests:
-            self.waiting.prepend_requests(skipped_waiting_requests)
-
-        # If fast path scheduled none, fall back to the original scheduling
-        if not num_scheduled_tokens:
-            return super().schedule()
-
-        # Compute common prefix blocks (aligned with v1)
-        num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
-        if self.running:
-            any_request = self.running[0]
-            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
-                any_request, len(self.running)
-            )
-
-        # Assemble SchedulerOutput
-        new_reqs_data = [
-            OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
-            for req in scheduled_new_reqs
-        ]
-        cached_reqs_data = self._make_cached_request_data(
-            scheduled_running_reqs,
-            scheduled_resumed_reqs,
-            num_scheduled_tokens,
-            scheduled_spec_decode_tokens,
-            req_to_new_blocks,
-        )
-        scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
-        structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
-            scheduled_requests, scheduled_spec_decode_tokens
-        )
-
-        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
-        scheduler_output = SchedulerOutput(
-            scheduled_new_reqs=new_reqs_data,
-            scheduled_cached_reqs=cached_reqs_data,
-            num_scheduled_tokens=num_scheduled_tokens,
-            total_num_scheduled_tokens=total_num_scheduled_tokens,
-            scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
-            scheduled_encoder_inputs=scheduled_encoder_inputs,
-            num_common_prefix_blocks=num_common_prefix_blocks,
-            finished_req_ids=self.finished_req_ids,
-            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
-            structured_output_request_ids=structured_output_request_ids,
-            grammar_bitmask=grammar_bitmask,
-        )
-
-        # KVTransfer: package metadata
-        if self.connector is not None:
-            meta = self.connector.build_connector_meta(scheduler_output)
-            scheduler_output.kv_connector_metadata = meta
-
-        # Update internal state (advance num_computed_tokens, free encoder inputs,
-        # etc.)
-        self._update_after_schedule(scheduler_output)
-        return scheduler_output
-
-    """
-    Scheduler for the diffusion model.
-    This scheduler is modified to stop the request immediately for the diffusion model.
-    This is because the diffusion model can generate the final image/audio in one step.
-    Note: This is just a minimal modification to the original scheduler,
-    and there should be some further efforts to optimize the scheduler.
-    The original scheduler is still used for the AR model.
-    """
-
-    def update_from_output(
-        self,
-        scheduler_output: SchedulerOutput,
-        model_runner_output: OmniModelRunnerOutput,
-    ) -> dict[int, EngineCoreOutputs]:
-        """Update the scheduler state based on the model runner output.
-
-        This method is modified to stop the request immediately for the diffusion model.
-        """
-        sampled_token_ids = model_runner_output.sampled_token_ids
-        logprobs = model_runner_output.logprobs
-        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
-        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
-        pooler_outputs = model_runner_output.pooler_output
-        num_nans_in_logits = model_runner_output.num_nans_in_logits
-        kv_connector_output = model_runner_output.kv_connector_output
-
-        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
-        spec_decoding_stats: Optional[SpecDecodingStats] = None
-        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
-
-        # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
-        # the below loop can be a performance bottleneck. We should do our best
-        # to avoid expensive operations inside the loop.
-        stopped_running_reqs: set[Request] = set()
-        stopped_preempted_reqs: set[Request] = set()
-        for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
-            assert num_tokens_scheduled > 0
-            request = self.requests.get(req_id)
-            if request is None:
-                # The request is already finished. This can happen if the
-                # request is aborted while the model is executing it (e.g.,
-                # in pipeline parallelism).
-                continue
-
-            req_index = model_runner_output.req_id_to_index[req_id]
-            generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
-
-            scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
-            if scheduled_spec_token_ids:
-                num_draft_tokens = len(scheduled_spec_token_ids)
-                num_accepted = len(generated_token_ids) - 1
-                num_rejected = num_draft_tokens - num_accepted
-                # num_computed_tokens represents the number of tokens
-                # processed in the current step, considering scheduled
-                # tokens and rejections. If some tokens are rejected,
-                # num_computed_tokens is decreased by the number of rejected
-                # tokens.
-                request.num_computed_tokens -= num_rejected
-                spec_decoding_stats = self.make_spec_decoding_stats(
-                    spec_decoding_stats,
-                    num_draft_tokens=num_draft_tokens,
-                    num_accepted_tokens=num_accepted,
-                )
-
-            new_logprobs = None
-            new_token_ids = generated_token_ids
-            kv_transfer_params = None
-            status_before_stop = request.status
-            pooler_output = None
-            if pooler_outputs:
-                pooler_output = pooler_outputs[req_index]
-
-            # Diffusion request: completes in one step; mark finished and free resources
-            request.status = RequestStatus.FINISHED_STOPPED
-            # Optional: set a stop_reason for front-end clarity
-            # (does not affect protocol)
-            request.stop_reason = request.stop_reason or "diffusion_done"
-            kv_transfer_params = self._free_request(request)
-            if status_before_stop == RequestStatus.RUNNING:
-                stopped_running_reqs.add(request)
-            else:
-                stopped_preempted_reqs.add(request)
-
-            # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
-                # NOTE: once we support N tokens per step (spec decode),
-                # the outer lists can be of length > 1.
-                new_logprobs = logprobs.slice(req_index, req_index + 1)
-
-            if new_token_ids and self.structured_output_manager.should_advance(request):
-                # NOTE: structured_output_request should not be None if
-                # use_structured_output, we have check above, so safe to ignore
-                # type warning
-                request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids)  # type: ignore[union-attr]  # noqa: E501
-
-            # spec_token_ids comes from the model runner output
-            if num_nans_in_logits is not None and req_id in num_nans_in_logits:
-                request.num_nans_in_logits = num_nans_in_logits[req_id]
-
-            # Get prompt logprobs for this request.
-            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
-            pooling_output_dict = None
-            if pooler_output is not None:
-                # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
-                # which expects Optional[dict[str, torch.Tensor]]
-                pooling_output_dict = {"model_outputs": pooler_output}
-            if new_token_ids or pooler_output is not None or kv_transfer_params:
-                # Add EngineCoreOutput for this Request.
-                outputs[request.client_index].append(
-                    EngineCoreOutput(
-                        request_id=req_id,
-                        new_token_ids=new_token_ids,
-                        finish_reason=request.get_finished_reason(),
-                        new_logprobs=new_logprobs,
-                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooling_output_dict,
-                        stop_reason=request.stop_reason,
-                        events=request.take_events(),
-                        kv_transfer_params=kv_transfer_params,
-                        trace_headers=request.trace_headers,
-                        num_cached_tokens=request.num_cached_tokens,
-                    )
-                )
-            else:
-                # Invariant: EngineCore returns no partial prefill outputs.
-                assert not prompt_logprobs_tensors
-
-        # Remove the stopped requests from the running and waiting queues.
-        if stopped_running_reqs:
-            if stopped_running_reqs:
-                self.running = remove_all(self.running, stopped_running_reqs)
-        if stopped_preempted_reqs:
-            # This is a rare case and unlikely to impact performance.
-            self.waiting.remove_requests(stopped_preempted_reqs)
-
-        # KV Connector: update state for finished KV Transfers.
-        if model_runner_output.kv_connector_output:
-            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
-
-        # Create EngineCoreOutputs for all clients that have requests with
-        # outputs in this step.
-        engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
-
-        finished_req_ids = self.finished_req_ids_dict
-        if finished_req_ids:
-            # Include ids of requests that finished since last outputs
-            # were sent.
-            for client_index, finished_set in finished_req_ids.items():
-                # Set finished request set in EngineCoreOutputs for this client.
-                if (eco := engine_core_outputs.get(client_index)) is not None:
-                    eco.finished_requests = finished_set
-                else:
-                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
-            finished_req_ids.clear()
-
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
-            # Return stats to only one of the front-ends.
-            if (eco := next(iter(engine_core_outputs.values()), None)) is None:
-                # We must return the stats even if there are no request
-                # outputs this step.
-                engine_core_outputs[0] = eco = EngineCoreOutputs()
-            eco.scheduler_stats = stats
-
-        return engine_core_outputs

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -12,7 +12,7 @@ from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 
-class OmniScheduler(VLLMScheduler):
+class OmniARScheduler(VLLMScheduler):
     """
     OmniScheduler: Scheduler for vLLM-Omni multimodal processing.
 

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -10,15 +10,15 @@ from vllm.v1.core.sched.scheduler import (
     SchedulerOutput,
     SpecDecodingStats,
 )
+from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 
 from vllm_omni.core.sched.output import OmniNewRequestData
-from vllm_omni.core.sched.scheduler import OmniScheduler
 from vllm_omni.outputs import OmniModelRunnerOutput
 
 
-class GenerationScheduler(OmniScheduler):
+class OmniGenerationScheduler(VLLMScheduler):
     def schedule(self) -> SchedulerOutput:
         """Diffusion fast path:
         - Feed all input tokens of the request at once

vllm_omni/diffusion/attention/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/attention/backends/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/attention/backends/abstract.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+import torch
+
+
+class AttentionBackend(ABC):
+    """Abstract class for diffusion attention backends."""
+
+    accept_output_buffer: bool = False
+
+    @staticmethod
+    @abstractmethod
+    def get_name() -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> type["AttentionImpl"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_supported_head_sizes() -> list[int]:
+        """Get the list of supported head sizes for this backend."""
+        raise NotImplementedError
+
+    @classmethod
+    def supports_head_size(cls, head_size: int) -> bool:
+        supported_head_sizes = cls.get_supported_head_sizes()
+        return (not supported_head_sizes) or head_size in supported_head_sizes
+
+
+@dataclass
+class AttentionMetadata:
+    attn_mask: torch.Tensor | None = None
+
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+
+class AttentionImpl(ABC, Generic[T]):
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    # only tested with flash_attn v3
+    # from flash_attn_interface import flash_attn_func as flash_attn_3_func  # not available in flash-attn 2.8.1
+    from flash_attn import flash_attn_func  # can be FA2 or FA3
+except ImportError:
+    logger.warning(
+        "FlashAttentionBackend is not available. You may install flash-attn "
+        "by running `uv pip install flash-attn==2.8.1 --no-build-isolation`"
+        " or install pre-built flash-attn from https://github.com/Dao-AILab/flash-attention/releases"
+    )
+    flash_attn_3_func = None
+
+
+class FlashAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+
+class FlashAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        # TODO: flash_attn_func does not support attn_mask.
+        out: torch.Tensor = flash_attn_func(
+            query,
+            key,
+            value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+        return out

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+
+class SDPABackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [x for x in range(1024)]  # todo
+
+    @staticmethod
+    def get_name() -> str:
+        return "SDPA"
+
+    @staticmethod
+    def get_impl_cls() -> type["SDPAImpl"]:
+        return SDPAImpl
+
+
+class SDPAImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        attention_mask = attn_metadata.attn_mask if attn_metadata else None
+
+        output = torch.nn.functional.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=self.causal,
+            scale=self.softmax_scale,
+        )
+        out = output.permute(0, 2, 1, 3)
+        return out

vllm_omni/diffusion/attention/layer.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.selector import get_attn_backend
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.attn_backend = get_attn_backend(-1)
+        self.attn_impl_cls = self.attn_backend.get_impl_cls()
+        self.attention = self.attn_impl_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            num_kv_heads=num_kv_heads,
+        )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        # shape: (batch_size, seq_len, num_heads, head_size)
+        attn_output = self.attention.forward(query, key, value, attn_metadata)
+        return attn_output

vllm_omni/diffusion/attention/selector.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+from functools import cache
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+
+logger = init_logger(__name__)
+
+# environment variable value -> backend class
+SUPPORTED_BACKENDS = {
+    "FLASH_ATTN": FlashAttentionBackend,
+    "TORCH_SDPA": SDPABackend,
+}
+
+
+@cache
+def get_attn_backend(head_size: int) -> type[AttentionBackend]:
+    """
+    Get attention backend for diffusion models.
+
+    The backend is selected based on the following priority:
+    1. DIFFUSION_ATTENTION_BACKEND environment variable (if set, e.g. export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN)
+    2. Default backend (SDPA)
+
+    Args:
+        head_size: Head size (currently not used for selection, but kept for API compatibility)
+
+    Returns:
+        The selected attention backend class
+    """
+    # Check environment variable
+    backend_name: str | None = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+
+    if backend_name is not None:
+        backend_name_upper = backend_name.upper()
+        if backend_name_upper not in SUPPORTED_BACKENDS:
+            valid_backends = list(SUPPORTED_BACKENDS.keys())
+            raise ValueError(
+                f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}"
+            )
+        logger.info(f"Using attention backend '{backend_name_upper}' for diffusion")
+        return SUPPORTED_BACKENDS[backend_name_upper]
+
+    # Default to SDPA
+    return SDPABackend

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -8,7 +8,6 @@ from typing import Any, Optional, Union
 import torch
 import torch.nn as nn
 from diffusers.models.attention import FeedForward
-from diffusers.models.attention_dispatch import dispatch_attention_fn
 
 # TODO replace this with vLLM implementation
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
@@ -18,6 +17,8 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
 
+from vllm_omni.diffusion.attention.layer import Attention
+
 logger = init_logger(__name__)
 
 
@@ -262,12 +263,12 @@ class QwenImageCrossAttention(nn.Module):
         self.norm_added_q = RMSNorm(head_dim, eps=eps)
         self.norm_added_k = RMSNorm(head_dim, eps=eps)
 
-        # self.attn = Attention(
-        #     num_heads=num_heads,
-        #     head_size=self.head_dim,
-        #     scale=1.0,  # softmax_scale
-        #     attn_type=AttentionType.ENCODER,
-        # )
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )
 
     def forward(
         self,
@@ -316,20 +317,11 @@ class QwenImageCrossAttention(nn.Module):
         joint_value = torch.cat([txt_value, img_value], dim=1)
 
         # Compute joint attention
-        # joint_hidden_states = self.attn(
-        #     joint_query,
-        #     joint_key,
-        #     joint_value,
-        # )
-        joint_hidden_states = dispatch_attention_fn(
+        joint_hidden_states = self.attn(
             joint_query,
             joint_key,
             joint_value,
-            dropout_p=0.0,
-            is_causal=False,
         )
-
-        # Reshape back
         joint_hidden_states = joint_hidden_states.flatten(2, 3)
         joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
 
@@ -368,19 +360,6 @@ class QwenImageTransformerBlock(nn.Module):
             nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
         )
         self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        # self.attn = Attention(
-        #     query_dim=dim,
-        #     cross_attention_dim=None,  # Enable cross attention for joint computation
-        #     added_kv_proj_dim=dim,  # Enable added KV projections for text stream
-        #     dim_head=attention_head_dim,
-        #     heads=num_attention_heads,
-        #     out_dim=dim,
-        #     context_pre_only=False,
-        #     bias=True,
-        #     processor=QwenDoubleStreamAttnProcessor2_0(),
-        #     qk_norm=qk_norm,
-        #     eps=eps,
-        # )
         self.attn = QwenImageCrossAttention(
             dim=dim,
             num_heads=num_attention_heads,

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -22,10 +22,11 @@ from typing import Optional
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from diffusers.models.attention_dispatch import dispatch_attention_fn
 from torch.nn.utils.rnn import pad_sequence
 from vllm.model_executor.layers.layernorm import RMSNorm
 
+from vllm_omni.diffusion.attention.layer import Attention
+
 ADALN_EMBED_DIM = 256
 SEQ_MULTI_OF = 32
 
@@ -105,6 +106,13 @@ class ZImageAttention(nn.Module):
 
         self.to_out = nn.ModuleList([nn.Linear(dim, dim, bias=False)])
 
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -143,13 +151,11 @@ class ZImageAttention(nn.Module):
             attention_mask = attention_mask[:, None, None, :]
 
         # Compute joint attention
-        hidden_states = dispatch_attention_fn(
+        hidden_states = self.attn(
             query,
             key,
             value,
-            attn_mask=attention_mask,
-            dropout_p=0.0,
-            is_causal=False,
+            # attn_mask=attention_mask, # we don't support multi prompts now.
         )
 
         # Reshape back

vllm_omni/diffusion/registry.py

@@ -3,6 +3,8 @@
 
 import importlib
 
+from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry
+
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 
 _DIFFUSION_MODELS = {
@@ -20,22 +22,28 @@ _DIFFUSION_MODELS = {
 }
 
 
+DiffusionModelRegistry = _ModelRegistry(
+    {
+        model_arch: _LazyRegisteredModel(
+            module_name=f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}",
+            class_name=cls_name,
+        )
+        for model_arch, (mod_folder, mod_relname, cls_name) in _DIFFUSION_MODELS.items()
+    }
+)
+
+
 def initialize_model(
     od_config: OmniDiffusionConfig,
 ):
-    if od_config.model_class_name in _DIFFUSION_MODELS:
-        mod_folder, mod_relname, cls_name = _DIFFUSION_MODELS[od_config.model_class_name]
-        module_name = f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}"
-        module = importlib.import_module(module_name)
-        model_class = getattr(module, cls_name)
-        model = model_class(od_config=od_config, prefix=mod_relname)
-
+    model_class = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
+    if model_class is not None:
+        model = model_class(od_config=od_config)
         # Configure VAE memory optimization settings from config
         if hasattr(model.vae, "use_slicing"):
             model.vae.use_slicing = od_config.vae_use_slicing
         if hasattr(model.vae, "use_tiling"):
             model.vae.use_tiling = od_config.vae_use_tiling
-
         return model
     else:
         raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")

vllm_omni/diffusion/utils/hf_utils.py

@@ -1,12 +1,13 @@
 from functools import lru_cache
 
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
 
 def load_diffusers_config(model_name) -> dict:
+    from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
     config = DiffusionPipeline.load_config(model_name)
     return config
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -1415,7 +1415,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             inputs_embeds = None
 
         if self.use_deepstack and inputs_embeds is not None and get_pp_group().is_first_rank:
-            deepstack_input_embeds = self._get_deepstack_input_embeds(inputs_embeds.size(0))
+            deepstack_input_embeds = self._get_deepstack_input_embeds(input_ids.size(1))
         else:
             deepstack_input_embeds = None
 
@@ -1431,7 +1431,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         )
 
         if inputs_embeds is not None and get_pp_group().is_first_rank:
-            self._clear_deepstack_input_embeds(inputs_embeds.size(0))
+            self._clear_deepstack_input_embeds(input_ids.size(1))
 
         return hidden_states.unsqueeze(0), captured_hidden_states
 

vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml

@@ -9,7 +9,7 @@ stage_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
       trust_remote_code: true
@@ -35,7 +35,7 @@ stage_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
       trust_remote_code: true
@@ -60,8 +60,8 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_diffusion_worker.NPUDiffusionWorker
-      scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
+      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
       trust_remote_code: true

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -11,12 +11,13 @@ stage_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
       trust_remote_code: true
       engine_output_type: latent
       enable_prefix_caching: false
+      max_num_batched_tokens: 32768
     is_comprehension: true
     final_output: true
     final_output_type: text
@@ -37,11 +38,12 @@ stage_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
+      max_num_batched_tokens: 32768
       engine_output_type: latent
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
@@ -62,12 +64,13 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
-      scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
+      max_num_batched_tokens: 32768
       engine_output_type: audio
     engine_input_source: [1]
     final_output: true

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -13,7 +13,7 @@ stage_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: true
       trust_remote_code: true
@@ -42,8 +42,8 @@ stage_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-       gpu_memory_utilization: 0.3
+       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+       gpu_memory_utilization: 0.5
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -72,7 +72,7 @@ stage_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false

vllm_omni/worker/gpu_diffusion_model_runner.py

@@ -1,576 +0,0 @@
-from __future__ import annotations
-
-import gc
-import logging
-from typing import Any
-
-import numpy as np
-import torch
-from vllm.config import CUDAGraphMode
-from vllm.forward_context import BatchDescriptor
-from vllm.multimodal.inputs import MultiModalKwargs
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.spec_decode.eagle import EagleProposer
-from vllm.v1.utils import record_function_or_nullcontext
-from vllm.v1.worker.gpu_model_runner import (
-    EMPTY_MODEL_RUNNER_OUTPUT,
-    AsyncGPUModelRunnerOutput,
-    IntermediateTensors,
-    PerLayerAttnMetadata,
-    get_pp_group,
-    set_forward_context,
-)
-from vllm.v1.worker.ubatch_utils import UBatchSlices
-from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
-
-from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-
-logger = logging.getLogger(__name__)
-
-
-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`.
-    """
-
-    def _preprocess(
-        self,
-        scheduler_output: SchedulerOutput,
-        intermediate_tensors: IntermediateTensors | None = None,
-        ubatch_slices: UBatchSlices | None = None,
-        num_tokens_after_padding: torch.Tensor | None = None,
-    ) -> tuple[
-        int,
-        int,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor,
-        IntermediateTensors | None,
-        dict[str, Any],
-    ]:
-        # 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_after_padding = self.get_dp_padding(num_input_tokens)
-        num_input_tokens += num_pad
-
-        # _prepare_inputs may reorder the batch, so we must gather multi
-        # modal outputs after that to ensure the correct order
-        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
-            # Run the multimodal encoder if any.
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-
-            # NOTE(woosuk): To unify token ids and soft tokens (vision
-            # embeddings), we always use embeddings (rather than token ids)
-            # as input to the multimodal model, even when the input is text.
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids.gpu[:num_input_tokens],
-                multimodal_embeddings=mm_embeds or None,
-            )
-
-            # TODO(woosuk): Avoid the copy. Optimize.
-            self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
-
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = {
-                **self._init_model_kwargs(num_input_tokens),
-                **self._extract_mm_kwargs(scheduler_output),
-            }
-        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
-            # Get the input embeddings for the tokens that are not input embeds,
-            # then put them into the appropriate positions.
-            # TODO(qthequartermasterman): Since even when prompt embeds are
-            # enabled, (a) not all requests will use prompt embeds, and (b)
-            # after the initial prompt is processed, the rest of the generated
-            # tokens will be token ids, it is not desirable to have the
-            # embedding layer outside of the CUDA graph all the time. The v0
-            # engine avoids this by "double compiling" the CUDA graph, once
-            # with input_ids and again with inputs_embeds, for all num_tokens.
-            # If a batch only has token ids, then including the embedding layer
-            # in the CUDA graph will be more performant (like in the else case
-            # below).
-            token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1)
-            # Some tokens ids may need to become embeds
-            if token_ids_idx.numel() > 0:
-                token_ids = self.input_ids.gpu[token_ids_idx]
-                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
-                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
-
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-        else:
-            # For text-only models, we use token ids as input.
-            # While it is possible to use embeddings as input just like the
-            # multimodal models, it is not desirable for performance since
-            # then the embedding layer is not included in the CUDA graph.
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = None
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-        if self.uses_mrope:
-            positions = self.mrope_positions.gpu[:, :num_input_tokens]
-        else:
-            positions = self.positions.gpu[:num_input_tokens]
-
-        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
-            )
-
-        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
-            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
-            model_kwargs.update(encoder_inputs)
-
-        return (
-            num_input_tokens,
-            num_input_tokens,
-            num_tokens_after_padding,
-            input_ids,
-            inputs_embeds,
-            positions,
-            intermediate_tensors,
-            model_kwargs,
-        )
-
-    @torch.inference_mode()
-    def execute_model(
-        self,
-        scheduler_output: SchedulerOutput,
-        intermediate_tensors: IntermediateTensors | None = None,
-    ) -> OmniModelRunnerOutput | IntermediateTensors:
-        with record_function_or_nullcontext("Preprocess"):
-            with self.synchronize_input_prep():
-                # Update persistent batch states.
-                self._update_states(scheduler_output)
-                if not scheduler_output.total_num_scheduled_tokens:
-                    return EMPTY_MODEL_RUNNER_OUTPUT
-
-                # Prepare the decoder inputs.
-                (
-                    attn_metadata,
-                    logits_indices,
-                    spec_decode_metadata,
-                    num_scheduled_tokens_np,
-                    spec_decode_common_attn_metadata,
-                    max_query_len,
-                    ubatch_slices,
-                    num_tokens_after_padding,
-                ) = self._prepare_inputs(scheduler_output)
-
-            # Input token count for this iteration (not used by diffusion, but
-            # retained to keep DP padding/ordering consistent)
-            (
-                num_scheduled_tokens,
-                num_input_tokens,
-                num_tokens_across_dp,
-                input_ids,
-                inputs_embeds,
-                positions,
-                intermediate_tensors,
-                model_kwargs,
-            ) = self._preprocess(
-                scheduler_output,
-                intermediate_tensors,
-                ubatch_slices,
-                num_tokens_after_padding,
-            )
-
-        cudagraph_runtime_mode = CUDAGraphMode.NONE
-        # Run the model.
-        # Use persistent buffers for CUDA graphs.
-        with (
-            set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=None,
-                ubatch_slices=ubatch_slices,
-            ),
-            record_function_or_nullcontext("Forward"),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
-        ):
-            outputs = self._run_diffusion(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                multimodal_kwargs=model_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[torch.Tensor | None] = []
-        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")
-
-        output = OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
-            sampled_token_ids=[],
-            logprobs=None,
-            prompt_logprobs_dict={},
-            pooler_output=pooler_output,
-            kv_connector_output=kv_connector_output,
-            num_nans_in_logits={},
-        )
-
-        if not self.use_async_scheduling:
-            return output
-
-        return AsyncGPUModelRunnerOutput(
-            model_runner_output=output,
-            sampled_token_ids=[],
-            invalid_req_indices=[],
-            async_output_copy_stream=self.async_output_copy_stream,
-        )
-
-    def _run_diffusion(
-        self,
-        *,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None,
-        inputs_embeds: torch.Tensor | None,
-        multimodal_kwargs: dict,
-        logits_indices: torch.Tensor,
-    ) -> torch.Tensor | list[torch.Tensor]:
-        """Runs the diffusion process and returns per-request tensors.
-
-        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,
-        )
-
-        # 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,
-        cudagraph_runtime_mode: CUDAGraphMode | None = None,
-        force_attention: bool = False,
-        uniform_decode: bool = False,
-        skip_eplb: bool = False,
-        is_profile: bool = False,
-        create_mixed_batch: bool = False,
-        remove_lora: bool = True,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Run a dummy forward pass to warm up/profile run or capture the
-        CUDA graph for the model.
-
-        Args:
-            num_tokens: Number of tokens to run the dummy forward pass.
-            cudagraph_runtime_mode: used to control the behavior.
-                - if not set will determine the cudagraph mode based on using
-                    the self.cudagraph_dispatcher.
-                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
-                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
-                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
-                    needed.
-            force_attention: If True, always create attention metadata. Used to
-                warm up attention backend when mode is NONE.
-            uniform_decode: If True, the batch is a uniform decode batch.
-            skip_eplb: If True, skip EPLB state update.
-            is_profile: If True, this is a profile run.
-            create_mixed_batch: If True, create a mixed batch with both decode
-                (1 token) and prefill (multiple tokens) requests.
-            remove_lora: If False, dummy LoRAs are not destroyed after the run
-        """
-        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
-            CUDAGraphMode.NONE,
-            CUDAGraphMode.PIECEWISE,
-            CUDAGraphMode.FULL,
-        }
-
-        # If cudagraph_mode.decode_mode() == FULL and
-        # cudagraph_mode.separate_routine(). This means that we are using
-        # different graphs and/or modes for mixed prefill-decode batches vs.
-        # uniform decode batches. A uniform decode batch means that all
-        # requests have identical query length, except a potential virtual
-        # request (shorter) in the batch account for padding.
-        # Uniform decode batch could either be common pure decode, where
-        # max_query_len == 1, or speculative decode, where
-        # max_query_len == 1 + num_spec_decode_tokens.
-
-        # When setting max_query_len = 1, we switch to and capture the optimized
-        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
-        # for GQA/MQA.
-        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
-
-        # 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)
-
-        num_tokens_after_padding = None
-
-        # If we failed to microbatch, currently need to resynchronize
-        # TODO(lucas,sage): we should be able to avoid this second sync by
-        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
-        #  a single `coordinate_batch_across_dp` function.
-        if num_tokens_after_padding is None:
-            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-            num_tokens_after_padding = num_tokens + num_pad
-        else:
-            num_tokens_across_dp = num_tokens_after_padding
-            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
-
-        attn_metadata: PerLayerAttnMetadata | None = None
-
-        # If force_attention is True, we always capture attention. Otherwise,
-        # it only happens for cudagraph_runtime_mode=FULL.
-        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-            attn_metadata = {}
-
-            seq_lens = max_query_len
-            self.seq_lens.np[:num_reqs] = seq_lens
-            self.seq_lens.np[num_reqs:] = 0
-            self.seq_lens.copy_to_gpu()
-
-            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
-            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
-            self.query_start_loc.copy_to_gpu()
-
-            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.gpu[: num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
-                    seq_lens=self.seq_lens.gpu[: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],  # noqa: E501
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=max_query_len,
-                    max_seq_len=self.max_model_len,
-                    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.gpu[:num_tokens],
-                    causal=True,
-                )
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-                    assert type(attn_metadata) is dict
-                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
-                        common_attn_metadata
-                    )  # noqa: E501
-                    for layer_name in attn_group.layer_names:
-                        attn_metadata[layer_name] = attn_metadata_i
-
-        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
-            model_kwargs = self._init_model_kwargs(num_tokens)
-            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = {
-                    **model_kwargs,
-                    **self._dummy_mm_kwargs(num_reqs),
-                }
-            elif self.enable_prompt_embeds:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = self._init_model_kwargs(num_tokens)
-            else:
-                input_ids = self.input_ids.gpu[:num_tokens]
-                inputs_embeds = None
-
-            if self.uses_mrope:
-                positions = self.mrope_positions.gpu[:, :num_tokens]
-            else:
-                positions = self.positions.gpu[: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)
-
-            # filter out the valid batch descriptor
-            _cg_mode, batch_descriptor = (
-                self.cudagraph_dispatcher.dispatch(
-                    BatchDescriptor(
-                        num_tokens=num_tokens_after_padding,
-                        uniform_decode=uniform_decode,
-                    )
-                )
-                if not is_profile
-                else (CUDAGraphMode.NONE, None)
-            )
-            if cudagraph_runtime_mode is not None:
-                # we allow forcing NONE when the dispatcher disagrees to support
-                # warm ups for cudagraph capture
-                assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
-                    f"Cudagraph runtime mode mismatch at dummy_run. "
-                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
-                )
-            else:
-                cudagraph_runtime_mode = _cg_mode
-
-            with (
-                self.maybe_randomize_inputs(input_ids),
-                set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens_after_padding,
-                    num_tokens_across_dp=num_tokens_across_dp,
-                    cudagraph_runtime_mode=cudagraph_runtime_mode,
-                    batch_descriptor=batch_descriptor,
-                ),
-            ):
-                outputs = self.model(
-                    input_ids=input_ids,
-                    positions=positions,
-                    intermediate_tensors=intermediate_tensors,
-                    inputs_embeds=inputs_embeds,
-                    **model_kwargs,
-                )
-
-            if self.use_aux_hidden_state_outputs:
-                hidden_states, _ = outputs
-            else:
-                hidden_states = outputs
-
-            if self.speculative_config and self.speculative_config.use_eagle():
-                assert isinstance(self.drafter, EagleProposer)
-                self.drafter.dummy_run(num_tokens)
-
-        # This is necessary to avoid blocking DP.
-        # For dummy runs, we typically skip EPLB since we don't have any real
-        # requests to process.
-        # However, in DP settings, there may be cases when some DP ranks do
-        # not have any requests to process, so they're executing dummy batches.
-        # In such cases, we still have to trigger EPLB to make sure
-        # ranks execute the rearrangement in synchronization.
-
-        # logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
-        return 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.
-        if self.supports_mm_inputs:
-            if self.model_config.multimodal_config.skip_mm_profiling:
-                logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
-            else:
-                mm_budget = self.mm_budget
-                assert mm_budget is not None
-
-                if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
-                    # NOTE: Currently model is profiled with a single non-text
-                    # modality with the max possible input tokens even when
-                    # it supports multiple.
-                    dummy_modality = mm_budget.get_modality_with_max_tokens()
-                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality]
-
-                    logger.info(
-                        "Encoder cache will be initialized with a budget "
-                        "of %s tokens, and profiled with %s %s items of the maximum feature size.",
-                        encoder_budget,
-                        max_mm_items_per_batch,
-                        dummy_modality,
-                    )
-
-                    # Create dummy batch of multimodal inputs.
-                    batched_dummy_mm_inputs = self._get_mm_dummy_batch(
-                        dummy_modality,
-                        max_mm_items_per_batch,
-                    )
-
-                    # Run multimodal encoder.
-                    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,
-                    )
-
-                    # NOTE: This happens when encoder cache needs to store
-                    # the embeddings that encoder outputs are scattered onto.
-                    # In this case we create dummy embeddings of size
-                    # (encode_budget, hidden_size) and scatter encoder
-                    # output into it.
-                    encoder_output_shape = dummy_encoder_outputs[0].shape
-                    if encoder_output_shape[0] < encoder_budget:
-                        expanded_outputs = []
-                        for output in dummy_encoder_outputs:
-                            expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1]))
-                            num_tokens = output.shape[0]
-                            expanded[:num_tokens].copy_(output)
-                            expanded_outputs.append(expanded)
-
-                        dummy_encoder_outputs = expanded_outputs
-
-                    # Cache the dummy encoder outputs.
-                    self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
-
-        # Add `is_profile` here to pre-allocate communication buffers
-        hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
-        if get_pp_group().is_last_rank:
-            pass
-        self._sync_device()
-        del hidden_states
-        self.encoder_cache.clear()
-        gc.collect()

vllm_omni/worker/gpu_diffusion_worker.py

@@ -1,86 +0,0 @@
-import gc
-import os
-
-import torch
-from vllm.model_executor import set_random_seed
-from vllm.platforms import current_platform
-from vllm.utils import GiB_bytes, MemorySnapshot
-from vllm.v1.utils import report_usage_stats
-from vllm.v1.worker.gpu_worker import Worker as GPUWorker
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-
-from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
-
-
-class GPUDiffusionWorker(GPUWorker):
-    """GPU worker for diffusion-based omni model stages.
-
-    Extends the base GPUWorker to initialize and manage diffusion model
-    runners for non-autoregressive generation stages (e.g., image or
-    audio generation stages).
-    """
-
-    def init_device(self):
-        """Initialize the GPU device and diffusion model runner.
-
-        Sets up CUDA device, initializes distributed environment,
-        sets random seed, and creates a GPUDiffusionModelRunner instance.
-
-        Raises:
-            ValueError: If free GPU memory is insufficient for the
-                requested memory utilization
-            RuntimeError: If device type is not supported
-        """
-        if self.device_config.device.type == "cuda":
-            # 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)
-
-            current_platform.check_if_supports_dtype(self.model_config.dtype)
-
-            # Initialize the distributed environment BEFORE taking
-            # memory snapshot
-            # This ensures NCCL buffers are allocated before we measure
-            # available memory
-            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)
-
-            # Now take memory snapshot after NCCL is initialized
-            gc.collect()
-            torch.cuda.empty_cache()
-
-            # 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:
-
-                def GiB(b):
-                    return 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}")
-
-        # 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)

vllm_omni/worker/npu/npu_diffusion_worker.py

@@ -1,15 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm_ascend.worker.worker_v1 import NPUWorker
-
-from vllm_omni.worker.npu.npu_diffusion_model_runner import NPUDiffusionModelRunner
-
-
-class NPUDiffusionWorker(NPUWorker):
-    """NPU diffusion worker for code2wav stage in Omni model."""
-
-    def init_device(self):
-        device = self._init_device()
-
-        self.model_runner: NPUDiffusionModelRunner = NPUDiffusionModelRunner(self.vllm_config, device)

vllm_omni/worker/npu/npu_generation_model_runner.py

@@ -53,8 +53,8 @@ if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
 
 
-class NPUDiffusionModelRunner(OmniNPUModelRunner):
-    """Diffusion model runner for vLLM-omni on NPU (non-autoregressive)."""
+class NPUGenerationModelRunner(OmniNPUModelRunner):
+    """Generation model runner for vLLM-omni on NPU (non-autoregressive)."""
 
     def _prepare_inputs(
         self,
@@ -455,7 +455,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner):
             ):
                 self.maybe_setup_kv_connector(scheduler_output)
 
-                outputs = self._run_diffusion(
+                outputs = self._run_generation(
                     input_ids=input_ids,
                     positions=positions,
                     intermediate_tensors=intermediate_tensors,
@@ -525,7 +525,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner):
             async_output_copy_stream=self.async_output_copy_stream,
         )
 
-    def _run_diffusion(
+    def _run_generation(
         self,
         *,
         input_ids: torch.Tensor,
@@ -535,7 +535,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner):
         multimodal_kwargs: dict,
         logits_indices: torch.Tensor,
     ) -> torch.Tensor | list[torch.Tensor]:
-        """Runs the diffusion process and returns per-request tensors.
+        """Runs the generation process and returns per-request tensors.
 
         Tries model interfaces in the following order for maximal compatibility:
         1) model.sample(condition=..., **kwargs)
@@ -558,7 +558,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner):
         # TODO: add the diffuse method for other models
 
         raise RuntimeError(
-            "The loaded model does not expose diffusion interfaces 'sample', "
+            "The loaded model does not expose generation interfaces 'sample', "
             "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
         )
 

vllm_omni/worker/npu/npu_generation_worker.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_ascend.worker.worker_v1 import NPUWorker
+
+from vllm_omni.worker.npu.npu_generation_model_runner import NPUGenerationModelRunner
+
+
+class NPUGenerationWorker(NPUWorker):
+    """NPU generation worker for code2wav stage in Omni model."""
+
+    def init_device(self):
+        device = self._init_device()
+
+        self.model_runner: NPUGenerationModelRunner = NPUGenerationModelRunner(self.vllm_config, device)

vllm_omni/worker/npu/npu_model_runner.py

@@ -66,17 +66,26 @@ class OmniNPUModelRunner(NPUModelRunner):
             if use_audio_in_video_value is not None:
                 use_audio_in_video = bool(use_audio_in_video_value.item())
 
-        assert supports_mrope(self.get_model()), "M-RoPE support is not implemented."
-
-        req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(
-            req_state.prompt_token_ids,
-            hf_config=self.model_config.hf_config,
-            image_grid_thw=image_grid_thw,
-            video_grid_thw=video_grid_thw,
-            second_per_grid_ts=second_per_grid_ts,
-            audio_feature_lengths=audio_feature_lengths,
-            use_audio_in_video=use_audio_in_video,
-        )
+        if supports_mrope(self.model):
+            req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
+                req_state.prompt_token_ids,
+                hf_config=self.model_config.hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                audio_feature_lengths=audio_feature_lengths,
+                use_audio_in_video=use_audio_in_video,
+            )
+        else:
+            req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(
+                req_state.prompt_token_ids,
+                hf_config=self.model_config.hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                audio_feature_lengths=audio_feature_lengths,
+                use_audio_in_video=use_audio_in_video,
+            )
 
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         # Remove finished requests from the cached states.