delete

Signed-off-by: Roger Wang <hey@rogerw.io>
update
2026-02-01 00:09:07 -08:00 · 2026-02-01 00:08:24 -08:00
54 changed files with 862 additions and 987 deletions
@@ -233,43 +233,43 @@ steps:
        volumes:
        - "/fsx/hf_cache:/fsx/hf_cache"

-  # - label: "Omni Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: image-build
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-  #     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-  #     - pytest -s -v tests/e2e/online_serving/test_async_omni.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 2
-  #               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
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 30
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_async_omni.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                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

  - label: "Diffusion Image Edit Test with H100 (1 GPU)"
    timeout_in_minutes: 20
@@ -306,37 +306,37 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: image-build
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   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
+  - label: "Bagel Text2Img Model Test with H100"
+    timeout_in_minutes: 30
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+    agents:
+      queue: "mithril-h100-pool"
+    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
@@ -15,7 +15,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---

 *Latest News* 🔥
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability.
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.

@@ -1,21 +1,20 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.15.0
+ARG VLLM_BASE_TAG=v0.14.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
+
 COPY . .
+
 # Install system dependencies
 RUN apt-get update && \
-  apt-get install -y ffmpeg && \
-  apt-get clean && \
-  rm -rf /var/lib/apt/lists/*
-
-# Install vLLM wheel for the specific commit
-# URL format: https://wheels.vllm.ai/<commit_hash>/
-# Current vLLM commit: f176443446f659dbab5315e056e605d8984fd976 (v0.15.0)
-RUN uv pip install --system --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/f176443446f659dbab5315e056e605d8984fd976
+    apt-get install -y ffmpeg && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*

 # Install vllm-omni into the same uv-managed Python environment used by the base image.
-RUN uv pip install --system --no-cache-dir ".[dev]"
+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 []
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0
 FROM ${BASE_IMAGE} AS final

 ARG COMMON_WORKDIR=/app
@@ -18,19 +18,20 @@ The following table shows which models are currently supported by parallelism me

 ### ImageGen

-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|
-| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |
-| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |
-| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ❌        |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |
+| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
+|-------|------------------|------------|---------|--------------|--------------------------|
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
+| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
+

 !!! note "TP Limitations for Diffusion Models"
    We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
@@ -410,7 +410,7 @@ def main(args):
        elif stage_outputs.final_output_type == "audio":
            for output in stage_outputs.request_output:
                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
+                audio_tensor = output.multimodal_output["audio"]
                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
@@ -418,7 +418,7 @@ def main(args):
        elif stage_outputs.final_output_type == "audio":
            for output in stage_outputs.request_output:
                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
+                audio_tensor = output.multimodal_output["audio"]
                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")

                # Convert to numpy array and ensure correct format
@@ -244,9 +244,9 @@ def main(args):
    for stage_outputs in omni_generator:
        for output in stage_outputs.request_output:
            request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            audio_tensor = output.multimodal_output["audio"]
            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            audio_samplerate = output.outputs[0].multimodal_output["sr"].item()
+            audio_samplerate = output.multimodal_output["sr"].item()
            # Convert to numpy array and ensure correct format
            audio_numpy = audio_tensor.float().detach().cpu().numpy()

@@ -111,10 +111,6 @@ def run_tts_generation(args) -> None:
        print(response.text)
        return

-    if response.content.decode("utf-8").startswith('{"error"'):
-        print(f"Error: {response.content.decode('utf-8')}")
-        return
-
    # Save audio response
    output_path = args.output or "tts_output.wav"
    with open(output_path, "wb") as f:
@@ -12,7 +12,7 @@ stage_args:
      model_arch: BagelForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.60
      enforce_eager: true
      trust_remote_code: true
      engine_output_type: text
@@ -44,7 +44,7 @@ stage_args:
      max_batch_size: 1
    engine_args:
      model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.4
      enforce_eager: true
      trust_remote_code: true
      engine_output_type: image
@@ -12,7 +12,7 @@ stage_args:
      model_arch: BagelForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.65
      enforce_eager: true
      trust_remote_code: true
      engine_output_type: text
@@ -43,7 +43,7 @@ stage_args:
      max_batch_size: 1
    engine_args:
      model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.4
      enforce_eager: true
      trust_remote_code: true
      engine_output_type: image
@@ -161,7 +161,7 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima
 def test_bagel_text2img_shared_memory_connector():
    """Test Bagel text2img with shared memory connector."""
    config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
-    omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
+    omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path)

    try:
        generated_image = _generate_bagel_image(omni)
@@ -288,7 +288,7 @@ def test_bagel_text2img_mooncake_connector():
            http_port=MOONCAKE_HTTP_PORT,
        )

-        omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
+        omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file)

        generated_image = _generate_bagel_image(omni)
        _validate_pixels(generated_image)
@@ -88,7 +88,7 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t
        assert len(audio_output.request_output) > 0

        # Verify audio tensor exists and has content
-        audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
+        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
        assert audio_tensor is not None
        assert audio_tensor.numel() > 0

@@ -75,6 +75,6 @@ def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
        assert len(audio_output.request_output) > 0

        # Verify audio tensor exists and has content
-        audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
+        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
        assert audio_tensor is not None
        assert audio_tensor.numel() > 0
@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 """
-System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.
+System test for Ulysses sequence parallel backend.

-Tests verify that SP inference produces correct outputs compared to baseline.
+This test verifies that Ulysses-SP (DeepSpeed Ulysses Sequence Parallel) works
+correctly with diffusion models. It uses minimal settings to keep test time
+short for CI.
 """

 import gc
@@ -12,7 +14,6 @@ import os
 import sys
 import time
 from pathlib import Path
-from typing import NamedTuple

 import numpy as np
 import pytest
@@ -31,22 +32,24 @@ from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.platforms import current_omni_platform

-# Test configuration
-MODELS = ["riverclouds/qwen_image_random"]
+models = ["riverclouds/qwen_image_random"]
+
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
-DEFAULT_HEIGHT = 256
-DEFAULT_WIDTH = 256
-DEFAULT_SEED = 42
-DEFAULT_STEPS = 4
-DIFF_MEAN_THRESHOLD = 2e-2
-DIFF_MAX_THRESHOLD = 2e-1


-class InferenceResult(NamedTuple):
-    """Result of an inference run."""
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)

-    images: list[Image.Image]
-    elapsed_ms: float
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    ta = _pil_to_float_rgb_tensor(a)
+    tb = _pil_to_float_rgb_tensor(b)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    return abs_diff.mean().item(), abs_diff.max().item()


 def _cleanup_distributed():
@@ -58,6 +61,7 @@ def _cleanup_distributed():
        os.environ.pop(key, None)

    gc.collect()
+
    if current_omni_platform.is_available():
        current_omni_platform.empty_cache()
        current_omni_platform.synchronize()
@@ -65,216 +69,199 @@ def _cleanup_distributed():
    time.sleep(5)


-def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
-    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
-    ta = torch.from_numpy(np.asarray(a.convert("RGB"), dtype=np.float32) / 255.0)
-    tb = torch.from_numpy(np.asarray(b.convert("RGB"), dtype=np.float32) / 255.0)
-    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
-    abs_diff = torch.abs(ta - tb)
-    return abs_diff.mean().item(), abs_diff.max().item()
-
-
-def _run_inference(
-    model_name: str,
-    dtype: torch.dtype,
-    attn_backend: str,
-    ulysses_degree: int = 1,
-    ring_degree: int = 1,
-    height: int = DEFAULT_HEIGHT,
-    width: int = DEFAULT_WIDTH,
-    seed: int = DEFAULT_SEED,
-    warmup: bool = True,
-) -> InferenceResult:
-    """Run inference with specified configuration.
-
-    Args:
-        warmup: If True, run one warmup iteration before the timed run.
-    """
-    parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-    omni = Omni(
+def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height: int, width: int, seed: int):
+    """Run baseline inference (no SP)."""
+    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
+    baseline = Omni(
        model=model_name,
-        parallel_config=parallel_config,
+        parallel_config=baseline_parallel_config,
        dtype=dtype,
        attention_backend=attn_backend,
    )

    try:
-        # Warmup run (not timed)
-        if warmup:
-            _ = omni.generate(
-                PROMPT,
-                OmniDiffusionSamplingParams(
-                    height=height,
-                    width=width,
-                    num_inference_steps=DEFAULT_STEPS,
-                    guidance_scale=0.0,
-                    generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed + 1000),
-                    num_outputs_per_prompt=1,
-                ),
-            )
-
-        # Timed run
-        start = time.time()
-        outputs = omni.generate(
+        outputs = baseline.generate(
            PROMPT,
            OmniDiffusionSamplingParams(
                height=height,
                width=width,
-                num_inference_steps=DEFAULT_STEPS,
+                num_inference_steps=4,
                guidance_scale=0.0,
                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
                num_outputs_per_prompt=1,
            ),
        )
-        elapsed_ms = (time.time() - start) * 1000
-
-        return InferenceResult(
-            images=outputs[0].request_output[0].images,
-            elapsed_ms=elapsed_ms,
-        )
+        return outputs[0].request_output[0].images
    finally:
-        omni.close()
+        baseline.close()
        _cleanup_distributed()


-# =============================================================================
-# Correctness & Performance Tests
-# =============================================================================
+def _run_sp(
+    model_name: str,
+    dtype: torch.dtype,
+    attn_backend: str,
+    height: int,
+    width: int,
+    seed: int,
+    ulysses_degree: int,
+    ring_degree: int,
+):
+    """Run SP inference."""
+    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    sp = Omni(
+        model=model_name,
+        parallel_config=sp_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )

-# SP configurations: (ulysses_degree, ring_degree, height, width, warmup, is_perf_test)
-# - warmup: whether to run warmup for this SP config
-# - is_perf_test: whether this is a performance test (show speedup metrics)
-SP_CONFIGS = [
-    (2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),  # Ulysses-2 - performance test
-    (1, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),  # Ring-2 - performance test
-    (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),  # Hybrid - correctness only
-    (4, 1, 272, 272, False, False),  # Ulysses-4 - shape and correctness
-]
-
-
-def _get_sp_mode(ulysses_degree: int, ring_degree: int) -> str:
-    """Get SP mode name for logging."""
-    if ulysses_degree > 1 and ring_degree == 1:
-        return f"ulysses-{ulysses_degree}"
-    elif ring_degree > 1 and ulysses_degree == 1:
-        return f"ring-{ring_degree}"
-    else:
-        return f"hybrid-{ulysses_degree}x{ring_degree}"
-
-
-@pytest.mark.parametrize("model_name", MODELS)
-def test_sp_correctness(model_name: str):
-    """Test that SP inference produces correct outputs and measure performance.
-
-    Runs baseline once per unique (height, width), then tests all SP configs.
-
-    Note: Run with `pytest -v -s` to see detailed output.
-    """
-    device_count = current_omni_platform.get_device_count()
-
-    # Cache baseline results by (height, width)
-    # Key: (height, width), Value: (result, warmup_used)
-    baseline_cache: dict[tuple[int, int], InferenceResult] = {}
-
-    # Collect results for summary
-    results: list[dict] = []
-
-    print("\n" + "=" * 70)
-    print(f"Sequence Parallel Test - Model: {model_name}")
-    print(f"Available GPUs: {device_count}")
-    print("=" * 70)
-
-    for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS:
-        sp_size = ulysses_degree * ring_degree
-        sp_mode = _get_sp_mode(ulysses_degree, ring_degree)
-
-        if device_count < sp_size:
-            print(f"\n[{sp_mode}] SKIPPED (requires {sp_size} GPUs)")
-            continue
-
-        # Determine baseline warmup: only for default size (performance tests)
-        cache_key = (height, width)
-        baseline_warmup = height == DEFAULT_HEIGHT and width == DEFAULT_WIDTH
-
-        # Get or compute baseline for this (height, width)
-        if cache_key not in baseline_cache:
-            print(f"\n--- Running baseline {height}x{width} (warmup={baseline_warmup}) ---")
-            baseline = _run_inference(
-                model_name,
-                torch.bfloat16,
-                "sdpa",
+    try:
+        outputs = sp.generate(
+            PROMPT,
+            OmniDiffusionSamplingParams(
                height=height,
                width=width,
-                warmup=baseline_warmup,
-            )
-            assert len(baseline.images) == 1
-            baseline_cache[cache_key] = baseline
-            print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms")
-        else:
-            baseline = baseline_cache[cache_key]
-
-        # Run SP
-        print(f"\n--- Running {sp_mode} (warmup={sp_warmup}) ---")
-        sp_result = _run_inference(
-            model_name,
-            torch.bfloat16,
-            "sdpa",
-            ulysses_degree=ulysses_degree,
-            ring_degree=ring_degree,
-            height=height,
-            width=width,
-            warmup=sp_warmup,
+                num_inference_steps=4,
+                guidance_scale=0.0,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
+                num_outputs_per_prompt=1,
+            ),
        )
-        assert len(sp_result.images) == 1
+        return outputs[0].request_output[0].images
+    finally:
+        sp.close()
+        _cleanup_distributed()

-        # Compare outputs (correctness)
-        mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])

-        # Build result entry
-        result = {
-            "mode": sp_mode,
-            "sp_size": sp_size,
-            "height": height,
-            "width": width,
-            "baseline_ms": baseline.elapsed_ms,
-            "sp_ms": sp_result.elapsed_ms,
-            "mean_diff": mean_diff,
-            "max_diff": max_diff,
-            "is_perf_test": is_perf_test,
-        }
-        results.append(result)
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test baseline inference only (no SP)."""
+    height = 256
+    width = 256
+    seed = 42

-        # Output based on test type
-        if is_perf_test:
-            speedup = baseline.elapsed_ms / sp_result.elapsed_ms if sp_result.elapsed_ms > 0 else 0
-            result["speedup"] = speedup
-            print(
-                f"[{sp_mode}] {sp_size} GPUs | "
-                f"baseline: {baseline.elapsed_ms:.0f}ms, sp: {sp_result.elapsed_ms:.0f}ms, "
-                f"speedup: {speedup:.2f}x"
-            )
-        else:
-            print(f"[{sp_mode}] {sp_size} GPUs | sp: {sp_result.elapsed_ms:.0f}ms (correctness only)")
+    images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)

-        print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}")
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height

-        # Assert correctness
-        assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
-            f"[{sp_mode}] SP output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP inference only (ulysses=2)."""
+    if current_omni_platform.get_device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=2, ring_degree=1)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP inference only (ring=2)."""
+    if current_omni_platform.get_device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=1, ring_degree=2)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("ulysses_degree", [1, 2])
+@pytest.mark.parametrize("ring_degree", [1, 2])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sequence_parallel(
+    model_name: str,
+    ulysses_degree: int,
+    ring_degree: int,
+    dtype: torch.dtype,
+    attn_backend: str,
+):
+    """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs."""
+    if ulysses_degree <= 1 and ring_degree <= 1:
+        pytest.skip(
+            "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; "
+            "provide ulysses_degree or ring_degree>1."
        )

-    # Summary
-    print("\n" + "=" * 70)
-    print("SUMMARY")
-    print("=" * 70)
-    print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'SP':<12} {'Speedup':<10} {'Status'}")
-    print("-" * 70)
-    for r in results:
-        speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A"
-        baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A"
-        status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL"
-        print(
-            f"{r['mode']:<15} {r['sp_size']:<6} {r['height']}x{r['width']:<5} "
-            f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
+    sp_size = ulysses_degree * ring_degree
+    if current_omni_platform.get_device_count() < sp_size:
+        pytest.skip(f"Test requires {sp_size} GPUs but only {current_omni_platform.get_device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
+
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
+
+    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
+
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
+
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
+    )
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP with ulysses_degree=4."""
+    ulysses_degree = 4
+    ring_degree = 1
+
+    if current_omni_platform.get_device_count() < ulysses_degree * ring_degree:
+        pytest.skip(
+            f"Test requires {ulysses_degree * ring_degree} GPUs but only {current_omni_platform.get_device_count()} available"
        )
-    print("=" * 70)
+
+    height = 272
+    width = 272
+    seed = 42
+
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
+
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
+
+    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
+
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
+
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
+    )
@@ -23,8 +23,6 @@ from vllm.assets.image import ImageAsset
 from vllm.utils.network_utils import get_open_port

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-# Increase timeout for downloading assets from S3 (default 5s is too short for CI)
-os.environ.setdefault("VLLM_IMAGE_FETCH_TIMEOUT", "60")

 models = ["Qwen/Qwen-Image-Edit-2509"]
 test_params = models
@@ -110,7 +110,6 @@ class OmniModelConfig(ModelConfig):
        mm_processor_cache_gb: float | None,
        mm_processor_cache_type: MMCacheType | None,
        mm_shm_cache_max_object_size_mb: int | None,
-        mm_encoder_only: bool | None,
        mm_encoder_tp_mode: MMEncoderTPMode | None,
        mm_encoder_attn_backend: AttentionBackendEnum | str | None,
        interleave_mm_strings: bool | None,
@@ -182,15 +181,6 @@ class OmniModelConfig(ModelConfig):
        )
        self.model_arch_config = self.get_model_arch_config()

-        if self.convert == "mm_encoder_only":
-            logger.warning_once(
-                "`--convert mm_encoder_only` is deprecated and "
-                "will be removed in v0.15. "
-                "Please use --mm-encoder-only` instead."
-            )
-            mm_encoder_only = True
-            self.convert = "none"
-
        architectures = self.architectures
        registry = self.registry
        is_generative_model = registry.is_text_generation_model(architectures, self)
@@ -272,7 +262,6 @@ class OmniModelConfig(ModelConfig):
                mm_processor_cache_gb=mm_processor_cache_gb,
                mm_processor_cache_type=mm_processor_cache_type,
                mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
-                mm_encoder_only=mm_encoder_only,
                mm_encoder_tp_mode=mm_encoder_tp_mode,
                mm_encoder_attn_backend=mm_encoder_attn_backend,
                interleave_mm_strings=interleave_mm_strings,
@@ -6,20 +6,20 @@ from dataclasses import asdict, dataclass
 from time import time
 from typing import Any

+import numpy as np
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
-from vllm.v1.core.sched.utils import remove_all
+from vllm.v1.core.sched.utils import check_stop, remove_all
 from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats

-from vllm_omni.core.sched.output import OmniSchedulerOutput
 from vllm_omni.distributed.omni_connectors.adapter import get_chunk, put_chunk
 from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
@@ -185,19 +185,13 @@ class OmniARScheduler(VLLMScheduler):
                get_chunk(self.omni_connector, scheduler_output)

            # Add information about requests needing KV cache transfer
-            finished_reqs = self.get_finished_requests_needing_kv_transfer()
+            scheduler_output.finished_requests_needing_kv_transfer = self.get_finished_requests_needing_kv_transfer()
        except Exception:
            # If anything goes wrong, leave the original output unchanged
            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
-            finished_reqs = {}
+            scheduler_output.finished_requests_needing_kv_transfer = {}

-        # Wrap in omni scheduler output to carry transfer metadata.
-        base_fields = SchedulerOutput.__dataclass_fields__.keys()
-        base_data = {name: getattr(scheduler_output, name) for name in base_fields}
-        return OmniSchedulerOutput(
-            **base_data,
-            finished_requests_needing_kv_transfer=finished_reqs,
-        )
+        return scheduler_output

    def update_from_output(
        self,
@@ -284,32 +278,47 @@ class OmniARScheduler(VLLMScheduler):
            pooler_output = pooler_outputs[req_index] if pooler_outputs else None
            kv_transfer_params = None
            status_before_stop = request.status
-            finish_reason = None
-            routed_experts = None

            # Check for stop and update request status.
            if new_token_ids:
                new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
-            elif request.pooling_params and pooler_output is not None:
-                # Pooling stops as soon as there is output.
-                request.status = RequestStatus.FINISHED_STOPPED
-                stopped = True

-            # If criteria returns True, it means we must STOP the request.
-            # If criteria returns False, it might have triggered a background
-            # transfer (e.g. prefill finished / special token) but continues decoding.
+            # If criteria returns True, it means we must STOP the request (e.g. legacy behavior, not used now)
+            # If criteria returns False, it might have triggered a background transfer
+            # (e.g. prefill finished / special token) but continues decoding
            if not stopped and self._process_kv_transfer_trigger(request, new_token_ids):
                stopped = True

+            # Stop checking for pooler models.
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+                if request.output_token_ids:
+                    stopped = check_stop(request, self.max_model_len)
+            routed_experts = None
            if stopped:
-                routed_experts = self._get_routed_experts(request)
+                # [Omni] Handle routed experts if enabled
+                if self.vllm_config.model_config.enable_return_routed_experts:
+                    kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
+                    block_ids = kv_blocks.get_block_ids()[0]
+                    num_tokens = request.num_tokens - 1

-                # Capture finish_reason BEFORE _handle_stopped_request, which may
-                # reset the status to WAITING for streaming requests that continue.
-                finish_reason = request.get_finished_reason()
-                finished = self._handle_stopped_request(request)
-                if finished:
-                    kv_transfer_params = self._free_request(request)
+                    # compute slot mapping
+                    block_ids_array = np.array(block_ids, dtype=np.int32)
+                    num_blocks = len(block_ids)
+                    block_size = self.block_size
+
+                    # generate block offsets
+                    block_offsets = np.arange(0, block_size)
+
+                    # compute slot mapping: slot = block_id * block_size + offset
+                    slot_mapping = (
+                        block_offsets.reshape((1, block_size)) + block_ids_array.reshape((num_blocks, 1)) * block_size
+                    ).flatten()[:num_tokens]
+
+                    routed_experts = self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
+
+                kv_transfer_params = self._free_request(request)
                if status_before_stop == RequestStatus.RUNNING:
                    stopped_running_reqs.add(request)
                else:
@@ -336,13 +345,13 @@ class OmniARScheduler(VLLMScheduler):

            # Get prompt logprobs for this request.
            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            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=finish_reason,
+                        finish_reason=request.get_finished_reason(),
                        new_logprobs=new_logprobs,
                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                        pooling_output=pooler_output,
@@ -1,9 +1,7 @@
 import time
 from collections import defaultdict

-from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -11,7 +9,6 @@ from vllm.v1.core.sched.request_queue import create_request_queue
 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.v1.metrics.perf import PerfStats
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats

@@ -59,21 +56,10 @@ class OmniGenerationScheduler(VLLMScheduler):
        # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
        skipped_waiting_requests = create_request_queue(self.policy)
        req_index = 0
-        # OMNI: Track requests that are already finished (e.g., marked by connector)
-        # These should be removed from running and not scheduled
-        already_finished_reqs: set[Request] = set()
        while req_index < len(self.running) and token_budget > 0:
            request = self.running[req_index]
            if self.omni_connector is not None:
                get_chunk_for_generation(self.omni_connector, request)
-
-            # OMNI: Skip requests that are already finished or not in self.requests
-            # This can happen when connector marks request as finished
-            if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests:
-                already_finished_reqs.add(request)
-                req_index += 1
-                continue
-
            num_computed_tokens = request.num_computed_tokens
            required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1)
            num_new_tokens = min(required_tokens, token_budget)
@@ -97,24 +83,12 @@ class OmniGenerationScheduler(VLLMScheduler):
            scheduled_running_reqs.append(request)
            req_index += 1

-        # OMNI: Remove already finished requests from running queue
-        if already_finished_reqs:
-            self.running = remove_all(self.running, already_finished_reqs)
-
        # Fast path selection and scheduling (treat all as diffusion requests,
        # independent of pooling_params)
        while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs:
            request = self.waiting.peek_request()
            if self.omni_connector is not None:
                get_chunk_for_generation(self.omni_connector, request)
-
-            # OMNI: Skip requests that are already finished or not in self.requests
-            # This can happen when connector marks request as finished
-            if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests:
-                # Pop the finished request from waiting queue and don't schedule it
-                self.waiting.pop_request()
-                continue
-
            # Uniformly treat as diffusion. A feature flag can be added later
            # via config or request tag.

@@ -283,16 +257,9 @@ class OmniGenerationScheduler(VLLMScheduler):
        num_nans_in_logits = model_runner_output.num_nans_in_logits
        kv_connector_output = model_runner_output.kv_connector_output

-        cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
-        perf_stats: PerfStats | None = None
-        if self.perf_metrics and self.perf_metrics.is_enabled():
-            perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
-
        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
        spec_decoding_stats: SpecDecodingStats | None = None
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
+        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
        # Merge connector-side stats (align with v0.14.0)
        if kv_connector_stats and self.connector:
            kv_stats = self.connector.get_kv_connector_stats()
@@ -341,14 +308,12 @@ class OmniGenerationScheduler(VLLMScheduler):
                    num_accepted_tokens=num_accepted,
                )

-            stopped = False
            new_logprobs = None
            new_token_ids = generated_token_ids
            kv_transfer_params = None
-            pooler_output = pooler_outputs[req_index] if pooler_outputs else None
-            status_before_stop = request.status
-            finish_reason = None
-            routed_experts = None
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]

            # Diffusion request: completes in one step; mark finished and free resources
            if request.status == RequestStatus.FINISHED_STOPPED or (
@@ -358,18 +323,8 @@ class OmniGenerationScheduler(VLLMScheduler):
                # Optional: set a stop_reason for front-end clarity
                # (does not affect protocol)
                request.stop_reason = request.stop_reason  # or "generation_done"
-                stopped = True
-
-            if stopped:
-                routed_experts = self._get_routed_experts(request)
-                finish_reason = request.get_finished_reason()
-                finished = self._handle_stopped_request(request)
-                if finished:
-                    kv_transfer_params = self._free_request(request)
-                if status_before_stop == RequestStatus.RUNNING:
-                    stopped_running_reqs.add(request)
-                else:
-                    stopped_preempted_reqs.add(request)
+                kv_transfer_params = self._free_request(request)
+                stopped_running_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:
@@ -389,13 +344,13 @@ class OmniGenerationScheduler(VLLMScheduler):

            # Get prompt logprobs for this request.
            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            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=finish_reason,
+                        finish_reason=request.get_finished_reason(),
                        new_logprobs=new_logprobs,
                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                        pooling_output=pooler_output,
@@ -404,7 +359,6 @@ class OmniGenerationScheduler(VLLMScheduler):
                        kv_transfer_params=kv_transfer_params,
                        trace_headers=request.trace_headers,
                        num_cached_tokens=request.num_cached_tokens,
-                        routed_experts=routed_experts,
                        num_nans_in_logits=request.num_nans_in_logits,
                    )
                )
@@ -452,7 +406,7 @@ class OmniGenerationScheduler(VLLMScheduler):
                    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, cudagraph_stats, perf_stats)) is not None:
+        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
@@ -1,6 +1,6 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass

-from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
+from vllm.v1.core.sched.output import CachedRequestData, NewRequestData
 from vllm.v1.request import Request

 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
@@ -66,11 +66,4 @@ class OmniCachedRequestData(CachedRequestData):
        prompt_token_ids: Mapping from request ID to list of prompt token IDs
    """

-    prompt_token_ids: dict[str, list[int]]
-
-
-@dataclass
-class OmniSchedulerOutput(SchedulerOutput):
-    """Scheduler output with omni-specific transfer metadata."""
-
-    finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
+    prompt_token_ids: dict[int, list[int]]
@@ -7,14 +7,11 @@


 import torch
-from vllm.logger import init_logger

 from vllm_omni.diffusion.attention.backends.ring.ring_kernels import pytorch_attn_forward
 from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse
 from vllm_omni.diffusion.distributed.comm import RingComm

-logger = init_logger(__name__)
-

 def ring_pytorch_attn_func(
    q,
@@ -3,16 +3,12 @@

 from __future__ import annotations

-from vllm.logger import init_logger
-
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
 from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group
 from vllm_omni.diffusion.forward_context import get_forward_context

-logger = init_logger(__name__)
-

 def build_parallel_attention_strategy(
    *,
@@ -30,8 +26,7 @@ def build_parallel_attention_strategy(
    try:
        cfg = get_forward_context().omni_diffusion_config
        p = cfg.parallel_config
-    except Exception as e:
-        logger.debug(f"No forward context available for parallel attention strategy: {e}")
+    except Exception:
        return NoParallelAttention()

    ulysses_degree = getattr(p, "ulysses_degree", 1)
@@ -42,18 +37,11 @@ def build_parallel_attention_strategy(
        # Ensure SP group is initialized and world size > 1
        if get_sequence_parallel_world_size() <= 1:
            return NoParallelAttention()
-    except Exception as e:
-        # Log warning if SP is configured but group is not available
-        if ulysses_degree > 1 or ring_degree > 1:
-            logger.warning(
-                f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}) but SP group not available: {e}. "
-                f"Falling back to NoParallelAttention. This may cause incorrect results."
-            )
+    except Exception:
        return NoParallelAttention()

    # Ulysses (or Hybrid Ulysses+Ring)
    if ulysses_degree > 1:
-        logger.debug(f"Using UlyssesParallelAttention (ulysses_degree={ulysses_degree})")
        return UlyssesParallelAttention(
            sp_group=sp_group,
            scatter_idx=scatter_idx,
@@ -63,7 +51,6 @@ def build_parallel_attention_strategy(

    # Pure Ring Attention
    if ring_degree > 1:
-        logger.debug(f"Using RingParallelAttention (ring_degree={ring_degree})")
        return RingParallelAttention(
            sp_group=sp_group,
        )
@@ -9,11 +9,10 @@ import torch.nn as nn
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
-from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
@@ -36,9 +35,9 @@ class FeedForward(nn.Module):
        inner_dim = int(dim * mult)
        dim_out = dim_out if dim_out is not None else dim

-        self.w_in = ColumnParallelLinear(dim, inner_dim, bias=bias, return_bias=False)
+        self.w_in = ReplicatedLinear(dim, inner_dim, bias=bias, return_bias=False)
        self.act = get_act_fn("gelu_pytorch_tanh")
-        self.w_out = RowParallelLinear(inner_dim, dim_out, bias=bias, return_bias=False)
+        self.w_out = ReplicatedLinear(inner_dim, dim_out, bias=bias, return_bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.w_in(hidden_states)
@@ -84,11 +83,12 @@ class LongCatImageAttention(nn.Module):
            hidden_size=query_dim,
            head_size=self.head_dim,
            total_num_heads=self.heads,
+            disable_tp=True,
            bias=bias,
        )

        if not self.pre_only:
-            self.to_out = RowParallelLinear(self.inner_dim, self.out_dim, bias=out_bias)
+            self.to_out = torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)

        if self.added_kv_proj_dim is not None:
            self.norm_added_q = RMSNorm(dim_head, eps=eps)
@@ -98,10 +98,11 @@ class LongCatImageAttention(nn.Module):
                hidden_size=self.added_kv_proj_dim,
                head_size=self.head_dim,
                total_num_heads=self.heads,
+                disable_tp=True,
                bias=added_proj_bias,
            )

-            self.to_add_out = RowParallelLinear(self.inner_dim, query_dim, bias=out_bias)
+            self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)

        self.attn = Attention(
            num_heads=heads,
@@ -184,26 +185,22 @@ class LongCatImageAttention(nn.Module):
        """
        qkv, _ = self.to_qkv(hidden_states)

-        q_size = self.to_qkv.num_heads * self.head_dim
-        kv_size = self.to_qkv.num_kv_heads * self.head_dim
-        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
+        query, key, value = qkv.chunk(3, dim=-1)

-        query = query.unflatten(-1, (self.to_qkv.num_heads, -1))
-        key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
-        value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
+        query = query.unflatten(-1, (self.heads, -1))
+        key = key.unflatten(-1, (self.heads, -1))
+        value = value.unflatten(-1, (self.heads, -1))

        query = self.norm_q(query)
        key = self.norm_k(key)

        if self.added_kv_proj_dim is not None:
            encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
-            q_size = self.add_kv_proj.num_heads * self.head_dim
-            kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
-            encoder_query, encoder_key, encoder_value = encoder_qkv.split([q_size, kv_size, kv_size], dim=-1)
+            encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1)

-            encoder_query = encoder_query.unflatten(-1, (self.add_kv_proj.num_heads, -1))
-            encoder_key = encoder_key.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
-            encoder_value = encoder_value.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
+            encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
+            encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
+            encoder_value = encoder_value.unflatten(-1, (self.heads, -1))

            # Apply RMSNorm to text Q/K
            encoder_query = self.norm_added_q(encoder_query)
@@ -287,14 +284,11 @@ class LongCatImageAttention(nn.Module):
            encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
                [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
            )
-            hidden_states, _ = self.to_out(hidden_states)
+            hidden_states = self.to_out(hidden_states)
            encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)

            return hidden_states, encoder_hidden_states
        else:
-            # For single-stream blocks, there's no to_out (RowParallelLinear) to handle the reduction
-            if get_tensor_model_parallel_world_size() > 1:
-                hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1)
            return hidden_states


@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Any

 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig
 from vllm.engine.arg_utils import EngineArgs
@@ -71,7 +70,7 @@ class OmniEngineArgs(EngineArgs):
    engine_output_type: str | None = None
    hf_config_name: str | None = None
    custom_process_next_stage_input_func: str | None = None
-    stage_connector_spec: dict[str, Any] = field(default_factory=dict)
+    stage_connector_spec: dict[str, any] = field(default_factory=dict)
    async_chunk: bool = False
    omni_kv_config: dict | None = None

@@ -173,7 +172,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
    engine_output_type: str | None = None
    hf_config_name: str | None = None
    custom_process_next_stage_input_func: str | None = None
-    stage_connector_spec: dict[str, Any] = field(default_factory=dict)
+    stage_connector_spec: dict[str, any] = field(default_factory=dict)
    async_chunk: bool = False
    omni_kv_config: dict | None = None

@@ -1,9 +1,7 @@
-import os
 import time
 from collections.abc import Mapping
 from typing import Any, cast

-import numpy as np
 import torch
 from vllm.config import VllmConfig
 from vllm.inputs import ProcessorInputs, PromptType
@@ -11,13 +9,12 @@ from vllm.inputs.parse import split_enc_dec_inputs
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict
-from vllm.multimodal.processing.context import set_request_id
 from vllm.multimodal.utils import argsort_mm_positions
 from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
+from vllm.tokenizers import TokenizerLike
 from vllm.utils import length_from_prompt_token_ids_or_embeds
-from vllm.utils.torch_utils import set_default_torch_num_threads
 from vllm.v1.engine.input_processor import InputProcessor

 from vllm_omni.engine import (
@@ -41,6 +38,7 @@ class OmniInputProcessor(InputProcessor):

    Args:
        vllm_config: Global vLLM configuration
+        tokenizer: Tokenizer instance for text processing
        mm_registry: Multi-modal registry for processing multimodal inputs
    """

@@ -77,12 +75,13 @@ class OmniInputProcessor(InputProcessor):
    def __init__(
        self,
        vllm_config: VllmConfig,
+        tokenizer: TokenizerLike,
        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
    ):
-        super().__init__(vllm_config, mm_registry)
+        super().__init__(vllm_config, tokenizer, mm_registry)
        self.input_preprocessor = OmniInputPreprocessor(
            self.model_config,
-            vllm_config.observability_config,
+            self.tokenizer,
            mm_registry,
            mm_processor_cache=self.mm_processor_cache,
        )
@@ -98,8 +97,7 @@ class OmniInputProcessor(InputProcessor):
        trace_headers: Mapping[str, str] | None = None,
        priority: int = 0,
        data_parallel_rank: int | None = None,
-        resumable: bool = False,
-    ) -> OmniEngineCoreRequest:
+    ) -> tuple[str | None, OmniEngineCoreRequest]:
        """Process input prompt into an engine core request.

        Converts a prompt (text, tokens, or multimodal) into an
@@ -132,12 +130,9 @@ class OmniInputProcessor(InputProcessor):
        self._validate_lora(lora_request)
        self._validate_params(params)

-        parallel_config = self.vllm_config.parallel_config
-        dp_size = parallel_config.data_parallel_size
-        dp_local_size = parallel_config.data_parallel_size_local
-        num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size
-        if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks):
-            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).")
+        data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
+        if data_parallel_rank is not None and not (0 <= data_parallel_rank < data_parallel_size):
+            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {data_parallel_size}).")

        if arrival_time is None:
            arrival_time = time.time()
@@ -159,7 +154,7 @@ class OmniInputProcessor(InputProcessor):
        else:
            # Otherwise, use user-provided uuids as multimodal hash overrides
            # if provided.
-            self._validate_mm_uuids(prompt)
+            self._validate_multi_modal_uuids(prompt)
            if isinstance(prompt, dict):
                mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids"))
            else:
@@ -169,19 +164,11 @@ class OmniInputProcessor(InputProcessor):
        # 1. Tokenize text prompt, with LoRA request if one exists.
        # 2. For multimodal models with a merged preprocessor, preprocess
        #   multimodal data and expand prompt token ids accordingly.
-        num_threads = int(os.environ.get("OMP_NUM_THREADS", "1"))
-        if "OMP_NUM_THREADS" not in os.environ:
-            logger.debug_once(
-                "OMP_NUM_THREADS is not set; defaulting Torch threads to %d for input preprocessing.",
-                num_threads,
-            )
-
-        with set_request_id(request_id), set_default_torch_num_threads(num_threads):
-            processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
-                prompt,
-                tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
-            )
+        processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
+            prompt,
+            tokenization_kwargs=tokenization_kwargs,
+            mm_uuids=mm_uuids,
+        )

        current_platform.validate_request(
            prompt=prompt,
@@ -194,13 +181,13 @@ class OmniInputProcessor(InputProcessor):
        encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
        self._validate_model_inputs(encoder_inputs, decoder_inputs)

-        # Normalize decoder prompt access across TypedDict variants.
-        if decoder_inputs["type"] == "embeds":
-            prompt_token_ids = None
-            prompt_embeds = decoder_inputs["prompt_embeds"]
-        else:
-            prompt_token_ids = decoder_inputs["prompt_token_ids"]
-            prompt_embeds = decoder_inputs.get("prompt_embeds")
+        # Mypy does not always properly infer the types of some elements of
+        # discriminated unions of TypedDicts, because of how it handles
+        # inheritance of TypedDict. If we explicitly extract the items we want
+        # we can avoid type errors from using `dict.get` later in the method.
+        _prompt_str: str | None = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
+        prompt_token_ids = decoder_inputs["prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None
+        prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs["type"] == "embeds" else None

        sampling_params = None
        pooling_params = None
@@ -232,26 +219,35 @@ class OmniInputProcessor(InputProcessor):

            mm_features = []
            for modality, idx in sorted_mm_idxs:
-                base_mm_hash = decoder_mm_hashes[modality][idx]
                mm_features.append(
                    MultiModalFeatureSpec(
                        data=decoder_mm_inputs[modality][idx],
                        modality=modality,
-                        identifier=self._get_mm_identifier(base_mm_hash, lora_request),
+                        identifier=decoder_mm_hashes[modality][idx],
                        mm_position=decoder_mm_positions[modality][idx],
-                        mm_hash=base_mm_hash,
                    )
                )

-        # Compatibility: decode serialized prompt embeds if provided.
-        if isinstance(prompt_embeds, PromptEmbedsPayload):
-            prompt_embeds = self._decode_prompt_embeds(prompt_embeds)
-
+        # Serialize prompt_embeds and additional_information if provided
+        # (direct-transfer path)
+        prompt_embeds_payload: PromptEmbedsPayload | None = None
        additional_information_payload: AdditionalInformationPayload | None = None
-        raw_info: dict[str, Any] | AdditionalInformationPayload | None = decoder_inputs.get("additional_information")
-        if isinstance(raw_info, AdditionalInformationPayload):
-            additional_information_payload = raw_info
-        elif raw_info is not None:
+        pe: torch.Tensor | None = decoder_inputs.get("prompt_embeds")  # type: ignore[operator]
+        if pe is not None:
+            if pe.ndim != 2:
+                raise ValueError("prompt_embeds must be of shape (seq_len, hidden_size)")
+            # Move to CPU and ensure contiguous memory for stable serialization
+            pe_cpu = pe.detach().to("cpu").contiguous()
+            seq_len, hidden_size = pe_cpu.shape
+            dtype_str = self._dtype_to_name(pe_cpu.dtype)
+            data_bytes = pe_cpu.numpy().tobytes()
+            prompt_embeds_payload = PromptEmbedsPayload(
+                data=data_bytes,
+                shape=[int(seq_len), int(hidden_size)],
+                dtype=dtype_str,
+            )
+        raw_info: dict[str, Any] | None = decoder_inputs.get("additional_information")  # type: ignore[operator]
+        if raw_info is not None:
            entries: dict[str, AdditionalInformationEntry] = {}
            for key, value in raw_info.items():
                if isinstance(value, torch.Tensor):
@@ -283,14 +279,6 @@ class OmniInputProcessor(InputProcessor):
            priority=priority,
            data_parallel_rank=data_parallel_rank,
            trace_headers=trace_headers,
-            prompt_embeds=prompt_embeds,
+            prompt_embeds=prompt_embeds_payload,
            additional_information=additional_information_payload,
-            resumable=resumable,
        )
-
-    @staticmethod
-    def _decode_prompt_embeds(payload: PromptEmbedsPayload) -> torch.Tensor:
-        dtype = getattr(np, payload.dtype)
-        arr = np.frombuffer(payload.data, dtype=dtype)
-        arr = arr.reshape(payload.shape)
-        return torch.from_numpy(arr)
@@ -2,7 +2,6 @@ from ast import Dict
 from collections.abc import Callable
 from typing import Any

-import numpy as np
 import torch
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput
@@ -10,11 +9,7 @@ from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
 from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
-from vllm.v1.engine.output_processor import (
-    OutputProcessorOutput,
-    RequestOutputCollector,
-    RequestState,
-)
+from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestOutputCollector, RequestState
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats

@@ -145,7 +140,6 @@ class OmniRequestState(RequestState):
        finish_reason: FinishReason | None,
        stop_reason: int | str | None,
        kv_transfer_params: dict[str, Any] | None = None,
-        routed_experts: np.ndarray | None = None,
    ) -> OmniRequestOutput | PoolingRequestOutput | None:
        """Create a request output from generation results.

@@ -164,27 +158,12 @@ class OmniRequestState(RequestState):
            OmniRequestOutput or PoolingRequestOutput if output should be
            emitted (based on finish status and output kind), None otherwise
        """
-        # Pooling-only requests should follow base behavior.
-        if self.detokenizer is None and pooling_output is not None:
-            return super().make_request_output(
-                new_token_ids,
-                pooling_output,
-                finish_reason,
-                stop_reason,
-                kv_transfer_params,
-                routed_experts,
-            )
-
        finished = finish_reason is not None
        final_only = self.output_kind == RequestOutputKind.FINAL_ONLY

        if not finished and final_only:
            return None

-        # Consolidate accumulated tensors when finishing.
-        if finished:
-            self._consolidate_multimodal_tensors()
-
        if self.stream_interval > 1:
            assert self.detokenizer is not None

@@ -206,7 +185,7 @@ class OmniRequestState(RequestState):
                self.sent_tokens_offset = len(self.detokenizer.output_token_ids)

        request_id = self.request_id
-        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts)
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)

        if self.parent_req is None:
            outputs = [output]
@@ -222,10 +201,9 @@ class OmniRequestState(RequestState):
        token_ids: list[int],
        finish_reason: FinishReason | None,
        stop_reason: int | str | None,
-        routed_experts: np.ndarray | None = None,
    ) -> Any:
        # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts)
+        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
        try:
            if self.mm_accumulated is not None:
                # Attach accumulated multimodal dict on the completion output
@@ -273,6 +251,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
        super().__init__(tokenizer=tokenizer, log_stats=log_stats)
        self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
        self._reqid_to_mm_type: dict[str, str] = {}
+        self.request_states: dict[str, OmniRequestState] = {}
        self.engine_core_output_type = engine_core_output_type

    def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
@@ -312,10 +291,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
            ValueError: If the request ID is already registered
        """
        request_id = request.request_id
-        req_state = self.request_states.get(request_id)
-        if req_state is not None:
-            self._update_streaming_request_state(req_state, request, prompt)
-            return
+        if request_id in self.request_states:
+            raise ValueError(f"Request id {request_id} already running.")

        req_state = OmniRequestState.from_new_request(
            tokenizer=self.tokenizer,
@@ -327,12 +304,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
            log_stats=self.log_stats,
            stream_interval=self.stream_interval,
        )
-        if self._requests_drained.is_set():
-            self._requests_drained.clear()
        self.request_states[request_id] = req_state
        if parent_req:
            self.parent_requests[parent_req.request_id] = parent_req
-        self.external_req_ids[req_state.external_req_id].append(request_id)

    def process_outputs(
        self,
@@ -340,27 +314,121 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
        engine_core_timestamp: float | None = None,
        iteration_stats: IterationStats | None = None,
    ) -> OutputProcessorOutput:
+        """Process engine core outputs into request outputs.
+
+        Converts EngineCoreOutput objects into RequestOutput objects,
+        handling multimodal outputs by routing them through appropriate
+        processors and accumulating tensors in request states.
+
+        Args:
+            engine_core_outputs: List of engine core outputs to process
+            engine_core_timestamp: Optional timestamp for the outputs
+            iteration_stats: Optional iteration statistics
+
+        Returns:
+            OutputProcessorOutput containing processed request outputs
+            and list of request IDs to abort
+        """
        self._reqid_to_mm_type.clear()
        for eco in engine_core_outputs:
            mm_type = (self.engine_core_output_type or "").lower()
            if mm_type:
                self._reqid_to_mm_type[eco.request_id] = mm_type
            self._route_and_normalize(eco)
-            req_state = self.request_states.get(eco.request_id)
-            if req_state is None or not isinstance(req_state, OmniRequestState):
-                continue
-            if eco.pooling_output is not None and req_state.detokenizer is not None:
-                req_state.add_multimodal_tensor(
-                    eco.pooling_output,
-                    (getattr(eco, "output_type", self.engine_core_output_type) or "").lower(),
-                )
-                # Force text path in base processor for multimodal outputs.
-                eco.pooling_output = None

-        return super().process_outputs(
-            engine_core_outputs,
-            engine_core_timestamp=engine_core_timestamp,
-            iteration_stats=iteration_stats,
+        # Build RequestOutputs without delegating to base, so we can keep ids
+        request_outputs: list[Any] = []
+        reqs_to_abort: list[str] = []
+        for eco in engine_core_outputs:
+            req_id = eco.request_id
+            req_state = self.request_states.get(req_id)
+            if req_state is None:
+                continue
+
+            # 1) Stats
+            self._update_stats_from_output(req_state, eco, engine_core_timestamp, iteration_stats)
+
+            new_token_ids = eco.new_token_ids
+            pooling_output = eco.pooling_output
+            finish_reason = eco.finish_reason
+            stop_reason = eco.stop_reason
+            kv_transfer_params = eco.kv_transfer_params
+            req_state.num_cached_tokens = eco.num_cached_tokens
+            req_state.is_prefilling = False
+
+            # 2) Detokenize and logprobs when text path
+            assert req_state.detokenizer is not None
+            assert req_state.logprobs_processor is not None
+            stop_string = req_state.detokenizer.update(new_token_ids, finish_reason == FinishReason.STOP)
+            if stop_string:
+                finish_reason = FinishReason.STOP
+                stop_reason = stop_string
+            req_state.logprobs_processor.update_from_output(eco)
+
+            # 2.5) Accumulate multimodal tensors in RequestState
+            try:
+                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
+                if pooling_output is not None and isinstance(req_state, OmniRequestState):
+                    req_state.add_multimodal_tensor(pooling_output, mm_type)
+            except Exception:
+                logger.debug(
+                    "Failed to accumulate multimodal tensor for request %s",
+                    req_id,
+                    exc_info=True,
+                )
+
+            # 3) Create RequestOutput objects, forcing combined mode to keep ids
+            pooling_for_make = pooling_output
+            if pooling_output is not None and new_token_ids:
+                # Do not consume pooling path now; keep ids and attach mm later
+                pooling_for_make = None
+
+            # Consolidate any accumulated tensor lists before creating output
+            if isinstance(req_state, OmniRequestState) and finish_reason is not None:
+                req_state._consolidate_multimodal_tensors()
+
+            ro = req_state.make_request_output(
+                new_token_ids,
+                pooling_for_make,
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+            )
+            if ro:
+                # Attach accumulated multimodal payload if any
+                try:
+                    if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None:
+                        if not hasattr(ro, "multimodal_output"):
+                            setattr(ro, "multimodal_output", {})
+                        ro.multimodal_output = req_state.mm_accumulated
+                except Exception:
+                    logger.exception("Error attaching multimodal payload in process_outputs")
+                if req_state.queue is not None:
+                    req_state.queue.put(ro)
+                else:
+                    request_outputs.append(ro)
+
+            # 4) Free completed
+            if finish_reason is not None:
+                self.request_states.pop(req_id)
+                parent_req = req_state.parent_req
+                if parent_req and not parent_req.child_requests:
+                    self.parent_requests.pop(parent_req.request_id, None)
+                if not self.request_states:
+                    self._requests_drained.set()
+                if not eco.finished:
+                    reqs_to_abort.append(req_id)
+                self._update_stats_from_finished(req_state, finish_reason, iteration_stats)
+                if self.tracer:
+                    self.do_tracing(eco, req_state, iteration_stats)
+                # Cleanup per-request mm state
+                if isinstance(req_state, OmniRequestState):
+                    req_state.mm_accumulated = None
+                    req_state.mm_type = None
+
+        return OutputProcessorOutput(
+            request_outputs=request_outputs,
+            reqs_to_abort=reqs_to_abort,
        )

    # ---- routing helpers ----
@@ -196,10 +196,11 @@ class AsyncOmni(OmniBase):
            if stage.vllm_config is not None and stage.tokenizer is not None:
                try:
                    vllm_config = stage.vllm_config
+                    tokenizer = stage.tokenizer
                    # Initialize input_processor
-                    # OMNI: OmniInputProcessor creates tokenizer internally from vllm_config
                    self.input_processor = OmniInputProcessor(
                        vllm_config=vllm_config,
+                        tokenizer=tokenizer,
                    )
                    # Initialize model_config
                    self.model_config = vllm_config.model_config
@@ -674,15 +675,6 @@ class AsyncOmni(OmniBase):
                return stage.is_tracing_enabled
        return False

-    @property
-    def renderer(self):
-        """Return the renderer from input_processor if available.
-
-        OMNI: Required by upstream OpenAIServingModels.__init__ which
-        accesses engine_client.renderer.
-        """
-        return self.input_processor.renderer
-
    async def do_log_stats(self) -> None:
        pass

@@ -116,6 +116,7 @@ class AsyncOmniLLM(AsyncLLM):
        # InputProcessor (converts Inputs --> EngineCoreRequests).
        self.input_processor = OmniInputProcessor(
            vllm_config=vllm_config,
+            tokenizer=tokenizer,
            mm_registry=mm_registry,
        )

@@ -11,11 +11,11 @@ from vllm.entrypoints.chat_utils import (
    BaseMultiModalItemTracker,
    ChatCompletionContentPartParam,
    ChatCompletionMessageParam,
-    ChatTemplateContentFormat,
    ConversationMessage,
    MultiModalDataDict,
    MultiModalUUIDDict,
    _AssistantParser,
+    _ChatTemplateContentFormat,
    _ContentPart,
    _get_full_multimodal_text_prompt,
    _parse_chat_message_content_part,
@@ -39,27 +39,17 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
        self._mm_processor_kwargs = mm_processor_kwargs

    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
-        # OMNI: Follow upstream async pattern - create coroutine that resolves to (data, uuid)
-        coro = self._video_with_uuid_async(video_url, uuid)
-        placeholder = self._tracker.add("video", coro)
+        video = self._connector.fetch_video_async(video_url=video_url) if video_url else None
+
+        placeholder = self._tracker.add("video", video, uuid)
        self._add_placeholder("video", placeholder)

        # Extract audio from video if use_audio_in_video is True
        if video_url and self._mm_processor_kwargs and self._mm_processor_kwargs.get("use_audio_in_video", False):
-            audio_coro = self._audio_from_video_with_uuid_async(video_url, uuid)
-            audio_placeholder = self._tracker.add("audio", audio_coro)
+            audio_coro = self._extract_audio_from_video_async(video_url)
+            audio_placeholder = self._tracker.add("audio", audio_coro, uuid)
            self._add_placeholder("audio", audio_placeholder)

-    async def _video_with_uuid_async(self, video_url: str | None, uuid: str | None):
-        """Fetch video and return (video, uuid) tuple."""
-        video = await self._connector.fetch_video_async(video_url=video_url) if video_url else None
-        return video, uuid
-
-    async def _audio_from_video_with_uuid_async(self, video_url: str, uuid: str | None):
-        """Extract audio from video and return (audio, uuid) tuple."""
-        audio = await self._extract_audio_from_video_async(video_url)
-        return audio, uuid
-
    async def _extract_audio_from_video_async(self, video_url: str) -> tuple[np.ndarray, int | float]:
        """
        Extract audio from video URL using librosa.
@@ -138,21 +128,13 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
 def parse_chat_messages_futures(
    messages: list[ChatCompletionMessageParam],
    model_config: ModelConfig,
-    content_format: ChatTemplateContentFormat,
+    content_format: _ChatTemplateContentFormat,
    mm_processor_kwargs: dict[str, Any] | None = None,
 ) -> tuple[
    list[ConversationMessage],
-    Awaitable[tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]],
+    Awaitable[MultiModalDataDict | None],
+    MultiModalUUIDDict | None,
 ]:
-    """Parse chat messages and return conversation with multimodal data future.
-
-    OMNI: Updated to use upstream vLLM v0.15.0 API where resolve_items()
-    returns both mm_data and mm_uuids together as a tuple.
-
-    Returns:
-        Tuple of (conversation, mm_future) where mm_future resolves to
-        (mm_data, mm_uuids) when awaited.
-    """
    conversation: list[ConversationMessage] = []
    mm_tracker = OmniAsyncMultiModalItemTracker(model_config)

@@ -173,14 +155,13 @@ def parse_chat_messages_futures(

    _postprocess_messages(conversation)

-    # OMNI: Use upstream resolve_items() which returns (mm_data, mm_uuids) tuple
-    return conversation, mm_tracker.resolve_items()
+    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()


 def _parse_chat_message_content(
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
-    content_format: ChatTemplateContentFormat,
+    content_format: _ChatTemplateContentFormat,
    interleave_strings: bool,
    mm_processor_kwargs: dict[str, Any] | None = None,
 ) -> list[ConversationMessage]:
@@ -161,7 +161,9 @@ class OmniLLM(LLM):
            log_stats=self.llm_engine.log_stats,
            engine_core_output_type=engine_args.engine_output_type,
        )
-        self.llm_engine.input_processor = OmniInputProcessor(vllm_config=self.llm_engine.vllm_config)
+        self.llm_engine.input_processor = OmniInputProcessor(
+            vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
+        )
        self.engine_class = type(self.llm_engine)

        self.request_counter = Counter()
@@ -670,22 +670,12 @@ def _stage_worker(
                    break

        lock_files = acquired_lock_fds
-
-        # Set FD_CLOEXEC on all lock file descriptors to prevent child processes
-        # (e.g., EngineCore) from inheriting them, which would cause deadlock
-        for lock_fd in acquired_lock_fds:
-            try:
-                flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD)
-                fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
-            except (OSError, ValueError):
-                pass
    except Exception as e:
        logger.debug(
            "[Stage-%s] Failed to set up sequential initialization lock: %s",
            stage_id,
            e,
        )
-
    # Init engine based on stage_type
    logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys()))
    if engine_args.get("async_chunk", False):
@@ -696,29 +686,30 @@ def _stage_worker(
            break
        engine_args["stage_connector_spec"] = stage_connector_spec
        engine_args["stage_id"] = stage_id
-    if stage_type == "diffusion":
-        engine_args.pop("model_stage", None)
-        engine_args.pop("model", None)
-        stage_engine = OmniDiffusion(
-            model=model,
-            stage_id=stage_id,
-            engine_input_source=stage_payload.get("engine_input_source", []),
-            **engine_args,
-        )
-    else:
-        # Default to LLM engine
-        stage_engine = OmniLLM(model=model, **engine_args)
-
-    # Release all locks AFTER engine initialization completes
-    for lock_fd in lock_files:
-        try:
-            fcntl.flock(lock_fd, fcntl.LOCK_UN)
-            _os.close(lock_fd)
-            logger.debug("Released initialization lock (fd=%s)", lock_fd)
-        except (OSError, ValueError):
-            pass
-    lock_files = []  # Clear after release
-
+    try:
+        if stage_type == "diffusion":
+            engine_args.pop("model_stage", None)
+            engine_args.pop("model", None)
+            stage_engine = OmniDiffusion(
+                model=model,
+                stage_id=stage_id,
+                engine_input_source=stage_payload.get("engine_input_source", []),
+                **engine_args,
+            )
+        else:
+            # Default to LLM engine
+            stage_engine = OmniLLM(model=model, **engine_args)
+    finally:
+        # Release all locks by closing file descriptors
+        # Locks are automatically released when file descriptors are closed
+        # or when process dies
+        for lock_fd in lock_files:
+            try:
+                fcntl.flock(lock_fd, fcntl.LOCK_UN)
+                _os.close(lock_fd)
+                logger.debug("Released initialization lock (fd=%s)", lock_fd)
+            except (OSError, ValueError):
+                pass
    logger.debug("Engine initialized")
    # Initialize OmniConnectors if configured
    connectors: dict[tuple[str, str], OmniConnectorBase] | None = {}
@@ -19,41 +19,40 @@ from typing import Annotated, Any, cast

 import httpx
 import vllm.envs as envs
-from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
+from fastapi import Depends, File, Form, HTTPException, Request, UploadFile
 from fastapi.responses import JSONResponse, StreamingResponse
 from PIL import Image
 from starlette.datastructures import State
 from starlette.routing import Route
 from vllm import SamplingParams
 from vllm.engine.protocol import EngineClient
-from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
-from vllm.entrypoints.chat_utils import load_chat_template
+from vllm.entrypoints.anthropic.serving_messages import AnthropicServingMessages
 from vllm.entrypoints.launcher import serve_http
 from vllm.entrypoints.logger import RequestLogger
-from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
-from vllm.entrypoints.openai.api_server import base, load_log_config
-from vllm.entrypoints.openai.api_server import build_app as build_openai_app
-from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
-from vllm.entrypoints.openai.chat_completion.protocol import (
+from vllm.entrypoints.openai.api_server import (
+    base,
+    build_app,
+    load_log_config,
+    router,
+    setup_server,
+)
+from vllm.entrypoints.openai.orca_metrics import metrics_header
+from vllm.entrypoints.openai.protocol import (
    ChatCompletionRequest,
    ChatCompletionResponse,
-)
-
-# yapf conflicts with isort for this block
-# yapf: disable
-# yapf: enable
-from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
-from vllm.entrypoints.openai.engine.protocol import (
    ErrorResponse,
    ModelCard,
    ModelList,
    ModelPermission,
 )
-from vllm.entrypoints.openai.models.protocol import BaseModelPath
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.openai.orca_metrics import metrics_header
-from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
-from vllm.entrypoints.openai.translations.serving import (
+
+# yapf conflicts with isort for this block
+# yapf: disable
+# yapf: enable
+from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
+from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
+from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses
+from vllm.entrypoints.openai.serving_transcription import (
    OpenAIServingTranscription,
    OpenAIServingTranslation,
 )
@@ -64,8 +63,10 @@ from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
 from vllm.entrypoints.pooling.score.serving import ServingScores
 from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.utils import (
    load_aware_call,
+    process_chat_template,
    process_lora_modules,
    with_cancellation,
 )
@@ -92,41 +93,29 @@ from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id

 logger = init_logger(__name__)
-router = APIRouter()
-
-
-def _remove_route_from_router(
-    router: APIRouter,
-    path: str,
-    methods: set[str] | None = None,
-) -> None:
-    methods_set = {method.upper() for method in methods} if methods else None
-    for route in list(router.routes):
-        if getattr(route, "path", None) != path:
-            continue
-        if methods_set is not None:
-            route_methods = {method.upper() for method in (getattr(route, "methods", None) or set())}
-            if not (route_methods & methods_set):
-                continue
-        router.routes.remove(route)
-

 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format"


-def _remove_route_from_app(app, path: str, methods: set[str] | None = None):
-    """Remove a route from the app by path and optionally by methods.
+def _remove_route_from_router(router_obj, path: str, methods: set[str] | None = None):
+    """Remove a route from the router by path and optionally by methods.

-    OMNI: used to override upstream /v1/chat/completions with omni behavior.
+    This is needed because vllm's api_server registers routes when imported,
+    and we need to override some routes (like /v1/chat/completions) with
+    omni-specific implementations.
    """
    routes_to_remove = []
-    for route in app.routes:
+    for route in router_obj.routes:
        if isinstance(route, Route) and route.path == path:
            if methods is None or (hasattr(route, "methods") and route.methods & methods):
                routes_to_remove.append(route)

    for route in routes_to_remove:
-        app.routes.remove(route)
+        router_obj.routes.remove(route)
+
+
+# Remove vllm's /v1/chat/completions route so we can register our own omni version
+_remove_route_from_router(router, "/v1/chat/completions", {"POST"})


 class _DiffusionServingModels:
@@ -173,7 +162,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
    # Add process-specific prefix to stdout and stderr.
    decorate_logs("APIServer")

-    listen_address, sock = setup_openai_server(args)
+    listen_address, sock = setup_server(args)

    # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models
    await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
@@ -190,7 +179,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
        ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin)

    # Load logging config for uvicorn if specified
-    log_config = load_log_config(getattr(args, "log_config_file", None))
+    log_config = load_log_config(args.log_config_file)
    if log_config is not None:
        uvicorn_kwargs["log_config"] = log_config

@@ -198,19 +187,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
        args,
        client_config=client_config,
    ) as engine_client:
-        supported_tasks: tuple[str, ...]
-        if hasattr(engine_client, "get_supported_tasks"):
-            supported_tasks = tuple(await engine_client.get_supported_tasks())
-        else:
-            supported_tasks = ("generate",)
-        if not supported_tasks:
-            supported_tasks = ("generate",)
-
-        app = build_openai_app(args)
-        # OMNI: Remove upstream routes that we override with omni-specific handlers
-        _remove_route_from_app(app, "/v1/chat/completions", {"POST"})
-        _remove_route_from_app(app, "/v1/models", {"GET"})  # Remove upstream /v1/models to use omni's handler
-        app.include_router(router)
+        app = build_app(args)

        await omni_init_app_state(engine_client, app.state, args)

@@ -397,8 +374,6 @@ async def omni_init_app_state(
        state.vllm_config = None
        state.diffusion_engine = engine_client
        state.openai_serving_models = _DiffusionServingModels(base_model_paths)
-        # OMNI: tokenization endpoints are not supported in pure diffusion mode.
-        state.openai_serving_tokenization = None

        # Use for_diffusion method to create chat handler
        state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion(
@@ -426,7 +401,11 @@ async def omni_init_app_state(
        supported_tasks = set(await engine_client.get_supported_tasks())
    logger.info("Supported tasks: %s", supported_tasks)

-    resolved_chat_template = load_chat_template(args.chat_template)
+    resolved_chat_template = await process_chat_template(
+        args.chat_template,
+        engine_client,
+        vllm_config.model_config if vllm_config is not None else None,
+    )

    if args.tool_server == "demo":
        tool_server: ToolServer | None = DemoToolServer()
@@ -465,10 +444,10 @@ async def omni_init_app_state(
                tokenizer = await engine_client.get_tokenizer()
                if tokenizer is not None:
                    # Initialize input_processor
-                    # OMNI: OmniInputProcessor creates tokenizer internally from vllm_config
                    if not hasattr(engine_client, "input_processor") or engine_client.input_processor is None:
                        engine_client.input_processor = OmniInputProcessor(
                            vllm_config=vllm_config,
+                            tokenizer=tokenizer,
                        )
                        logger.info("Initialized input_processor for AsyncOmni")

@@ -712,13 +691,7 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
    metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "")
    handler = Omnichat(raw_request)
    if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
-        if base_server is None:
-            raise HTTPException(
-                status_code=HTTPStatus.NOT_FOUND.value,
-                detail="The model does not support Chat Completions API",
-            )
-        return base_server.create_error_response(message="The model does not support Chat Completions API")
+        return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
    try:
        generator = await handler.create_chat_completion(request, raw_request)
    except Exception as e:
@@ -787,13 +760,7 @@ _remove_route_from_router(router, "/v1/audio/speech", {"POST"})
 async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request):
    handler = Omnispeech(raw_request)
    if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
-        if base_server is None:
-            raise HTTPException(
-                status_code=HTTPStatus.NOT_FOUND.value,
-                detail="The model does not support Speech API",
-            )
-        return base_server.create_error_response(message="The model does not support Speech API")
+        return base(raw_request).create_error_response(message="The model does not support Speech API")
    try:
        return await handler.create_speech(request, raw_request)
    except Exception as e:
@@ -1,4 +1,6 @@
-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionStreamResponse
+from vllm.entrypoints.openai.protocol import (
+    ChatCompletionStreamResponse,
+)


 class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse):
@@ -13,7 +13,6 @@ import torch
 from fastapi import Request
 from PIL import Image
 from pydantic import TypeAdapter
-from vllm.renderers import RendererLike

 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
@@ -29,10 +28,17 @@ from vllm.entrypoints.chat_utils import (
    ChatCompletionMessageParam,
    ChatTemplateContentFormatOption,
    ConversationMessage,
+    apply_hf_chat_template,
+    apply_mistral_chat_template,
    get_history_tool_calls_cnt,
    make_tool_call_id,
+    resolve_chat_template_content_format,
 )
-from vllm.entrypoints.openai.chat_completion.protocol import (
+from vllm.entrypoints.openai.parser.harmony_utils import (
+    get_streamable_parser_for_assistant,
+    parse_chat_output,
+)
+from vllm.entrypoints.openai.protocol import (
    ChatCompletionNamedToolChoiceParam,
    ChatCompletionRequest,
    ChatCompletionResponse,
@@ -40,9 +46,6 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse,
    ChatMessage,
-)
-from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
-from vllm.entrypoints.openai.engine.protocol import (
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
@@ -52,29 +55,20 @@ from vllm.entrypoints.openai.engine.protocol import (
    FunctionDefinition,
    PromptTokenUsageInfo,
    RequestResponseMetadata,
+    ResponsesRequest,
    ToolCall,
    UsageInfo,
 )
-from vllm.entrypoints.openai.engine.serving import ChatLikeRequest, clamp_prompt_logprobs
-from vllm.entrypoints.openai.parser.harmony_utils import (
-    get_streamable_parser_for_assistant,
-    parse_chat_output,
+from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
+from vllm.entrypoints.openai.serving_engine import (
+    ChatLikeRequest,
+    clamp_prompt_logprobs,
 )
-from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
 from vllm.entrypoints.utils import should_include_usage
 from vllm.inputs.data import PromptType, TokensPrompt
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
-from vllm.renderers.hf import (
-    resolve_chat_template_content_format,
-)
-from vllm.renderers.hf import (
-    safe_apply_chat_template as apply_hf_chat_template,
-)
-from vllm.renderers.mistral import (
-    safe_apply_chat_template as apply_mistral_chat_template,
-)
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.tokenizers.mistral import (
@@ -177,10 +171,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):

            model_name = self.models.model_name(lora_request)

-            renderer = self.renderer
-            tokenizer = renderer.get_tokenizer()
-            if tokenizer is None:
-                tokenizer = await self.engine_client.get_tokenizer()
+            tokenizer = await self.engine_client.get_tokenizer()

            tool_parser = self.tool_parser

@@ -237,7 +228,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    engine_prompts,
                ) = await self._preprocess_chat(
                    request,
-                    renderer,
+                    tokenizer,
                    request.messages,
                    chat_template=request.chat_template or self.chat_template,
                    chat_template_content_format=self.chat_template_content_format,
@@ -296,7 +287,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):

                generators.append(generator)
        except ValueError as e:
-            return self.create_error_response(e)
+            # TODO: Use a vllm-specific Validation Error
+            return self.create_error_response(str(e))

        assert len(generators) == 1
        (result_generator,) = generators
@@ -324,12 +316,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                request_metadata,
            )
        except ValueError as e:
-            return self.create_error_response(e)
+            # TODO: Use a vllm-specific Validation Error
+            return self.create_error_response(str(e))

    async def _preprocess_chat(
        self,
        request: ChatLikeRequest | ResponsesRequest,
-        renderer: RendererLike,
+        tokenizer: TokenizerLike,
        messages: list[ChatCompletionMessageParam],
        chat_template: str | None,
        chat_template_content_format: ChatTemplateContentFormatOption,
@@ -347,22 +340,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        list[TokensPrompt],
    ]:
        model_config = self.model_config
-        tokenizer = renderer.get_tokenizer() if renderer is not None else None

-        if tokenizer is None or isinstance(tokenizer, MistralTokenizer):
-            resolved_content_format = (
-                chat_template_content_format if chat_template_content_format != "auto" else "string"
-            )
-        else:
-            resolved_content_format = resolve_chat_template_content_format(
-                chat_template,
-                tool_dicts,
-                chat_template_content_format,
-                tokenizer,
-                model_config=model_config,
-            )
-        # OMNI: Updated for vLLM v0.15.0 API - resolve_items() returns (mm_data, mm_uuids) tuple
-        conversation, mm_future = parse_chat_messages_futures(
+        resolved_content_format = resolve_chat_template_content_format(
+            chat_template,
+            tool_dicts,
+            chat_template_content_format,
+            tokenizer,
+            model_config=model_config,
+        )
+        conversation, mm_data_future, mm_uuids = parse_chat_messages_futures(
            messages,
            model_config,
            content_format=resolved_content_format,
@@ -396,21 +382,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                **_chat_template_kwargs,
            )
        else:
-            hf_chat_template_kwargs = dict(_chat_template_kwargs)
-            hf_chat_template_kwargs.pop("tools", None)
-            hf_chat_template_kwargs.pop("chat_template", None)
            request_prompt = apply_hf_chat_template(
-                model_config=model_config,
                tokenizer=tokenizer,
                conversation=conversation,
-                tools=tool_dicts,
-                chat_template=chat_template,
-                tokenize=False,
-                **hf_chat_template_kwargs,
+                model_config=model_config,
+                **_chat_template_kwargs,
            )

-        # OMNI: Await the combined future to get both mm_data and mm_uuids
-        mm_data, mm_uuids = await mm_future
+        mm_data = await mm_data_future

        # tool parsing is done only if a tool_parser has been set and if
        # tool_choice is not "none" (if tool_choice is "none" but a tool_parser
@@ -655,7 +634,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                )
        except RuntimeError as e:
            logger.exception("Error in reasoning parser creation.")
-            data = self.create_streaming_error_response(e)
+            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
            yield "data: [DONE]\n\n"
            return
@@ -667,7 +646,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                tool_parsers = [None] * num_choices
        except Exception as e:
            logger.exception("Error in tool parser creation.")
-            data = self.create_streaming_error_response(e)
+            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
            yield "data: [DONE]\n\n"
            return
@@ -1307,8 +1286,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    )

        except Exception as e:
+            # TODO: Use a vllm-specific Validation Error
            logger.exception("Error in chat completion stream generator.")
-            data = self.create_streaming_error_response(e)
+            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"
@@ -1333,7 +1313,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")
        except ValueError as e:
-            return self.create_error_response(e)
+            # TODO: Use a vllm-specific Validation Error
+            return self.create_error_response(str(e))

        assert final_outputs is not None

@@ -1506,7 +1487,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    )
                except RuntimeError as e:
                    logger.exception("Error in reasoning parser creation.")
-                    return self.create_error_response(e)
+                    return self.create_error_response(str(e))
                # If the reasoning parser is enabled,
                # tool calls are extracted exclusively from the content.
                reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request)
@@ -1591,7 +1572,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    tool_parser = self.tool_parser(tokenizer)
                except RuntimeError as e:
                    logger.exception("Error in tool parser creation.")
-                    return self.create_error_response(e)
+                    return self.create_error_response(str(e))

                tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request)
                # In the OpenAI API the finish_reason is "tools_called"
@@ -1676,9 +1657,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
    ):
        choices: list[ChatCompletionResponseChoice] = []
        final_res = omni_outputs.request_output
-        # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
-        # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421
-        audio_data = final_res.outputs[0].multimodal_output.get("audio")
+        audio_data = final_res.multimodal_output.get("audio")
        if stream:
            audio_tensor = audio_data[-1].float().detach().cpu().numpy()
        else:
@@ -1764,11 +1743,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        if omni_outputs.images:
            images = omni_outputs.images
        # Fall back to request_output for pipeline mode
-        # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
-        elif final_res is not None and final_res.outputs:
-            completion_output = final_res.outputs[0]
-            if hasattr(completion_output, "multimodal_output") and completion_output.multimodal_output:
-                image_data = completion_output.multimodal_output.get("image")
+        elif final_res is not None:
+            if hasattr(final_res, "multimodal_output") and final_res.multimodal_output:
+                image_data = final_res.multimodal_output.get("image")
                if image_data is not None:
                    if isinstance(image_data, Image.Image):
                        images.append(image_data)
@@ -3,7 +3,7 @@ from typing import Any

 from fastapi import Request
 from fastapi.responses import Response
-from vllm.entrypoints.openai.engine.serving import OpenAIServing
+from vllm.entrypoints.openai.serving_engine import OpenAIServing
 from vllm.logger import init_logger
 from vllm.utils import random_uuid

@@ -258,26 +258,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):

            # Extract audio from output
            # Audio can be in final_output.multimodal_output or final_output.request_output.multimodal_output
-            # Support both "audio" and "model_outputs" keys for compatibility with different models
            audio_output = None
            if hasattr(final_output, "multimodal_output") and final_output.multimodal_output:
                audio_output = final_output.multimodal_output
-            if not audio_output and hasattr(final_output, "request_output"):
+            if (not audio_output or "audio" not in audio_output) and hasattr(final_output, "request_output"):
                if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"):
                    audio_output = final_output.request_output.multimodal_output

-            # Check for audio data using either "audio" or "model_outputs" key
-            audio_key = None
-            if audio_output:
-                if "audio" in audio_output:
-                    audio_key = "audio"
-                elif "model_outputs" in audio_output:
-                    audio_key = "model_outputs"
-
-            if not audio_output or audio_key is None:
+            if not audio_output or "audio" not in audio_output:
                return self.create_error_response("TTS model did not produce audio output.")

-            audio_tensor = audio_output[audio_key]
+            audio_tensor = audio_output["audio"]
            sample_rate = audio_output.get("sr", 24000)
            if hasattr(sample_rate, "item"):
                sample_rate = sample_rate.item()
@@ -305,7 +296,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")
        except ValueError as e:
-            return self.create_error_response(e)
+            return self.create_error_response(str(e))
        except Exception as e:
            logger.exception("Speech generation failed: %s", e)
            return self.create_error_response(f"Speech generation failed: {e}")
@@ -104,9 +104,6 @@ class OmniEmbedsPrompt(EmbedsPrompt):
 # Must ensure that all additional prompt types are inherited from vLLM prompt types
 # Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime.
 # Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate()
-OmniSingletonPrompt: TypeAlias = str | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt
-"""Omni singleton prompt type extending vLLM's SingletonPrompt with additional fields."""
-
 OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt


@@ -1,10 +1,10 @@
+from vllm.inputs.data import SingletonPrompt
 from vllm.inputs.parse import (
    ParsedEmbedsPrompt,
    ParsedSingletonPrompt,
    ParsedStrPrompt,
    ParsedTextPrompt,
    ParsedTokensPrompt,
-    SingletonPrompt,
 )


@@ -1,17 +1,12 @@
 from typing import Any

 from typing_extensions import assert_never
-from vllm.inputs.data import SingletonInputs, SingletonPrompt
+from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
 from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict

-from vllm_omni.inputs.data import (
-    OmniTextPrompt,
-    OmniTokenInputs,
-    OmniTokensPrompt,
-    token_inputs_omni,
-)
+from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni
 from vllm_omni.inputs.parse import parse_singleton_prompt_omni

 logger = init_logger(__name__)
@@ -27,7 +22,7 @@ class OmniInputPreprocessor(InputPreprocessor):

    def _process_text(
        self,
-        parsed_content: OmniTextPrompt,
+        parsed_content: TextPrompt,
        tokenization_kwargs: dict[str, Any] | None = None,
        *,
        mm_uuids: MultiModalUUIDDict | None = None,
@@ -67,7 +62,7 @@ class OmniInputPreprocessor(InputPreprocessor):

    def _process_tokens(
        self,
-        parsed_content: OmniTokensPrompt,
+        parsed_content: TokensPrompt,
        tokenization_kwargs: dict[str, Any] | None = None,
        *,
        mm_uuids: MultiModalUUIDDict | None = None,
@@ -85,10 +80,6 @@ class OmniInputPreprocessor(InputPreprocessor):
                tokenization_kwargs=tokenization_kwargs,
                mm_uuids=mm_uuids,
            )
-            if prompt_embeds is not None:
-                inputs["prompt_embeds"] = prompt_embeds
-            if additional_information is not None:
-                inputs["additional_information"] = additional_information
        else:
            inputs = token_inputs_omni(
                prompt_token_ids=prompt_token_ids,
@@ -123,6 +114,8 @@ class OmniInputPreprocessor(InputPreprocessor):
        """
        parsed = parse_singleton_prompt_omni(prompt)

+        if parsed["type"] == "embeds":
+            return self._process_embeds(parsed["content"])
        if parsed["type"] == "tokens":
            return self._process_tokens(
                parsed["content"],
@@ -136,7 +129,7 @@ class OmniInputPreprocessor(InputPreprocessor):
            )
        if parsed["type"] == "str":
            return self._process_text(
-                OmniTextPrompt(prompt=parsed["content"]),
+                TextPrompt(prompt=parsed["content"]),
                tokenization_kwargs=tokenization_kwargs,
                mm_uuids=mm_uuids,
            )
@@ -1,32 +1,136 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Omni-extended MRotaryEmbedding with multimodal position computation methods.
-
-This module extends the upstream vLLM MRotaryEmbedding with additional methods
-for computing input positions for various multimodal scenarios including:
- Image/Video inputs (Qwen2.5-VL style)
- Audio inputs (Qwen2.5-Omni style)
- Audio-in-video interleaved inputs
- GLM4V style inputs
-"""
-
 import itertools

+import numpy as np
 import torch
 from transformers import PretrainedConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
+from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding

 logger = init_logger(__name__)


-class OmniMRotaryEmbedding(MRotaryEmbedding):
-    """Omni-extended MRotaryEmbedding with multimodal position computation.
-
-    Extends the upstream MRotaryEmbedding with additional class methods for
-    computing input positions for various multimodal scenarios.
+def _apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    is_neox_style: bool,
+) -> torch.Tensor:
    """
+    Args:
+        x: [num_tokens, num_heads, head_size]
+        cos: [num_tokens, head_size // 2]
+        sin: [num_tokens, head_size // 2]
+        is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
+            positional embeddings.
+    """
+    cos = cos.unsqueeze(-2).to(x.dtype)
+    sin = sin.unsqueeze(-2).to(x.dtype)
+    if is_neox_style:
+        x1, x2 = torch.chunk(x, 2, dim=-1)
+    else:
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    if is_neox_style:
+        return torch.cat((o1, o2), dim=-1)
+    else:
+        return torch.stack((o1, o2), dim=-1).flatten(-2)
+
+
+class MRotaryEmbedding(RotaryEmbedding):
+    """Rotary Embedding with Multimodal Sections."""
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: float,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+        mrope_section: list[int] | None = None,
+        mrope_interleaved: bool = False,
+        # YaRN parameters.
+        *,
+        scaling_factor: float | None = None,
+        extrapolation_factor: float = 1,
+        attn_factor: float = 1,
+        beta_fast: int = 32,
+        beta_slow: int = 1,
+    ) -> None:
+        # In Qwen2.5-VL, the maximum index value is related to the duration of
+        # the input video. We enlarge max_position_embeddings to 4 times to get
+        # a larger the cos and sin cache.
+        self.cache_max_position_num = max_position_embeddings * 4
+        super().__init__(
+            head_size,
+            rotary_dim,
+            self.cache_max_position_num,
+            base,
+            is_neox_style,
+            dtype,
+        )
+
+        self.mrope_section = mrope_section
+        if self.mrope_section:
+            print(
+                "Warning: mrope_section check is disabled in Qwen2.5-Omni, "
+                "this may cause errors, and should be restored in the future."
+            )
+            # assert sum(self.mrope_section) == rotary_dim // 2
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """PyTorch-native implementation equivalent to forward().
+
+        Args:
+            positions:
+                [num_tokens,] (text only) or
+                [3, num_tokens] (T/H/W positions with multimodal inputs)
+            query: [num_tokens, num_heads * head_size]
+            key: [num_tokens, num_kv_heads * head_size]
+        """
+        assert positions.ndim == 1 or positions.ndim == 2
+        assert key is not None
+
+        num_tokens = positions.shape[-1]
+        cos_sin = self.cos_sin_cache[positions]
+        cos, sin = cos_sin.chunk(2, dim=-1)
+        if positions.ndim == 2:
+            assert self.mrope_section
+
+            cos = torch.cat(
+                [m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
+            sin = torch.cat(
+                [m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
+
+        query_shape = query.shape
+        query = query.view(num_tokens, -1, self.head_size)
+        query_rot = query[..., : self.rotary_dim]
+        query_pass = query[..., self.rotary_dim :]
+        # query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,
+        #                                       self.is_neox_style)
+        query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
+        query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
+
+        key_shape = key.shape
+        key = key.view(num_tokens, -1, self.head_size)
+        key_rot = key[..., : self.rotary_dim]
+        key_pass = key[..., self.rotary_dim :]
+        # key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,
+        #                                     self.is_neox_style)
+        key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
+        key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
+        return query, key

    @classmethod
    def get_input_positions(
@@ -336,6 +440,9 @@ class OmniMRotaryEmbedding(MRotaryEmbedding):
            |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
        """

+        # TODO(fyabc): refactor and share more code with
+        #  _vl_get_input_positions_tensor.
+
        thinker_config = hf_config.thinker_config
        try:
            audio_token_id = thinker_config.audio_token_index
@@ -506,6 +613,29 @@ class OmniMRotaryEmbedding(MRotaryEmbedding):
            ranges[index].append(num)
        return ranges

+    @staticmethod
+    def get_next_input_positions(
+        mrope_position_delta: int,
+        context_len: int,
+        seq_len: int,
+    ) -> list[list[int]]:
+        return [list(range(context_len + mrope_position_delta, seq_len + mrope_position_delta)) for _ in range(3)]
+
+    @staticmethod
+    def get_next_input_positions_tensor(
+        out: np.ndarray,
+        out_offset: int,
+        mrope_position_delta: int,
+        context_len: int,
+        num_new_tokens: int,
+    ):
+        values = np.arange(
+            mrope_position_delta + context_len,
+            mrope_position_delta + context_len + num_new_tokens,
+            dtype=out.dtype,
+        )
+        out[:, out_offset : out_offset + num_new_tokens] = values
+
    @classmethod
    def omni_get_updates_use_audio_in_video(
        cls,
@@ -1,6 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from .mrope import OmniMRotaryEmbedding
-
-__all__ = ["OmniMRotaryEmbedding"]
@@ -22,9 +22,6 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
-from vllm.model_executor.models.vision import (
-    get_llm_pos_ids_for_vision,
-)

 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -38,6 +35,7 @@ from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
+from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
 from vllm_omni.platforms import current_omni_platform

 TALKER_CODEC_EOS_TOKEN_ID = 8294
@@ -286,6 +286,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "visual"),
+                multimodal_config=multimodal_config,
            )
        else:
            self.visual = None
@@ -82,7 +82,7 @@ from vllm.model_executor.models.vision import (
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems
 from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems
-from vllm.multimodal.processing.processor import (
+from vllm.multimodal.processing import (
    MultiModalPromptUpdates,
    PlaceholderFeaturesInfo,
    PromptReplacement,
@@ -609,8 +609,7 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
            feature_lens=audio_feature_lengths,
            aftercnn_lens=audio_output_lengths,
        )
-        # OMNI: audio_tower.forward() returns hidden_states tensor directly
-        audio_features = audio_outputs
+        audio_features = audio_outputs.last_hidden_state
        return audio_features.split(audio_output_lengths.tolist())


@@ -675,6 +674,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "visual"),
+            multimodal_config=multimodal_config,
        )
        self.quant_config = quant_config

@@ -504,19 +504,12 @@ class Qwen3TTSConfig(PretrainedConfig):
        self.tts_bos_token_id = tts_bos_token_id
        self.tts_eos_token_id = tts_eos_token_id

-        # TODO: remove these dummy values after
-        self.image_token_id = 0  # dummy image token id
-        self.video_token_id = 0  # dummy video token id
-        self.vision_start_token_id = 0  # dummy vision start token id
-        self.vision_config = PretrainedConfig()  # dummy vision config
-        self.vision_config.spatial_merge_size = 1
-
    def get_text_config(self, **kwargs):
        # vLLM expects text config to expose hidden_size/num_attention_heads.
        # For Qwen3 TTS, the talker config is the text model config.
        config = self.talker_config
-        # if hasattr(config, "rope_parameters"):
-        #     delattr(config, "rope_parameters")
+        if hasattr(config, "rope_parameters"):
+            delattr(config, "rope_parameters")
        return config


@@ -430,22 +430,6 @@ def mel_spectrogram(
    return mel_spec


-def _compute_default_rope_parameters(
-    config,
-    device,
-):
-    base = config.rope_theta
-    partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
-    head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
-    dim = int(head_dim * partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
-    return inv_freq, attention_factor
-
-
 class Qwen3TTSPreTrainedModel(PreTrainedModel):
    config_class = Qwen3TTSConfig
    base_model_prefix = "model"
@@ -518,9 +502,7 @@ class Qwen3TTSTalkerRotaryEmbedding(nn.Module):
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
-        self.rope_init_fn: Callable = _compute_default_rope_parameters
-        if self.rope_type != "default":
-            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -556,9 +538,7 @@ class Qwen3TTSRotaryEmbedding(nn.Module):
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
-        self.rope_init_fn: Callable = _compute_default_rope_parameters
-        if self.rope_type != "default":
-            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -1409,6 +1389,7 @@ class Qwen3TTSTalkerModel(Qwen3TTSTalkerTextPreTrainedModel):

    def __init__(self, config):
        super().__init__(config)
+        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.layers = nn.ModuleList(
            [Qwen3TTSTalkerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
@@ -0,0 +1,23 @@
+import torch
+
+
+def get_llm_pos_ids_for_vision(
+    start_idx: int,
+    vision_idx: int,
+    spatial_merge_size: int,
+    t_index: list[int],
+    grid_hs: torch.Tensor,
+    grid_ws: torch.Tensor,
+) -> torch.Tensor:
+    llm_pos_ids_list = []
+    llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
+    llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
+    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten()
+    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten()
+    t_index_tensor = (
+        torch.Tensor(t_index).to(llm_grid_h.device).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
+    )
+    _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
+    llm_pos_ids_list.append(_llm_pos_ids + start_idx)
+    llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
+    return llm_pos_ids
@@ -59,7 +59,7 @@ class OmniRequestOutput:
    prompt: OmniPromptType | None = None
    latents: torch.Tensor | None = None
    metrics: dict[str, Any] = field(default_factory=dict)
-    _multimodal_output: dict[str, Any] = field(default_factory=dict)
+    multimodal_output: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_pipeline(
@@ -116,21 +116,10 @@ class OmniRequestOutput:
            prompt=prompt,
            latents=latents,
            metrics=metrics or {},
-            _multimodal_output=multimodal_output or {},
+            multimodal_output=multimodal_output or {},
            finished=True,
        )

-    @property
-    def multimodal_output(self) -> dict[str, Any]:
-        """Return multimodal output from the underlying request output or local field.
-
-        For pipeline outputs, this proxies to request_output.multimodal_output.
-        For diffusion outputs, this returns the local _multimodal_output field.
-        """
-        if self.request_output is not None:
-            return getattr(self.request_output, "multimodal_output", {})
-        return self._multimodal_output
-
    @property
    def num_images(self) -> int:
        """Return the number of generated images."""
@@ -240,7 +229,7 @@ class OmniRequestOutput:
            f"prompt={self.prompt!r}",
            f"latents={self.latents}",
            f"metrics={self.metrics}",
-            f"multimodal_output={self._multimodal_output}",
+            f"multimodal_output={self.multimodal_output}",
        ]

        return f"OmniRequestOutput({', '.join(parts)})"
@@ -12,7 +12,7 @@ from vllm.v1.request import Request as _OriginalRequest
 import vllm_omni.logger  # noqa: F401
 from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
-from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding
+from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 from vllm_omni.request import OmniRequest

 for module_name, module in sys.modules.items():
@@ -26,7 +26,7 @@ for module_name, module in sys.modules.items():
    if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
        module.TokensPrompt = OmniTokensPrompt
    if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:
-        module.MRotaryEmbedding = OmniMRotaryEmbedding
+        module.MRotaryEmbedding = MRotaryEmbedding
    if hasattr(module, "Request") and module.Request == _OriginalRequest:
        module.Request = OmniRequest
    if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
@@ -1,8 +1,6 @@
 from collections.abc import Callable
 from typing import TYPE_CHECKING

-import numpy as np
-import torch
 from vllm.v1.request import Request

 if TYPE_CHECKING:
@@ -27,35 +25,21 @@ class OmniRequest(Request):

    def __init__(
        self,
-        prompt_embeds: PromptEmbedsPayload | torch.Tensor | None = None,
+        prompt_embeds: PromptEmbedsPayload | None = None,
        # Optional external request ID for tracking
        external_req_id: str | None = None,
        additional_information: AdditionalInformationPayload | None = None,
        *args,
        **kwargs,
    ):
-        prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds)
-        super().__init__(prompt_embeds=prompt_embeds_tensor, *args, **kwargs)
-        # Preserve serialized prompt embeddings payload (optional)
-        self.prompt_embeds_payload: PromptEmbedsPayload | None = (
-            prompt_embeds if isinstance(prompt_embeds, PromptEmbedsPayload) else None
-        )
+        super().__init__(*args, **kwargs)
+        # Serialized prompt embeddings payload (optional)
+        self.prompt_embeds: PromptEmbedsPayload | None = prompt_embeds
        # Optional external request ID for tracking
        self.external_req_id: str | None = external_req_id
        # Serialized additional information payload (optional)
        self.additional_information: AdditionalInformationPayload | None = additional_information

-    @staticmethod
-    def _maybe_decode_prompt_embeds(
-        prompt_embeds: PromptEmbedsPayload | torch.Tensor | None,
-    ) -> torch.Tensor | None:
-        if isinstance(prompt_embeds, PromptEmbedsPayload):
-            dtype = getattr(np, prompt_embeds.dtype)
-            arr = np.frombuffer(prompt_embeds.data, dtype=dtype)
-            arr = arr.reshape(prompt_embeds.shape)
-            return torch.from_numpy(arr)
-        return prompt_embeds
-
    @classmethod
    def from_engine_core_request(
        cls,
@@ -90,5 +74,4 @@ class OmniRequest(Request):
            trace_headers=request.trace_headers,
            block_hasher=block_hasher,
            additional_information=request.additional_information,
-            resumable=request.resumable,
        )
@@ -13,8 +13,7 @@ import numpy as np
 import torch
 from vllm.config import CUDAGraphMode
 from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
-from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
-from vllm.distributed.parallel_state import get_pp_group, get_tp_group
+from vllm.distributed.kv_transfer import get_kv_transfer_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
@@ -29,6 +28,9 @@ from vllm.v1.worker.gpu_model_runner import (
    EMPTY_MODEL_RUNNER_OUTPUT,
    AsyncGPUModelRunnerOutput,
    IntermediateTensors,
+    get_pp_group,
+    get_tp_group,
+    has_kv_transfer_group,
 )
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
@@ -219,24 +221,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
            use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
            ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices

-            # True if any attention backend handles KV cache update separately
-            # from forward() (i.e., forward_includes_kv_cache_update=False). When true,
-            # slot_mappings must use padded dimensions to match the key/value tensors.
-            from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec
-
-            has_separate_kv_update = not all(
-                all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id])
-                for id, spec in enumerate(self.kv_cache_config.kv_cache_groups)
-                if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec)
-            )
-
-            slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
-                num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded,
-                num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs),
-                num_tokens_unpadded=num_tokens_unpadded,
-                ubatch_slices=ubatch_slices_padded,
-            )
-
            attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
                num_tokens=num_tokens_unpadded,
                num_tokens_padded=num_tokens_padded if pad_attn else None,
@@ -248,7 +232,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
                use_spec_decode=use_spec_decode,
                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
                cascade_attn_prefix_lens=cascade_attn_prefix_lens,
-                slot_mappings=slot_mappings_by_group,
            )

            (
@@ -279,7 +262,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
                cudagraph_runtime_mode=cudagraph_mode,
                batch_descriptor=batch_desc,
                ubatch_slices=ubatch_slices_padded,
-                slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
            ),
            record_function_or_nullcontext("gpu_model_runner: forward"),
            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
@@ -91,11 +91,6 @@ class GPUARWorker(OmniWorkerMixin, GPUWorker):
        num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
        init_workspace_manager(self.device, num_ubatches)

-        if self.use_v2_model_runner:
-            # OMNI: v2 model runner does not yet include omni hooks.
-            logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.")
-            self.use_v2_model_runner = False
-
        # Construct the model runner
        self.model_runner = GPUARModelRunner(self.vllm_config, self.device)

@@ -15,11 +15,10 @@ import torch
 from vllm.config import CUDAGraphMode
 from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
-from vllm.distributed.parallel_state import get_pp_group
-from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
    RoutedExpertsCapturer,
 )
+from vllm.model_executor.models.interfaces import supports_mm_encoder_only
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
@@ -30,6 +29,8 @@ from vllm.v1.worker.gpu_model_runner import (
    AsyncGPUModelRunnerOutput,
    IntermediateTensors,
    PerLayerAttnMetadata,
+    get_pp_group,
+    set_forward_context,
 )
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
@@ -189,24 +190,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
            use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
            ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices

-            # OMNI: True if any attention backend handles KV cache update separately
-            # from forward() (i.e., forward_includes_kv_cache_update=False). When true,
-            # slot_mappings must use padded dimensions to match the key/value tensors.
-            from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec
-
-            has_separate_kv_update = not all(
-                all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id])
-                for id, spec in enumerate(self.kv_cache_config.kv_cache_groups)
-                if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec)
-            )
-
-            slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
-                num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded,
-                num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs),
-                num_tokens_unpadded=num_tokens_unpadded,
-                ubatch_slices=ubatch_slices_padded,
-            )
-
            attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
                num_tokens=num_tokens_unpadded,
                num_tokens_padded=num_tokens_padded if pad_attn else None,
@@ -218,7 +201,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                use_spec_decode=use_spec_decode,
                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
                cascade_attn_prefix_lens=cascade_attn_prefix_lens,
-                slot_mappings=slot_mappings_by_group,
            )

            (
@@ -253,7 +235,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                cudagraph_runtime_mode=cudagraph_mode,
                batch_descriptor=batch_desc,
                ubatch_slices=ubatch_slices_padded,
-                slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
            ),
            record_function_or_nullcontext("Forward"),
            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
@@ -454,8 +435,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
            remove_lora: If False, dummy LoRAs are not destroyed after the run
            activate_lora: If False, dummy_run is performed without LoRAs.
        """
-        mm_config = self.vllm_config.model_config.multimodal_config
-        if mm_config and mm_config.mm_encoder_only:
+        if supports_mm_encoder_only(self.model):
            # The current dummy run only covers LM execution, so we can skip it.
            # mm encoder dummy run may need to add in the future.
            return torch.tensor([]), torch.tensor([])
@@ -559,14 +539,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):

        attn_metadata: PerLayerAttnMetadata | None = None

-        # OMNI: Get slot mappings before building attention metadata
-        slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
-            num_tokens_padded=num_tokens,
-            num_reqs_padded=num_reqs_padded,
-            num_tokens_unpadded=num_tokens_unpadded,
-            ubatch_slices=ubatch_slices_padded,
-        )
-
        # 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:
@@ -592,7 +564,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                max_query_len=max_query_len,
                ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
                for_cudagraph_capture=is_graph_capturing,
-                slot_mappings=slot_mappings_by_group,
            )

        with self.maybe_dummy_run_with_lora(
@@ -657,7 +628,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                    cudagraph_runtime_mode=cudagraph_runtime_mode,
                    batch_descriptor=batch_desc,
                    ubatch_slices=ubatch_slices_padded,
-                    slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
                ),
            ):
                outputs = self.model(
@@ -91,11 +91,6 @@ class GPUGenerationWorker(OmniWorkerMixin, GPUWorker):
        num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
        init_workspace_manager(self.device, num_ubatches)

-        if self.use_v2_model_runner:
-            # OMNI: v2 model runner does not yet include omni hooks.
-            logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.")
-            self.use_v2_model_runner = False
-
        self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device)

        if self.rank == 0:
@@ -8,7 +8,7 @@ from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
-from vllm.model_executor.models.interfaces import supports_mrope
+from vllm.model_executor.models.interfaces import supports_mm_encoder_only, supports_mrope
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
 from vllm.utils.import_utils import LazyLoader
@@ -448,8 +448,7 @@ class OmniGPUModelRunner(GPUModelRunner):
            remove_lora: If False, dummy LoRAs are not destroyed after the run
            activate_lora: If False, dummy_run is performed without LoRAs.
        """
-        mm_config = self.vllm_config.model_config.multimodal_config
-        if mm_config and mm_config.mm_encoder_only:
+        if supports_mm_encoder_only(self.model):
            # The current dummy run only covers LM execution, so we can skip it.
            # mm encoder dummy run may need to add in the future.
            return torch.tensor([]), torch.tensor([])
@@ -553,13 +552,6 @@ class OmniGPUModelRunner(GPUModelRunner):

        attn_metadata: PerLayerAttnMetadata | None = None

-        slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
-            num_tokens_padded=num_tokens,
-            num_reqs_padded=num_reqs_padded,
-            num_tokens_unpadded=num_tokens_unpadded,
-            ubatch_slices=ubatch_slices_padded,
-        )
-
        # 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:
@@ -585,7 +577,6 @@ class OmniGPUModelRunner(GPUModelRunner):
                max_query_len=max_query_len,
                ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
                for_cudagraph_capture=is_graph_capturing,
-                slot_mappings=slot_mappings_by_group,
            )

        with self.maybe_dummy_run_with_lora(
@@ -650,7 +641,6 @@ class OmniGPUModelRunner(GPUModelRunner):
                    cudagraph_runtime_mode=cudagraph_runtime_mode,
                    batch_descriptor=batch_desc,
                    ubatch_slices=ubatch_slices_padded,
-                    slot_mapping=slot_mappings,
                ),
            ):
                if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
@@ -698,7 +688,6 @@ class OmniGPUModelRunner(GPUModelRunner):
                    num_tokens,
                    use_cudagraphs=use_cudagraphs,
                    is_graph_capturing=is_graph_capturing,
-                    slot_mappings=slot_mappings,
                )

        # We register layerwise NVTX hooks here after the first dynamo tracing is
Author	SHA1	Message	Date
Roger Wang	f3c9b8e420	delete Signed-off-by: Roger Wang <hey@rogerw.io>	2026-02-01 00:09:07 -08:00
Roger Wang	c72e2fa0e5	update Signed-off-by: Roger Wang <hey@rogerw.io>	2026-02-01 00:08:24 -08:00