From c4220f0979cb07df4443e2688392171fd53f150f Mon Sep 17 00:00:00 2001
From: lif <1835304752@qq.com>
Date: Wed, 28 Jan 2026 13:11:15 +0800
Subject: [PATCH 01/32] [BugFix] Add /health and /v1/models endpoints for
 diffusion mode, fix streaming compatibility (#454)

Signed-off-by: majiayu000 <1835304752@qq.com>
---
 .../openai_api/test_image_server.py           |  59 +++++-
 tests/test_outputs.py                         | 171 ++++++++++++++++++
 vllm_omni/entrypoints/openai/api_server.py    |  88 ++++++++-
 vllm_omni/outputs.py                          |  52 ++++++
 4 files changed, 358 insertions(+), 12 deletions(-)
 create mode 100644 tests/test_outputs.py

diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 45f4b33ad..86331e085 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -124,6 +124,8 @@ class FakeAsyncOmni:
 def mock_async_diffusion():
     """Mock AsyncOmniDiffusion instance that returns fake images"""
     mock = Mock()
+    mock.is_running = True  # For health endpoint
+    mock.check_health = AsyncMock()  # For LLM mode health check
 
     async def generate(**kwargs):
         # Return n PIL images wrapped in result object
@@ -147,8 +149,9 @@ def test_client(mock_async_diffusion):
 
     # Set up app state with diffusion engine
     app.state.engine_client = mock_async_diffusion
+    app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
     app.state.stage_configs = [{"stage_type": "diffusion"}]
-    app.state.served_model_names = "Qwen/Qwen-Image"
+    app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
 
     return TestClient(app)
 
@@ -169,10 +172,58 @@ def async_omni_test_client():
     return TestClient(app)
 
 
-@pytest.mark.skip(reason="Async API server uses different health check mechanism")
 def test_health_endpoint(test_client):
-    """Test health check endpoint - skipped for async server"""
-    pass
+    """Test health check endpoint for diffusion mode"""
+    response = test_client.get("/health")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "healthy"
+
+
+def test_health_endpoint_no_engine():
+    """Test health check endpoint when no engine is initialized"""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+    # Don't set any engine
+
+    client = TestClient(app)
+    response = client.get("/health")
+    assert response.status_code == 503
+    data = response.json()
+    assert data["status"] == "unhealthy"
+
+
+def test_models_endpoint(test_client):
+    """Test /v1/models endpoint for diffusion mode"""
+    response = test_client.get("/v1/models")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["object"] == "list"
+    assert len(data["data"]) == 1
+    assert data["data"][0]["id"] == "Qwen/Qwen-Image"
+    assert data["data"][0]["object"] == "model"
+
+
+def test_models_endpoint_no_engine():
+    """Test /v1/models endpoint when no engine is initialized"""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+    # Don't set any engine
+
+    client = TestClient(app)
+    response = client.get("/v1/models")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["object"] == "list"
+    assert len(data["data"]) == 0
 
 
 def test_generate_single_image(test_client):
diff --git a/tests/test_outputs.py b/tests/test_outputs.py
new file mode 100644
index 000000000..c0c6201f6
--- /dev/null
+++ b/tests/test_outputs.py
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for OmniRequestOutput class."""
+
+from unittest.mock import Mock
+
+from PIL import Image
+
+from vllm_omni.outputs import OmniRequestOutput
+
+
+class TestOmniRequestOutput:
+    """Tests for OmniRequestOutput class."""
+
+    def test_from_diffusion(self):
+        """Test creating output from diffusion model."""
+        images = [Image.new("RGB", (64, 64), color="red")]
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=images,
+            prompt="a cat",
+            metrics={"steps": 50},
+        )
+        assert output.request_id == "test-123"
+        assert output.images == images
+        assert output.prompt == "a cat"
+        assert output.metrics == {"steps": 50}
+        assert output.is_diffusion_output
+        assert output.num_images == 1
+
+    def test_from_pipeline(self):
+        """Test creating output from pipeline stage."""
+        mock_request_output = Mock()
+        mock_request_output.request_id = "pipeline-123"
+        mock_request_output.prompt_token_ids = [1, 2, 3]
+        mock_request_output.outputs = [Mock()]
+        mock_request_output.encoder_prompt_token_ids = None
+        mock_request_output.prompt_logprobs = None
+        mock_request_output.num_cached_tokens = 10
+        mock_request_output.kv_transfer_params = None
+        mock_request_output.multimodal_output = {"image": Mock()}
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.request_id == "pipeline-123"
+        assert output.stage_id == 0
+        assert output.final_output_type == "text"
+        assert output.is_pipeline_output
+
+    def test_prompt_token_ids_property(self):
+        """Test prompt_token_ids property for streaming compatibility."""
+        mock_request_output = Mock()
+        mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.prompt_token_ids == [1, 2, 3, 4, 5]
+
+    def test_prompt_token_ids_none_when_no_request_output(self):
+        """Test prompt_token_ids returns None when no request_output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[],
+            prompt="a cat",
+        )
+        assert output.prompt_token_ids is None
+
+    def test_outputs_property(self):
+        """Test outputs property for chat completion compatibility."""
+        mock_output = Mock()
+        mock_request_output = Mock()
+        mock_request_output.outputs = [mock_output]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.outputs == [mock_output]
+
+    def test_outputs_empty_when_no_request_output(self):
+        """Test outputs returns empty list when no request_output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[],
+            prompt="a cat",
+        )
+        assert output.outputs == []
+
+    def test_encoder_prompt_token_ids_property(self):
+        """Test encoder_prompt_token_ids property."""
+        mock_request_output = Mock()
+        mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.encoder_prompt_token_ids == [10, 20, 30]
+
+    def test_num_cached_tokens_property(self):
+        """Test num_cached_tokens property."""
+        mock_request_output = Mock()
+        mock_request_output.num_cached_tokens = 42
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.num_cached_tokens == 42
+
+    def test_multimodal_output_property(self):
+        """Test multimodal_output property."""
+        mock_request_output = Mock()
+        mock_audio = Mock()
+        expected_output = {"audio": mock_audio}
+        mock_request_output.multimodal_output = expected_output
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="audio",
+            request_output=mock_request_output,
+        )
+
+        assert output.multimodal_output is expected_output
+
+    def test_to_dict_diffusion(self):
+        """Test to_dict for diffusion output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[Image.new("RGB", (64, 64), color="red")],
+            prompt="a cat",
+            metrics={"steps": 50},
+        )
+        result = output.to_dict()
+
+        assert result["request_id"] == "test-123"
+        assert result["finished"] is True
+        assert result["final_output_type"] == "image"
+        assert result["num_images"] == 1
+        assert result["prompt"] == "a cat"
+
+    def test_to_dict_pipeline(self):
+        """Test to_dict for pipeline output."""
+        mock_request_output = Mock()
+        mock_request_output.request_id = "pipeline-123"
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+        result = output.to_dict()
+
+        assert result["request_id"] == "pipeline-123"
+        assert result["finished"] is True
+        assert result["final_output_type"] == "text"
+        assert result["stage_id"] == 0
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 3c44da8b3..a47a864cd 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -666,14 +666,6 @@ def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
     return request.app.state.openai_serving_speech
 
 
-# Remove the original /v1/chat/completions route before registering our own
-# This prevents duplicate route registration warnings in FastAPI logs.
-for route in router.routes[:]:
-    if hasattr(route, "path") and route.path == "/v1/chat/completions":
-        router.routes.remove(route)
-        break
-
-
 @router.post(
     "/v1/chat/completions",
     dependencies=[Depends(validate_json_request)],
@@ -763,6 +755,86 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
 
 
+# Health and Model endpoints for diffusion mode
+
+
+# Remove existing health endpoint if present (from vllm imports)
+# to ensure our handler takes precedence
+_remove_route_from_router(router, "/health")
+
+
+@router.get("/health")
+async def health(raw_request: Request) -> JSONResponse:
+    """Health check endpoint that works for both LLM and diffusion modes.
+
+    Returns 200 OK if the server is healthy.
+    For LLM mode: delegates to engine_client health check
+    For diffusion mode: checks if diffusion_engine is running
+    """
+    # Check if we're in diffusion mode
+    diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None)
+    if diffusion_engine is not None:
+        # Diffusion mode health check
+        if hasattr(diffusion_engine, "is_running") and diffusion_engine.is_running:
+            return JSONResponse(content={"status": "healthy"})
+        return JSONResponse(
+            content={"status": "unhealthy", "reason": "Diffusion engine is not running"},
+            status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+        )
+
+    # LLM mode - delegate to engine_client
+    engine_client = getattr(raw_request.app.state, "engine_client", None)
+    if engine_client is not None:
+        await engine_client.check_health()
+        return JSONResponse(content={"status": "healthy"})
+
+    return JSONResponse(
+        content={"status": "unhealthy", "reason": "No engine initialized"},
+        status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+    )
+
+
+# Remove existing models endpoint if present (from vllm imports)
+# to ensure our handler takes precedence
+_remove_route_from_router(router, "/v1/models")
+
+
+@router.get("/v1/models")
+async def show_available_models(raw_request: Request) -> JSONResponse:
+    """Show available models endpoint that works for both LLM and diffusion modes.
+
+    Returns model information in OpenAI-compatible format.
+    """
+    # Check if we're in diffusion mode
+    diffusion_model_name = getattr(raw_request.app.state, "diffusion_model_name", None)
+    if diffusion_model_name is not None:
+        # Diffusion mode - return the loaded model
+        return JSONResponse(
+            content={
+                "object": "list",
+                "data": [
+                    {
+                        "id": diffusion_model_name,
+                        "object": "model",
+                        "created": 0,
+                        "owned_by": "vllm-omni",
+                        "permission": [],
+                    }
+                ],
+            }
+        )
+
+    # LLM mode - delegate to openai_serving_models
+    openai_serving_models = getattr(raw_request.app.state, "openai_serving_models", None)
+    if openai_serving_models is not None:
+        models = await openai_serving_models.show_available_models()
+        return JSONResponse(content=models.model_dump())
+
+    return JSONResponse(
+        content={"object": "list", "data": []},
+    )
+
+
 # Image generation API endpoints
 
 
diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py
index 24b168287..958cebd3e 100644
--- a/vllm_omni/outputs.py
+++ b/vllm_omni/outputs.py
@@ -123,6 +123,58 @@ class OmniRequestOutput:
         """Return the number of generated images."""
         return len(self.images)
 
+    # Pass-through properties keep vLLM serving codepaths compatible with
+    # OmniRequestOutput for pipeline outputs (Issue #345).
+    @property
+    def prompt_token_ids(self) -> list[int] | None:
+        """Return prompt token IDs from the underlying request output.
+
+        This property is required for compatibility with vLLM's streaming
+        chat completion generator which checks res.prompt_token_ids.
+        """
+        if self.request_output is not None:
+            return getattr(self.request_output, "prompt_token_ids", None)
+        return None
+
+    @property
+    def outputs(self) -> list[Any]:
+        """Return outputs from the underlying request output.
+
+        This property is required for compatibility with vLLM's streaming
+        and non-streaming chat completion generators.
+        """
+        if self.request_output is not None:
+            return getattr(self.request_output, "outputs", [])
+        return []
+
+    @property
+    def encoder_prompt_token_ids(self) -> list[int] | None:
+        """Return encoder prompt token IDs from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "encoder_prompt_token_ids", None)
+        return None
+
+    @property
+    def prompt_logprobs(self) -> Any:
+        """Return prompt logprobs from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "prompt_logprobs", None)
+        return None
+
+    @property
+    def num_cached_tokens(self) -> int | None:
+        """Return number of cached tokens from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "num_cached_tokens", None)
+        return None
+
+    @property
+    def kv_transfer_params(self) -> Any:
+        """Return KV transfer params from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "kv_transfer_params", None)
+        return None
+
     @property
     def is_diffusion_output(self) -> bool:
         """Check if this is a diffusion model output."""

From 9ece725fc52d86ee48c0e4f4afa1bd6b0cdd8201 Mon Sep 17 00:00:00 2001
From: zhou zhuoxin <zhouzhuoxin1508@outlook.com>
Date: Wed, 28 Jan 2026 15:20:07 +0800
Subject: [PATCH 02/32] [Bugfix] Add missing dependencies (onnxruntime, sox)
 for Qwen3-TTS support (#981)

Signed-off-by: zhou zhuoxin <zhouzhuoxin1508@outlook.com>
---
 pyproject.toml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 306733063..809d1f53c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,6 +42,8 @@ dependencies = [
     "fa3-fwd", # flash attention 3, maintained by @ZJY0516
     "openai-whisper>=20250625",
     "imageio[ffmpeg]>=2.37.2",
+    "onnxruntime>=1.19.0",
+    "sox>=1.5.0",
     # "vllm==0.14.0",  # TODO: fix the entrypoints overwrite problem
 ]
 

From d64bbde5616495006310a34c469baeab16d8a925 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?=
 <11222265+fhfuih@users.noreply.github.com>
Date: Wed, 28 Jan 2026 16:07:42 +0800
Subject: [PATCH 03/32] =?UTF-8?q?[Frontend][Model]=20Support=20batch=20req?=
 =?UTF-8?q?uest=20with=20refined=20OmniDiffusionReq=E2=80=A6=20(#797)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
---
 docs/getting_started/quickstart.md            |   40 +
 .../diffusion/cache_dit_acceleration.md       |    3 +-
 .../diffusion/parallelism_acceleration.md     |   42 +-
 docs/user_guide/diffusion/teacache.md         |   16 +-
 docs/user_guide/diffusion_acceleration.md     |   48 +-
 .../offline_inference/text_to_image.md        |   37 +-
 examples/offline_inference/bagel/end2end.py   |   48 +-
 .../image_to_image/image_edit.py              |   32 +-
 .../image_to_video/image_to_video.py          |   25 +-
 .../lora_inference/lora_inference.py          |   18 +-
 .../text_to_audio/text_to_audio.py            |   23 +-
 .../offline_inference/text_to_image/README.md |   37 +-
 .../text_to_image/gradio_demo.py              |   15 +-
 .../text_to_image/text_to_image.py            |   23 +-
 .../text_to_video/text_to_video.py            |   23 +-
 .../omni_connectors/test_kv_flow.py           |   21 +-
 tests/e2e/offline_inference/conftest.py       |   21 +-
 tests/e2e/offline_inference/test_cache_dit.py |   16 +-
 .../test_diffusion_cpu_offload.py             |   13 +-
 .../offline_inference/test_diffusion_lora.py  |   33 +-
 .../e2e/offline_inference/test_ovis_image.py  |   39 +-
 .../test_sequence_parallel.py                 |   30 +-
 .../test_stable_audio_model.py                |   23 +-
 tests/e2e/offline_inference/test_t2i_model.py |   15 +-
 tests/e2e/offline_inference/test_t2v_model.py |   18 +-
 tests/e2e/offline_inference/test_teacache.py  |   16 +-
 .../test_zimage_tensor_parallel.py            |   16 +-
 tests/e2e/online_serving/test_async_omni.py   |   19 +-
 .../openai_api/test_image_server.py           |   46 +-
 tests/entrypoints/test_omni_diffusion.py      | 1103 +++++++++++++++++
 tests/entrypoints/test_omni_llm.py            |   22 +-
 vllm_omni/diffusion/diffusion_engine.py       |  100 +-
 vllm_omni/diffusion/executor/abstract.py      |    4 +-
 .../diffusion/executor/multiproc_executor.py  |    6 +-
 .../diffusion/models/bagel/pipeline_bagel.py  |   36 +-
 .../flux2_klein/pipeline_flux2_klein.py       |   65 +-
 .../models/glm_image/pipeline_glm_image.py    |   92 +-
 .../longcat_image/pipeline_longcat_image.py   |   58 +-
 .../pipeline_longcat_image_edit.py            |  105 +-
 .../models/ovis_image/pipeline_ovis_image.py  |   35 +-
 .../models/qwen_image/pipeline_qwen_image.py  |   37 +-
 .../qwen_image/pipeline_qwen_image_edit.py    |   96 +-
 .../pipeline_qwen_image_edit_plus.py          |  104 +-
 .../qwen_image/pipeline_qwen_image_layered.py |  112 +-
 .../diffusion/models/sd3/pipeline_sd3.py      |   33 +-
 .../stable_audio/pipeline_stable_audio.py     |   30 +-
 .../models/wan2_2/pipeline_wan2_2.py          |  132 +-
 .../models/wan2_2/pipeline_wan2_2_i2v.py      |  138 ++-
 .../models/wan2_2/pipeline_wan2_2_ti2v.py     |  134 +-
 .../models/z_image/pipeline_z_image.py        |   38 +-
 vllm_omni/diffusion/request.py                |  193 +--
 vllm_omni/diffusion/scheduler.py              |    4 +-
 .../worker/gpu_diffusion_model_runner.py      |   40 +-
 .../diffusion/worker/gpu_diffusion_worker.py  |   36 +-
 vllm_omni/entrypoints/async_omni.py           |   78 +-
 vllm_omni/entrypoints/async_omni_diffusion.py |  120 +-
 vllm_omni/entrypoints/omni.py                 |  132 +-
 vllm_omni/entrypoints/omni_diffusion.py       |   63 +-
 vllm_omni/entrypoints/omni_stage.py           |  201 ++-
 vllm_omni/entrypoints/openai/api_server.py    |   58 +-
 vllm_omni/entrypoints/openai/serving_chat.py  |   59 +-
 vllm_omni/inputs/data.py                      |  177 ++-
 .../model_executor/stage_configs/bagel.yaml   |    6 -
 vllm_omni/outputs.py                          |    6 +-
 64 files changed, 2995 insertions(+), 1414 deletions(-)
 create mode 100644 tests/entrypoints/test_omni_diffusion.py

diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md
index b64bf1423..d4087621a 100644
--- a/docs/getting_started/quickstart.md
+++ b/docs/getting_started/quickstart.md
@@ -46,6 +46,46 @@ if __name__ == "__main__":
     images[0].save("coffee.png")
 ```
 
+You can pass a list of prompts and wait for them to process altogether, shown below.
+
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        # stage_configs_path="./stage-config.yaml",  # See below
+    )
+    prompts = [
+        "a cup of coffee on a table",
+        "a toy dinosaur on a sandy beach",
+        "a fox waking up in bed and yawning",
+    ]
+    omni_outputs = omni.generate(prompts)
+    for i_prompt, prompt_output in enumerate(omni_outputs):
+        this_request_output = prompt_output.request_output[0]
+        this_images = this_request_output.images
+        for i_image, image in enumerate(this_images):
+            image.save(f"p{i_prompt}-img{i_image}.jpg")
+            print("saved to", f"p{i_prompt}-img{i_image}.jpg")
+            # saved to p0-img0.jpg
+            # saved to p1-img0.jpg
+            # saved to p2-img0.jpg
+```
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
 For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
 
 ## Online Serving with OpenAI-Completions API
diff --git a/docs/user_guide/diffusion/cache_dit_acceleration.md b/docs/user_guide/diffusion/cache_dit_acceleration.md
index 57f6a94e6..fd1bd522a 100644
--- a/docs/user_guide/diffusion/cache_dit_acceleration.md
+++ b/docs/user_guide/diffusion/cache_dit_acceleration.md
@@ -18,6 +18,7 @@ Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cac
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 # Simplest way: just enable cache-dit with default parameters
 omni = Omni(
@@ -27,7 +28,7 @@ omni = Omni(
 
 images = omni.generate(
     "a beautiful landscape",
-    num_inference_steps=50,
+    OmniDiffusionSamplingParams(num_inference_steps=50),
 )
 ```
 
diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md
index 09e4651ae..45e53294e 100644
--- a/docs/user_guide/diffusion/parallelism_acceleration.md
+++ b/docs/user_guide/diffusion/parallelism_acceleration.md
@@ -67,10 +67,12 @@ omni = Omni(
 )
 
 outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=512,
-    height=512,
+    "a cat reading a book",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=9,
+        width=512,
+        height=512,
+    ),
 )
 ```
 
@@ -83,6 +85,7 @@ outputs = omni.generate(
 An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2
 
@@ -91,7 +94,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -133,6 +139,7 @@ Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the inpu
 An example of offline inference script using Ring-Attention is shown below:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ring_degree = 2
 
@@ -141,7 +148,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ring_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -183,6 +193,7 @@ You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism.
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 
 # Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
@@ -191,7 +202,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 ##### Online Serving
@@ -374,11 +388,15 @@ omni = Omni(
 )
 
 outputs = omni.generate(
-    prompt="turn this cat to a dog",
-    negative_prompt="low quality, blurry",
-    true_cfg_scale=4.0,
-    pil_image=input_image,
-    num_inference_steps=50,
+    {
+        "prompt": "turn this cat to a dog",
+        "negative_prompt": "low quality, blurry",
+    },
+    OmniDiffusionSamplingParams(
+        true_cfg_scale=4.0,
+        pil_image=input_image,
+        num_inference_steps=50,
+    ),
 )
 ```
 
diff --git a/docs/user_guide/diffusion/teacache.md b/docs/user_guide/diffusion/teacache.md
index 3fe961405..40dafeb88 100644
--- a/docs/user_guide/diffusion/teacache.md
+++ b/docs/user_guide/diffusion/teacache.md
@@ -8,6 +8,7 @@ Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 # Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
 omni = Omni(
@@ -17,7 +18,12 @@ omni = Omni(
         "rel_l1_thresh": 0.2  # Optional, defaults to 0.2
     }
 )
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ### Using Environment Variable
@@ -68,13 +74,19 @@ Controls the balance between speed and quality. Lower values prioritize quality,
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 omni = Omni(
     model="Qwen/Qwen-Image",
     cache_backend="tea_cache",
     cache_config={"rel_l1_thresh": 0.2}
 )
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ## Performance Tuning
diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md
index 42202d8d7..8d31747d2 100644
--- a/docs/user_guide/diffusion_acceleration.md
+++ b/docs/user_guide/diffusion_acceleration.md
@@ -96,6 +96,7 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 omni = Omni(
     model="Qwen/Qwen-Image",
@@ -103,13 +104,19 @@ omni = Omni(
     cache_config={"rel_l1_thresh": 0.2}  # Optional, defaults to 0.2
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ### Using Cache-DiT
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 omni = Omni(
     model="Qwen/Qwen-Image",
@@ -123,7 +130,12 @@ omni = Omni(
     }
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ### Using Ulysses-SP
@@ -131,6 +143,7 @@ outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_st
 Run text-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2
 
@@ -139,13 +152,17 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 
 Run image-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2
 
@@ -154,8 +171,13 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )
 
-outputs = omni.generate(prompt="turn this cat to a dog",
-        pil_image=input_image, num_inference_steps=50)
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "multi_modal_data": {"image": input_image}
+    },
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
 ```
 
 ### Using Ring-Attention
@@ -163,6 +185,7 @@ outputs = omni.generate(prompt="turn this cat to a dog",
 Run text-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ring_degree = 2
 
@@ -171,7 +194,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ring_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 ### Using CFG-Parallel
@@ -182,6 +208,7 @@ CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 cfg_parallel_size = 2
 
@@ -190,8 +217,13 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
 )
 
-outputs = omni.generate(prompt="turn this cat to a dog",
-        pil_image=input_image, num_inference_steps=50, true_cfg_scale=4.0)
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "multi_modal_data": {"image": input_image}
+    },
+    OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
+)
 ```
 
 ## Documentation
diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md
index 5c6f167e1..2dd0f7b7c 100644
--- a/docs/user_guide/examples/offline_inference/text_to_image.md
+++ b/docs/user_guide/examples/offline_inference/text_to_image.md
@@ -23,7 +23,7 @@ if __name__ == "__main__":
     images[0].save("coffee.png")
 ```
 
-Or put more than one prompt in a request, processing them sequentially.
+Or put more than one prompt in a request.
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
@@ -40,6 +40,41 @@ if __name__ == "__main__":
       image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
+Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
+This is useful for models that support negative prompts.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    outputs = omni.generate([
+      {
+        "prompt": "a cup of coffee on a table"，
+        "negative_prompt": "low resolution"
+      },
+      {
+        "prompt": "a toy dinosaur on a sandy beach"，
+        "negative_prompt": "cinematic, realistic"
+      }
+    ])
+    for i, output in enumerate(outputs):
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
+```
+
 ## Local CLI Usage
 
 ```bash
diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py
index 93000662d..397fd333e 100644
--- a/examples/offline_inference/bagel/end2end.py
+++ b/examples/offline_inference/bagel/end2end.py
@@ -1,5 +1,8 @@
 import argparse
 import os
+from typing import cast
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 
 def parse_args():
@@ -50,21 +53,24 @@ def parse_args():
 def main():
     args = parse_args()
     model_name = args.model
+    prompts: list[OmniPromptType] = []
     try:
         # Preferred: load from txt file (one prompt per line)
         if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
             with open(args.txt_prompts, encoding="utf-8") as f:
                 lines = [ln.strip() for ln in f.readlines()]
-            args.prompts = [ln for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+            prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+        else:
+            prompts = args.prompts
     except Exception as e:
         print(f"[Error] Failed to load prompts: {e}")
         raise
 
-    if args.prompts is None:
+    if not prompts:
         # Default prompt for text2img test if none provided
-        args.prompts = ["<|im_start|>A cute cat<|im_end|>"]
-        print(f"[Info] No prompts provided, using default: {args.prompts}")
+        prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        print(f"[Info] No prompts provided, using default: {prompts}")
     omni_outputs = []
 
     from PIL import Image
@@ -77,21 +83,27 @@ def main():
         print("[Info] Running in img2img mode (Stage 1 only)")
         client = OmniDiffusion(model=model_name)
 
-        generate_kwargs = {
-            "prompt": args.prompts,
-            "seed": 52,
-            "need_kv_receive": False,
-            "num_inference_steps": args.steps,
-        }
-
         if args.image_path:
             if os.path.exists(args.image_path):
                 loaded_image = Image.open(args.image_path).convert("RGB")
-                generate_kwargs["pil_image"] = loaded_image
+                prompts = [
+                    {
+                        "prompt": cast(str, p),
+                        "multi_modal_data": {"image": loaded_image},
+                    }
+                    for p in prompts
+                ]
             else:
                 print(f"[Warning] Image path {args.image_path} does not exist.")
 
-        result = client.generate(**generate_kwargs)
+        result = client.generate(
+            prompts,
+            OmniDiffusionSamplingParams(
+                seed=52,
+                need_kv_receive=False,
+                num_inference_steps=args.steps,
+            ),
+        )
 
         # Ensure result is a list for iteration
         if not isinstance(result, list):
@@ -100,8 +112,6 @@ def main():
             omni_outputs = result
 
     else:
-        import copy
-
         from vllm_omni.entrypoints.omni import Omni
 
         omni_kwargs = {}
@@ -144,11 +154,11 @@ def main():
                 prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
                 formatted_prompts.append(prompt_dict)
 
-        params_list = copy.deepcopy(omni.default_sampling_params_list)
+        params_list = omni.default_sampling_params_list
         if args.modality == "text2img":
-            params_list[0]["max_tokens"] = 1
+            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
             if len(params_list) > 1:
-                params_list[1]["num_inference_steps"] = args.steps
+                params_list[1].num_inference_steps = args.steps  # type: ignore # The second stage is an OmniDiffusionSamplingParam
 
         omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py
index 405a0cafa..f2065e6b2 100644
--- a/examples/offline_inference/image_to_image/image_edit.py
+++ b/examples/offline_inference/image_to_image/image_edit.py
@@ -79,6 +79,7 @@ from PIL import Image
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -374,25 +375,28 @@ def main():
     print(f"{'=' * 60}\n")
 
     generation_start = time.perf_counter()
-    # Generate edited image
-    generate_kwargs = {
-        "prompt": args.prompt,
-        "pil_image": input_image,
-        "negative_prompt": args.negative_prompt,
-        "generator": generator,
-        "true_cfg_scale": args.cfg_scale,
-        "guidance_scale": args.guidance_scale,
-        "num_inference_steps": args.num_inference_steps,
-        "num_outputs_per_prompt": args.num_outputs_per_prompt,
-        "layers": args.layers,
-        "resolution": args.resolution,
-    }
 
     if profiler_enabled:
         print("[Profiler] Starting profiling...")
         omni.start_profile()
 
-    outputs = omni.generate(**generate_kwargs)
+    # Generate edited image
+    outputs = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": input_image},
+        },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_outputs_per_prompt,
+            layers=args.layers,
+            resolution=args.resolution,
+        ),
+    )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
 
diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py
index 2104ec0fd..06c13fbee 100644
--- a/examples/offline_inference/image_to_video/image_to_video.py
+++ b/examples/offline_inference/image_to_video/image_to_video.py
@@ -27,6 +27,7 @@ import PIL.Image
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -119,16 +120,20 @@ def main():
 
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        pil_image=image,
-        height=height,
-        width=width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        guidance_scale_2=args.guidance_scale_high,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": image},
+        },
+        OmniDiffusionSamplingParams(
+            height=height,
+            width=width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
     )
 
     # Extract video frames from OmniRequestOutput
diff --git a/examples/offline_inference/lora_inference/lora_inference.py b/examples/offline_inference/lora_inference/lora_inference.py
index 17e9d6196..5e4299edb 100644
--- a/examples/offline_inference/lora_inference/lora_inference.py
+++ b/examples/offline_inference/lora_inference/lora_inference.py
@@ -5,6 +5,7 @@ import argparse
 from pathlib import Path
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 
@@ -95,18 +96,17 @@ def main():
         )
         print(f"Activating pre-loaded LoRA: id={lora_request_id}, scale={args.lora_scale}")
 
-    gen_kwargs = {
-        "prompt": args.prompt,
-        "height": args.height,
-        "width": args.width,
-        "num_inference_steps": args.num_inference_steps,
-    }
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+    )
 
     if lora_request:
-        gen_kwargs["lora_request"] = lora_request
-        gen_kwargs["lora_scale"] = args.lora_scale
+        sampling_params.lora_request = lora_request
+        sampling_params.lora_scale = args.lora_scale
 
-    outputs = omni.generate(**gen_kwargs)
+    outputs = omni.generate(args.prompt, sampling_params)
 
     if not outputs or len(outputs) == 0:
         raise ValueError("No output generated from omni.generate()")
diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py
index ca1a455d1..1124dbdb4 100644
--- a/examples/offline_inference/text_to_audio/text_to_audio.py
+++ b/examples/offline_inference/text_to_audio/text_to_audio.py
@@ -21,6 +21,7 @@ import numpy as np
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.utils.platform_utils import detect_device_type
 
 
@@ -143,16 +144,20 @@ def main():
 
     # Generate audio
     outputs = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_waveforms,
-        extra={
-            "audio_start_in_s": args.audio_start,
-            "audio_end_in_s": audio_end_in_s,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
         },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_waveforms,
+            extra_args={
+                "audio_start_in_s": args.audio_start,
+                "audio_end_in_s": audio_end_in_s,
+            },
+        ),
     )
 
     generation_end = time.perf_counter()
diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md
index 0f4351ac2..f4de19db8 100644
--- a/examples/offline_inference/text_to_image/README.md
+++ b/examples/offline_inference/text_to_image/README.md
@@ -20,7 +20,7 @@ if __name__ == "__main__":
     images[0].save("coffee.png")
 ```
 
-Or put more than one prompt in a request, processing them sequentially.
+Or put more than one prompt in a request.
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
@@ -37,6 +37,41 @@ if __name__ == "__main__":
       image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
+Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
+This is useful for models that support negative prompts.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    outputs = omni.generate([
+      {
+        "prompt": "a cup of coffee on a table"，
+        "negative_prompt": "low resolution"
+      },
+      {
+        "prompt": "a toy dinosaur on a sandy beach"，
+        "negative_prompt": "cinematic, realistic"
+      }
+    ])
+    for i, output in enumerate(outputs):
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
+```
+
 ## Local CLI Usage
 
 ```bash
diff --git a/examples/offline_inference/text_to_image/gradio_demo.py b/examples/offline_inference/text_to_image/gradio_demo.py
index 8d4ff2b20..30446f50e 100644
--- a/examples/offline_inference/text_to_image/gradio_demo.py
+++ b/examples/offline_inference/text_to_image/gradio_demo.py
@@ -5,6 +5,7 @@ import gradio as gr
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -101,12 +102,14 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
         generator = torch.Generator(device=device).manual_seed(seed)
         outputs = omni.generate(
             prompt.strip(),
-            height=height,
-            width=width,
-            generator=generator,
-            true_cfg_scale=float(cfg_scale_value),
-            num_inference_steps=num_steps,
-            num_outputs_per_prompt=num_images,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                generator=generator,
+                true_cfg_scale=float(cfg_scale_value),
+                num_inference_steps=num_steps,
+                num_outputs_per_prompt=num_images,
+            ),
         )
         images_outputs = []
         for output in outputs:
diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py
index 428c2437e..7fc6ec832 100644
--- a/examples/offline_inference/text_to_image/text_to_image.py
+++ b/examples/offline_inference/text_to_image/text_to_image.py
@@ -10,6 +10,7 @@ import torch
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -189,15 +190,19 @@ def main():
 
     generation_start = time.perf_counter()
     outputs = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        true_cfg_scale=args.cfg_scale,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_images_per_prompt,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_images_per_prompt,
+        ),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index c655fff14..054db95b3 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -11,6 +11,7 @@ import torch
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -108,15 +109,19 @@ def main():
 
     generation_start = time.perf_counter()
     frames = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        guidance_scale_2=args.guidance_scale_high,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py
index a8d2f03b2..2bb06d4e0 100644
--- a/tests/distributed/omni_connectors/test_kv_flow.py
+++ b/tests/distributed/omni_connectors/test_kv_flow.py
@@ -1,10 +1,12 @@
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import torch
 
 from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 
 
@@ -109,11 +111,12 @@ class TestKVFlow(unittest.TestCase):
         transfer_data.metadata = {"kv_lens": [self.seq_len], "ropes": [0]}
 
         # 2. Setup Request with Injected Data
-        req = OmniDiffusionRequest(prompt="test")
-        from types import SimpleNamespace
+        sp = OmniDiffusionSamplingParams(
+            past_key_values=SimpleNamespace(**layer_blocks),
+            kv_metadata=transfer_data.metadata,
+        )
 
-        req.past_key_values = SimpleNamespace(**layer_blocks)
-        req.kv_metadata = transfer_data.metadata
+        req = OmniDiffusionRequest(["test"], sp)
 
         # 3. Setup Pipeline
         pipeline = MockBagelPipeline()
@@ -144,7 +147,7 @@ class TestKVFlow(unittest.TestCase):
         current_cache = RealNaiveCache(self.num_layers)
 
         # --- Logic from Source Code ---
-        injected_kv = req.past_key_values
+        injected_kv = req.sampling_params.past_key_values
         if isinstance(current_cache, RealNaiveCache) and hasattr(injected_kv, "key_cache"):
             # Assuming injected_kv is SimpleNamespace or object with list attrs
             for layer_idx in range(len(injected_kv.key_cache)):
@@ -176,9 +179,11 @@ class TestKVFlow(unittest.TestCase):
         data_dict = runner_test_result.to_dict()
 
         # 3. Receiver (Request Setup)
-        req = OmniDiffusionRequest(prompt="integration_test")
-        req.past_key_values = data_dict["layer_blocks"]
-        req.kv_metadata = data_dict["metadata"]
+        sp = OmniDiffusionSamplingParams(
+            past_key_values=data_dict["layer_blocks"],
+            kv_metadata=data_dict["metadata"],
+        )
+        req = OmniDiffusionRequest(["integration_test"], sp)  # noqa: F841
 
         # 4. Receiver (Injection Simulation)
         # Use the logic verification again
diff --git a/tests/e2e/offline_inference/conftest.py b/tests/e2e/offline_inference/conftest.py
index b3a4ade92..9ee0e4b74 100644
--- a/tests/e2e/offline_inference/conftest.py
+++ b/tests/e2e/offline_inference/conftest.py
@@ -7,11 +7,12 @@ Pytest configuration and fixtures for vllm-omni tests.
 from typing import Any
 
 import pytest
+from vllm import TextPrompt
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from vllm.sampling_params import SamplingParams
 
 from tests.conftest import _run_post_test_cleanup
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
 PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
@@ -64,7 +65,7 @@ class OmniRunner:
             **kwargs,
         )
 
-    def get_default_sampling_params_list(self) -> list[SamplingParams]:
+    def get_default_sampling_params_list(self) -> list[OmniSamplingParams]:
         """
         Get a list of default sampling parameters for all stages.
 
@@ -82,7 +83,7 @@ class OmniRunner:
         videos: PromptVideoInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
         modalities: list[str] | None = None,
-    ) -> list[dict[str, Any]]:
+    ) -> list[TextPrompt]:
         """
         Construct Omni input format from prompts and multimodal data.
 
@@ -175,7 +176,7 @@ class OmniRunner:
                 f"<|im_start|>assistant\n"
             )
 
-            input_dict: dict[str, Any] = {"prompt": full_prompt}
+            input_dict: TextPrompt = {"prompt": full_prompt}
             if multi_modal_data:
                 input_dict["multi_modal_data"] = multi_modal_data
             if modalities:
@@ -189,8 +190,8 @@ class OmniRunner:
 
     def generate(
         self,
-        prompts: list[dict[str, Any]],
-        sampling_params_list: list[SamplingParams] | None = None,
+        prompts: list[TextPrompt],
+        sampling_params_list: list[OmniSamplingParams] | None = None,
     ) -> list[OmniRequestOutput]:
         """
         Generate outputs for the given prompts.
@@ -212,7 +213,7 @@ class OmniRunner:
     def generate_multimodal(
         self,
         prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
         system_prompt: str | None = None,
         audios: PromptAudioInput = None,
         images: PromptImageInput = None,
@@ -249,7 +250,7 @@ class OmniRunner:
     def generate_audio(
         self,
         prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
         system_prompt: str | None = None,
         audios: PromptAudioInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
@@ -276,7 +277,7 @@ class OmniRunner:
     def generate_video(
         self,
         prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
         system_prompt: str | None = None,
         videos: PromptVideoInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
@@ -303,7 +304,7 @@ class OmniRunner:
     def generate_image(
         self,
         prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
         system_prompt: str | None = None,
         images: PromptImageInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
diff --git a/tests/e2e/offline_inference/test_cache_dit.py b/tests/e2e/offline_inference/test_cache_dit.py
index 9eaee00c3..18d3988e0 100644
--- a/tests/e2e/offline_inference/test_cache_dit.py
+++ b/tests/e2e/offline_inference/test_cache_dit.py
@@ -15,6 +15,8 @@ from pathlib import Path
 import pytest
 import torch
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
@@ -55,12 +57,14 @@ def test_cache_dit(model_name: str):
 
         outputs = m.generate(
             "a photo of a cat sitting on a laptop keyboard",
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator("cuda").manual_seed(42),
-            num_outputs_per_prompt=1,  # Single output for speed
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=0.0,
+                generator=torch.Generator("cuda").manual_seed(42),
+                num_outputs_per_prompt=1,  # Single output for speed
+            ),
         )
         # Extract images from request_output[0]['images']
         first_output = outputs[0]
diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py
index 5b43400e8..f6606dce6 100644
--- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py
@@ -6,6 +6,7 @@ import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
 from tests.utils import GPUMemoryMonitor
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.utils.platform_utils import is_npu, is_rocm
 
 # ruff: noqa: E402
@@ -30,11 +31,13 @@ def inference(model_name: str, offload: bool = True):
 
     m.generate(
         "a photo of a cat sitting on a laptop keyboard",
-        height=height,
-        width=width,
-        num_inference_steps=9,
-        guidance_scale=0.0,
-        generator=torch.Generator("cuda").manual_seed(42),
+        OmniDiffusionSamplingParams(
+            height=height,
+            width=width,
+            num_inference_steps=9,
+            guidance_scale=0.0,
+            generator=torch.Generator("cuda").manual_seed(42),
+        ),
     )
     peak = monitor.peak_used_mb
     monitor.stop()
diff --git a/tests/e2e/offline_inference/test_diffusion_lora.py b/tests/e2e/offline_inference/test_diffusion_lora.py
index 1761f253c..ef734da61 100644
--- a/tests/e2e/offline_inference/test_diffusion_lora.py
+++ b/tests/e2e/offline_inference/test_diffusion_lora.py
@@ -7,6 +7,7 @@ import pytest
 import torch
 from safetensors.torch import save_file
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import is_npu
 
@@ -88,12 +89,14 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
 
         outputs = m.generate(
             prompt,
-            height=height,
-            width=width,
-            num_inference_steps=2,
-            guidance_scale=0.0,
-            generator=torch.Generator("cuda").manual_seed(42),
-            num_outputs_per_prompt=1,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=2,
+                guidance_scale=0.0,
+                generator=torch.Generator("cuda").manual_seed(42),
+                num_outputs_per_prompt=1,
+            ),
         )
         images = _extract_images(outputs)
 
@@ -116,14 +119,16 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
             )
             outputs_lora = m.generate(
                 prompt,
-                height=height,
-                width=width,
-                num_inference_steps=2,
-                guidance_scale=0.0,
-                generator=torch.Generator("cuda").manual_seed(42),
-                num_outputs_per_prompt=1,
-                lora_request=lora_request,
-                lora_scale=2.0,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=2,
+                    guidance_scale=0.0,
+                    generator=torch.Generator("cuda").manual_seed(42),
+                    num_outputs_per_prompt=1,
+                    lora_request=lora_request,
+                    lora_scale=2.0,
+                ),
             )
             images_lora = _extract_images(outputs_lora)
             assert len(images_lora) == 1
diff --git a/tests/e2e/offline_inference/test_ovis_image.py b/tests/e2e/offline_inference/test_ovis_image.py
index 33e1cb880..f1bc73817 100644
--- a/tests/e2e/offline_inference/test_ovis_image.py
+++ b/tests/e2e/offline_inference/test_ovis_image.py
@@ -25,6 +25,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image import OvisImagePipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 
 @pytest.fixture
@@ -162,11 +163,13 @@ def test_basic_generation(ovis_pipeline):
     """Test the forward pass logic."""
     # Setup request
     req = OmniDiffusionRequest(
-        prompt="A photo of a cat",
-        height=256,
-        width=256,
-        num_inference_steps=2,
-        guidance_scale=1.0,
+        prompts=["A photo of a cat"],
+        sampling_params=OmniDiffusionSamplingParams(
+            height=256,
+            width=256,
+            num_inference_steps=2,
+            guidance_scale=1.0,
+        ),
     )
 
     output = ovis_pipeline(req)
@@ -184,12 +187,18 @@ def test_basic_generation(ovis_pipeline):
 def test_guidance_scale(ovis_pipeline):
     """Test that classifier-free guidance path is taken when scale > 1.0."""
     req = OmniDiffusionRequest(
-        prompt="A photo of a cat",
-        negative_prompt="bad quality",
-        height=256,
-        width=256,
-        num_inference_steps=1,
-        guidance_scale=2.0,  # Trigger CFG
+        prompts=[
+            {
+                "prompt": "A photo of a cat",
+                "negative_prompt": "bad quality",
+            }
+        ],
+        sampling_params=OmniDiffusionSamplingParams(
+            height=256,
+            width=256,
+            num_inference_steps=1,
+            guidance_scale=2.0,  # Trigger CFG
+        ),
     )
 
     ovis_pipeline(req)
@@ -200,9 +209,11 @@ def test_resolution_check(ovis_pipeline):
     """Test resolution divisible validation logic if present."""
     # Pass odd resolution
     req = OmniDiffusionRequest(
-        prompt="test",
-        height=250,  # Not divisible by 16 (8*2)
-        width=250,
+        prompts=["test"],
+        sampling_params=OmniDiffusionSamplingParams(
+            height=250,  # Not divisible by 16 (8*2)
+            width=250,
+        ),
     )
 
     # Should warn but proceed (as per code I read earlier) or resize?
diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py
index 9de9fd981..367e2e59e 100644
--- a/tests/e2e/offline_inference/test_sequence_parallel.py
+++ b/tests/e2e/offline_inference/test_sequence_parallel.py
@@ -21,6 +21,8 @@ import torch
 import torch.distributed as dist
 from PIL import Image
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[3]
 if str(REPO_ROOT) not in sys.path:
@@ -81,12 +83,14 @@ def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height
     try:
         outputs = baseline.generate(
             PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=4,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=4,
+                guidance_scale=0.0,
+                generator=torch.Generator(get_device_name()).manual_seed(seed),
+                num_outputs_per_prompt=1,
+            ),
         )
         return outputs[0].request_output[0].images
     finally:
@@ -116,12 +120,14 @@ def _run_sp(
     try:
         outputs = sp.generate(
             PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=4,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=4,
+                guidance_scale=0.0,
+                generator=torch.Generator(get_device_name()).manual_seed(seed),
+                num_outputs_per_prompt=1,
+            ),
         )
         return outputs[0].request_output[0].images
     finally:
diff --git a/tests/e2e/offline_inference/test_stable_audio_model.py b/tests/e2e/offline_inference/test_stable_audio_model.py
index 729b501e4..df2ca5e42 100644
--- a/tests/e2e/offline_inference/test_stable_audio_model.py
+++ b/tests/e2e/offline_inference/test_stable_audio_model.py
@@ -5,6 +5,7 @@ import numpy as np
 import pytest
 import torch
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
 # ruff: noqa: E402
@@ -29,16 +30,20 @@ def test_stable_audio_model(model_name: str):
     sample_rate = 44100  # Stable Audio uses 44100 Hz
 
     outputs = m.generate(
-        "The sound of a dog barking",
-        negative_prompt="Low quality.",
-        num_inference_steps=4,  # Minimal steps for speed
-        guidance_scale=7.0,
-        generator=torch.Generator("cuda").manual_seed(42),
-        num_outputs_per_prompt=1,
-        extra={
-            "audio_start_in_s": audio_start_in_s,
-            "audio_end_in_s": audio_end_in_s,
+        prompts={
+            "prompt": "The sound of a dog barking",
+            "negative_prompt": "Low quality.",
         },
+        sampling_params_list=OmniDiffusionSamplingParams(
+            num_inference_steps=4,  # Minimal steps for speed
+            guidance_scale=7.0,
+            generator=torch.Generator("cuda").manual_seed(42),
+            num_outputs_per_prompt=1,
+            extra_args={
+                "audio_start_in_s": audio_start_in_s,
+                "audio_end_in_s": audio_end_in_s,
+            },
+        ),
     )
 
     # Extract audio from OmniRequestOutput
diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py
index 09fadc93a..ba6cac34c 100644
--- a/tests/e2e/offline_inference/test_t2i_model.py
+++ b/tests/e2e/offline_inference/test_t2i_model.py
@@ -5,6 +5,7 @@ from pathlib import Path
 import pytest
 import torch
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import is_npu, is_rocm
 
@@ -42,12 +43,14 @@ def test_diffusion_model(model_name: str):
         width = 256
         outputs = m.generate(
             "a photo of a cat sitting on a laptop keyboard",
-            height=height,
-            width=width,
-            num_inference_steps=2,
-            guidance_scale=0.0,
-            generator=torch.Generator("cuda").manual_seed(42),
-            num_outputs_per_prompt=2,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=2,
+                guidance_scale=0.0,
+                generator=torch.Generator("cuda").manual_seed(42),
+                num_outputs_per_prompt=2,
+            ),
         )
         # Extract images from request_output[0]['images']
         first_output = outputs[0]
diff --git a/tests/e2e/offline_inference/test_t2v_model.py b/tests/e2e/offline_inference/test_t2v_model.py
index 0263310e3..a378291ac 100644
--- a/tests/e2e/offline_inference/test_t2v_model.py
+++ b/tests/e2e/offline_inference/test_t2v_model.py
@@ -5,6 +5,8 @@ from pathlib import Path
 import pytest
 import torch
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
@@ -33,13 +35,15 @@ def test_video_diffusion_model(model_name: str):
     width = 640
     num_frames = 5
     outputs = m.generate(
-        "A cat sitting on a table",
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        num_inference_steps=2,
-        guidance_scale=1.0,
-        generator=torch.Generator("cuda").manual_seed(42),
+        prompts="A cat sitting on a table",
+        sampling_params_list=OmniDiffusionSamplingParams(
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=2,
+            guidance_scale=1.0,
+            generator=torch.Generator("cuda").manual_seed(42),
+        ),
     )
     first_output = outputs[0]
     assert first_output.final_output_type == "image"
diff --git a/tests/e2e/offline_inference/test_teacache.py b/tests/e2e/offline_inference/test_teacache.py
index 7ce915b96..7d6261388 100644
--- a/tests/e2e/offline_inference/test_teacache.py
+++ b/tests/e2e/offline_inference/test_teacache.py
@@ -15,6 +15,8 @@ from pathlib import Path
 import pytest
 import torch
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
@@ -51,12 +53,14 @@ def test_teacache(model_name: str):
 
         outputs = m.generate(
             "a photo of a cat sitting on a laptop keyboard",
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator("cuda").manual_seed(42),
-            num_outputs_per_prompt=1,  # Single output for speed
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=0.0,
+                generator=torch.Generator("cuda").manual_seed(42),
+                num_outputs_per_prompt=1,  # Single output for speed
+            ),
         )
         # Extract images from request_output[0]['images']
         first_output = outputs[0]
diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py
index a9fdec4dc..ea8f0acfe 100644
--- a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py
+++ b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py
@@ -12,6 +12,8 @@ import torch
 from PIL import Image
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
@@ -90,12 +92,14 @@ def _run_zimage_generate(
         num_requests = 4  # 1 warmup + 3 timed
         gen = m.generate(
             [PROMPT] * num_requests,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            seed=seed,
-            num_outputs_per_prompt=1,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=0.0,
+                seed=seed,
+                num_outputs_per_prompt=1,
+            ),
             py_generator=True,
         )
 
diff --git a/tests/e2e/online_serving/test_async_omni.py b/tests/e2e/online_serving/test_async_omni.py
index 112b68ff4..d90727e1b 100644
--- a/tests/e2e/online_serving/test_async_omni.py
+++ b/tests/e2e/online_serving/test_async_omni.py
@@ -5,6 +5,7 @@ from contextlib import ExitStack
 from pathlib import Path
 
 import pytest
+from vllm import SamplingParams
 from vllm.inputs import PromptType
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni, ClientRequestState
@@ -25,15 +26,15 @@ async def generate(
 ) -> tuple[int, str]:
     # Ensure generate doesn't complete too fast for cancellation test.
     await asyncio.sleep(0.2)
-    thinker_sampling_params = {
-        "temperature": 0.4,  # Deterministic
-        "top_p": 0.9,
-        "top_k": 1,
-        "max_tokens": max_tokens,
-        "repetition_penalty": 1.05,
-        "stop_token_ids": [151645],  # Qwen EOS token <|im_end|>
-        "seed": SEED,
-    }
+    thinker_sampling_params = SamplingParams(
+        temperature=0.4,  # Deterministic
+        top_p=0.9,
+        top_k=1,
+        max_tokens=max_tokens,
+        repetition_penalty=1.05,
+        stop_token_ids=[151645],  # Qwen EOS token <|im_end|>
+        seed=SEED,
+    )
 
     sampling_params_list = [
         thinker_sampling_params,
diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 86331e085..575dfaa33 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -14,11 +14,13 @@ from unittest.mock import AsyncMock, Mock
 import pytest
 from fastapi.testclient import TestClient
 from PIL import Image
+from vllm import SamplingParams
 
 from vllm_omni.entrypoints.openai.image_api_utils import (
     encode_image_base64,
     parse_size,
 )
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 # Unit Tests
 
@@ -111,7 +113,7 @@ class FakeAsyncOmni:
 
     def __init__(self):
         self.stage_list = ["llm", "diffusion"]
-        self.default_sampling_params_list = [{"temperature": 0.1}, {"top_p": 0.9}]
+        self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
 
     async def generate(self, prompt, request_id, sampling_params_list):
@@ -129,7 +131,8 @@ def mock_async_diffusion():
 
     async def generate(**kwargs):
         # Return n PIL images wrapped in result object
-        n = kwargs.get("num_outputs_per_prompt", 1)
+        print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
+        n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
         images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
         return MockGenerationResult(images)
 
@@ -268,11 +271,11 @@ def test_generate_images_async_omni_sampling_params(async_omni_test_client):
     captured = engine.captured_sampling_params_list
     assert captured is not None
     assert len(captured) == 2
-    assert captured[0] == {"temperature": 0.1}
-    assert captured[1]["num_outputs_per_prompt"] == 2
-    assert captured[1]["height"] == 256
-    assert captured[1]["width"] == 256
-    assert captured[1]["seed"] == 7
+    assert captured[0].temperature == 0.1
+    assert captured[1].num_outputs_per_prompt == 2
+    assert captured[1].height == 256
+    assert captured[1].width == 256
+    assert captured[1].seed == 7
 
 
 def test_generate_multiple_images(test_client):
@@ -505,30 +508,11 @@ def test_parameters_passed_through(test_client, mock_async_diffusion):
 
     # Ensure generate() was called exactly once
     mock_async_diffusion.generate.assert_awaited_once()
-    call_kwargs = mock_async_diffusion.generate.call_args[1]
-    assert call_kwargs["num_inference_steps"] == 100
-    assert call_kwargs["guidance_scale"] == 7.5
-    assert call_kwargs["true_cfg_scale"] == 3.0
-    assert call_kwargs["seed"] == 42
-
-
-def test_optional_parameters_omitted(test_client, mock_async_diffusion):
-    """Verify optional parameters not passed when omitted"""
-    response = test_client.post(
-        "/v1/images/generations",
-        json={
-            "prompt": "test",
-            "size": "512x512",
-        },
-    )
-    assert response.status_code == 200
-
-    # Ensure generate() was called exactly once
-    mock_async_diffusion.generate.assert_awaited_once()
-    call_kwargs = mock_async_diffusion.generate.call_args[1]
-    assert "num_inference_steps" not in call_kwargs
-    assert "guidance_scale" not in call_kwargs
-    assert "true_cfg_scale" not in call_kwargs
+    call_kwargs = mock_async_diffusion.generate.call_args[1]["sampling_params_list"][0]
+    assert call_kwargs.num_inference_steps == 100
+    assert call_kwargs.guidance_scale == 7.5
+    assert call_kwargs.true_cfg_scale == 3.0
+    assert call_kwargs.seed == 42
 
 
 def test_model_field_omitted_works(test_client):
diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py
new file mode 100644
index 000000000..c4884e3ab
--- /dev/null
+++ b/tests/entrypoints/test_omni_diffusion.py
@@ -0,0 +1,1103 @@
+import uuid
+import warnings
+from queue import Empty, Queue
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies.
+warnings.filterwarnings(
+    "ignore",
+    message=r"builtin type SwigPy.*has no __module__ attribute",
+    category=DeprecationWarning,
+)
+
+
+class _FakeEngineArgs(dict):
+    """Fake engine args that can be used both as object attributes and as **kwargs."""
+
+    def __init__(self, args_dict: dict[str, Any]):
+        super().__init__(args_dict)
+        # Add required attributes if not present
+        if "model_stage" not in self:
+            self["model_stage"] = None
+        if "engine_output_type" not in self:
+            self["engine_output_type"] = None
+        # Also set as attributes for object-style access
+        for key, value in self.items():
+            setattr(self, key, value)
+
+
+class _FakeStageConfig:
+    """Fake stage config object that mimics the real stage config structure."""
+
+    def __init__(self, config_dict: dict[str, Any]):
+        # engine_args needs to work both as object (for OmniStage) and as dict (for **kwargs)
+        engine_args_dict = config_dict.get("engine_args", {})
+        self.engine_args = _FakeEngineArgs(engine_args_dict)
+        self.final_output = config_dict.get("final_output", False)
+        self.final_output_type = config_dict.get("final_output_type", None)
+        self.stage_id = config_dict.get("stage_id", 0)
+        # Store original dict for reference
+        self._config_dict = config_dict
+
+
+class _FakeQueue:
+    """Fake queue using standard library Queue to replace mp.Queue."""
+
+    def __init__(self, maxsize=0):
+        self._queue = Queue(maxsize=maxsize)
+
+    def put(self, item):
+        self._queue.put(item)
+
+    def put_nowait(self, item):
+        self._queue.put_nowait(item)
+
+    def get(self):
+        return self._queue.get()
+
+    def get_nowait(self):
+        return self._queue.get_nowait()
+
+    def empty(self):
+        return self._queue.empty()
+
+
+class _FakeStage:
+    """Lightweight Stage stub for multi-process pipeline version with queue support."""
+
+    def __init__(self, config, stage_init_timeout: int = 300):
+        # Handle both dict and object configs
+        if isinstance(config, dict):
+            config = _FakeStageConfig(config)
+        self.config = config
+        self.stage_config = config
+        self.engine = None
+        self.engine_outputs = None
+        # Set attributes that OmniStage expects
+        self.stage_id = getattr(config, "stage_id", 0)
+        self.engine_args = config.engine_args
+        self.model_stage = getattr(config.engine_args, "model_stage", None)
+        self.stage_type = "diffusion"
+        # set default sampling params
+        self.default_sampling_params = OmniDiffusionSamplingParams(num_inference_steps=1)
+        # Allow configuring final_output and final_output_type
+        self.final_output = config.final_output if hasattr(config, "final_output") else False
+        self.final_output_type = getattr(config, "final_output_type", None)
+        # Configurable processing logic, default returns placeholder
+        processed_input = getattr(config, "_config_dict", {}).get("processed_input", ["processed"])
+        self._processed_input = processed_input
+        # Queue references (set by attach_queues)
+        self._in_q = None
+        self._out_q = None
+        self._proc = None  # Mock process reference
+        self._stage_init_timeout = max(0, int(stage_init_timeout))
+
+    def attach_queues(self, in_q, out_q):
+        """Attach input and output queues."""
+        self._in_q = in_q
+        self._out_q = out_q
+
+    def init_stage_worker(
+        self,
+        model: str,
+        *,
+        is_async: bool = False,
+        shm_threshold_bytes: int = 65536,
+        ctx=None,
+        batch_timeout: int = 10,
+        **kwargs,
+    ):
+        """Mock init_stage_worker: don't start real process, just send stage_ready message."""
+        # Create a mock process object
+        self._proc = MagicMock()
+        self._proc.start = MagicMock()
+        self._proc.join = MagicMock()
+        self._proc.is_alive = MagicMock(return_value=False)
+        self._proc.terminate = MagicMock()
+        # Send stage_ready message to output queue
+        if self._out_q is not None:
+            try:
+                self._out_q.put_nowait({"type": "stage_ready", "stage_id": self.stage_id})
+            except Exception:
+                pass
+
+    def stop_stage_worker(self):
+        """Mock stop_stage_worker: clean up queue references."""
+        if self._in_q is not None:
+            try:
+                self._in_q.put_nowait(SHUTDOWN_TASK)
+            except Exception:
+                pass
+
+    def submit(self, payload: dict[str, Any]):
+        """Submit task to input queue."""
+        if self._in_q is not None:
+            self._in_q.put(payload)
+
+    def try_collect(self) -> Any:
+        """Non-blocking collect from output queue."""
+        if self._out_q is None:
+            return None
+        try:
+            return self._out_q.get_nowait()
+        except Empty:
+            return None
+
+    def set_engine_outputs(self, outputs):
+        """Set engine outputs for the stage."""
+        self.engine_outputs = outputs
+
+    def process_engine_inputs(self, stage_list, prompts):
+        """Process engine inputs: return preset processed result."""
+        return self._processed_input
+
+
+class _FakeEngine:
+    """Lightweight Engine stub: provides generate iterator output."""
+
+    def __init__(self, outputs: list[Any]):
+        self._outputs = outputs
+
+    def generate(self, prompts, sampling_params):
+        # Record the most recent prompts for outer assertions
+        self._last_prompts = prompts
+        # Simplified: return preset list at once, ensuring iterability
+        yield from self._outputs
+
+
+@pytest.fixture
+def fake_stage_config():
+    return {
+        # Don't include 'model' in engine_args since it's passed separately
+        "engine_args": {},
+        "final_output": True,
+        "final_output_type": "text",
+        # Second stage will use processed_input to verify the chain
+        "processed_input": ["processed-by-stage"],
+    }
+
+
+def _setup_engine_mocks(monkeypatch):
+    """Helper function to set up common engine mocks."""
+    fake_engine = MagicMock()
+    # Add necessary attributes to fake_engine
+    fake_engine.tokenizer = MagicMock()
+    fake_engine.log_stats = False
+    fake_engine.vllm_config = MagicMock()
+    fake_engine.vllm_config.model_config = MagicMock()
+    fake_engine.vllm_config.model_config.io_processor_plugin = None
+    fake_engine.get_supported_tasks = MagicMock(return_value=[])
+    fake_engine.model_config = MagicMock()
+    fake_engine.model_config.io_processor_plugin = None
+    # Add registry with resolve_model_cls method
+    fake_registry = MagicMock()
+    fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch"))
+    fake_engine.model_config.registry = fake_registry
+    fake_engine.vllm_config.model_config.registry = fake_registry
+
+    monkeypatch.setattr(
+        "vllm.v1.engine.llm_engine.LLMEngine.from_engine_args",
+        lambda **kw: fake_engine,
+        raising=False,
+    )
+
+    # Mock model_config.registry.resolve_model_cls to return a tuple
+    # Use a real class instead of MagicMock to avoid inspect.getsource issues
+    class FakeModelClass:
+        pass
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.utils.get_model_architecture",
+        lambda model_config: (FakeModelClass, "test_arch"),
+        raising=False,
+    )
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.utils._get_model_architecture",
+        lambda model_config: (FakeModelClass, "test_arch"),
+        raising=False,
+    )
+
+    # Mock try_create_mm_pooling_model_cls to return the class as-is
+    monkeypatch.setattr(
+        "vllm.model_executor.models.adapters.try_create_mm_pooling_model_cls",
+        lambda model_cls: model_cls,
+        raising=False,
+    )
+
+    # Mock _enable_processor_cache to return False
+    monkeypatch.setattr(
+        "vllm.multimodal.cache._enable_processor_cache",
+        lambda model_config, mm_registry: False,
+        raising=False,
+    )
+
+    # Mock get_io_processor to return None
+    monkeypatch.setattr(
+        "vllm.plugins.io_processors.get_io_processor",
+        lambda vllm_config, io_processor_plugin: None,
+        raising=False,
+    )
+
+
+def _setup_multiprocessing_mocks(monkeypatch):
+    """Helper function to set up multiprocessing mocks."""
+    import multiprocessing as mp
+
+    # Mock Process
+    fake_process_class = MagicMock()
+    fake_process_instance = MagicMock()
+    fake_process_instance.start = MagicMock()
+    fake_process_instance.join = MagicMock()
+    fake_process_instance.is_alive = MagicMock(return_value=False)
+    fake_process_instance.terminate = MagicMock()
+    fake_process_class.return_value = fake_process_instance
+
+    # Mock get_context to return a context with Queue that returns _FakeQueue
+    fake_ctx = MagicMock()
+    fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize)
+    fake_ctx.Process = fake_process_class
+
+    def _mock_get_context(method):
+        return fake_ctx
+
+    monkeypatch.setattr(mp, "get_context", _mock_get_context, raising=False)
+    monkeypatch.setattr(mp, "Process", fake_process_class, raising=False)
+
+
+def _setup_ipc_mocks(monkeypatch):
+    """Helper function to set up IPC function mocks."""
+
+    # Mock _encode: simple serialization
+    def _fake_encode(obj, threshold, obj_key, shm_key):
+        return {obj_key: obj}
+
+    # Mock _load: extract object from result
+    def _fake_load(result, obj_key, shm_key):
+        return result.get(obj_key)
+
+    # Mock _set: calculate serialization size
+    def _fake_set(obj):
+        return str(obj).encode()
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni._encode", _fake_encode, raising=False)
+    monkeypatch.setattr("vllm_omni.entrypoints.omni._load", _fake_load, raising=False)
+    monkeypatch.setattr("vllm_omni.entrypoints.omni._set", _fake_set, raising=False)
+
+
+def _setup_log_mocks(monkeypatch):
+    """Helper function to set up logging and stats mocks."""
+    # Mock OrchestratorMetrics to be a simple class that doesn't require file operations
+
+    class _FakeOrchestratorMetrics:
+        def __init__(self, num_stages, enable_stats, wall_start_ts):
+            self.num_stages = num_stages
+            self.enable_stats = enable_stats
+            self.stage_first_ts = [None] * num_stages
+            self.stage_last_ts = [None] * num_stages
+            self.e2e_done = set()
+
+        def on_stage_metrics(self, stage_id, req_id, metrics):
+            pass
+
+        def on_finalize_request(self, stage_id, req_id, start_ts):
+            self.e2e_done.add(req_id)
+
+        def on_forward(self, from_stage, to_stage, req_id, size_bytes, tx_ms, use_shm):
+            pass
+
+        def build_and_log_summary(self, final_stage_id):
+            return "Fake summary"
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni.OrchestratorMetrics",
+        _FakeOrchestratorMetrics,
+        raising=False,
+    )
+
+
+@pytest.fixture(autouse=True)
+def mock_get_config(monkeypatch):
+    """Auto-mock get_config and related model loading functions to avoid model path validation."""
+    # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports
+    # This prevents ImportError when async_omni is imported (which happens via omni_stage)
+    import sys
+
+    fake_tokenizer = MagicMock()
+    fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
+    fake_tokenizer.decode = MagicMock(return_value="test")
+
+    # Mock init_tokenizer_from_configs (used in async_omni)
+    def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
+        return fake_tokenizer
+
+    # Strategy 1: Mock in the original location (vllm.transformers_utils.tokenizer)
+    # This works if the module hasn't been imported yet
+    monkeypatch.setattr(
+        "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs",
+        _mock_init_tokenizer_from_configs,
+        raising=False,
+    )
+
+    # Strategy 2: If the module is already in sys.modules, patch it directly
+    tokenizer_module_path = "vllm.transformers_utils.tokenizer"
+    if tokenizer_module_path in sys.modules:
+        tokenizer_module = sys.modules[tokenizer_module_path]
+        setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+
+    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni
+
+    # This is because async_omni imports processor.py, which imports this function at module level
+    # Mock length_from_prompt_token_ids_or_embeds (used in processor.py)
+    def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
+        # Return a reasonable default length
+        if prompt_token_ids is not None:
+            if isinstance(prompt_token_ids, list):
+                return len(prompt_token_ids)
+            elif hasattr(prompt_token_ids, "shape"):
+                return prompt_token_ids.shape[-1] if len(prompt_token_ids.shape) > 0 else 1
+        if prompt_embeds is not None:
+            if hasattr(prompt_embeds, "shape"):
+                return prompt_embeds.shape[-2] if len(prompt_embeds.shape) > 1 else 1
+        return 10  # Default length
+
+    # Mock in vllm.utils
+    monkeypatch.setattr(
+        "vllm.utils.length_from_prompt_token_ids_or_embeds",
+        _mock_length_from_prompt_token_ids_or_embeds,
+        raising=False,
+    )
+    # Also mock in processor module if it's imported
+    monkeypatch.setattr(
+        "vllm_omni.engine.input_processor.length_from_prompt_token_ids_or_embeds",
+        _mock_length_from_prompt_token_ids_or_embeds,
+        raising=False,
+    )
+    # If processor module is already imported, patch it directly
+    processor_module_path = "vllm_omni.engine.input_processor"
+    if processor_module_path in sys.modules:
+        processor_module = sys.modules[processor_module_path]
+        setattr(
+            processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds
+        )
+
+    # Strategy 3: Now mock async_omni AFTER length_from_prompt_token_ids_or_embeds is mocked
+    # This prevents ImportError when async_omni imports processor.py
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs",
+        _mock_init_tokenizer_from_configs,
+        raising=False,
+    )
+
+    # Strategy 4: If async_omni is already imported, patch it directly
+    async_omni_path = "vllm_omni.entrypoints.async_omni"
+    if async_omni_path in sys.modules:
+        async_omni_module = sys.modules[async_omni_path]
+        setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+
+    # Now mock get_config and other functions
+    fake_hf_config = MagicMock()
+    fake_hf_config.model_type = "qwen2_5_omni"
+
+    def _mock_get_config(model, **kwargs):
+        return fake_hf_config
+
+    monkeypatch.setattr(
+        "vllm.transformers_utils.config.get_config",
+        _mock_get_config,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.get_config",
+        _mock_get_config,
+        raising=False,
+    )
+
+    # Mock transformers' cached_file to avoid downloading model configs
+    def _mock_cached_file(path_or_repo_id, *args, **kwargs):
+        import os
+        import tempfile
+
+        fake_config_file = os.path.join(tempfile.gettempdir(), "fake_config.json")
+        if not os.path.exists(fake_config_file):
+            with open(fake_config_file, "w") as f:
+                f.write('{"model_type": "qwen2_5_omni"}')
+        return fake_config_file
+
+    monkeypatch.setattr(
+        "transformers.utils.hub.cached_file",
+        _mock_cached_file,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "transformers.utils.hub.cached_files",
+        lambda path_or_repo_id, filenames, **kwargs: (
+            [_mock_cached_file(path_or_repo_id, filenames[0])] if filenames else None
+        ),
+        raising=False,
+    )
+
+
+def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
+    """Test that stage configs are auto-loaded when stage_configs_path is None."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [
+            _FakeStageConfig(fake_stage_config),
+            _FakeStageConfig(fake_stage_config),
+        ]
+
+    # Remove modules from cache BEFORE setting mocks
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    # Set up mocks
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    # Mock load_stage_configs_from_model
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+
+    # Replace OmniStage
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    # Import the module after mocks are set
+    import vllm_omni.entrypoints.omni as omni_module
+
+    # Patch the imported function and class in the module
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+    # Verify: auto-loaded stage_configs and stage_list have consistent count
+    assert isinstance(omni.stage_configs, list)
+    assert len(omni.stage_configs) == 2
+    assert len(omni.stage_list) == 2
+    # Verify: each Stage is _FakeStage instance
+    for st in omni.stage_list:
+        assert isinstance(st, _FakeStage)
+    # Verify: queues are attached
+    for st in omni.stage_list:
+        assert st._in_q is not None
+        assert st._out_q is not None
+    # Verify: all stages are ready
+    assert len(omni._stages_ready) == 2
+
+
+def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
+    """Test that generate raises ValueError when sampling_params_list length doesn't match."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(fake_stage_config)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+    with pytest.raises(ValueError):
+        omni.generate(prompts=["hi"], sampling_params_list=[])
+
+
+def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
+    """Test multi-stage generation pipeline with queue polling."""
+    stage_cfg0 = dict(fake_stage_config)
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["processed_input"] = ["processed-for-stage-1"]
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_module, "uuid", uuid)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
+    # Simulate worker behavior: manually put results into output queues
+    # Note: We put results before calling generate, which simulates worker processes
+    # that have already completed. The polling loop will collect them in stage order.
+    # Stage 0 output (will be collected first)
+    omni.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 0, "text": "s0"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    # Stage 1 output (will be collected after stage 0 forwards to it)
+    # Note: In real flow, stage 1 result would appear after stage 0 forwards,
+    # but for testing we pre-populate it. The polling loop processes stages
+    # in order, so stage 0 result will be collected first, then forwarded,
+    # then stage 1 result will be collected.
+    omni.stage_list[1]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 1, "text": "s1"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+
+    sampling_params_list = [
+        OmniDiffusionSamplingParams(num_inference_steps=1),
+        OmniDiffusionSamplingParams(num_inference_steps=1, max_sequence_length=10),
+    ]
+    prompts = ["hi"]
+    outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
+
+    # Both stages have final_output=True, so should aggregate two OmniRequestOutput
+    assert len(outputs) == 2
+    # Verify stage outputs are set
+    assert omni.stage_list[0].engine_outputs == [{"stage": 0, "text": "s0"}]
+    assert omni.stage_list[1].engine_outputs == [{"stage": 1, "text": "s1"}]
+    # Verify stage 0 input queue received the task
+    assert not omni.stage_list[0]._in_q.empty()
+    # Verify stage 1 received forwarded task (process_engine_inputs was called)
+    assert omni.stage_list[1].process_engine_inputs([], []) is not None
+
+
+def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
+    """Test single-stage generation pipeline with multiple inputs in one batch."""
+    stage_cfg0 = dict(fake_stage_config)
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["final_output"] = False
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_module, "uuid", uuid)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
+    # Simulate worker behavior: manually put results into output queues
+    # Note: We put results before calling generate, which simulates worker processes
+    # that have already completed. The polling loop will collect them in stage order.
+    omni.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 0, "text": "s0"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    omni.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 0, "text": "s0"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    omni.stage_list[1]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 1}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    omni.stage_list[1]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 1}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+
+    outputs = omni.generate(
+        prompts=[
+            {
+                "prompt": "hi",
+                "negative_prompt": "hi",
+                "multi_modal_data": {"image": ["dog.jpg", "cat.jpg"]},
+            },
+            {
+                "prompt": "hi",
+                "negative_prompt": "hi",
+                "multi_modal_data": {"image": ["dog.jpg", "cat.jpg"]},
+            },
+        ],
+        sampling_params_list=[
+            OmniDiffusionSamplingParams(num_inference_steps=1),
+            OmniDiffusionSamplingParams(num_inference_steps=1),
+        ],
+    )
+
+    assert len(outputs) == 2
+
+
+def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
+    """Test that generate returns empty list when all stages have final_output=False."""
+    stage_cfg0 = dict(fake_stage_config)
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["final_output"] = False
+    stage_cfg1["final_output"] = False
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_module, "uuid", uuid)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
+    # Simulate worker behavior: put results into output queues
+    omni.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 0}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    omni.stage_list[1]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 1}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+
+    outputs = omni.generate(
+        prompts=["p"],
+        sampling_params_list=[
+            OmniDiffusionSamplingParams(num_inference_steps=1),
+            OmniDiffusionSamplingParams(num_inference_steps=1, max_sequence_length=10),
+        ],
+    )
+    assert outputs == []
+
+
+def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config):
+    """Test that generate uses default sampling params when sampling_params_list is None."""
+    stage_cfg0 = dict(fake_stage_config)
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["final_output"] = False
+    stage_cfg1["final_output"] = False
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_module, "uuid", uuid)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
+    # Simulate worker behavior: put results into output queues
+    omni.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 0}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    omni.stage_list[1]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 1}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    # Use the default sampling params
+    omni.generate(prompts=["p"], sampling_params_list=None)
+
+
+def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
+    """Test that _wait_for_stages_ready handles timeout correctly."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(fake_stage_config)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+
+    # Create a stage that doesn't send stage_ready message
+    class _FakeStageNoReady(_FakeStage):
+        def init_stage_worker(self, *args, **kwargs):
+            # Don't send stage_ready message
+            self._proc = MagicMock()
+            self._proc.start = MagicMock()
+            self._proc.join = MagicMock()
+            self._proc.is_alive = MagicMock(return_value=False)
+            self._proc.terminate = MagicMock()
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs))
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    # Use very short timeout
+    omni = Omni(model="any", init_timeout=0.01)
+    # Verify that no stages are ready
+    assert len(omni._stages_ready) == 0
+
+
+def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
+    """Test that generate handles error messages from stages correctly."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(fake_stage_config)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    # Mock uuid.uuid4() to return a predictable value for request ID generation
+    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid)
+    monkeypatch.setattr(omni_module, "uuid", uuid)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+
+    # Generate the expected request ID format: "0_<uuid>"
+    expected_request_id = f"0_{test_uuid}"
+
+    # Put error message in output queue
+    omni.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "error": "test error",
+        }
+    )
+    # Also put a valid result after error to allow the loop to complete
+    # (error handling continues the loop, so we need a valid result to finish)
+    omni.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": expected_request_id,
+            "engine_outputs": [{"stage": 0, "text": "result"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+
+    # Generate should handle error gracefully (log but continue)
+    sampling_params_list = [OmniDiffusionSamplingParams(num_inference_steps=1)]
+    outputs = omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
+    # Should return final output (error was logged but didn't stop processing)
+    assert isinstance(outputs, list)
+    # Since final_output=True, should have one output
+    assert len(outputs) == 1
+
+
+def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
+    """Test that close() sends shutdown signal to all input queues."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        return [_FakeStageConfig(fake_stage_config)]
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        raising=False,
+    )
+
+    import vllm_omni.entrypoints.omni as omni_module
+
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(model="any", init_timeout=1)
+
+    # Call close
+    omni.close()
+
+    # Verify shutdown signal (None) was sent to input queue
+    # Use get_nowait to avoid blocking (close() uses put_nowait, so should be safe)
+    try:
+        shutdown_signal = omni.stage_list[0]._in_q.get_nowait()
+        assert shutdown_signal == SHUTDOWN_TASK
+    except Empty:
+        # If queue was already empty or only had stage_ready, that's also acceptable
+        # The important thing is that close() was called without error
+        pass
+
+    # Verify stop_stage_worker was called (process should be set)
+    assert omni.stage_list[0]._proc is not None
diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py
index fd2ba2673..f99c6d833 100644
--- a/tests/entrypoints/test_omni_llm.py
+++ b/tests/entrypoints/test_omni_llm.py
@@ -5,6 +5,7 @@ from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
+from vllm import SamplingParams
 
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
 
@@ -82,8 +83,9 @@ class _FakeStage:
         self.stage_id = getattr(config, "stage_id", 0)
         self.engine_args = config.engine_args
         self.model_stage = getattr(config.engine_args, "model_stage", None)
+        self.stage_type = "llm"
         # set default sampling params
-        self.default_sampling_params = {"temperature": 1.0}
+        self.default_sampling_params = SamplingParams(temperature=1.0)
         # Allow configuring final_output and final_output_type
         self.final_output = config.final_output if hasattr(config, "final_output") else False
         self.final_output_type = getattr(config, "final_output_type", None)
@@ -637,8 +639,10 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
         }
     )
 
-    # Use dicts instead of object() for serializable sampling params
-    sampling_params_list = [{"temperature": 0.7}, {"temperature": 0.8}]
+    sampling_params_list = [
+        SamplingParams(temperature=0.7),
+        SamplingParams(temperature=0.8),
+    ]
     prompts = ["hi"]
     outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
 
@@ -722,8 +726,13 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
         }
     )
 
-    # Use dicts instead of object() for serializable sampling params
-    outputs = omni.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}])
+    outputs = omni.generate(
+        prompts=["p"],
+        sampling_params_list=[
+            SamplingParams(temperature=0.7),
+            SamplingParams(temperature=0.8),
+        ],
+    )
     assert outputs == []
 
 
@@ -922,8 +931,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     )
 
     # Generate should handle error gracefully (log but continue)
-    # Use dict instead of object() for serializable sampling params
-    sampling_params_list = [{"temperature": 0.7}]
+    sampling_params_list = [SamplingParams(temperature=0.7)]
     outputs = omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
     # Should return final output (error was logged but didn't stop processing)
     assert isinstance(outputs, list)
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 2e00f41b9..b0b9b89db 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -17,6 +17,7 @@ from vllm_omni.diffusion.registry import (
     get_diffusion_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -60,36 +61,31 @@ class DiffusionEngine:
             self.close()
             raise e
 
-    def step(self, requests: list[OmniDiffusionRequest]):
+    def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
         # Apply pre-processing if available
         if self.pre_process_func is not None:
             preprocess_start_time = time.time()
-            requests = self.pre_process_func(requests)
+            request = self.pre_process_func(request)
             preprocess_time = time.time() - preprocess_start_time
             logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
 
-        output = self.add_req_and_wait_for_response(requests)
+        output = self.add_req_and_wait_for_response(request)
         if output.error:
             raise Exception(f"{output.error}")
         logger.info("Generation completed successfully.")
 
         if output.output is None:
             logger.warning("Output is None, returning empty OmniRequestOutput")
-            # Return empty output for the first request
-            if len(requests) > 0:
-                request = requests[0]
-                request_id = request.request_id or ""
-                prompt = request.prompt
-                if isinstance(prompt, list):
-                    prompt = prompt[0] if prompt else None
-                return OmniRequestOutput.from_diffusion(
-                    request_id=request_id,
+            return [
+                OmniRequestOutput.from_diffusion(
+                    request_id=request.request_ids[i] if i < len(request.request_ids) else "",
                     images=[],
                     prompt=prompt,
                     metrics={},
                     latents=None,
                 )
-            return None
+                for i, prompt in enumerate(request.prompts)
+            ]
 
         postprocess_start_time = time.time()
         outputs = self.post_process_func(output.output) if self.post_process_func is not None else output.output
@@ -102,13 +98,10 @@ class DiffusionEngine:
             outputs = [outputs] if outputs is not None else []
 
         # Handle single request or multiple requests
-        if len(requests) == 1:
+        if len(request.prompts) == 1:
             # Single request: return single OmniRequestOutput
-            request = requests[0]
-            request_id = request.request_id or ""
-            prompt = request.prompt
-            if isinstance(prompt, list):
-                prompt = prompt[0] if prompt else None
+            prompt = request.prompts[0]
+            request_id = request.request_ids[0] if request.request_ids else ""
 
             metrics = {}
             if output.trajectory_timesteps is not None:
@@ -116,37 +109,38 @@ class DiffusionEngine:
 
             if supports_audio_output(self.od_config.model_class_name):
                 audio_payload = outputs[0] if len(outputs) == 1 else outputs
-                return OmniRequestOutput.from_diffusion(
-                    request_id=request_id,
-                    images=[],
-                    prompt=prompt,
-                    metrics=metrics,
-                    latents=output.trajectory_latents,
-                    multimodal_output={"audio": audio_payload},
-                    final_output_type="audio",
-                )
+                return [
+                    OmniRequestOutput.from_diffusion(
+                        request_id=request_id,
+                        images=[],
+                        prompt=prompt,
+                        metrics=metrics,
+                        latents=output.trajectory_latents,
+                        multimodal_output={"audio": audio_payload},
+                        final_output_type="audio",
+                    ),
+                ]
             else:
-                return OmniRequestOutput.from_diffusion(
-                    request_id=request_id,
-                    images=outputs,
-                    prompt=prompt,
-                    metrics=metrics,
-                    latents=output.trajectory_latents,
-                )
+                return [
+                    OmniRequestOutput.from_diffusion(
+                        request_id=request_id,
+                        images=outputs,
+                        prompt=prompt,
+                        metrics=metrics,
+                        latents=output.trajectory_latents,
+                    ),
+                ]
         else:
             # Multiple requests: return list of OmniRequestOutput
             # Split images based on num_outputs_per_prompt for each request
             results = []
             output_idx = 0
 
-            for request in requests:
-                request_id = request.request_id or ""
-                prompt = request.prompt
-                if isinstance(prompt, list):
-                    prompt = prompt[0] if prompt else None
+            for i, prompt in enumerate(request.prompts):
+                request_id = request.request_ids[i] if i < len(request.request_ids) else ""
 
                 # Get images for this request
-                num_outputs = request.num_outputs_per_prompt
+                num_outputs = request.sampling_params.num_outputs_per_prompt
                 request_outputs = outputs[output_idx : output_idx + num_outputs] if output_idx < len(outputs) else []
                 output_idx += num_outputs
 
@@ -192,8 +186,8 @@ class DiffusionEngine:
         """
         return DiffusionEngine(config)
 
-    def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
-        return self.executor.add_req(requests)
+    def add_req_and_wait_for_response(self, request: OmniDiffusionRequest):
+        return self.executor.add_req(request)
 
     def start_profile(self, trace_filename: str | None = None) -> None:
         """
@@ -316,8 +310,6 @@ class DiffusionEngine:
 
     def _dummy_run(self):
         """A dummy run to warm up the model."""
-        prompt = "dummy run"
-        # note that num_inference_steps=1 will cause timestep and temb None in the pipeline
         num_inference_steps = 1
         height = 1024
         width = 1024
@@ -327,17 +319,19 @@ class DiffusionEngine:
             dummy_image = PIL.Image.new("RGB", (width, height), color=(0, 0, 0))
         else:
             dummy_image = None
+        prompt: OmniTextPrompt = {"prompt": "dummy run", "multi_modal_data": {"image": dummy_image}}
         req = OmniDiffusionRequest(
-            prompt=prompt,
-            height=height,
-            width=width,
-            pil_image=dummy_image,
-            num_inference_steps=num_inference_steps,
-            num_outputs_per_prompt=1,
+            prompts=[prompt],
+            sampling_params=OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                num_outputs_per_prompt=1,
+            ),
         )
         logger.info("dummy run to warm up the model")
-        requests = self.pre_process_func([req]) if self.pre_process_func is not None else [req]
-        self.add_req_and_wait_for_response(requests)
+        request = self.pre_process_func(req) if self.pre_process_func is not None else req
+        self.add_req_and_wait_for_response(request)
 
     def collective_rpc(
         self,
diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py
index 021c210ef..e41f41d11 100644
--- a/vllm_omni/diffusion/executor/abstract.py
+++ b/vllm_omni/diffusion/executor/abstract.py
@@ -3,7 +3,7 @@ from typing import Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 
@@ -59,7 +59,7 @@ class DiffusionExecutor(ABC):
         pass
 
     @abstractmethod
-    def add_req(self, requests: list[OmniDiffusionRequest]):
+    def add_req(self, requests: OmniDiffusionRequest) -> DiffusionOutput:
         """Add requests to the execution queue."""
         pass
 
diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py
index a04a40e86..21a94acd3 100644
--- a/vllm_omni/diffusion/executor/multiproc_executor.py
+++ b/vllm_omni/diffusion/executor/multiproc_executor.py
@@ -6,7 +6,7 @@ from typing import Any
 
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE
+from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import Scheduler
@@ -130,8 +130,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
         return processes, result_handle
 
-    def add_req(self, requests: list[OmniDiffusionRequest]):
-        return self.scheduler.add_req(requests)
+    def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
+        return self.scheduler.add_req(request)
 
     def collective_rpc(
         self,
diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
index 5d3bcdd11..bdb9f1f5c 100644
--- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
+++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
@@ -272,16 +272,22 @@ class BagelPipeline(nn.Module):
 
     @torch.inference_mode()
     def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
-        prompt = req.prompt or ""
-        if isinstance(prompt, list):
-            # vllm-omni request supports list; Bagel pipeline currently supports first prompt.
-            prompt = prompt[0] if prompt else ""
+        if len(req.prompts) > 1:
+            logger.warning(
+                """This model only supports a single prompt, not a batched request.""",
+                """Taking only the first image for now.""",
+            )
+        # TODO: In online mode, sometimes it receives [{"prompts": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(req.prompts[0], str) else (req.prompts[0].get("prompt") or "")
+
         max_hw = int(self.bagel.max_latent_size * self.bagel.latent_downsample)
-        if req.height is None and req.width is None:
+        if req.sampling_params.height is None and req.sampling_params.width is None:
             height = width = max_hw
         else:
-            height = int(req.height) if req.height is not None else max_hw
-            width = int(req.width) if req.width is not None else max_hw
+            height = int(req.sampling_params.height) if req.sampling_params.height is not None else max_hw
+            width = int(req.sampling_params.width) if req.sampling_params.width is not None else max_hw
         if height > max_hw or width > max_hw:
             raise ValueError(
                 f"Requested resolution {height}x{width} exceeds Bagel checkpoint limit "
@@ -292,7 +298,7 @@ class BagelPipeline(nn.Module):
 
         # Map request params to Bagel gen params (defaults follow Bagel inferencer)
         gen_params = BagelGenParams(
-            num_timesteps=int(req.num_inference_steps or 50),
+            num_timesteps=int(req.sampling_params.num_inference_steps or 50),
             timestep_shift=3.0,
         )
 
@@ -304,7 +310,7 @@ class BagelPipeline(nn.Module):
 
         # Add text prompt (prefill) on gen context.
         # [Omni] Check for injected KV Cache from remote transfer
-        injected_kv = getattr(req, "past_key_values", None)
+        injected_kv = req.sampling_params.past_key_values
         if injected_kv is not None:
             logger.info("Using injected KV Cache (direct)")
             gen_context["past_key_values"] = injected_kv
@@ -316,9 +322,13 @@ class BagelPipeline(nn.Module):
             gen_context["ropes"] = [seq_len]
 
         else:
-            image_input = getattr(req, "pil_image", None)
+            image_input = (
+                None if isinstance(first_prompt, str) else (first_prompt.get("multi_modal_data") or {}).get("image")
+            )
             if image_input and not isinstance(image_input, list):
                 image_input = [image_input]
+            if image_input:
+                image_input = [Image.open(image) if isinstance(image, str) else image for image in image_input]
 
             if image_input:
                 # If we have an image, we prefill with it
@@ -414,10 +424,10 @@ class BagelPipeline(nn.Module):
             gen_context["kv_lens"] = newlens
             gen_context["ropes"] = new_rope
 
-        if req.seed is not None:
-            torch.manual_seed(req.seed)
+        if req.sampling_params.seed is not None:
+            torch.manual_seed(req.sampling_params.seed)
             if self.device.type == "cuda":
-                torch.cuda.manual_seed(req.seed)
+                torch.cuda.manual_seed(req.sampling_params.seed)
 
         # Prepare latent query and run flow
         generation_input = self.bagel.prepare_vae_latent(
diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
index 7d1b6578d..2eac7b2ec 100644
--- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
+++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
@@ -19,7 +19,7 @@ import json
 import math
 import os
 from collections.abc import Callable, Iterable
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import PIL.Image
@@ -742,27 +742,54 @@ class Flux2KleinPipeline(nn.Module, SupportImageInput):
             `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
             generated images.
         """
+        if len(req.prompts) > 1:
+            logger.warning(
+                """This model only supports a single prompt, not a batched request.""",
+                """Taking only the first image for now.""",
+            )
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
 
-        prompt = req.prompt if req.prompt is not None else prompt
-        image = req.pil_image if req.pil_image is not None else image
-        height = req.height or height
-        width = req.width or width
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        guidance_scale = req.guidance_scale if req.guidance_scale is not None else guidance_scale
-        generator = req.generator or generator
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        if (
+            raw_image := None
+            if isinstance(first_prompt, str)
+            else first_prompt.get("multi_modal_data", {}).get("image")
+        ) is None:
+            pass  # use image from param list
+        elif isinstance(raw_image, list):
+            image = [PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image, im) for im in raw_image]
+        else:
+            image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image)
 
-        if isinstance(req.prompt_embeds, torch.Tensor):
-            prompt_embeds = req.prompt_embeds
-        if isinstance(req.negative_prompt_embeds, torch.Tensor):
-            negative_prompt_embeds = req.negative_prompt_embeds
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
+        )
+        generator = req.sampling_params.generator or generator
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers)
 
-        if req.max_sequence_length is not None:
-            max_sequence_length = req.max_sequence_length
-        if getattr(req, "text_encoder_out_layers", None) is not None:
-            text_encoder_out_layers = req.text_encoder_out_layers
+        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
+        if any(p is not None for p in req_prompt_embeds):
+            # If at list one prompt is provided as an embedding,
+            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
+            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
+            # And `torch.stack` automatically raises an exception for us
+            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
+
+        req_negative_prompt_embeds = [
+            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
+        ]
+        if any(p is not None for p in req_negative_prompt_embeds):
+            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
index 9a03a9349..d222342b5 100644
--- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
+++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
@@ -17,6 +17,7 @@ import logging
 import os
 import re
 from collections.abc import Iterable
+from typing import cast
 
 import numpy as np
 import PIL.Image
@@ -48,6 +49,7 @@ from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
     GlmImageTransformer2DModel,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -77,16 +79,26 @@ def get_glm_image_pre_process_func(od_config: OmniDiffusionConfig):
     # GLM-Image uses patch_size=2 for transformer
     patch_size = 2
 
-    def pre_process_func(requests: list[OmniDiffusionRequest]):
+    def pre_process_func(request: OmniDiffusionRequest):
         """Pre-process condition images for Image Edit mode."""
-        for req in requests:
-            images = req.pil_image
-            if images is None:
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
+
+            if raw_image is None:
                 # Text-to-image mode, no preprocessing needed
                 continue
 
-            if not isinstance(images, list):
-                images = [images]
+            if not isinstance(raw_image, list):
+                raw_image = [raw_image]
+            images = [
+                PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im)
+                for im in raw_image
+            ]
 
             preprocessed = []
             height, width = None, None
@@ -110,14 +122,19 @@ def get_glm_image_pre_process_func(od_config: OmniDiffusionConfig):
                     height, width = img_h, img_w
 
             # Store in request
-            req.preprocessed_image = preprocessed
-            req.prompt_image = images  # Keep original PIL images
-            if req.height is None:
-                req.height = height
-            if req.width is None:
-                req.width = width
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt, additional_information={})
+            elif "additional_information" not in prompt:
+                prompt["additional_information"] = {}
+            prompt["additional_information"]["preprocessed_image"] = processed  # type: ignore
+            prompt["additional_information"]["prompt_image"] = images  # type: ignore
+            request.prompts[i] = prompt
+            if request.sampling_params.height is None:
+                request.sampling_params.height = height
+            if request.sampling_params.width is None:
+                request.sampling_params.width = width
 
-        return requests
+        return request
 
     return pre_process_func
 
@@ -821,27 +838,44 @@ class GlmImagePipeline(nn.Module):
         Returns:
             DiffusionOutput containing generated image
         """
-        prompt = req.prompt or ""
-        if isinstance(prompt, list):
-            prompt = prompt[0] if prompt else ""
+        if len(req.prompts) > 1:
+            logger.warning(
+                """This model only supports a single prompt, not a batched request.""",
+                """Taking only the first image for now.""",
+            )
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
 
         # Get pre-computed prompt embeddings if provided
-        prompt_embeds = req.prompt_embeds if isinstance(req.prompt_embeds, torch.Tensor) else None
+        if isinstance(first_prompt, str):
+            prompt_embeds = None
+        else:
+            prompt_embeds = first_prompt.get("prompt_embeds")
+            if not isinstance(prompt_embeds, torch.Tensor):
+                prompt_embeds = None
 
         # Get condition images for Image Edit mode
         # Use pre-processed images from pre_process_func
-        preprocessed_images = req.preprocessed_image
-        condition_images = getattr(req, "prompt_image", None)
-        img_height = req.height
-        img_width = req.width
+        preprocessed_images = (
+            None
+            if isinstance(first_prompt, str)
+            else first_prompt.get("additional_information", {}).get("preprocessed_image")
+        )
+        condition_images = (
+            None
+            if isinstance(first_prompt, str)
+            else first_prompt.get("additional_information", {}).get("prompt_image")
+        )
+        img_height = req.sampling_params.height
+        img_width = req.sampling_params.width
 
         is_image_edit = preprocessed_images is not None
 
         # Use image dimensions as default if available
-        height = req.height or img_height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or img_width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.num_inference_steps or 50
-        guidance_scale = req.guidance_scale or 1.5
+        height = req.sampling_params.height or img_height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or img_width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        guidance_scale = req.sampling_params.guidance_scale or 1.5
 
         # 0. Validate inputs
         self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds)
@@ -851,13 +885,13 @@ class GlmImagePipeline(nn.Module):
 
         # Set seed if provided
         generator = None
-        if req.seed is not None:
-            generator = torch.Generator(device=self.device).manual_seed(req.seed)
+        if req.sampling_params.seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
 
         # 1. Get prior tokens - either from external source (multistage) or generate internally
         # Check if prior_token_ids are provided externally (from AR stage in multistage mode)
-        external_prior_tokens = req.extra.get("prior_token_ids") if req.extra else None
-        external_prior_image_ids = req.extra.get("prior_token_image_ids") if req.extra else None
+        external_prior_tokens = req.sampling_params.extra_args.get("prior_token_ids")
+        external_prior_image_ids = req.sampling_params.extra_args.get("prior_token_image_ids")
 
         if external_prior_tokens is not None:
             # Multistage mode: use externally provided prior tokens from vLLM AR stage
diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py
index e918afe0b..8b616ec45 100644
--- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py
+++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py
@@ -451,7 +451,7 @@ class LongCatImagePipeline(
             raise ValueError(
                 "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
             )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+        elif prompt is not None and not isinstance(prompt, (str, list)):
             raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
 
         if negative_prompt is not None and negative_prompt_embeds is not None:
@@ -470,11 +470,11 @@ class LongCatImagePipeline(
         num_inference_steps: int = 50,
         sigmas: list[float] | None = None,
         guidance_scale: float = 4.5,
-        num_images_per_prompt: int | None = 1,
+        num_images_per_prompt: int = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
         latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
         output_type: str | None = "pil",
         return_dict: bool = True,
         joint_attention_kwargs: dict[str, Any] | None = None,
@@ -482,20 +482,44 @@ class LongCatImagePipeline(
         cfg_renorm_min: float | None = 0.0,
         enable_prompt_rewrite: bool | None = True,
     ) -> DiffusionOutput:
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
+            negative_prompt = None
+        elif req.prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        height = req.height or height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        generator = req.generator or generator
-        guidance_scale = req.guidance_scale if getattr(req, "guidance_scale", None) is not None else guidance_scale
-        num_images_per_prompt = getattr(req, "num_outputs_per_prompt", None) or num_images_per_prompt
-        enable_prompt_rewrite = getattr(req, "enable_prompt_rewrite", None) or enable_prompt_rewrite
-        enable_cfg_renorm = getattr(req, "enable_cfg_renorm", None) or enable_cfg_renorm
-        cfg_renorm_min = getattr(req, "cfg_renorm_min", None) or cfg_renorm_min
-        prompt_embeds = getattr(req, "prompt_embeds", None) or prompt_embeds
-        negative_prompt_embeds = getattr(req, "negative_prompt_embeds", None) or negative_prompt_embeds
+        height = req.sampling_params.height or height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        generator = req.sampling_params.generator or generator
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
+        )
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt is not None
+            else num_images_per_prompt
+        )
+        enable_prompt_rewrite = req.sampling_params.extra_args.get("enable_prompt_rewrite", enable_prompt_rewrite)
+        enable_cfg_renorm = req.sampling_params.extra_args.get("enable_cfg_renorm", enable_cfg_renorm)
+        cfg_renorm_min = req.sampling_params.extra_args.get("cfg_renorm_min", cfg_renorm_min)
+
+        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
+        if any(p is not None for p in req_prompt_embeds):
+            # If at list one prompt is provided as an embedding,
+            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
+            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
+            # And `torch.stack` automatically raises an exception for us
+            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
+
+        req_negative_prompt_embeds = [
+            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
+        ]
+        if any(p is not None for p in req_negative_prompt_embeds):
+            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
 
         self.check_inputs(
             prompt,
diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py
index 607553aee..f2c3fd648 100644
--- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py
+++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py
@@ -6,7 +6,7 @@ import math
 import os
 import re
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import PIL.Image
@@ -34,6 +34,7 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import (
 )
 from vllm_omni.diffusion.models.longcat_image.pipeline_longcat_image import calculate_shift
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -59,22 +60,38 @@ def get_longcat_image_edit_pre_process_func(
     latent_channels = vae_config.get("latent_channels", 16)
 
     def pre_process_func(
-        requests: list[OmniDiffusionRequest],
+        request: OmniDiffusionRequest,
     ):
-        """Pre-process requests for QwenImageEditPipeline."""
-        for req in requests:
-            image = req.pil_image
+        """Pre-process requests for LongCatImageEditPipeline."""
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
 
-            image_size = image[0].size if isinstance(image, list) else image.size
+            if raw_image is None or isinstance(raw_image, list):
+                raise ValueError(
+                    """Received no image or a list of image. Only a single image is supported by this model."""
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
+                )
+
+            if isinstance(raw_image, str):
+                image = PIL.Image.open(raw_image)
+            else:
+                image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image)
+
+            image_size = image.size
             calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1])
-            height = req.height or calculated_height
-            width = req.width or calculated_width
+            height = request.sampling_params.height or calculated_height
+            width = request.sampling_params.width or calculated_width
 
             # Store calculated dimensions in request
-            req.calculated_height = calculated_height
-            req.calculated_width = calculated_width
-            req.height = height
-            req.width = width
+            prompt["additional_information"]["calculated_height"] = calculated_height
+            prompt["additional_information"]["calculated_width"] = calculated_width
+            request.sampling_params.height = height
+            request.sampling_params.width = width
 
             # Preprocess image
             if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels):
@@ -83,9 +100,10 @@ def get_longcat_image_edit_pre_process_func(
                 image = image_processor.preprocess(image, calculated_height, calculated_width)
 
                 # Store preprocessed image and prompt image in request
-                req.preprocessed_image = image
-                req.prompt_image = prompt_image
-        return requests
+                prompt["additional_information"]["preprocessed_image"] = image
+                prompt["additional_information"]["prompt_image"] = prompt_image
+            request.prompts[i] = prompt
+        return request
 
     return pre_process_func
 
@@ -500,42 +518,59 @@ class LongCatImageEditPipeline(nn.Module, SupportImageInput):
         self,
         req: OmniDiffusionRequest,
         image: PIL.Image.Image | torch.Tensor | None = None,
-        prompt: str | list[str] = None,
-        negative_prompt: str | list[str] = None,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
         num_inference_steps: int = 50,
         sigmas: list[float] | None = None,
         guidance_scale: float = 3.5,
         num_images_per_prompt: int | None = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
         latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
         output_type: str | None = "pil",
         return_dict: bool = True,
         joint_attention_kwargs: dict[str, Any] | None = None,
     ):
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
-        negative_prompt = "" if negative_prompt is None else negative_prompt
-        guidance_scale = req.guidance_scale if req.guidance_scale is not None else guidance_scale
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        num_images_per_prompt = getattr(req, "num_outputs_per_prompt", None) or num_images_per_prompt
-        generator = req.generator or generator
-        prompt_embeds = getattr(req, "prompt_embeds", None) or prompt_embeds
-        negative_prompt_embeds = getattr(req, "negative_prompt_embeds", None) or negative_prompt_embeds
-        height = req.height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or self.default_sample_size * self.vae_scale_factor
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        if len(req.prompts) > 1:
+            logger.warning(
+                """This model only supports a single prompt, not a batched request.""",
+                """Taking only the first image for now.""",
+            )
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+        negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+        prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("prompt_embeds")
+        negative_prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt_embeds")  # type: ignore # Why it is list[torch.Tensor] in OmniTokenInputs or OmniEmbedsPrompt? Doesn't make sense
+
+        sigmas = req.sampling_params.sigmas or sigmas
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
+        )
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt is not None
+            else num_images_per_prompt
+        )
+        generator = req.sampling_params.generator or generator
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
 
         if prompt is not None:
             batch_size = 1 if isinstance(prompt, str) else len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
 
-        if hasattr(req, "preprocessed_image"):
-            prompt_image = req.prompt_image
-            image = req.preprocessed_image
-            calculated_height = req.calculated_height if hasattr(req, "calculated_height") else height
-            calculated_width = req.calculated_width if hasattr(req, "calculated_width") else width
+        if not isinstance(first_prompt, str) and "preprocessed_image" in (
+            additional_information := first_prompt.get("additional_information", {})
+        ):
+            prompt_image = additional_information.get("prompt_image")
+            image = additional_information.get("preprocessed_image")
+            calculated_height = additional_information.get("calculated_height", height)
+            calculated_width = additional_information.get("calculated_width", width)
         else:
             image_size = image[0].size if isinstance(image, list) else image.size
             calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1])
diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
index 6bf4ca1b5..a87410a6c 100644
--- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
+++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
@@ -519,8 +519,8 @@ class OvisImagePipeline(
     def forward(
         self,
         req: OmniDiffusionRequest,
-        prompt: str | list[str] = None,
-        negative_prompt: str | list[str] = None,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
         guidance_scale: float = 5.0,
         height: int | None = None,
         width: int | None = None,
@@ -606,16 +606,27 @@ class OvisImagePipeline(
             [`~pipelines.ovis_image.OvisImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
             returning a tuple, the first element is a list with the generated images.
         """
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
-        height = req.height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        guidance_scale = req.guidance_scale if req.guidance_scale is not None else guidance_scale
-        generator = req.generator or generator
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
+            negative_prompt = None
+        elif req.prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
+        )
+        generator = req.sampling_params.generator or generator
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
 
         # Steps:
         # 1. Check Inputs
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
index 2b9d906d3..e0a37b8bc 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
@@ -662,21 +662,28 @@ class QwenImagePipeline(
         callback_on_step_end_tensor_inputs: list[str] = ["latents"],
         max_sequence_length: int = 512,
     ) -> DiffusionOutput:
-        # # TODO: only support single prompt now
-        # if req.prompt is not None:
-        #     prompt = req.prompt[0] if isinstance(req.prompt, list) else req.prompt
-        prompt = req.prompt
-        negative_prompt = req.negative_prompt
-        height = req.height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        generator = req.generator or generator
-        true_cfg_scale = req.true_cfg_scale or true_cfg_scale
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
+            negative_prompt = None
+        elif req.prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        generator = req.sampling_params.generator or generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
         # 1. check inputs
         # 2. encode prompts
         # 3. prepare latents and timesteps
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
index ac74507c7..ba65bf07e 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import PIL.Image
@@ -39,6 +39,7 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -64,16 +65,33 @@ def get_qwen_image_edit_pre_process_func(
     latent_channels = vae_config.get("z_dim", 16)
 
     def pre_process_func(
-        requests: list[OmniDiffusionRequest],
+        request: OmniDiffusionRequest,
     ):
         """Pre-process requests for QwenImageEditPipeline."""
-        for req in requests:
-            image = req.pil_image
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
 
-            image_size = image[0].size if isinstance(image, list) else image.size
+            # Only handles single image
+            if raw_image is None or isinstance(raw_image, list):
+                raise ValueError(
+                    """Received no image or a list of image. Only a single image is supported by this model."""
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
+                )
+
+            if isinstance(raw_image, str):
+                image = PIL.Image.open(raw_image)
+            else:
+                image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image)
+
+            image_size = image.size
             calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
-            height = req.height or calculated_height
-            width = req.width or calculated_width
+            height = request.sampling_params.height or calculated_height
+            width = request.sampling_params.width or calculated_width
 
             # Ensure dimensions are multiples of vae_scale_factor * 2
             multiple_of = vae_scale_factor * 2
@@ -81,10 +99,10 @@ def get_qwen_image_edit_pre_process_func(
             width = width // multiple_of * multiple_of
 
             # Store calculated dimensions in request
-            req.calculated_height = calculated_height
-            req.calculated_width = calculated_width
-            req.height = height
-            req.width = width
+            prompt["additional_information"]["calculated_height"] = calculated_height
+            prompt["additional_information"]["calculated_width"] = calculated_width
+            request.sampling_params.height = height
+            request.sampling_params.width = width
 
             # Preprocess image
             if image is not None and not (
@@ -96,10 +114,10 @@ def get_qwen_image_edit_pre_process_func(
                 image = image.unsqueeze(2)
 
                 # Store preprocessed image and prompt image in request
-                req.preprocessed_image = image
-                req.prompt_image = prompt_image
-
-        return requests
+                prompt["additional_information"]["preprocessed_image"] = image
+                prompt["additional_information"]["prompt_image"] = prompt_image
+            request.prompts[i] = prompt
+        return request
 
     return pre_process_func
 
@@ -717,17 +735,27 @@ class QwenImageEditPipeline(
         max_sequence_length: int = 512,
     ) -> DiffusionOutput:
         """Forward pass for image editing."""
-        prompt = req.prompt
-        negative_prompt = req.negative_prompt
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        if len(req.prompts) > 1:
+            logger.warning(
+                """This model only supports a single prompt, not a batched request.""",
+                """Taking only the first image for now.""",
+            )
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+        negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
         # Get preprocessed image from request (pre-processing is done in DiffusionEngine)
-        if hasattr(req, "preprocessed_image"):
-            prompt_image = req.prompt_image
-            image = req.preprocessed_image
-            calculated_height = req.calculated_height
-            calculated_width = req.calculated_width
-            height = req.height
-            width = req.width
+        if not isinstance(first_prompt, str) and "preprocessed_image" in (
+            additional_information := first_prompt.get("additional_information", {})
+        ):
+            prompt_image = additional_information.get("prompt_image")
+            image = additional_information.get("preprocessed_image")
+            calculated_height = additional_information.get("calculated_height")
+            calculated_width = additional_information.get("calculated_width")
+            height = req.sampling_params.height
+            width = req.sampling_params.width
         else:
             # fallback to run pre-processing in pipeline (debug only)
             image_size = image[0].size if isinstance(image, list) else image.size
@@ -745,14 +773,18 @@ class QwenImageEditPipeline(
                 image = self.image_processor.preprocess(image, calculated_height, calculated_width)
                 image = image.unsqueeze(2)
 
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        generator = req.generator or generator
-        true_cfg_scale = req.true_cfg_scale or true_cfg_scale
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        generator = req.sampling_params.generator or generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
 
         # 1. check inputs
         # 2. encode prompts
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
index 4e2fe35f1..52f1d64ec 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
@@ -5,7 +5,7 @@ import json
 import logging
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import PIL.Image
@@ -42,6 +42,7 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -70,24 +71,33 @@ def get_qwen_image_edit_plus_pre_process_func(
     latent_channels = vae_config.get("z_dim", 16)
 
     def pre_process_func(
-        requests: list[OmniDiffusionRequest],
+        request: OmniDiffusionRequest,
     ):
         """Pre-process requests for QwenImageEditPlusPipeline."""
-        for req in requests:
-            image = req.pil_image
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
 
             # Handle single image or list of images
-            if image is None:
+            if raw_image is None:
                 continue
 
-            if not isinstance(image, list):
-                image = [image]
+            if not isinstance(raw_image, list):
+                raw_image = [raw_image]
+            image = [
+                PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im)
+                for im in raw_image
+            ]
 
             # Calculate dimensions based on first image
             image_size = image[0].size
             calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1])
-            height = req.height or calculated_height
-            width = req.width or calculated_width
+            height = request.sampling_params.height or calculated_height
+            width = request.sampling_params.width or calculated_width
 
             # Ensure dimensions are multiples of vae_scale_factor * 2
             multiple_of = vae_scale_factor * 2
@@ -95,10 +105,10 @@ def get_qwen_image_edit_plus_pre_process_func(
             width = width // multiple_of * multiple_of
 
             # Store calculated dimensions in request
-            req.calculated_height = calculated_height
-            req.calculated_width = calculated_width
-            req.height = height
-            req.width = width
+            prompt["additional_information"]["calculated_height"] = calculated_height
+            prompt["additional_information"]["calculated_width"] = calculated_width
+            request.sampling_params.height = height
+            request.sampling_params.width = width
 
             # Preprocess images into condition_images (for prompt encoding) and vae_images (for VAE encoding)
             condition_images = []
@@ -124,12 +134,12 @@ def get_qwen_image_edit_plus_pre_process_func(
                 vae_images.append(image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
 
             # Store preprocessed images in request
-            req.condition_images = condition_images
-            req.vae_images = vae_images
-            req.condition_image_sizes = condition_image_sizes
-            req.vae_image_sizes = vae_image_sizes
-
-        return requests
+            prompt["additional_information"]["condition_images"] = condition_images
+            prompt["additional_information"]["vae_images"] = vae_images
+            prompt["additional_information"]["condition_image_sizes"] = condition_image_sizes
+            prompt["additional_information"]["vae_image_sizes"] = vae_image_sizes
+            request.prompts[i] = prompt
+        return request
 
     return pre_process_func
 
@@ -281,7 +291,7 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput):
 
     def _get_qwen_prompt_embeds(
         self,
-        prompt: str | list[str] = None,
+        prompt: str | list[str],
         image: list[torch.Tensor] | torch.Tensor | None = None,
         dtype: torch.dtype | None = None,
     ):
@@ -655,19 +665,31 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput):
         max_sequence_length: int = 512,
     ) -> DiffusionOutput:
         """Forward pass for image editing with support for multiple images."""
-        prompt = req.prompt
-        negative_prompt = req.negative_prompt
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        if len(req.prompts) > 1:
+            logger.warning(
+                """This model only supports a single prompt, not a batched request.""",
+                """Taking only the first image for now.""",
+            )
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+        negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
         # Get preprocessed images from request (pre-processing is done in DiffusionEngine)
-        if hasattr(req, "vae_images") and hasattr(req, "condition_images"):
-            condition_images = req.condition_images
-            vae_images = req.vae_images
-            condition_image_sizes = req.condition_image_sizes
-            vae_image_sizes = req.vae_image_sizes
-            calculated_height = req.calculated_height
-            calculated_width = req.calculated_width
-            height = req.height
-            width = req.width
+        if (
+            not isinstance(first_prompt, str)
+            and "vae_images" in (additional_information := first_prompt.get("additional_information", {}))
+            and "condition_images" in additional_information
+        ):
+            condition_images = additional_information.get("condition_images")
+            vae_images = additional_information.get("vae_images")
+            condition_image_sizes = additional_information.get("condition_image_sizes")
+            vae_image_sizes = additional_information.get("vae_image_sizes")
+            calculated_height = additional_information.get("calculated_height")
+            calculated_width = additional_information.get("calculated_width")
+            height = req.sampling_params.height
+            width = req.sampling_params.width
         else:
             # fallback to run pre-processing in pipeline (debug only)
             if image is None:
@@ -701,14 +723,18 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput):
                 condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
                 vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
 
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        generator = req.generator or generator
-        true_cfg_scale = req.true_cfg_scale or true_cfg_scale
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        generator = req.sampling_params.generator or generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
 
         # 1. check inputs
         # 2. encode prompts
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
index ae2d039e3..ae27ba708 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import PIL.Image
@@ -38,6 +38,7 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -69,17 +70,35 @@ def get_qwen_image_layered_pre_process_func(
     image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
 
     def pre_process_func(
-        requests: list[OmniDiffusionRequest],
+        request: OmniDiffusionRequest,
     ):
         """Pre-process requests for QwenImageLayeredPipeline."""
-        for req in requests:
-            image = req.pil_image
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
+
+            if raw_image is None or isinstance(raw_image, list):
+                raise ValueError(
+                    """Received no image or a list of image. Only a single image is supported by this model."""
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
+                )
+
+            if isinstance(raw_image, str):
+                image = PIL.Image.open(raw_image)
+            else:
+                image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image)
 
             # 1. calculate dimensions
-            image_size = image[0].size if isinstance(image, list) else image.size
-            assert req.resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {req.resolution}"
+            image_size = image.size
+            assert request.sampling_params.resolution in [640, 1024], (
+                f"resolution must be either 640 or 1024, but got {request.sampling_params.resolution}"
+            )
             calculated_width, calculated_height = calculate_dimensions(
-                req.resolution * req.resolution, image_size[0] / image_size[1]
+                request.sampling_params.resolution * request.sampling_params.resolution, image_size[0] / image_size[1]
             )
             height = calculated_height
             width = calculated_width
@@ -89,10 +108,10 @@ def get_qwen_image_layered_pre_process_func(
             height = height // multiple_of * multiple_of
 
             # Store calculated dimensions in request
-            req.calculated_height = calculated_height
-            req.calculated_width = calculated_width
-            req.height = height
-            req.width = width
+            prompt["additional_information"]["calculated_height"] = calculated_height
+            prompt["additional_information"]["calculated_width"] = calculated_width
+            request.sampling_params.height = height
+            request.sampling_params.width = width
 
             # 2. Preprocess image
             if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels):
@@ -103,10 +122,10 @@ def get_qwen_image_layered_pre_process_func(
                 # image = image.to(dtype=self.text_encoder.dtype)  # do it later
 
                 # Store preprocessed image and prompt image in request
-                req.preprocessed_image = image
-                req.prompt_image = prompt_image
-
-        return requests
+                prompt["additional_information"]["preprocessed_image"] = image
+                prompt["additional_information"]["prompt_image"] = prompt_image
+            request.prompts[i] = prompt
+        return request
 
     return pre_process_func
 
@@ -689,7 +708,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         req: OmniDiffusionRequest,
         image: PIL.Image.Image | torch.Tensor | None = None,
         prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] = None,
+        negative_prompt: str | list[str] | None = None,
         true_cfg_scale: float = 4.0,
         layers: int | None = 4,
         num_inference_steps: int = 50,
@@ -713,29 +732,48 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
 
         # 1. Get preprocessed image from request (pre-processing is done in DiffusionEngine)
         # Override parameters from request if provided
-        prompt = req.prompt if req.prompt is not None else prompt
-        layers = req.layers if req.layers is not None else layers
-        resolution = req.resolution if req.resolution is not None else resolution
-        cfg_normalize = req.cfg_normalize if req.cfg_normalize is not None else cfg_normalize
-        use_en_prompt = req.use_en_prompt if req.use_en_prompt is not None else use_en_prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        generator = req.generator or generator
-        true_cfg_scale = req.true_cfg_scale or true_cfg_scale
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        if len(req.prompts) > 1:
+            logger.warning(
+                """This model only supports a single prompt, not a batched request.""",
+                """Taking only the first image for now.""",
+            )
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+        negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        if hasattr(req, "preprocessed_image"):
-            prompt_image = req.prompt_image
-            image = req.preprocessed_image
+        layers = req.sampling_params.layers if req.sampling_params.layers is not None else layers
+        resolution = req.sampling_params.resolution if req.sampling_params.resolution is not None else resolution
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        cfg_normalize = (
+            req.sampling_params.cfg_normalize if req.sampling_params.cfg_normalize is not None else cfg_normalize
+        )
+        use_en_prompt = (
+            req.sampling_params.use_en_prompt if req.sampling_params.use_en_prompt is not None else use_en_prompt
+        )
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        generator = req.sampling_params.generator or generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
+
+        if not isinstance(first_prompt, str) and "preprocessed_image" in (
+            additional_information := first_prompt.get("additional_information", {})
+        ):
+            prompt_image = additional_information.get("prompt_image")
+            image = additional_information.get("preprocessed_image")
             image = image.to(dtype=self.text_encoder.dtype)  # Now we get the type
-            calculated_height = req.calculated_height
-            calculated_width = req.calculated_width
-            height = req.height
-            width = req.width
+            calculated_height = additional_information.get("calculated_height")
+            calculated_width = additional_information.get("calculated_width")
+            height = req.sampling_params.height
+            width = req.sampling_params.width
         else:
             # fallback to run pre-processing in pipeline (debug only)
             image_size = image[0].size if isinstance(image, list) else image.size
diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py
index dd17af883..34a0eb6c1 100644
--- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py
+++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py
@@ -562,19 +562,24 @@ class StableDiffusion3Pipeline(
         negative_pooled_prompt_embeds: torch.Tensor | None = None,
         max_sequence_length: int = 256,
     ) -> DiffusionOutput:
-        # # TODO: only support single prompt now
-        # if req.prompt is not None:
-        #     prompt = req.prompt[0] if isinstance(req.prompt, list) else req.prompt
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
-        height = req.height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or self.default_sample_size * self.vae_scale_factor
-        sigmas = req.sigmas or sigmas
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        generator = req.generator or generator
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        negative_prompt = [
+            "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts
+        ] or negative_prompt
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        generator = req.sampling_params.generator or generator
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
         # 1. check inputs
         # 2. encode prompts
         # 3. prepare latents and timesteps
@@ -595,7 +600,7 @@ class StableDiffusion3Pipeline(
             max_sequence_length=max_sequence_length,
         )
 
-        self._guidance_scale = req.guidance_scale
+        self._guidance_scale = req.sampling_params.guidance_scale
         self._current_timestep = None
         self._interrupt = False
 
diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py
index c947bd934..c48d68efd 100644
--- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py
+++ b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py
@@ -332,7 +332,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput):
         sample_size: int,
         dtype: torch.dtype,
         device: torch.device,
-        generator: torch.Generator | None,
+        generator: torch.Generator | list[torch.Generator] | None,
         latents: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Prepare initial latent noise."""
@@ -357,7 +357,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput):
         num_inference_steps: int = 100,
         guidance_scale: float = 7.0,
         num_waveforms_per_prompt: int = 1,
-        generator: torch.Generator | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
         latents: torch.Tensor | None = None,
         prompt_embeds: torch.Tensor | None = None,
         negative_prompt_embeds: torch.Tensor | None = None,
@@ -385,20 +385,26 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput):
             DiffusionOutput containing generated audio
         """
         # Extract from request
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
+            negative_prompt = None
+        elif req.prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
 
         if generator is None:
-            generator = req.generator
-        if generator is None and req.seed is not None:
-            generator = torch.Generator(device=self.device).manual_seed(req.seed)
+            generator = req.sampling_params.generator
+        if generator is None and req.sampling_params.seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
 
         # Get audio duration from request extra params or defaults
-        audio_start_in_s = req.extra.get("audio_start_in_s", audio_start_in_s)
-        audio_end_in_s = req.extra.get("audio_end_in_s", audio_end_in_s)
+        audio_start_in_s = req.sampling_params.extra_args.get("audio_start_in_s", audio_start_in_s)
+        audio_end_in_s = req.sampling_params.extra_args.get("audio_end_in_s", audio_end_in_s)
 
         # Calculate audio length
         downsample_ratio = self.vae.hop_length
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
index b691cfedb..712c05400 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
@@ -4,8 +4,10 @@
 from __future__ import annotations
 
 import json
+import logging
 import os
 from collections.abc import Iterable
+from typing import cast
 
 import PIL.Image
 import torch
@@ -21,6 +23,9 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniTextPrompt
+
+logger = logging.getLogger(__name__)
 
 
 def retrieve_latents(
@@ -127,39 +132,54 @@ def get_wan22_pre_process_func(
 
     video_processor = VideoProcessor(vae_scale_factor=8)
 
-    def pre_process_func(requests: list[OmniDiffusionRequest]) -> list[OmniDiffusionRequest]:
-        for req in requests:
-            # Load image if path is provided
-            if req.image_path is not None and req.pil_image is None:
-                req.pil_image = PIL.Image.open(req.image_path).convert("RGB")
+    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
 
-            if req.pil_image is not None:
-                image = req.pil_image
+            if raw_image is None:
+                continue
 
-                # Calculate dimensions based on aspect ratio if not provided
-                if req.height is None or req.width is None:
-                    # Default max area for 720P
-                    max_area = 720 * 1280
-                    aspect_ratio = image.height / image.width
+            if not isinstance(raw_image, (str, PIL.Image.Image)):
+                raise TypeError(
+                    f"""Unsupported image format {raw_image.__class__}.""",
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
+                )
+            image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image
 
-                    # Calculate dimensions maintaining aspect ratio
-                    mod_value = 16  # Must be divisible by 16
-                    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-                    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+            # Calculate dimensions based on aspect ratio if not provided
+            if request.sampling_params.height is None or request.sampling_params.width is None:
+                # Default max area for 720P
+                max_area = 720 * 1280
+                aspect_ratio = image.height / image.width
 
-                    if req.height is None:
-                        req.height = height
-                    if req.width is None:
-                        req.width = width
+                # Calculate dimensions maintaining aspect ratio
+                mod_value = 16  # Must be divisible by 16
+                height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+                width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
 
-                # Resize image to target dimensions
-                image = image.resize((req.width, req.height), PIL.Image.Resampling.LANCZOS)
-                req.pil_image = image
+                if request.sampling_params.height is None:
+                    request.sampling_params.height = height
+                if request.sampling_params.width is None:
+                    request.sampling_params.width = width
 
-                # Preprocess for VAE
-                req.preprocessed_image = video_processor.preprocess(image, height=req.height, width=req.width)
+            # Resize image to target dimensions
+            image = image.resize(
+                (request.sampling_params.width, request.sampling_params.height),  # type: ignore # Above has ensured that width & height are not None
+                PIL.Image.Resampling.LANCZOS,
+            )
+            prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above
 
-        return requests
+            # Preprocess for VAE
+            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
+                image, height=request.sampling_params.height, width=request.sampling_params.width
+            )
+            request.prompts[i] = prompt
+        return request
 
     return pre_process_func
 
@@ -289,20 +309,27 @@ class Wan22Pipeline(nn.Module):
         guidance_scale: float | tuple[float, float] = 4.0,
         frame_num: int = 81,
         output_type: str | None = "np",
-        generator: torch.Generator | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
         prompt_embeds: torch.Tensor | None = None,
         negative_prompt_embeds: torch.Tensor | None = None,
         attention_kwargs: dict | None = None,
         **kwargs,
     ) -> DiffusionOutput:
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+        # Get parameters from request or arguments
+        if len(req.prompts) > 1:
+            raise ValueError(
+                """This model only supports a single prompt, not a batched request.""",
+                """Please pass in a single prompt object or string, or a single-item list.""",
+            )
+        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
+            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
+            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
         if prompt is None and prompt_embeds is None:
             raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
 
-        height = req.height or height
-        width = req.width or width
-        num_frames = req.num_frames if req.num_frames else frame_num
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
 
         # Ensure dimensions are compatible with VAE and patch size
         # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size)
@@ -310,16 +337,16 @@ class Wan22Pipeline(nn.Module):
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
-        num_steps = req.num_inference_steps or num_inference_steps
+        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
 
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
         guidance_high = (
-            req.guidance_scale_2
-            if req.guidance_scale_2 is not None
+            req.sampling_params.guidance_scale_2
+            if req.sampling_params.guidance_scale_2 is not None
             else (
                 guidance_scale[1]
                 if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
@@ -351,9 +378,9 @@ class Wan22Pipeline(nn.Module):
 
         # Seed / generator
         if generator is None:
-            generator = req.generator
-        if generator is None and req.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.seed)
+            generator = req.sampling_params.generator
+        if generator is None and req.sampling_params.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
         # Encode prompts
         if prompt_embeds is None:
@@ -361,8 +388,8 @@ class Wan22Pipeline(nn.Module):
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.num_outputs_per_prompt or 1,
-                max_sequence_length=req.max_sequence_length or 512,
+                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+                max_sequence_length=req.sampling_params.max_sequence_length or 512,
                 device=device,
                 dtype=dtype,
             )
@@ -384,7 +411,22 @@ class Wan22Pipeline(nn.Module):
             boundary_timestep = self.boundary_ratio * self.scheduler.config.num_train_timesteps
 
         # Handle I2V mode when expand_timesteps=True and image is provided
-        image = req.pil_image
+        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        if isinstance(raw_image, list):
+            if len(raw_image) > 1:
+                logger.warning(
+                    """Received a list of image. Only a single image is supported by this model."""
+                    """Taking only the first image for now."""
+                )
+            raw_image = raw_image[0]
+        if raw_image is None:
+            image = None
+        elif isinstance(raw_image, str):
+            image = PIL.Image.open(raw_image)
+        else:
+            image = cast(PIL.Image.Image | torch.Tensor, raw_image)
+
         latent_condition = None
         first_frame_mask = None
 
@@ -415,7 +457,7 @@ class Wan22Pipeline(nn.Module):
                 dtype=torch.float32,
                 device=device,
                 generator=generator,
-                latents=req.latents,
+                latents=req.sampling_params.latents,
             )
 
             # Encode image condition
@@ -457,7 +499,7 @@ class Wan22Pipeline(nn.Module):
                 dtype=torch.float32,
                 device=device,
                 generator=generator,
-                latents=req.latents,
+                latents=req.sampling_params.latents,
             )
 
         if attention_kwargs is None:
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
index ed8c0a987..ea045e5d6 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
@@ -3,8 +3,10 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from collections.abc import Iterable
+from typing import cast
 
 import numpy as np
 import PIL.Image
@@ -26,6 +28,9 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     retrieve_latents,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniTextPrompt
+
+logger = logging.getLogger(__name__)
 
 
 def _load_model_index(model: str, local_files_only: bool) -> dict:
@@ -77,39 +82,56 @@ def get_wan22_i2v_pre_process_func(
 
     video_processor = VideoProcessor(vae_scale_factor=8)
 
-    def pre_process_func(requests: list[OmniDiffusionRequest]) -> list[OmniDiffusionRequest]:
-        for req in requests:
-            # Load image if path is provided
-            if req.image_path is not None and req.pil_image is None:
-                req.pil_image = PIL.Image.open(req.image_path).convert("RGB")
+    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
 
-            if req.pil_image is not None:
-                image = req.pil_image
+            if raw_image is None:
+                raise ValueError(
+                    """No image is provided. This model requires an image to run.""",
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
+                )
+            if not isinstance(raw_image, (str, PIL.Image.Image)):
+                raise TypeError(
+                    f"""Unsupported image format {raw_image.__class__}.""",
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
+                )
+            image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image
 
-                # Calculate dimensions based on aspect ratio if not provided
-                if req.height is None or req.width is None:
-                    # Default max area for 480P
-                    max_area = 480 * 832
-                    aspect_ratio = image.height / image.width
+            # Calculate dimensions based on aspect ratio if not provided
+            if request.sampling_params.height is None or request.sampling_params.width is None:
+                # Default max area for 480P
+                max_area = 480 * 832
+                aspect_ratio = image.height / image.width
 
-                    # Calculate dimensions maintaining aspect ratio
-                    mod_value = 16  # Must be divisible by 16
-                    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-                    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+                # Calculate dimensions maintaining aspect ratio
+                mod_value = 16  # Must be divisible by 16
+                height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+                width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
 
-                    if req.height is None:
-                        req.height = height
-                    if req.width is None:
-                        req.width = width
+                if request.sampling_params.height is None:
+                    request.sampling_params.height = height
+                if request.sampling_params.width is None:
+                    request.sampling_params.width = width
 
-                # Resize image to target dimensions
-                image = image.resize((req.width, req.height), PIL.Image.Resampling.LANCZOS)
-                req.pil_image = image
+            # Resize image to target dimensions
+            image = image.resize(
+                (request.sampling_params.width, request.sampling_params.height),  # type: ignore # Above has ensured that width & height are not None
+                PIL.Image.Resampling.LANCZOS,
+            )
+            prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above
 
-                # Preprocess for VAE
-                req.preprocessed_image = video_processor.preprocess(image, height=req.height, width=req.width)
-
-        return requests
+            # Preprocess for VAE
+            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
+                image, height=request.sampling_params.height, width=request.sampling_params.width
+            )
+            request.prompts[i] = prompt
+        return request
 
     return pre_process_func
 
@@ -266,7 +288,7 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput):
         guidance_scale: float | tuple[float, float] = 5.0,
         frame_num: int = 81,
         output_type: str | None = "np",
-        generator: torch.Generator | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
         prompt_embeds: torch.Tensor | None = None,
         negative_prompt_embeds: torch.Tensor | None = None,
         image_embeds: torch.Tensor | None = None,
@@ -275,31 +297,51 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput):
         **kwargs,
     ) -> DiffusionOutput:
         # Get parameters from request or arguments
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+        if len(req.prompts) > 1:
+            raise ValueError(
+                """This model only supports a single prompt, not a batched request.""",
+                """Please pass in a single prompt object or string, or a single-item list.""",
+            )
+        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
+            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
+            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
         if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 I2V generation.")
+            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
 
         # Get image from request
         if image is None:
-            image = req.pil_image
-        if image is None:
-            raise ValueError("Image is required for I2V generation.")
+            multi_modal_data = (
+                req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+            )
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if raw_image is None:
+                raise ValueError("Image is required for I2V generation.")
+            if isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    logger.warning(
+                        """Received a list of image. Only a single image is supported by this model."""
+                        """Taking only the first image for now."""
+                    )
+                raw_image = raw_image[0]
+            if isinstance(raw_image, str):
+                image = PIL.Image.open(raw_image)
+            else:
+                image = cast(PIL.Image.Image | torch.Tensor, raw_image)
 
-        height = req.height or height
-        width = req.width or width
-        num_frames = req.num_frames if req.num_frames else frame_num
-        num_steps = req.num_inference_steps or num_inference_steps
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_frames = req.sampling_params.num_frames or frame_num
+        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
 
         # Handle guidance scales
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
         guidance_high = (
-            req.guidance_scale_2
-            if req.guidance_scale_2 is not None
+            req.sampling_params.guidance_scale_2
+            if req.sampling_params.guidance_scale_2 is not None
             else (
                 guidance_scale[1]
                 if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
@@ -333,9 +375,9 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput):
 
         # Generator setup
         if generator is None:
-            generator = req.generator
-        if generator is None and req.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.seed)
+            generator = req.sampling_params.generator
+        if generator is None and req.sampling_params.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
         # Encode prompts
         if prompt_embeds is None:
@@ -343,8 +385,8 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput):
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.num_outputs_per_prompt or 1,
-                max_sequence_length=req.max_sequence_length or 512,
+                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+                max_sequence_length=req.sampling_params.max_sequence_length or 512,
                 device=device,
                 dtype=dtype,
             )
@@ -410,7 +452,7 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput):
             dtype=torch.float32,
             device=device,
             generator=generator,
-            latents=req.latents,
+            latents=req.sampling_params.latents,
             last_image=last_image_tensor,
         )
 
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
index b943000cf..5944b678a 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
@@ -16,8 +16,10 @@ The key difference from the MoE-based I2V pipeline is:
 
 from __future__ import annotations
 
+import logging
 import os
 from collections.abc import Iterable
+from typing import cast
 
 import numpy as np
 import PIL.Image
@@ -39,6 +41,9 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     retrieve_latents,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniTextPrompt
+
+logger = logging.getLogger(__name__)
 
 
 def get_wan22_ti2v_post_process_func(
@@ -67,39 +72,56 @@ def get_wan22_ti2v_pre_process_func(
 
     video_processor = VideoProcessor(vae_scale_factor=8)
 
-    def pre_process_func(requests: list[OmniDiffusionRequest]) -> list[OmniDiffusionRequest]:
-        for req in requests:
-            # Load image if path is provided
-            if req.image_path is not None and req.pil_image is None:
-                req.pil_image = PIL.Image.open(req.image_path).convert("RGB")
+    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        for i, prompt in enumerate(request.prompts):
+            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(prompt, str):
+                prompt = OmniTextPrompt(prompt=prompt)
+            if "additional_information" not in prompt:
+                prompt["additional_information"] = {}
 
-            if req.pil_image is not None:
-                image = req.pil_image
+            if raw_image is None:
+                raise ValueError(
+                    """No image is provided. This model requires an image to run.""",
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
+                )
+            if not isinstance(raw_image, (str, PIL.Image.Image)):
+                raise TypeError(
+                    f"""Unsupported image format {raw_image.__class__}.""",
+                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
+                )
+            image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image
 
-                # Calculate dimensions based on aspect ratio if not provided
-                if req.height is None or req.width is None:
-                    # Default max area for 720P (TI2V-5B default)
-                    max_area = 720 * 1280
-                    aspect_ratio = image.height / image.width
+            # Calculate dimensions based on aspect ratio if not provided
+            if request.sampling_params.height is None or request.sampling_params.width is None:
+                # Default max area for 720P (TI2V-5B default)
+                max_area = 720 * 1280
+                aspect_ratio = image.height / image.width
 
-                    # Calculate dimensions maintaining aspect ratio
-                    mod_value = 16  # Must be divisible by 16
-                    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-                    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+                # Calculate dimensions maintaining aspect ratio
+                mod_value = 16  # Must be divisible by 16
+                height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+                width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
 
-                    if req.height is None:
-                        req.height = height
-                    if req.width is None:
-                        req.width = width
+                if request.sampling_params.height is None:
+                    request.sampling_params.height = height
+                if request.sampling_params.width is None:
+                    request.sampling_params.width = width
 
-                # Resize image to target dimensions
-                image = image.resize((req.width, req.height), PIL.Image.Resampling.LANCZOS)
-                req.pil_image = image
+            # Resize image to target dimensions
+            image = image.resize(
+                (request.sampling_params.width, request.sampling_params.height),  # type: ignore # Above has ensured that width & height are not None
+                PIL.Image.Resampling.LANCZOS,
+            )
+            prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above
 
-                # Preprocess for VAE
-                req.preprocessed_image = video_processor.preprocess(image, height=req.height, width=req.width)
-
-        return requests
+            # Preprocess for VAE
+            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
+                image, height=request.sampling_params.height, width=request.sampling_params.width
+            )
+            request.prompts[i] = prompt
+        return request
 
     return pre_process_func
 
@@ -205,31 +227,53 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput):
         guidance_scale: float = 5.0,
         frame_num: int = 81,
         output_type: str | None = "np",
-        generator: torch.Generator | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
         prompt_embeds: torch.Tensor | None = None,
         negative_prompt_embeds: torch.Tensor | None = None,
         attention_kwargs: dict | None = None,
         **kwargs,
     ) -> DiffusionOutput:
         # Get parameters from request or arguments
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+        if len(req.prompts) > 1:
+            raise ValueError(
+                """This model only supports a single prompt, not a batched request.""",
+                """Please pass in a single prompt object or string, or a single-item list.""",
+            )
+        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
+            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
+            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
         if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 TI2V generation.")
+            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
 
         # Get image from request (optional for TI2V)
         if image is None:
-            image = req.pil_image
+            multi_modal_data = (
+                req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+            )
+            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+            if isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    logger.warning(
+                        """Received a list of image. Only a single image is supported by this model."""
+                        """Taking only the first image for now."""
+                    )
+                raw_image = raw_image[0]
+            if raw_image is None:
+                image = None
+            elif isinstance(raw_image, str):
+                image = PIL.Image.open(raw_image)
+            else:
+                image = cast(PIL.Image.Image | torch.Tensor, raw_image)
 
         # Default dimensions for TI2V-5B (720P)
-        height = req.height or height
-        width = req.width or width
-        num_frames = req.num_frames if req.num_frames else frame_num
-        num_steps = req.num_inference_steps or num_inference_steps
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
+        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
-        if req.guidance_scale_provided:
-            guidance_scale = req.guidance_scale
+        if req.sampling_params.guidance_scale_provided:
+            guidance_scale = req.sampling_params.guidance_scale
 
         self._guidance_scale = guidance_scale
 
@@ -254,9 +298,9 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput):
 
         # Generator setup
         if generator is None:
-            generator = req.generator
-        if generator is None and req.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.seed)
+            generator = req.sampling_params.generator
+        if generator is None and req.sampling_params.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
         # Encode prompts
         if prompt_embeds is None:
@@ -264,8 +308,8 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput):
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 do_classifier_free_guidance=guidance_scale > 1.0,
-                num_videos_per_prompt=req.num_outputs_per_prompt or 1,
-                max_sequence_length=req.max_sequence_length or 512,
+                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+                max_sequence_length=req.sampling_params.max_sequence_length or 512,
                 device=device,
                 dtype=dtype,
             )
@@ -310,7 +354,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput):
                 dtype=torch.float32,
                 device=device,
                 generator=generator,
-                latents=req.latents,
+                latents=req.sampling_params.latents,
             )
         else:
             # T2V mode: prepare random latents
@@ -323,7 +367,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput):
                 dtype=torch.float32,
                 device=device,
                 generator=generator,
-                latents=req.latents,
+                latents=req.sampling_params.latents,
             )
             latent_condition = None
             first_frame_mask = torch.ones(
diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
index b92d3de88..62d92900a 100644
--- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
+++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
@@ -316,15 +316,15 @@ class ZImagePipeline(nn.Module):
         self,
         req: OmniDiffusionRequest,
         prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
+        height: int = 1024,
+        width: int = 1024,
         num_inference_steps: int = 50,
         sigmas: list[float] | None = None,
         guidance_scale: float = 5.0,
         cfg_normalization: bool = False,
         cfg_truncation: float = 1.0,
         negative_prompt: str | list[str] | None = None,
-        num_images_per_prompt: int | None = 1,
+        num_images_per_prompt: int = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
         latents: torch.FloatTensor | None = None,
         prompt_embeds: list[torch.FloatTensor] | None = None,
@@ -413,16 +413,28 @@ class ZImagePipeline(nn.Module):
             `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
             generated images.
         """
-        prompt = req.prompt
-        negative_prompt = req.negative_prompt
-        height: int = req.height or 1024
-        width: int = req.width or 1024
-        num_inference_steps = req.num_inference_steps or 50
-        generator = req.generator
-        guidance_scale = req.guidance_scale if req.guidance_rescale is not None else guidance_scale
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
+            negative_prompt = None
+        elif req.prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        generator = req.sampling_params.generator
+        sigmas = req.sampling_params.sigmas or sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_rescale is not None else guidance_scale
+        )
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
 
         vae_scale = self.vae_scale_factor * 2
         if height % vae_scale != 0:
diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py
index 89c0a79f1..a6005290c 100644
--- a/vllm_omni/diffusion/request.py
+++ b/vllm_omni/diffusion/request.py
@@ -2,14 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import pprint
-from dataclasses import asdict, dataclass, field
-from typing import Any
+from dataclasses import dataclass, field
 
-import PIL.Image
-import torch
-
-from vllm_omni.lora.request import LoRARequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 
 @dataclass
@@ -17,9 +12,8 @@ class OmniDiffusionRequest:
     """
     Complete state passed through the pipeline execution.
 
-    This dataclass contains all information needed during the diffusion pipeline
-    execution, allowing methods to update specific components without needing
-    to manage numerous individual parameters.
+    This dataclass contains the prompts and sampling parameters for the diffusion pipeline
+    execution. It also contains a request_id for other components to trace this request and its outputs.
     """
 
     # TODO(will): double check that args are separate from server_args
@@ -27,173 +21,24 @@ class OmniDiffusionRequest:
     # specific arguments.
     # data_type: DataType
 
-    request_id: str | None = None
+    prompts: list[OmniPromptType]  # Actually supporting str-based prompts
+    sampling_params: OmniDiffusionSamplingParams
 
-    generator: torch.Generator | list[torch.Generator] | None = None
-
-    # Image inputs
-    image_path: str | None = None
-    # Image encoder hidden states
-    image_embeds: list[torch.Tensor] = field(default_factory=list)
-    pil_image: torch.Tensor | PIL.Image.Image | None = None
-    pixel_values: torch.Tensor | PIL.Image.Image | None = None
-    preprocessed_image: torch.Tensor | None = None
-
-    # Text inputs
-    prompt: str | list[str] | None = None
-    negative_prompt: str | list[str] | None = None
-    prompt_path: str | None = None
-    output_path: str = "outputs/"
-    # without extension
-    output_file_name: str | None = None
-    output_file_ext: str | None = None
-    # Primary encoder embeddings
-    prompt_embeds: list[torch.Tensor] | torch.Tensor = field(default_factory=list)
-    negative_prompt_embeds: list[torch.Tensor] | None = None
-    prompt_attention_mask: list[torch.Tensor] | None = None
-    negative_attention_mask: list[torch.Tensor] | None = None
-    clip_embedding_pos: list[torch.Tensor] | None = None
-    clip_embedding_neg: list[torch.Tensor] | None = None
-
-    pooled_embeds: list[torch.Tensor] = field(default_factory=list)
-    neg_pooled_embeds: list[torch.Tensor] = field(default_factory=list)
-
-    # Additional text-related parameters
-    max_sequence_length: int | None = None
-    prompt_template: dict[str, Any] | None = None
-    do_classifier_free_guidance: bool = False
-
-    # Batch info
-    num_outputs_per_prompt: int = 1
-    seed: int | None = None
-    seeds: list[int] | None = None
-
-    # layered info
-    layers: int = 4
-
-    # cfg info
-    cfg_normalize: bool = False
-
-    # caption language
-    use_en_prompt: bool = False
-
-    # different bucket in (640, 1024) to determine the condition and output resolution
-    resolution: int = 640
-
-    # Tracking if embeddings are already processed
-    is_prompt_processed: bool = False
-
-    # Latent tensors
-    latents: torch.Tensor | None = None
-    raw_latent_shape: torch.Tensor | None = None
-    noise_pred: torch.Tensor | None = None
-    image_latent: torch.Tensor | None = None
-
-    # Latent dimensions
-    height_latents: list[int] | int | None = None
-    width_latents: list[int] | int | None = None
-    num_frames: list[int] | int = 1  # Default for image models
-    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
-
-    # Original dimensions (before VAE scaling)
-    height: list[int] | int | None = None
-    width: list[int] | int | None = None
-    fps: list[int] | int | None = None
-    height_not_provided: bool = False
-    width_not_provided: bool = False
-
-    # Timesteps
-    timesteps: torch.Tensor | None = None
-    timestep: torch.Tensor | float | int | None = None
-    step_index: int | None = None
-    boundary_ratio: float | None = None
-
-    # Scheduler parameters
-    num_inference_steps: int = 50
-    guidance_scale: float = 1.0
-    guidance_scale_provided: bool = False
-    guidance_scale_2: float | None = None
-    guidance_rescale: float = 0.0
-    eta: float = 0.0
-    sigmas: list[float] | None = None
-
-    true_cfg_scale: float | None = None  # qwen-image specific now
-
-    n_tokens: int | None = None
-
-    # Other parameters that may be needed by specific schedulers
-    extra_step_kwargs: dict[str, Any] = field(default_factory=dict)
-
-    # [Omni] KV Cache Transfer, for bagel model now
-    past_key_values: Any | None = None  # Injected KV Cache
-    kv_metadata: dict[str, Any] | None = None  # Metadata for KV Cache (e.g., kv_lens, ropes)
-    need_kv_receive: bool = True  # Flag to indicate if this request expects KV transfer
-
-    # Component modules (populated by the pipeline)
-    modules: dict[str, Any] = field(default_factory=dict)
-
-    return_trajectory_latents: bool = False
-    return_trajectory_decoded: bool = False
-    trajectory_timesteps: list[torch.Tensor] | None = None
-    trajectory_latents: torch.Tensor | None = None
-
-    # Extra parameters that might be needed by specific pipeline implementations
-    extra: dict[str, Any] = field(default_factory=dict)
-
-    # Misc
-    save_output: bool = True
-    return_frames: bool = False
-
-    # LoRA
-    lora_request: LoRARequest | None = None
-    lora_scale: float = 1.0
-
-    # STA parameters
-    STA_param: list | None = None
-    is_cfg_negative: bool = False
-    mask_search_final_result_pos: list[list] | None = None
-    mask_search_final_result_neg: list[list] | None = None
-
-    # VSA parameters
-    VSA_sparsity: float = 0.0
-    # perf_logger: PerformanceLogger | None = None
-
-    # stage logging
-    # logging_info: PipelineLoggingInfo = field(default_factory=PipelineLoggingInfo)
-
-    # profile
-    profile: bool = False
-    num_profiled_timesteps: int = 8
-
-    # debugging
-    debug: bool = False
-
-    # results
-    output: torch.Tensor | None = None
-
-    @property
-    def batch_size(self):
-        # Determine batch size
-        if isinstance(self.prompt, list):
-            batch_size = len(self.prompt)
-        elif self.prompt is not None:
-            batch_size = 1
-        else:
-            batch_size = self.prompt_embeds[0].shape[0]
-
-        # Adjust batch size for number of videos per prompt
-        batch_size *= self.num_outputs_per_prompt
-        return batch_size
+    request_ids: list[str] = field(default_factory=list)
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
         # Set do_classifier_free_guidance based on guidance scale and negative prompt
-        if self.guidance_scale > 1.0 and self.negative_prompt is not None:
-            self.do_classifier_free_guidance = True
-        if self.negative_prompt_embeds is None:
-            self.negative_prompt_embeds = []
-        if self.guidance_scale_2 is None:
-            self.guidance_scale_2 = self.guidance_scale
+        if self.sampling_params.guidance_scale > 1.0 and any(
+            (not isinstance(p, str) and p.get("negative_prompt")) for p in self.prompts
+        ):
+            self.sampling_params.do_classifier_free_guidance = True
+        if self.sampling_params.guidance_scale_2 is None:
+            self.sampling_params.guidance_scale_2 = self.sampling_params.guidance_scale
 
-    def __str__(self):
-        return pprint.pformat(asdict(self), indent=2, width=120)
+        # The dataclass default value is 0 (false-like), used to detect whether user explicitly provides this value
+        # After this check is done, reset this value to old default 1
+        if self.sampling_params.guidance_scale:
+            self.sampling_params.guidance_scale_provided = True
+        else:
+            self.sampling_params.guidance_scale = 1.0
diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py
index 21b0d4914..f104d0522 100644
--- a/vllm_omni/diffusion/scheduler.py
+++ b/vllm_omni/diffusion/scheduler.py
@@ -41,14 +41,14 @@ class Scheduler:
     def get_broadcast_handle(self):
         return self.mq.export_handle()
 
-    def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
+    def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         """Sends a request to the scheduler and waits for the response."""
         try:
             # Prepare RPC request for generation
             rpc_request = {
                 "type": "rpc",
                 "method": "generate",
-                "args": (requests,),
+                "args": (request,),
                 "kwargs": {},
                 "output_rank": 0,
                 "exec_all_ranks": True,
diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py b/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
index 54327be81..bf56df155 100644
--- a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
@@ -173,12 +173,13 @@ class GPUDiffusionModelRunner:
     def _receive_kv_cache_for_request(self, req: OmniDiffusionRequest) -> None:
         """Receive KV cache for a request via OmniConnector."""
         # TODO(wzliu)! must get control info from stage queue instead of hardcode
-        if not req.request_id:
+        if not req.request_ids:
             logger.warning("Request has no ID, cannot receive KV cache")
             return
+        request_id = req.request_ids[0]
 
         try:
-            logger.info(f"Attempting to receive KV cache for request {req.request_id}")
+            logger.info(f"Attempting to receive KV cache for request {request_id}")
 
             # TODO: Key used for transfer (must match sender side)
             # key = f"kv_cache_{req.request_id}"
@@ -197,7 +198,7 @@ class GPUDiffusionModelRunner:
                 from_stage = stage_id - 1
             else:
                 raise ValueError("Invalid stage id")
-            logger.info(f"Wait for KV cache for request {req.request_id} from stage {from_stage} to {to_stage}...")
+            logger.info(f"Wait for KV cache for request {request_id} from stage {from_stage} to {to_stage}...")
 
             # Check if we should receive KV cache based on config
             need_recv_cache = omni_kv_config.get("need_recv_cache", False)
@@ -207,7 +208,7 @@ class GPUDiffusionModelRunner:
                 start_time = time.time()
 
                 while True:
-                    get_key = f"omni_{from_stage}_to_{to_stage}_kv_cache_{req.request_id}"
+                    get_key = f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"
                     result = self.connector.get(
                         from_stage=from_stage,
                         to_stage=to_stage,
@@ -217,18 +218,18 @@ class GPUDiffusionModelRunner:
                         break
 
                     if time.time() - start_time > timeout:
-                        logger.error(f"Timeout waiting for KV cache for request {req.request_id} after {timeout}s")
+                        logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s")
                         result = None
                         break
 
                     time.sleep(0.5)
             else:
-                logger.info(f"Skip receiving KV cache for {req.request_id} (need_recv_cache=False)")
+                logger.info(f"Skip receiving KV cache for {request_id} (need_recv_cache=False)")
                 result = None
 
             if result:
                 data, size = result
-                logger.info(f"Successfully received KV cache for {req.request_id}")
+                logger.info(f"Successfully received KV cache for {request_id}")
 
                 # Assume data structure matches KVCacheTransferData.to_dict()
                 if isinstance(data, dict) and "layer_blocks" in data:
@@ -242,16 +243,16 @@ class GPUDiffusionModelRunner:
                                 cache_list[i] = tensor.to(self.pipeline.device).contiguous()
                     from types import SimpleNamespace
 
-                    req.past_key_values = SimpleNamespace(**layer_blocks)
+                    req.sampling_params.past_key_values = SimpleNamespace(**layer_blocks)
 
                 if "metadata" in data:
-                    req.kv_metadata = data["metadata"]
+                    req.sampling_params.kv_metadata = data["metadata"]
 
             else:
-                logger.warning(f"No KV cache received for {req.request_id} (timeout or empty)")
+                logger.warning(f"No KV cache received for {request_id} (timeout or empty)")
 
         except Exception as e:
-            logger.error(f"Error receiving KV cache for {req.request_id}: {e}")
+            logger.error(f"Error receiving KV cache for {request_id}: {e}")
             import traceback
 
             traceback.print_exc()
@@ -261,33 +262,30 @@ class GPUDiffusionModelRunner:
         return self.pipeline.load_weights(weights)
 
     @torch.inference_mode()
-    def execute_model(self, reqs: list[OmniDiffusionRequest]) -> DiffusionOutput:
+    def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
 
         Args:
-            reqs: List of diffusion requests to process.
+            req: A diffusion request containing a list of prompts to process.
 
         Returns:
             DiffusionOutput with generated results.
         """
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
-        if not reqs or len(reqs) == 0:
+        if len(req.prompts) == 0:
             raise ValueError("Cannot execute model with empty request list")
 
-        # TODO: dealing with first req for now
-        req = reqs[0]
-
         # [Omni] KV Cache Receiving Logic
-        if getattr(req, "need_kv_receive", False) and self.connector is not None:
+        if req.sampling_params.need_kv_receive and self.connector is not None:
             self._receive_kv_cache_for_request(req)
 
-        if req.generator is None and req.seed is not None:
-            req.generator = torch.Generator(device=self.device).manual_seed(req.seed)
+        if req.sampling_params.generator is None and req.sampling_params.seed is not None:
+            req.sampling_params.generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
 
         # Refresh cache context if needed
         if self.cache_backend is not None and self.cache_backend.is_enabled():
-            self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
+            self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps)
 
         with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
             with record_function("pipeline_forward"):
diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py b/vllm_omni/diffusion/worker/gpu_diffusion_worker.py
index cadb6b64d..d8b962a58 100644
--- a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/gpu_diffusion_worker.py
@@ -124,9 +124,9 @@ class GPUDiffusionWorker:
         )
         logger.info(f"Worker {self.rank}: Initialization complete.")
 
-    def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
+    def generate(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         """Generate output for the given requests."""
-        return self.execute_model(requests, self.od_config)
+        return self.execute_model(request, self.od_config)
 
     @classmethod
     def start_profile(cls, trace_path_template: str) -> str:
@@ -138,37 +138,17 @@ class GPUDiffusionWorker:
         """Stop profiling and return the result dictionary."""
         return CurrentProfiler.stop()
 
-    def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
+    def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput:
         """Execute a forward pass by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"
-        if self.lora_manager is not None and reqs:
-            req = reqs[0]
-
-            if len(reqs) > 1:
-                # This worker (and the current diffusion model runner) applies
-                # a single LoRA to the whole batch. Reject inconsistent LoRA
-                # settings to avoid silently applying the wrong adapter.
-                def _lora_key(r: OmniDiffusionRequest):
-                    if r.lora_request is None:
-                        return None
-                    lr = r.lora_request
-                    return (lr.lora_name, lr.lora_int_id, lr.lora_path, lr.tensorizer_config_dict)
-
-                key0 = _lora_key(req)
-                scale0 = req.lora_scale if key0 is not None else None
-                for other in reqs[1:]:
-                    if _lora_key(other) != key0:
-                        raise ValueError("All requests in a diffusion batch must share the same LoRARequest.")
-                    if key0 is not None and other.lora_scale != scale0:
-                        raise ValueError("All requests in a diffusion batch must share the same lora_scale.")
-
+        if self.lora_manager is not None:
             try:
-                self.lora_manager.set_active_adapter(req.lora_request, req.lora_scale)
+                self.lora_manager.set_active_adapter(req.sampling_params.lora_request, req.sampling_params.lora_scale)
             except Exception as exc:
-                if req.lora_request is not None:
+                if req.sampling_params.lora_request is not None:
                     raise
                 logger.warning("LoRA activation skipped: %s", exc)
-        return self.model_runner.execute_model(reqs)
+        return self.model_runner.execute_model(req)
 
     def load_weights(self, weights) -> set[str]:
         """Load weights by delegating to the model runner."""
@@ -360,7 +340,7 @@ class WorkerProc:
                 except Exception as e:
                     logger.error(f"Error processing RPC: {e}", exc_info=True)
                     if self.result_mq is not None:
-                        self.return_result({"status": "error", "error": str(e)})
+                        self.return_result(DiffusionOutput(error=str(e)))
 
             elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                 logger.info("Worker %s: Received shutdown message", self.gpu_id)
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 3dd687a8b..a78710fa2 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -3,10 +3,10 @@
 import asyncio
 import time
 import weakref
-from collections.abc import AsyncGenerator, Iterable
+from collections.abc import AsyncGenerator, Iterable, Sequence
 from dataclasses import asdict
 from pprint import pformat
-from typing import Any
+from typing import Any, cast
 
 from vllm.config import VllmConfig
 from vllm.inputs.preprocess import InputPreprocessor
@@ -32,6 +32,7 @@ from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
 )
+from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams, OmniTokensPrompt
 
 # Internal imports (our code)
 from vllm_omni.lora.request import LoRARequest
@@ -66,10 +67,8 @@ class AsyncOmni(OmniBase):
     asynchronous LLM and Diffusion models.
 
     Args:
-        *args: Variable length argument list.
-            - args[0]: Model name or path to load.
+        model: Model name or path to load.
         **kwargs: Arbitrary keyword arguments.
-            - model: Model name or path to load (if not in args).
             - stage_configs_path: Optional path to YAML file containing stage
               configurations. If None, configurations are loaded from the model.
             - log_stats: Whether to enable statistics logging
@@ -95,7 +94,7 @@ class AsyncOmni(OmniBase):
         ...     print(output)
     """
 
-    def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None:
+    def __init__(self, model: str, **kwargs: dict[str, Any]) -> None:
         # Pause/resume control attributes
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
@@ -104,7 +103,7 @@ class AsyncOmni(OmniBase):
         self.request_states: dict[str, ClientRequestState] = {}
         self.output_handler: asyncio.Task | None = None
 
-        super().__init__(*args, **kwargs)
+        super().__init__(model, **kwargs)
 
         # Register weak reference cleanup (called on garbage collection)
         self._weak_finalizer = weakref.finalize(
@@ -233,7 +232,14 @@ class AsyncOmni(OmniBase):
         if hasattr(self, "_weak_finalizer"):
             self._weak_finalizer()
 
-    async def generate(self, *args: Any, **kwargs: dict[str, Any]) -> AsyncGenerator[OmniRequestOutput, None]:
+    async def generate(
+        self,
+        prompt: OmniPromptType,
+        request_id: str,
+        sampling_params_list: Sequence[OmniSamplingParams] | None = None,
+        *,
+        output_modalities: list[str] | None = None,
+    ) -> AsyncGenerator[OmniRequestOutput, None]:
         """Generate outputs for the given prompt asynchronously.
 
         Coordinates multi-stage pipeline through YAML configuration.
@@ -243,21 +249,13 @@ class AsyncOmni(OmniBase):
         sampling parameters from the sampling_params_list.
 
         Args:
-            *args: Arguments for generation.
-                - prompt: Prompt to process. Can be a text string, token IDs,
-                    or multimodal prompt.
-                - request_id: Unique identifier for this request
-                - sampling_params_list: List of SamplingParams, one for each stage.
-                    Must have the same length as the number of stages.
-                    If None, uses default sampling params for each stage.
-            **kwargs: Additional arguments for generation.
-                - prompt: Prompt to process. Can be a text string, token IDs,
-                    or multimodal prompt.
-                - request_id: Unique identifier for this request
-                - sampling_params_list: List of SamplingParams, one for each stage.
-                    Must have the same length as the number of stages.
-                    If None, uses default sampling params for each stage.
-                - output_modalities: Optional list of output modalities.
+            prompt: Prompt to process. Can be a text string, token IDs,
+                or multimodal prompt.
+            request_id: Unique identifier for this request
+            sampling_params_list: List of SamplingParams, one for each stage.
+                Must have the same length as the number of stages.
+                If None, uses default sampling params for each stage.
+            output_modalities: Optional list of output modalities.
 
         Yields:
             OmniRequestOutput objects as they are produced by each stage.
@@ -276,33 +274,9 @@ class AsyncOmni(OmniBase):
             # Start output handler on the first call to generate()
             self._run_output_handler()
 
-            prompt = args[0] if args else kwargs.get("prompt")
-            request_id = args[1] if len(args) > 1 else kwargs.get("request_id")
-            sampling_params_list = args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
-            output_modalities = kwargs.get("output_modalities", None)
             # TODO: lora_request, trace_headers, priority are not supported yet
-
             if sampling_params_list is None:
-                # For Omni LLM, the params are parsed via the yaml file. For the current version,
-                # diffusion params can parsed via the command line.
-                omni_params_kwargs = {
-                    k: v for k, v in kwargs.items() if k not in ["prompt", "request_id", "output_modalities"]
-                }
-
-                per_stage_params: list[Any] = []
-                for stage_id, stage in enumerate(self.stage_list):
-                    stage_type = getattr(stage, "stage_type", "llm")
-                    if stage_type == "diffusion":
-                        default_dict = self.default_sampling_params_list[stage_id]
-                        # Merge user-provided kwargs
-                        merged = {**default_dict, **omni_params_kwargs}
-                        # Diffusion only needs to keep diff params, will be used via OmniDiffusionRequest
-                        per_stage_params.append(merged)
-                    else:
-                        # LLM directly constructs SamplingParams, don't use the merged params
-                        per_stage_params.append(self.default_sampling_params_list[stage_id])
-
-                sampling_params_list = per_stage_params
+                sampling_params_list = self.default_sampling_params_list
 
             if len(sampling_params_list) != len(self.stage_list):
                 raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
@@ -334,11 +308,11 @@ class AsyncOmni(OmniBase):
                 stage_queues = {stage_id: asyncio.Queue() for stage_id in range(num_stages)}
                 req_state.stage_queues = stage_queues
                 for i in range(num_stages):
-                    sp: SamplingParams = sampling_params_list[i]
-                    engine_inputs = prompt
+                    sp: SamplingParams = cast(SamplingParams, sampling_params_list[i])
+                    engine_inputs = cast(OmniTokensPrompt, prompt)
                     if i != 0:
-                        prompt_token_ids = prompt["prompt_token_ids"]
-                        prompt_1 = prompt.copy()
+                        prompt_token_ids = engine_inputs["prompt_token_ids"]
+                        prompt_1 = engine_inputs.copy()
                         prompt_1["prompt_token_ids"] = [0] * compute_talker_prompt_ids_length(prompt_token_ids)
                         prompt_1["multi_modal_data"] = prompt_1["mm_processor_kwargs"] = None
                         engine_inputs = prompt_1
diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py
index 6d34e3445..535f04f7d 100644
--- a/vllm_omni/entrypoints/async_omni_diffusion.py
+++ b/vllm_omni/entrypoints/async_omni_diffusion.py
@@ -12,16 +12,15 @@ import asyncio
 import uuid
 from collections.abc import AsyncGenerator, Iterable
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import fields
 from typing import Any
 
-from PIL import Image
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -109,68 +108,19 @@ class AsyncOmniDiffusion:
 
         logger.info("AsyncOmniDiffusion initialized with model: %s", model)
 
-    def _prepare_request(
-        self,
-        prompt: str,
-        request_id: str | None = None,
-        **kwargs: Any,
-    ) -> OmniDiffusionRequest:
-        """Prepare a diffusion request from prompt and parameters.
-
-        Args:
-            prompt: Text prompt for image generation
-            request_id: Optional unique identifier for the request
-            **kwargs: Additional generation parameters
-
-        Returns:
-            OmniDiffusionRequest ready for processing
-        """
-        if request_id is None:
-            request_id = f"diff-{uuid.uuid4().hex[:16]}"
-
-        field_names = {f.name for f in fields(OmniDiffusionRequest)}
-
-        init_kwargs = {
-            "prompt": prompt,
-            "request_id": request_id,
-        }
-
-        for key, value in kwargs.items():
-            if key in field_names:
-                init_kwargs[key] = value
-
-        if "guidance_scale" in kwargs:
-            init_kwargs["guidance_scale_provided"] = True
-
-        return OmniDiffusionRequest(**init_kwargs)
-
     async def generate(
         self,
-        prompt: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
         request_id: str | None = None,
-        num_inference_steps: int = 50,
-        guidance_scale: float | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        negative_prompt: str | None = None,
-        num_outputs_per_prompt: int = 1,
-        seed: int | None = None,
-        lora_request=None,
-        **kwargs: Any,
+        lora_request: LoRARequest | None = None,
     ) -> OmniRequestOutput:
         """Generate images asynchronously from a text prompt.
 
         Args:
             prompt: Text prompt describing the desired image
+            sampling_params: Sampling parameters
             request_id: Optional unique identifier for tracking the request
-            num_inference_steps: Number of denoising steps (default: 50)
-            guidance_scale: Classifier-free guidance scale (optional, uses model defaults if omitted)
-            height: Optional image height in pixels
-            width: Optional image width in pixels
-            negative_prompt: Optional negative prompt for guidance
-            num_outputs_per_prompt: Number of images to generate (default: 1)
-            seed: Optional random seed for reproducibility
-            **kwargs: Additional generation parameters
 
         Returns:
             OmniRequestOutput containing generated images
@@ -181,64 +131,38 @@ class AsyncOmniDiffusion:
         if request_id is None:
             request_id = f"diff-{uuid.uuid4().hex[:16]}"
 
-        # Prepare request
-        request_kwargs = {
-            "prompt": prompt,
-            "request_id": request_id,
-            "num_inference_steps": num_inference_steps,
-            "height": height,
-            "width": width,
-            "negative_prompt": negative_prompt,
-            "num_outputs_per_prompt": num_outputs_per_prompt,
-            "seed": seed,
-            "lora_request": lora_request,
-            **kwargs,
-        }
-        if guidance_scale is not None:
-            request_kwargs["guidance_scale"] = guidance_scale
+        if sampling_params.guidance_scale:
+            sampling_params.guidance_scale_provided = True
 
-        request = self._prepare_request(**request_kwargs)
+        if lora_request is not None:
+            sampling_params.lora_request = lora_request
+
+        request = OmniDiffusionRequest(
+            prompts=[prompt],
+            sampling_params=sampling_params,
+            request_ids=[request_id],
+        )
 
         logger.debug("Starting generation for request %s", request_id)
 
         # Run engine in thread pool
         loop = asyncio.get_event_loop()
         try:
+            # In async mode, only a single request is submitted at a time
             result = await loop.run_in_executor(
                 self._executor,
                 self.engine.step,
-                [request],
+                request,
             )
+            result = result[0]
         except Exception as e:
             logger.error("Generation failed for request %s: %s", request_id, e)
             raise RuntimeError(f"Diffusion generation failed: {e}") from e
 
-        # Check if result is already OmniRequestOutput
-        if isinstance(result, OmniRequestOutput):
-            # Update request_id if needed
-            if not result.request_id:
-                result.request_id = request_id
-            return result
-
-        # Process results if not OmniRequestOutput
-        images: list[Image.Image] = []
-        if result is not None:
-            if isinstance(result, list):
-                for item in result:
-                    if isinstance(item, Image.Image):
-                        images.append(item)
-            elif isinstance(result, Image.Image):
-                images.append(result)
-
-        return OmniRequestOutput.from_diffusion(
-            request_id=request_id,
-            images=images,
-            prompt=prompt,
-            metrics={
-                "num_inference_steps": num_inference_steps,
-                "guidance_scale": request.guidance_scale,
-            },
-        )
+        # Update request_id if needed
+        if not result.request_id:
+            result.request_id = request_id
+        return result
 
     async def generate_stream(
         self,
diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py
index 90b344c6f..97357dc3b 100644
--- a/vllm_omni/entrypoints/omni.py
+++ b/vllm_omni/entrypoints/omni.py
@@ -10,14 +10,13 @@ from collections.abc import Callable, Generator, Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import asdict
 from pprint import pformat
-from typing import Any
+from typing import Any, Literal, overload
 
 from omegaconf import OmegaConf
 from tqdm.auto import tqdm
-from vllm.inputs import PromptType
+from vllm import SamplingParams
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.distributed.omni_connectors import (
     get_stage_connector_config,
     initialize_orchestrator_connectors,
@@ -42,6 +41,7 @@ from vllm_omni.entrypoints.utils import (
     load_stage_configs_from_yaml,
     resolve_model_config_path,
 )
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -82,10 +82,8 @@ class OmniBase:
     """Base class for serving Omni models.
 
     Args:
-        *args: Variable length argument list.
-            - args[0]: Model name or path to load.
+        model: Model name or path to load.
         **kwargs: Arbitrary keyword arguments.
-            - model: Model name or path to load (if not in args).
             - stage_configs_path: Optional path to YAML file containing stage
               configurations. If None, configurations are loaded from the model.
             - log_stats: Whether to enable statistics logging
@@ -102,14 +100,9 @@ class OmniBase:
             - Additional keyword arguments passed to stage engines.
     """
 
-    def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None:
-        model = args[0] if args else kwargs.get("model", "")
-        assert model != "", "Null model id detected, please specify a model id."
+    def __init__(self, model: str, **kwargs: Any) -> None:
         model = omni_snapshot_download(model)
-        if args:
-            args[0] = model
-        elif kwargs.get("model", "") != "":
-            kwargs["model"] = model
+        kwargs["model"] = model
 
         # Stage management attributes
         self.stage_list: list[OmniStage] = []
@@ -502,10 +495,8 @@ class Omni(OmniBase):
     """Unified entrypoint for both LLM and Diffusion models for better usability.
 
     Args:
-        *args: Variable length argument list.
-            - args[0]: Model name or path to load.
+        model: Model name or path to load.
         **kwargs: Arbitrary keyword arguments.
-            - model: Model name or path to load (if not in args).
             - stage_configs_path: Optional path to YAML file containing stage
               configurations. If None, configurations are loaded from the model.
             - log_stats: Whether to enable statistics logging
@@ -527,8 +518,8 @@ class Omni(OmniBase):
         >>> print(outputs)
     """
 
-    def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(self, model: str, **kwargs: Any) -> None:
+        super().__init__(model, **kwargs)
 
         # Register weak reference cleanup (called on garbage collection)
         self._weak_finalizer = weakref.finalize(
@@ -539,8 +530,31 @@ class Omni(OmniBase):
             self._ray_pg,
         )
 
+    @overload
     def generate(
-        self, *args: Any, **kwargs: dict[str, Any]
+        self,
+        prompts: OmniPromptType | Sequence[OmniPromptType],
+        sampling_params_list: OmniSamplingParams | Sequence[OmniSamplingParams] | None = None,
+        *,
+        py_generator: Literal[True],
+    ) -> Generator[OmniRequestOutput, None, None]: ...
+
+    @overload
+    def generate(
+        self,
+        prompts: OmniPromptType | Sequence[OmniPromptType],
+        sampling_params_list: OmniSamplingParams | Sequence[OmniSamplingParams] | None = None,
+        *,
+        py_generator: Literal[False] = False,
+    ) -> list[OmniRequestOutput]: ...
+
+    def generate(
+        self,
+        prompts: OmniPromptType | Sequence[OmniPromptType],
+        sampling_params_list: OmniSamplingParams | Sequence[OmniSamplingParams] | None = None,
+        *,
+        py_generator: bool = False,
+        use_tqdm: bool | Callable[..., tqdm] = True,
     ) -> Generator[OmniRequestOutput, None, None] | list[OmniRequestOutput]:
         """Generate outputs for the given prompts.
 
@@ -548,12 +562,10 @@ class Omni(OmniBase):
         Each stage will use OmniLLM or OmniDiffusion based on stage_type.
 
         Args:
-            *args: Variable length argument list.
-                - args[0]: Input prompts for generation.
-                - args[1]: Optional list of per-stage parameters.
-            **kwargs: Arbitrary keyword arguments.
-                - prompt: Input prompts for generation (if not in args).
-                - sampling_params_list: Optional list of per-stage parameters (if not in args).
+            prompts: Input prompt(s) for generation.
+            sampling_params_list: Optional list of per-stage parameters.
+            py_generator: Whether the returned result(s) are wrapped in a generator instead of a list.
+            use_tqdm: Whether to use tqdm progress bar
 
         Returns:
             List of OmniRequestOutput objects, one for each input prompt.
@@ -563,40 +575,26 @@ class Omni(OmniBase):
         Raises:
             ValueError: If sampling_params_list is None or has incorrect length.
         """
-        prompts = args[0] if args else kwargs.get("prompts")
-        sampling_params_list = args[1] if len(args) > 1 else kwargs.get("sampling_params_list")
-        py_generator = kwargs.get("py_generator", False)
-        if prompts is None:
-            if kwargs.get("prompt") is None:
-                raise ValueError("prompts is required for generation")
-            prompts = kwargs.get("prompt")
-
         if sampling_params_list is None:
-            # For Omni LLM, the params are parsed via the yaml file. For the current version,
-            # diffusion params can parsed via the command line.
-            omni_params_kwargs = {
-                k: v for k, v in kwargs.items() if k not in ["prompt", "request_id", "output_modalities"]
-            }
-
-            per_stage_params: list[Any] = []
-            for stage_id, stage in enumerate(self.stage_list):
-                stage_type = getattr(stage, "stage_type", "llm")
-                if stage_type == "diffusion":
-                    default_dict = self.default_sampling_params_list[stage_id]
-                    # Merge user-provided kwargs
-                    merged = {**default_dict, **omni_params_kwargs}
-                    # Diffusion only needs to keep diff params, will be used via OmniDiffusionRequest
-                    per_stage_params.append(merged)
+            sampling_params_list = self.default_sampling_params_list
+        elif not isinstance(sampling_params_list, Sequence):
+            # TODO: After the recent introduction of BAGEL model (one LLM and one Diffusion),
+            # expect the text_to_image example code to run when only passing one OmniDiffusionSamplingParams
+            # This behavior may be confusing, and future PR can improve it.
+            per_stage_params: list[OmniSamplingParams] = []
+            for default_stage_sp in self.default_sampling_params_list:
+                default_sp_type = default_stage_sp.__class__
+                if default_sp_type == sampling_params_list.__class__:
+                    per_stage_params.append(sampling_params_list)
                 else:
-                    # LLM directly constructs SamplingParams, don't use the merged params
-                    per_stage_params.append(self.default_sampling_params_list[stage_id])
-
+                    per_stage_params.append(default_stage_sp)
             sampling_params_list = per_stage_params
+
         try:
             if py_generator:
                 return self._run_generation_with_generator(prompts, sampling_params_list)
             else:
-                outputs = list(self._run_generation(prompts, sampling_params_list))
+                outputs = list(self._run_generation(prompts, sampling_params_list, use_tqdm))
                 return outputs
         except Exception as e:
             logger.exception("[Orchestrator] Failed to run generation: %s", e)
@@ -606,8 +604,8 @@ class Omni(OmniBase):
 
     def _run_generation_with_generator(
         self,
-        prompts: PromptType | Sequence[PromptType] | OmniDiffusionRequest | Sequence[OmniDiffusionRequest],
-        sampling_params_list: Any | Sequence[Any] | None,
+        prompts: OmniPromptType | Sequence[OmniPromptType],
+        sampling_params_list: Sequence[OmniSamplingParams],
     ) -> Generator[OmniRequestOutput, None, None]:
         """Run generation through all stages in the pipeline and return a generator."""
         gen = self._run_generation(prompts, sampling_params_list)
@@ -622,8 +620,8 @@ class Omni(OmniBase):
 
     def _run_generation(
         self,
-        prompts: PromptType | Sequence[PromptType] | OmniDiffusionRequest | Sequence[OmniDiffusionRequest],
-        sampling_params_list: Any | Sequence[Any] | None = None,
+        prompts: OmniPromptType | Sequence[OmniPromptType],
+        sampling_params_list: Sequence[OmniSamplingParams],
         use_tqdm: bool | Callable[..., tqdm] = True,
     ) -> Generator[OmniRequestOutput, None, None]:
         """Run generation through all stages in the pipeline."""
@@ -631,18 +629,20 @@ class Omni(OmniBase):
         if sampling_params_list is None:
             raise ValueError("sampling_params_list is required for pipelined generation")
 
-        # Normalize sampling_params_list to a list
-        if not isinstance(sampling_params_list, (list, tuple)):
-            sampling_params_list = [sampling_params_list]
-        else:
-            sampling_params_list = list(sampling_params_list)
-
         if len(sampling_params_list) != len(self.stage_list):
             raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
 
+        for i, (stage, sp) in enumerate(zip(self.stage_list, sampling_params_list)):
+            ExpectedSPType = OmniDiffusionSamplingParams if stage.stage_type == "diffusion" else SamplingParams
+            if not isinstance(sp, ExpectedSPType):
+                raise ValueError(
+                    f"Expected sampling parameters with type {ExpectedSPType} in stage {i}, got {sp.__class__}"
+                )
+
         # Normalize prompts to a list for per-request iteration
-        if not isinstance(prompts, (list, tuple)):
-            request_prompts: list[PromptType] = [prompts]
+        # str is also Sequence but only test list-like containers here
+        if isinstance(prompts, str) or not isinstance(prompts, Sequence):
+            request_prompts: list[OmniPromptType] = [prompts]
         else:
             request_prompts = list(prompts)
 
@@ -650,8 +650,8 @@ class Omni(OmniBase):
         num_stages = len(self.stage_list)
 
         # Generate globally unique request IDs and map them to original prompts
-        request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
-        request_id_to_prompt: dict[str, PromptType] = {rid: p for rid, p in zip(request_ids, request_prompts)}
+        request_ids = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
+        request_id_to_prompt = {rid: p for rid, p in zip(request_ids, request_prompts)}
 
         # Track per-request start time for end-to-end timing
         _req_start_ts: dict[str, float] = {}
diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py
index f5b0a0a40..5ad9a91c8 100644
--- a/vllm_omni/entrypoints/omni_diffusion.py
+++ b/vllm_omni/entrypoints/omni_diffusion.py
@@ -3,7 +3,7 @@
 
 import logging
 import uuid
-from dataclasses import fields
+from collections.abc import Sequence
 
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
@@ -11,6 +11,8 @@ from vllm.transformers_utils.config import get_hf_file_to_dict
 from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+from vllm_omni.outputs import OmniRequestOutput
 
 # TODO configure logging properly
 logging.basicConfig(level=logging.INFO)
@@ -18,21 +20,6 @@ logging.basicConfig(level=logging.INFO)
 logger = init_logger(__name__)
 
 
-def prepare_requests(prompt: str | list[str], **kwargs):
-    field_names = {f.name for f in fields(OmniDiffusionRequest)}
-
-    init_kwargs = {"prompt": prompt}
-
-    for key, value in kwargs.items():
-        if key in field_names:
-            init_kwargs[key] = value
-
-    if "guidance_scale" in kwargs:
-        init_kwargs["guidance_scale_provided"] = True
-
-    return OmniDiffusionRequest(**init_kwargs)
-
-
 class OmniDiffusion:
     """
     It is the main class to interact with vLLM-Omni diffusion models.
@@ -100,42 +87,24 @@ class OmniDiffusion:
 
     def generate(
         self,
-        prompt: str | list[str],
-        **kwargs,
-    ):
-        prompts = []
-        if isinstance(prompt, str):
-            prompts.append(prompt)
-        elif isinstance(prompt, list):
-            prompts.extend(prompt)
+        prompts: OmniPromptType | Sequence[OmniPromptType],
+        sampling_params: OmniDiffusionSamplingParams,
+        request_ids: list[str] = [],
+    ) -> list[OmniRequestOutput]:
+        if isinstance(prompts, (str, dict)):
+            prompts = [prompts]
         else:
-            raise ValueError("Prompt must be a string or a list of strings")
-
-        requests: list[OmniDiffusionRequest] = []
+            prompts = list(prompts)
 
         # Check if request_id is provided in kwargs
-        request_id = kwargs.get("request_id")
-        request_ids = kwargs.pop("request_ids", None)
+        if len(request_ids) < len(prompts):
+            request_ids.extend(f"{i + len(request_ids)}_{uuid.uuid4()}" for i in range(len(prompts) - len(request_ids)))
 
-        for i, p in enumerate(prompts):
-            req_kwargs = kwargs.copy()
-            if request_ids and isinstance(request_ids, list) and i < len(request_ids):
-                req_kwargs["request_id"] = request_ids[i]
-            elif request_id is None:
-                # Generate default ID consistent with OmniLLM: "{i}_{uuid}"
-                req_kwargs["request_id"] = f"{i}_{uuid.uuid4()}"
+        request = OmniDiffusionRequest(prompts, sampling_params, request_ids)
+        return self._run_engine(request)
 
-            requests.append(
-                prepare_requests(
-                    p,
-                    **req_kwargs,
-                )
-            )
-        logger.info(f"Prepared {len(requests)} requests for generation.")
-        return self._run_engine(requests)
-
-    def _run_engine(self, requests: list[OmniDiffusionRequest]):
-        return self.engine.step(requests)
+    def _run_engine(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
+        return self.engine.step(request)
 
     def close(self) -> None:
         self.engine.close()
diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index e0e699c9c..0a5ee55be 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -15,9 +15,11 @@ import queue
 import sys
 import time
 import traceback
+from collections.abc import Sequence
 from dataclasses import fields
-from typing import Any
+from typing import Any, Literal, cast
 
+from vllm import PromptType, RequestOutput
 from vllm.inputs import TextPrompt
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
@@ -31,6 +33,7 @@ from vllm.v1.engine.llm_engine import LLMEngine
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.distributed.omni_connectors import build_stage_connectors
 from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
+from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase
 from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor
 from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
@@ -46,7 +49,8 @@ from vllm_omni.entrypoints.stage_utils import (
     maybe_dump_to_shm,
     set_stage_devices,
 )
-from vllm_omni.inputs.data import OmniTokensPrompt
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt
+from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils import detect_device_type
 
 logger = init_logger(__name__)
@@ -64,35 +68,6 @@ def _build_od_config(engine_args: dict[str, Any], model: str) -> dict[str, Any]:
     return od_config
 
 
-def prepare_sampling_params(sampling_params: Any, stage_type: str) -> Any:
-    """Prepare sampling parameters for the given stage type.
-
-    Args:
-        sampling_params: Raw sampling parameters (dict or SamplingParams)
-        stage_type: Either "llm" or "diffusion"
-
-    Returns:
-        Processed sampling parameters ready for engine consumption
-    """
-    if stage_type == "diffusion":
-        # For diffusion stages: extract kwargs, handling different input types
-        if isinstance(sampling_params, dict):
-            diffusion_kwargs = dict(sampling_params)
-        else:
-            diffusion_kwargs = getattr(sampling_params, "__dict__", {}) or {}
-
-        # Remove 'prompt' and 'request_id' to avoid conflict with explicit arguments
-        diffusion_kwargs.pop("prompt", None)
-        diffusion_kwargs.pop("request_id", None)
-        return diffusion_kwargs
-
-    else:  # stage_type == "llm"
-        # For LLM stages: ensure we have a SamplingParams object
-        if isinstance(sampling_params, dict):
-            return SamplingParams(**sampling_params)
-        return sampling_params
-
-
 class OmniStage:
     """Stage manager for orchestrating a single stage in the omni pipeline.
 
@@ -123,7 +98,7 @@ class OmniStage:
         self.engine_outputs = None
         self.is_comprehension = getattr(stage_config, "is_comprehension", False)
         # Support for different stage types: "llm" (default) or "diffusion"
-        self.stage_type = getattr(stage_config, "stage_type", "llm")
+        self.stage_type: Literal["llm", "diffusion"] = getattr(stage_config, "stage_type", "llm")
         if hasattr(stage_config, "custom_process_input_func"):
             # Import the module specified in the config (already a full module path)
             module_path, func_name = stage_config.custom_process_input_func.rsplit(".", 1)
@@ -137,7 +112,14 @@ class OmniStage:
         default_sampling_params = getattr(stage_config, "default_sampling_params", {})
         # For LLM stage, this can directly be a SamplingParams-compatible dict;
         # For diffusion stage, this only serves as default values for diffusion kwargs.
-        self.default_sampling_params = _to_dict(default_sampling_params)
+        default_sampling_params = _to_dict(default_sampling_params)
+        # Further convert it to dataclass to check fields
+        try:
+            self.default_sampling_params = (
+                SamplingParams if self.stage_type == "llm" else OmniDiffusionSamplingParams
+            )(**default_sampling_params)
+        except TypeError as error:
+            raise TypeError(f"Invalid default_sampling_params for stage {self.stage_id}: {error}") from error
         # Runtime orchestration state (added)
         self._in_q: mp.Queue | None = None
         self._out_q: mp.Queue | None = None
@@ -519,7 +501,7 @@ def _stage_worker(
     runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
     connectors_config = stage_payload.get("connectors_config", {})
-    stage_type = stage_payload.get("stage_type", "llm")
+    stage_type: Literal["llm", "diffusion"] = stage_payload.get("stage_type", "llm")
 
     # Aggregates for running average
     _agg_total_tokens = 0
@@ -708,15 +690,14 @@ def _stage_worker(
                 pass
     logger.debug("Engine initialized")
     # Initialize OmniConnectors if configured
-    connectors = {}
+    connectors: dict[tuple[str, str], OmniConnectorBase] | None = {}
     if connectors_config:
-        built_connectors = build_stage_connectors(
+        connectors = build_stage_connectors(
             stage_id=stage_id,
             connectors_config=connectors_config,
         )
-        if built_connectors is None:
+        if connectors is None:
             return
-        connectors = built_connectors
 
     # Signal readiness to orchestrator
     try:
@@ -785,6 +766,7 @@ def _stage_worker(
             continue
 
         batch_tasks: list[dict[str, Any]] = [task]
+        tasks_failed_to_add_to_batch: list[dict[str, Any]] = []
         start_time = _time.time()
         if max_batch_size > 1:
             while len(batch_tasks) < max_batch_size:
@@ -800,7 +782,20 @@ def _stage_worker(
                         if extra_type == OmniStageTaskType.PROFILER_STOP:
                             out_q.put({"type": "profiler_result", "data": p_data})
                         continue
-                    batch_tasks.append(extra)
+                    # Ensure that all tasks have the same sampling params
+                    # If no, put them in a temporary container and add back to queue
+                    # This should be always true, because user only calls omni.generate() once and it blocks
+                    # User can only pass one sampling param object, but the list of prompts are separated.
+                    if task.get("sampling_params") != extra.get("sampling_params"):
+                        logger.warning(
+                            """In offline mode, expect all prompts in one `omni.generate()` call to share same sampling params"""  # noqa: E501 # line too long
+                            f"""However, prompt {task.get("engine_inputs")} has sampling params {task.get("sampling_params")}, """  # noqa: E501 # line too long
+                            f"""whereas the prompt {extra.get("engine_inputs")} has sampling params {extra.get("sampling_params")}."""  # noqa: E501 # line too long
+                            """The two tasks cannot be combined in one batch request."""
+                        )
+                        tasks_failed_to_add_to_batch.append(extra)
+                    else:
+                        batch_tasks.append(extra)
                     end_time = _time.time()
                     duration = end_time - start_time
                     if duration > batch_timeout:
@@ -815,9 +810,13 @@ def _stage_worker(
                         break
                     else:
                         continue
+        for task_to_readd in tasks_failed_to_add_to_batch:
+            in_q.put(task_to_readd)
+        # Ensure that the popped tasks are with identical sampling params. Take one of them.
+        batch_engine_sampling_params: OmniSamplingParams = batch_tasks[0]["sampling_params"]
 
         batch_request_ids: list[Any] = []
-        batch_engine_inputs: list[Any] = []
+        batch_engine_inputs: list[OmniPromptType] = []
         _rx_bytes_by_rid: dict[Any, int] = {}
         _rx_decode_ms_by_rid: dict[Any, float] = {}
         _in_flight_ms_by_rid: dict[Any, float] = {}
@@ -839,6 +838,9 @@ def _stage_worker(
                 connectors=connectors,
                 stage_id=stage_id,
             )
+            # TODO: hack type annotation for now.
+            # A better way is to refine type annotation of connection and task/payloads, maybe using template types.
+            ein = cast(OmniPromptType | Sequence[OmniPromptType] | None, ein)
 
             if ein is None or _rx_metrics is None:
                 raise RuntimeError(
@@ -850,17 +852,14 @@ def _stage_worker(
             _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
 
             batch_request_ids.append(rid)
-            if isinstance(ein, list):
+            if isinstance(ein, (str, dict)):
+                # Types like OmniTextPrompt, TextPrompt are TypedDict, essentially dict and enters this branch
+                batch_engine_inputs.append(ein)
+            elif isinstance(ein, Sequence):
                 batch_engine_inputs.extend(ein)
-            elif isinstance(ein, dict):
-                batch_engine_inputs.append(ein)
-            elif isinstance(ein, str):
-                # For diffusion stage-0, ein might be a string prompt directly
-                batch_engine_inputs.append(ein)
             else:
-                # For other types (e.g., OmniTokensPrompt, TextPrompt), append as-is
+                # Other unknown types, append as-is
                 batch_engine_inputs.append(ein)
-        sampling_params = batch_tasks[0]["sampling_params"]
         logger.debug(
             "Received batch size=%d, request_ids=%s",
             len(batch_tasks),
@@ -868,65 +867,30 @@ def _stage_worker(
         )
         try:
             _batch_seq += 1
-            gen_outputs: list[Any] = []
+            gen_outputs: list[OmniRequestOutput | RequestOutput] = []
             _gen_t0 = _time.time()
             if stage_type == "diffusion":
-                # For diffusion, batch_engine_inputs should be prompts (strings)
-                # Convert to list of strings if needed
-                prompts = []
-                for ein in batch_engine_inputs:
-                    if isinstance(ein, str):
-                        prompts.append(ein)
-                    elif isinstance(ein, dict) and "prompt" in ein:
-                        prompts.append(ein["prompt"])
-                    elif hasattr(ein, "prompt"):
-                        prompts.append(ein.prompt)
-                    else:
-                        prompts.append(str(ein))
-                # Prepare diffusion kwargs from sampling parameters
-                diffusion_kwargs = prepare_sampling_params(sampling_params, "diffusion")
-
-                # Pass batch_request_ids to ensure correct ID mapping
-                diffusion_kwargs["request_ids"] = batch_request_ids
-
+                stage_engine = cast(OmniDiffusion, stage_engine)
+                batch_engine_sampling_params = cast(OmniDiffusionSamplingParams, batch_engine_sampling_params)
                 # Diffusion generate returns results directly, not an iterator
-                diffusion_results = stage_engine.generate(prompts, **diffusion_kwargs)
-                # Convert to list format compatible with LLM outputs
-                # Ensure each result has a request_id for proper mapping
-                if isinstance(diffusion_results, list):
-                    gen_outputs = diffusion_results
-                    # Assign request_ids if not present
-                    for idx, result in enumerate(gen_outputs):
-                        if not hasattr(result, "request_id") or result.request_id is None:
-                            if idx < len(batch_request_ids):
-                                if hasattr(result, "request_id"):
-                                    result.request_id = batch_request_ids[idx]
-                                else:
-                                    # Create a wrapper object if result doesn't support request_id
-                                    from types import SimpleNamespace
-
-                                    wrapped = SimpleNamespace()
-                                    wrapped.request_id = batch_request_ids[idx]
-                                    wrapped.output = result
-                                    gen_outputs[idx] = wrapped
-                else:
-                    gen_outputs = [diffusion_results]
-                    # Assign request_id to single result
-                    if len(batch_request_ids) > 0:
-                        if hasattr(gen_outputs[0], "request_id"):
-                            gen_outputs[0].request_id = batch_request_ids[0]
-                        else:
-                            from types import SimpleNamespace
-
-                            wrapped = SimpleNamespace()
-                            wrapped.request_id = batch_request_ids[0]
-                            wrapped.output = gen_outputs[0]
-                            gen_outputs[0] = wrapped
+                diffusion_results = stage_engine.generate(
+                    batch_engine_inputs, batch_engine_sampling_params, batch_request_ids
+                )
+                gen_outputs.extend(diffusion_results)
+                # Assign request_ids if not present
+                for idx, result in enumerate(gen_outputs):
+                    if not hasattr(result, "request_id") or result.request_id is None:
+                        if idx < len(batch_request_ids):
+                            result.request_id = batch_request_ids[idx]
             else:
-                # LLM engine: use vLLM native SamplingParams
-                llm_sampling_params = prepare_sampling_params(sampling_params, "llm")
-                for ro in stage_engine.generate(batch_engine_inputs, llm_sampling_params, use_tqdm=False):
-                    gen_outputs.append(ro)
+                stage_engine = cast(OmniLLM, stage_engine)
+                batch_engine_sampling_params = cast(SamplingParams, batch_engine_sampling_params)
+                results = stage_engine.generate(
+                    batch_engine_inputs,  # type: ignore # silent complaints about list of subclassed TypedDict
+                    batch_engine_sampling_params,
+                    use_tqdm=False,
+                )
+                gen_outputs.extend(results)
             _gen_t1 = _time.time()
             _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
             logger.debug(f"Generate done: batch={len(batch_tasks)}, req_ids={batch_request_ids}, gen_ms={_gen_ms:.1f}")
@@ -935,7 +899,7 @@ def _stage_worker(
             req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
             unmapped: list[Any] = []
             for ro in gen_outputs:
-                rid = getattr(ro, "request_id", None)
+                rid = ro.request_id
                 if rid in req_to_outputs:
                     req_to_outputs[rid].append(ro)
                 else:
@@ -1354,6 +1318,10 @@ async def _stage_worker_async(
                 connectors=connectors,
                 stage_id=stage_id,
             )
+            # TODO: hack type annotation for now.
+            # A better way is to refine type annotation of connection and task/payloads, maybe using template types.
+            ein = cast(OmniPromptType | Sequence[OmniPromptType] | None, ein)
+
             if ein is None or _rx_metrics is None:
                 raise RuntimeError(
                     f"[Stage-{stage_id}] Missing connector payload for request {rid}. "
@@ -1362,36 +1330,23 @@ async def _stage_worker_async(
             _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
             _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
 
-            sampling_params = task["sampling_params"]
             logger.debug("Received batch size=1, request_ids=%s", rid)
             _gen_t0 = _time.time()
-            if isinstance(ein, list):
+            if isinstance(ein, Sequence) and not isinstance(ein, str):
                 ein = ein[0]
 
             if stage_type == "diffusion":
-                # For diffusion, ein should be prompts (strings)
-                # Convert to string if needed
-                if isinstance(ein, str):
-                    prompt = ein
-                elif isinstance(ein, dict) and "prompt" in ein:
-                    prompt = ein["prompt"]
-                elif hasattr(ein, "prompt"):
-                    prompt = ein.prompt
-                else:
-                    prompt = str(ein)
-
-                # Prepare diffusion kwargs from sampling parameters
-                diffusion_kwargs = prepare_sampling_params(sampling_params, "diffusion")
+                diffusion_sampling_params = cast(OmniDiffusionSamplingParams, task["sampling_params"])
                 # AsyncOmniDiffusion.generate returns a single result, not an async generator
-                gen_output = await stage_engine.generate(prompt=prompt, request_id=rid, **diffusion_kwargs)
+                gen_output = await cast(AsyncOmniDiffusion, stage_engine).generate(ein, diffusion_sampling_params, rid)
                 _gen_t1 = _time.time()
                 _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
                 await generation_out_q.put((rid, gen_output, _gen_ms))
             else:
-                # LLM stages: ensure using SamplingParams
-                llm_sampling_params = prepare_sampling_params(sampling_params, "llm")
+                ein = cast(PromptType, ein)
+                llm_sampling_params: SamplingParams = task["sampling_params"]
                 gen_output = None
-                async for res in stage_engine.generate(ein, llm_sampling_params, rid):
+                async for res in cast(AsyncLLM, stage_engine).generate(ein, llm_sampling_params, rid):
                     gen_output = res
                     _gen_t1 = _time.time()
                     _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index a47a864cd..f80fe95fa 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -10,13 +10,14 @@ from argparse import Namespace
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from http import HTTPStatus
-from typing import Any
+from typing import Any, cast
 
 import vllm.envs as envs
 from fastapi import Depends, HTTPException, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 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_messages import AnthropicServingMessages
 from vllm.entrypoints.launcher import serve_http
@@ -80,6 +81,7 @@ from vllm_omni.entrypoints.openai.protocol.images import (
 )
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 
@@ -865,7 +867,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         HTTPException: For validation errors, missing engine, or generation failures
     """
     # Get engine client (AsyncOmni) from app state
-    engine_client: EngineClient | None = getattr(raw_request.app.state, "engine_client", None)
+    engine_client: EngineClient | AsyncOmni | None = getattr(raw_request.app.state, "engine_client", None)
     if engine_client is None or not hasattr(engine_client, "stage_list"):
         raise HTTPException(
             status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
@@ -925,10 +927,8 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
 
     try:
         # Build params - pass through user values directly
-        gen_params = {
-            "prompt": request.prompt,
-            "num_outputs_per_prompt": request.n,
-        }
+        prompt: OmniTextPrompt = {"prompt": request.prompt}
+        gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
 
         # Parse per-request LoRA (compatible with chat's extra_body.lora shape).
         if request.lora is not None:
@@ -960,62 +960,72 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
                     detail="Invalid lora object: both name and path are required.",
                 )
 
-            gen_params["lora_request"] = LoRARequest(str(lora_name), int(lora_int_id), str(lora_path))
+            gen_params.lora_request = LoRARequest(str(lora_name), int(lora_int_id), str(lora_path))
             if lora_scale is not None:
-                gen_params["lora_scale"] = float(lora_scale)
+                gen_params.lora_scale = float(lora_scale)
 
         # Parse and add size if provided
         if request.size:
             width, height = parse_size(request.size)
-            gen_params["height"] = height
-            gen_params["width"] = width
+            gen_params.height = height
+            gen_params.width = width
             size_str = f"{width}x{height}"
         else:
             size_str = "model default"
 
         # Add optional parameters ONLY if provided
         if request.num_inference_steps is not None:
-            gen_params["num_inference_steps"] = request.num_inference_steps
+            gen_params.num_inference_steps = request.num_inference_steps
         if request.negative_prompt is not None:
-            gen_params["negative_prompt"] = request.negative_prompt
+            prompt["negative_prompt"] = request.negative_prompt
         if request.guidance_scale is not None:
-            gen_params["guidance_scale"] = request.guidance_scale
+            gen_params.guidance_scale = request.guidance_scale
         if request.true_cfg_scale is not None:
-            gen_params["true_cfg_scale"] = request.true_cfg_scale
+            gen_params.true_cfg_scale = request.true_cfg_scale
         if request.seed is not None:
-            gen_params["seed"] = request.seed
-        gen_params["request_id"] = f"img_gen_{int(time.time())}"
+            gen_params.seed = request.seed
+        request_id = f"img_gen_{int(time.time())}"
 
         logger.info(f"Generating {request.n} image(s) {size_str}")
 
         # Generate images using AsyncOmni (multi-stage mode)
+        engine_client = cast(AsyncOmni, engine_client)
         result = None
         stage_list = getattr(engine_client, "stage_list", None)
         if isinstance(stage_list, list):
-            default_params_list = getattr(engine_client, "default_sampling_params_list", None)
+            default_params_list: list[OmniSamplingParams] | None = getattr(
+                engine_client, "default_sampling_params_list", None
+            )
             if not isinstance(default_params_list, list):
-                default_params_list = [{} for _ in stage_types]
+                default_params_list = [
+                    OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
+                ]
             else:
                 default_params_list = list(default_params_list)
             if len(default_params_list) != len(stage_types):
-                default_params_list = (default_params_list + [{} for _ in stage_types])[: len(stage_types)]
+                default_params_list = (
+                    default_params_list
+                    + [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
+                )[: len(stage_types)]
 
-            sampling_params_list: list[dict[str, Any]] = []
+            sampling_params_list: list[OmniSamplingParams] = []
             for idx, stage_type in enumerate(stage_types):
                 if stage_type == "diffusion":
                     sampling_params_list.append(gen_params)
                 else:
                     base_params = default_params_list[idx]
-                    sampling_params_list.append(dict(base_params) if isinstance(base_params, dict) else base_params)
+                    sampling_params_list.append(base_params)
 
             async for output in engine_client.generate(
-                prompt=gen_params["prompt"],
-                request_id=gen_params["request_id"],
+                prompt=prompt,
+                request_id=request_id,
                 sampling_params_list=sampling_params_list,
             ):
                 result = output
         else:
-            result = await engine_client.generate(**gen_params)
+            result = await engine_client.generate(
+                prompt=prompt, request_id=request_id, sampling_params_list=[gen_params]
+            )
 
         if result is None:
             raise HTTPException(
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index f35008260..db8c972a1 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -6,13 +6,16 @@ import uuid
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence
 from datetime import datetime, timedelta, timezone
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Final, Optional
+from typing import TYPE_CHECKING, Any, Final, Optional, cast
 
 import jinja2
 from fastapi import Request
 from PIL import Image
 from pydantic import TypeAdapter
 
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+
 try:
     import soundfile
 except ImportError:
@@ -274,16 +277,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     lora_request=lora_request,
                 )
 
-                trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers)
-
                 generator = self.engine_client.generate(
                     prompt=engine_prompt,
                     request_id=request_id,
                     sampling_params_list=sampling_params_list,
                     output_modalities=output_modalities,
-                    lora_request=lora_request,
-                    trace_headers=trace_headers,
-                    priority=request.priority,
                 )
 
                 generators.append(generator)
@@ -1904,29 +1902,30 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     logger.warning("Failed to decode reference image: %s", e)
 
             # Build generation kwargs
-            gen_kwargs: dict[str, Any] = {
+            gen_prompt: OmniTextPrompt = {
                 "prompt": prompt,
-                "request_id": request_id,
-                "num_inference_steps": num_inference_steps,
-                "height": height,
-                "width": width,
                 "negative_prompt": negative_prompt,
-                "num_outputs_per_prompt": num_outputs_per_prompt,
-                "seed": seed,
             }
+            gen_params = OmniDiffusionSamplingParams(
+                num_inference_steps=num_inference_steps,
+                height=height,
+                width=width,
+                num_outputs_per_prompt=num_outputs_per_prompt,
+                seed=seed,
+            )
 
             if guidance_scale is not None:
-                gen_kwargs["guidance_scale"] = guidance_scale
+                gen_params.guidance_scale = guidance_scale
 
             # Add Qwen-Image specific parameter
             if true_cfg_scale is not None:
-                gen_kwargs["true_cfg_scale"] = true_cfg_scale
+                gen_params.true_cfg_scale = true_cfg_scale
 
             # Add video generation parameters if set
             if num_frames is not None:
-                gen_kwargs["num_frames"] = num_frames
+                gen_params.num_frames = num_frames
             if guidance_scale_2 is not None:
-                gen_kwargs["guidance_scale_2"] = guidance_scale_2
+                gen_params.guidance_scale_2 = guidance_scale_2
 
             # Parse per-request LoRA (works for both AsyncOmniDiffusion and AsyncOmni).
             if lora_body and isinstance(lora_body, dict):
@@ -1949,16 +1948,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         lora_int_id = stable_lora_int_id(str(lora_path))
                     if lora_name and lora_path:
                         lora_req = LoRARequest(str(lora_name), int(lora_int_id), str(lora_path))
-                        gen_kwargs["lora_request"] = lora_req
+                        gen_params.lora_request = lora_req
                         if lora_scale is not None:
-                            gen_kwargs["lora_scale"] = float(lora_scale)
+                            gen_params.lora_scale = float(lora_scale)
                 except Exception as e:  # pragma: no cover - safeguard
                     logger.warning("Failed to parse LoRA request: %s", e)
 
             # Add reference image if provided
             if pil_images:
                 if len(pil_images) == 1:
-                    gen_kwargs["pil_image"] = pil_images[0]
+                    gen_prompt["multi_modal_data"] = {}
+                    gen_prompt["multi_modal_data"]["image"] = pil_images[0]
                 else:
                     od_config = getattr(self._diffusion_engine, "od_config", None)
                     supports_multimodal_inputs = getattr(od_config, "supports_multimodal_inputs", False)
@@ -1966,7 +1966,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         # TODO: entry is asyncOmni. We hack the od config here.
                         supports_multimodal_inputs = True
                     if supports_multimodal_inputs:
-                        gen_kwargs["pil_image"] = pil_images
+                        gen_prompt["multi_modal_data"] = {}
+                        gen_prompt["multi_modal_data"]["image"] = pil_images
                     else:
                         return self._create_error_response(
                             "Multiple input images are not supported by the current diffusion model. "
@@ -1979,18 +1980,24 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # Handle both AsyncOmniDiffusion (returns OmniRequestOutput) and AsyncOmni (returns AsyncGenerator)
             if hasattr(self._diffusion_engine, "stage_list"):
                 # AsyncOmni: iterate through async generator to get final output
+                diffusion_engine = cast(AsyncOmni, self._diffusion_engine)
                 result = None
-                async for output in self._diffusion_engine.generate(
-                    prompt=gen_kwargs["prompt"],
-                    request_id=gen_kwargs.get("request_id"),
-                    sampling_params_list=[gen_kwargs],  # Pass as single-stage params
+                async for output in diffusion_engine.generate(
+                    prompt=gen_prompt,
+                    sampling_params_list=[gen_params],  # Pass as single-stage params
+                    request_id=request_id,
                 ):
                     result = output
                 if result is None:
                     return self._create_error_response("No output generated from AsyncOmni")
             else:
                 # AsyncOmniDiffusion: direct call
-                result = await self._diffusion_engine.generate(**gen_kwargs)
+                diffusion_engine = cast(AsyncOmniDiffusion, self._diffusion_engine)
+                result = await diffusion_engine.generate(
+                    prompt=gen_prompt,
+                    sampling_params=gen_params,
+                    request_id=request_id,
+                )
             # Extract images from result
             # Handle nested OmniRequestOutput structure where images might be in request_output
             images = getattr(result.request_output, "images", [])
diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py
index ec291e43e..01de61b96 100644
--- a/vllm_omni/inputs/data.py
+++ b/vllm_omni/inputs/data.py
@@ -1,4 +1,11 @@
-from typing import Any
+import copy
+import pprint
+from dataclasses import asdict, dataclass, field
+from typing import Any, TypeAlias
+
+from vllm import PromptType, SamplingParams
+
+from vllm_omni.lora.request import LoRARequest
 
 try:
     from typing import NotRequired
@@ -7,7 +14,25 @@ except ImportError:
     from typing_extensions import NotRequired
 
 import torch
-from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
+from vllm.inputs.data import EmbedsPrompt, TextPrompt, TokenInputs, TokensPrompt
+
+
+class OmniTextPrompt(TextPrompt):
+    """Text prompt with optional embeddings and additional information.
+
+    Extends TextPrompt to support prompt embeddings and additional
+    information payloads for direct transfer between pipeline stages.
+
+    Attributes:
+        prompt_embeds: Optional tensor containing prompt embeddings
+        additional_information: Optional dictionary containing additional
+            information (tensors or lists) to pass along with the prompt
+    """
+
+    negative_prompt: NotRequired[str]
+    prompt_embeds: NotRequired[torch.Tensor]
+    negative_prompt_embeds: NotRequired[torch.Tensor]
+    additional_information: NotRequired[dict[str, Any]]
 
 
 class OmniTokensPrompt(TokensPrompt):
@@ -22,7 +47,9 @@ class OmniTokensPrompt(TokensPrompt):
             information (tensors or lists) to pass along with the prompt
     """
 
+    negative_prompt: NotRequired[str]
     prompt_embeds: NotRequired[torch.Tensor]
+    negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]
     """The embeddings of the prompt."""
 
     # New: optional additional information dictionary
@@ -44,7 +71,9 @@ class OmniTokenInputs(TokenInputs):
     """
 
     # New: optional prompt embeddings aligned with token ids
+    negative_prompt: NotRequired[str]
     prompt_embeds: NotRequired[torch.Tensor]
+    negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]
 
     # New: optional additional information dictionary
     # Values may be torch.Tensor or list
@@ -65,12 +94,19 @@ class OmniEmbedsPrompt(EmbedsPrompt):
 
     # New: optional prompt embeddings aligned with token ids
     prompt_embeds: NotRequired[torch.Tensor]
+    negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]
 
     # New: optional additional information dictionary
     # Values may be torch.Tensor or list
     additional_information: NotRequired[dict[str, Any]]
 
 
+# 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()
+OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt
+
+
 def token_inputs_omni(
     prompt_token_ids: list[int],
     prompt: str | None = None,
@@ -106,3 +142,140 @@ def token_inputs_omni(
         inputs["additional_information"] = additional_information
 
     return inputs
+
+
+@dataclass
+class OmniDiffusionSamplingParams:
+    """
+    The collection of sampling parameters passed to diffusion pipelines.
+
+    This dataclass contains all information needed during the diffusion pipeline
+    execution, allowing methods to update specific components without needing
+    to manage numerous individual parameters.
+    """
+
+    # Additional text-related parameters
+    max_sequence_length: int | None = None
+    prompt_template: dict[str, Any] | None = None
+    do_classifier_free_guidance: bool = False
+
+    # Batch info
+    num_outputs_per_prompt: int = 1
+    seed: int | None = None
+    generator: torch.Generator | list[torch.Generator] | None = None
+
+    # layered info
+    layers: int = 4
+
+    # cfg info
+    cfg_normalize: bool = False
+
+    # caption language
+    use_en_prompt: bool = False
+
+    # different bucket in (640, 1024) to determine the condition and output resolution
+    resolution: int = 640
+
+    # Tracking if embeddings are already processed
+    is_prompt_processed: bool = False
+
+    # Latent tensors
+    latents: torch.Tensor | None = None
+    raw_latent_shape: torch.Tensor | None = None
+    noise_pred: torch.Tensor | None = None
+    image_latent: torch.Tensor | None = None
+
+    # Latent dimensions
+    height_latents: list[int] | int | None = None
+    width_latents: list[int] | int | None = None
+    num_frames: int = 1  # Default for image models
+    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
+
+    # Original dimensions (before VAE scaling)
+    height: int | None = None
+    width: int | None = None
+    fps: int | None = None
+    height_not_provided: bool = False
+    width_not_provided: bool = False
+
+    # Timesteps
+    timesteps: torch.Tensor | None = None
+    timestep: torch.Tensor | float | int | None = None
+    step_index: int | None = None
+    boundary_ratio: float | None = None
+
+    # Scheduler parameters
+    num_inference_steps: int = 50
+    guidance_scale: float = 0.0
+    guidance_scale_provided: bool = False
+    guidance_scale_2: float | None = None
+    guidance_rescale: float = 0.0
+    eta: float = 0.0
+    sigmas: list[float] | None = None
+
+    true_cfg_scale: float | None = None  # qwen-image specific now
+
+    n_tokens: int | None = None
+    extra_step_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    # [Omni] KV Cache Transfer, for bagel model now
+    past_key_values: Any | None = None  # Injected KV Cache
+    kv_metadata: dict[str, Any] | None = None  # Metadata for KV Cache (e.g., kv_lens, ropes)
+    need_kv_receive: bool = True  # Flag to indicate if this request expects KV transfer
+
+    # Component modules
+    modules: dict[str, Any] = field(default_factory=dict)
+
+    return_trajectory_latents: bool = False
+    return_trajectory_decoded: bool = False
+    trajectory_timesteps: list[torch.Tensor] | None = None
+    trajectory_latents: torch.Tensor | None = None
+
+    # Extra parameters that might be needed by specific pipeline implementations
+    extra_args: dict[str, Any] = field(default_factory=dict)
+
+    # Misc
+    save_output: bool = True
+    return_frames: bool = False
+
+    # LoRA
+    lora_request: LoRARequest | None = None
+    lora_scale: float = 1.0
+
+    # STA parameters
+    STA_param: list | None = None
+    is_cfg_negative: bool = False
+    mask_search_final_result_pos: list[list] | None = None
+    mask_search_final_result_neg: list[list] | None = None
+
+    # VSA parameters
+    VSA_sparsity: float = 0.0
+    # perf_logger: PerformanceLogger | None = None
+
+    # stage logging
+    # logging_info: PipelineLoggingInfo = field(default_factory=PipelineLoggingInfo)
+
+    # profile
+    profile: bool = False
+    num_profiled_timesteps: int = 8
+
+    # debugging
+    debug: bool = False
+
+    # results
+    output: torch.Tensor | None = None
+
+    @property
+    def batch_size(self):
+        # This class is changed to only represent a single prompt request
+        # Only adjust batch size for number of videos per prompt
+        return self.num_outputs_per_prompt
+
+    def __str__(self):
+        return pprint.pformat(asdict(self), indent=2, width=120)
+
+    def clone(self) -> "OmniDiffusionSamplingParams":
+        return copy.deepcopy(self)
+
+
+OmniSamplingParams: TypeAlias = SamplingParams | OmniDiffusionSamplingParams
diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml
index cf5dff50d..96bdc01b6 100644
--- a/vllm_omni/model_executor/stage_configs/bagel.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel.yaml
@@ -58,13 +58,7 @@ stage_args:
     final_output_type: image
     is_comprehension: false
     default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 2048
       seed: 52
-      detokenize: True
-      repetition_penalty: 1.0
 
 # Runtime edges
 runtime:
diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py
index 958cebd3e..d67922000 100644
--- a/vllm_omni/outputs.py
+++ b/vllm_omni/outputs.py
@@ -6,6 +6,8 @@ from PIL import Image
 from vllm.outputs import RequestOutput
 from vllm.v1.outputs import ModelRunnerOutput
 
+from vllm_omni.inputs.data import OmniPromptType
+
 
 class OmniModelRunnerOutput(ModelRunnerOutput):
     """Model runner output for omni models.
@@ -54,7 +56,7 @@ class OmniRequestOutput:
 
     # Diffusion model fields
     images: list[Image.Image] = field(default_factory=list)
-    prompt: str | None = None
+    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)
@@ -89,7 +91,7 @@ class OmniRequestOutput:
         cls,
         request_id: str,
         images: list[Image.Image],
-        prompt: str | None = None,
+        prompt: OmniPromptType | None = None,
         metrics: dict[str, Any] | None = None,
         latents: torch.Tensor | None = None,
         multimodal_output: dict[str, Any] | None = None,

From ae96aa8a8bede85bd0699e1108c1c38db170a7d7 Mon Sep 17 00:00:00 2001
From: JohnJan <wuzhongjian_yewu@cmss.chinamobile.com>
Date: Wed, 28 Jan 2026 16:28:57 +0800
Subject: [PATCH 04/32] [Model]: add FLUX.1-dev model (#853)

---
 docs/models/supported_models.md               |   1 +
 .../diffusion/parallelism_acceleration.md     |   1 +
 .../diffusion/cache/cache_dit_backend.py      |   2 +-
 vllm_omni/diffusion/models/flux/__init__.py   |  17 +
 .../diffusion/models/flux/flux_transformer.py | 638 ++++++++++++++++
 .../diffusion/models/flux/pipeline_flux.py    | 712 ++++++++++++++++++
 vllm_omni/diffusion/registry.py               |   6 +
 7 files changed, 1376 insertions(+), 1 deletion(-)
 create mode 100644 vllm_omni/diffusion/models/flux/__init__.py
 create mode 100644 vllm_omni/diffusion/models/flux/flux_transformer.py
 create mode 100644 vllm_omni/diffusion/models/flux/pipeline_flux.py

diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index e713d7934..269c9261e 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -33,6 +33,7 @@ th {
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
 |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md
index 45e53294e..1a9f540ad 100644
--- a/docs/user_guide/diffusion/parallelism_acceleration.md
+++ b/docs/user_guide/diffusion/parallelism_acceleration.md
@@ -30,6 +30,7 @@ The following table shows which models are currently supported by parallelism me
 | **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"
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index 0c43659c9..ff2c98002 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -791,7 +791,7 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int],
 CUSTOM_DIT_ENABLERS.update(
     {
         "WanPipeline": enable_cache_for_wan22,
-        "FluxPipeline": enable_cache_for_flux,
+        # "FluxPipeline": enable_cache_for_flux,
         "LongCatImagePipeline": enable_cache_for_longcat_image,
         "LongCatImageEditPipeline": enable_cache_for_longcat_image,
         "StableDiffusion3Pipeline": enable_cache_for_sd3,
diff --git a/vllm_omni/diffusion/models/flux/__init__.py b/vllm_omni/diffusion/models/flux/__init__.py
new file mode 100644
index 000000000..cad3508a1
--- /dev/null
+++ b/vllm_omni/diffusion/models/flux/__init__.py
@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FLUX.1-dev diffusion model components."""
+
+from vllm_omni.diffusion.models.flux.flux_transformer import (
+    FluxTransformer2DModel,
+)
+from vllm_omni.diffusion.models.flux.pipeline_flux import (
+    FluxPipeline,
+    get_flux_post_process_func,
+)
+
+__all__ = [
+    "FluxPipeline",
+    "FluxTransformer2DModel",
+    "get_flux_post_process_func",
+]
diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py
new file mode 100644
index 000000000..faf6d08d3
--- /dev/null
+++ b/vllm_omni/diffusion/models/flux/flux_transformer.py
@@ -0,0 +1,638 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.embeddings import (
+    CombinedTimestepGuidanceTextProjEmbeddings,
+    CombinedTimestepTextProjEmbeddings,
+    get_1d_rotary_pos_embed,
+)
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from diffusers.utils import is_torch_npu_available
+from torch import nn
+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.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+logger = init_logger(__name__)
+
+
+class ColumnParallelApproxGELU(nn.Module):
+    def __init__(self, dim_in: int, dim_out: int, *, approximate: str, bias: bool = True):
+        super().__init__()
+        self.proj = ColumnParallelLinear(
+            dim_in,
+            dim_out,
+            bias=bias,
+            gather_output=False,
+            return_bias=False,
+        )
+        self.approximate = approximate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        return F.gelu(x, approximate=self.approximate)
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int | None = None,
+        mult: int = 4,
+        activation_fn: str = "gelu-approximate",
+        inner_dim: int | None = None,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        assert activation_fn == "gelu-approximate", "Only gelu-approximate is supported."
+
+        inner_dim = inner_dim or int(dim * mult)
+        dim_out = dim_out or dim
+
+        layers: list[nn.Module] = [
+            ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias),
+            nn.Identity(),  # placeholder for weight loading
+            RowParallelLinear(
+                inner_dim,
+                dim_out,
+                input_is_parallel=True,
+                return_bias=False,
+            ),
+        ]
+
+        self.net = nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+class FluxAttention(torch.nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        added_kv_proj_dim: int | None = None,
+        added_proj_bias: bool | None = True,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        out_dim: int = None,
+        context_pre_only: bool | None = None,
+        pre_only: bool = False,
+    ):
+        super().__init__()
+
+        self.head_dim = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.use_bias = bias
+        self.dropout = dropout
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.context_pre_only = context_pre_only
+        self.pre_only = pre_only
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.added_kv_proj_dim = added_kv_proj_dim
+        self.added_proj_bias = added_proj_bias
+
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=query_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.heads,
+            bias=bias,
+        )
+
+        if not self.pre_only:
+            self.to_out = nn.ModuleList(
+                [
+                    RowParallelLinear(
+                        self.inner_dim,
+                        self.out_dim,
+                        bias=out_bias,
+                        input_is_parallel=True,
+                        return_bias=False,
+                    ),
+                    nn.Dropout(dropout),
+                ]
+            )
+
+        if added_kv_proj_dim is not None:
+            self.norm_added_q = RMSNorm(dim_head, eps=eps)
+            self.norm_added_k = RMSNorm(dim_head, eps=eps)
+
+            self.add_kv_proj = QKVParallelLinear(
+                hidden_size=self.added_kv_proj_dim,
+                head_size=self.head_dim,
+                total_num_heads=self.heads,
+                bias=added_proj_bias,
+            )
+
+            self.to_add_out = RowParallelLinear(
+                self.inner_dim,
+                query_dim,
+                bias=out_bias,
+                input_is_parallel=True,
+                return_bias=False,
+            )
+
+        self.rope = RotaryEmbedding(is_neox_style=False)
+        self.attn = Attention(
+            num_heads=self.to_qkv.num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.to_qkv.num_kv_heads,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        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 = 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 = 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)
+            add_q_size = self.add_kv_proj.num_heads * self.head_dim
+            add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
+            encoder_query, encoder_key, encoder_value = encoder_qkv.split(
+                [add_q_size, add_kv_size, add_kv_size], 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 = self.norm_added_q(encoder_query)
+            encoder_key = self.norm_added_k(encoder_key)
+
+            query = torch.cat([encoder_query, query], dim=1)
+            key = torch.cat([encoder_key, key], dim=1)
+            value = torch.cat([encoder_value, value], dim=1)
+
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb  # [S, D/2]
+            cos = cos.to(query.dtype)
+            sin = sin.to(query.dtype)
+            query = self.rope(query, cos, sin)
+            key = self.rope(key, cos, sin)
+
+        hidden_states = self.attn(
+            query,
+            key,
+            value,
+        )
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            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[0](hidden_states)
+            hidden_states = self.to_out[1](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
+
+
+class FluxTransformerBlock(nn.Module):
+    def __init__(
+        self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
+    ):
+        super().__init__()
+
+        self.norm1 = AdaLayerNormZero(dim)
+        self.norm1_context = AdaLayerNormZero(dim)
+
+        self.attn = FluxAttention(
+            query_dim=dim,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            eps=eps,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim)
+
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(dim=dim, dim_out=dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+        joint_attention_kwargs = joint_attention_kwargs or {}
+
+        # Attention.
+        attention_outputs = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        if len(attention_outputs) == 2:
+            attn_output, context_attn_output = attention_outputs
+        elif len(attention_outputs) == 3:
+            attn_output, context_attn_output, ip_attn_output = attention_outputs
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = hidden_states + ff_output
+        if len(attention_outputs) == 3:
+            hidden_states = hidden_states + ip_attn_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class FluxSingleTransformerBlock(nn.Module):
+    def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.norm = AdaLayerNormZeroSingle(dim)
+        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+        self.attn = FluxAttention(
+            query_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            eps=1e-6,
+            pre_only=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        text_seq_len = encoder_hidden_states.shape[1]
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
+        return encoder_hidden_states, hidden_states
+
+
+class FluxPosEmbed(nn.Module):
+    # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
+    def __init__(self, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        n_axes = ids.shape[-1]
+        cos_out = []
+        sin_out = []
+        pos = ids.float()
+        is_mps = ids.device.type == "mps"
+        is_npu = ids.device.type == "npu"
+        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+        for i in range(n_axes):
+            freqs_cis = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[:, i],
+                theta=self.theta,
+                use_real=False,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(freqs_cis.real)
+            sin_out.append(freqs_cis.imag)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class FluxTransformer2DModel(nn.Module):
+    """
+    The Transformer model introduced in Flux.
+
+    Args:
+        od_config (`OmniDiffusionConfig`):
+            The configuration for the model.
+        patch_size (`int`, defaults to `1`):
+            Patch size to turn the input data into small patches.
+        in_channels (`int`, defaults to `64`):
+            The number of channels in the input.
+        out_channels (`int`, *optional*, defaults to `None`):
+            The number of channels in the output. If not specified, it defaults to `in_channels`.
+        num_layers (`int`, defaults to `19`):
+            The number of layers of dual stream DiT blocks to use.
+        num_single_layers (`int`, defaults to `38`):
+            The number of layers of single stream DiT blocks to use.
+        attention_head_dim (`int`, defaults to `128`):
+            The number of dimensions to use for each attention head.
+        num_attention_heads (`int`, defaults to `24`):
+            The number of attention heads to use.
+        joint_attention_dim (`int`, defaults to `4096`):
+            The number of dimensions to use for the joint attention (embedding/channel dimension of
+            `encoder_hidden_states`).
+        pooled_projection_dim (`int`, defaults to `768`):
+            The number of dimensions to use for the pooled projection.
+        guidance_embeds (`bool`, defaults to `False`):
+            Whether to use guidance embeddings for guidance-distilled variant of the model.
+        axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+            The dimensions to use for the rotary positional embeddings.
+    """
+
+    # the small and frequently-repeated block(s) of a model
+    # -- typically a transformer layer
+    # used for torch compile optimizations
+    _repeated_blocks = ["FluxTransformerBlock"]
+
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        out_channels: int = None,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        guidance_embeds: bool = True,
+        axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
+    ):
+        super().__init__()
+        model_config = od_config.tf_model_config
+        num_layers = model_config.num_layers
+        self.parallel_config = od_config.parallel_config
+        self.in_channels = in_channels
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.guidance_embeds = guidance_embeds
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor | Transformer2DModelOutput:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            img_ids: (`torch.Tensor`):
+                The position ids for image tokens.
+            txt_ids (`torch.Tensor`):
+                The position ids for text tokens.
+            guidance (`torch.Tensor`):
+                Guidance embeddings for guidance-distilled variant of the model.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+
+        hidden_states = self.x_embedder(hidden_states)
+        timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
+
+        if guidance is not None:
+            guidance = guidance.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
+
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if txt_ids.ndim == 3:
+            logger.warning(
+                "Passing `txt_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            logger.warning(
+                "Passing `img_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            img_ids = img_ids[0]
+
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        if is_torch_npu_available():
+            freqs_cos, freqs_sin = self.pos_embed(ids.cpu())
+            image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu())
+        else:
+            image_rotary_emb = self.pos_embed(ids)
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            # self-attn
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+            # cross-attn
+            (".add_kv_proj", ".add_q_proj", "q"),
+            (".add_kv_proj", ".add_k_proj", "k"),
+            (".add_kv_proj", ".add_v_proj", "v"),
+        ]
+
+        params_dict = dict(self.named_parameters())
+
+        # we need to load the buffers for beta and eps (XIELU)
+        for name, buffer in self.named_buffers():
+            if name.endswith(".beta") or name.endswith(".eps"):
+                params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            original_name = name
+            lookup_name = name
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in original_name:
+                    continue
+                lookup_name = original_name.replace(weight_name, param_name)
+                param = params_dict[lookup_name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if lookup_name not in params_dict and ".to_out.0." in lookup_name:
+                    lookup_name = lookup_name.replace(".to_out.0.", ".to_out.")
+                param = params_dict[lookup_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(original_name)
+            loaded_params.add(lookup_name)
+        return loaded_params
diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py
new file mode 100644
index 000000000..01e0bfba9
--- /dev/null
+++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py
@@ -0,0 +1,712 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+import json
+import logging
+import os
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import TextualInversionLoaderMixin
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.flux import FluxTransformer2DModel
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = logging.getLogger(__name__)
+
+
+def get_flux_post_process_func(
+    od_config: OmniDiffusionConfig,
+):
+    if od_config.output_type == "latent":
+        return lambda x: x
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+
+    def post_process_func(images: torch.Tensor):
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, int]:
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class FluxPipeline(
+    nn.Module,
+):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+
+        self.device = get_local_device()
+        model = od_config.model
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            model, subfolder="text_encoder", local_files_only=local_files_only
+        )
+        self.text_encoder_2 = T5EncoderModel.from_pretrained(
+            model, subfolder="text_encoder_2", local_files_only=local_files_only
+        )
+        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
+            self.device
+        )
+        self.transformer = FluxTransformer2DModel(od_config=od_config)
+
+        self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+        self.tokenizer_2 = T5TokenizerFast.from_pretrained(
+            model, subfolder="tokenizer_2", local_files_only=local_files_only
+        )
+
+        self.stage = None
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        # Flux latents are turned into 2x2 patches and packed.
+        # This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        # self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = 128
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        height,
+        width,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} "
+                f"but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. "
+                "Make sure to generate `pooled_prompt_embeds` from the same text encoder "
+                "that was used to generate `prompt_embeds`."
+            )
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. "
+                "Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder "
+                "that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: str | list[str] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 512,
+        dtype: torch.dtype | None = None,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer_2)
+
+        text_inputs = self.tokenizer_2(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(self.device), output_hidden_states=False)[0]
+
+        dtype = self.text_encoder_2.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: str | list[str],
+        num_images_per_prompt: int = 1,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = self.text_encoder(text_input_ids.to(self.device), output_hidden_states=False)
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds.pooler_output
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=self.device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        prompt_2: str | list[str],
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.FloatTensor | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        max_sequence_length: int = 512,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+        """
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            # We only use the pooled prompt output from the CLIPTextModel
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype
+        text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=self.device, dtype=dtype)
+
+        return prompt_embeds, pooled_prompt_embeds, text_ids
+
+    @staticmethod
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+        return latents
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is not None:
+            latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+            return latents.to(device=device, dtype=dtype), latent_image_ids
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+        latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+
+        return latents, latent_image_ids
+
+    def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len):
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        return timesteps, num_inference_steps
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def diffuse(
+        self,
+        prompt_embeds,
+        pooled_prompt_embeds,
+        negative_prompt_embeds,
+        negative_pooled_prompt_embeds,
+        latents,
+        latent_image_ids,
+        text_ids,
+        negative_text_ids,
+        timesteps,
+        do_true_cfg,
+        guidance,
+        true_cfg_scale,
+    ):
+        """Diffusion loop with optional image conditioning."""
+        self.scheduler.set_begin_index(0)
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            self._current_timestep = t
+            # broadcast to batch dimension and place on same device/dtype as latents
+            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+
+            self.transformer.do_true_cfg = do_true_cfg  # used in teacache hook
+            # Forward pass for positive prompt (or unconditional if no CFG)
+            noise_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                pooled_projections=pooled_prompt_embeds,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+                joint_attention_kwargs=self.joint_attention_kwargs,
+                return_dict=False,
+            )[0]
+
+            # Forward pass for negative prompt (CFG)
+            if do_true_cfg:
+                neg_noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=negative_pooled_prompt_embeds,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    txt_ids=negative_text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+                noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        return latents
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] = None,
+        prompt_2: str | list[str] = None,
+        negative_prompt: str | list[str] = None,
+        negative_prompt_2: str | list[str] = None,
+        true_cfg_scale: float = 1.0,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 28,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.FloatTensor | None = None,
+        prompt_embeds: torch.FloatTensor | None = None,
+        pooled_prompt_embeds: torch.FloatTensor | None = None,
+        negative_prompt_embeds: torch.FloatTensor | None = None,
+        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
+        output_type: str | None = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ):
+        """Forward pass for flux."""
+        prompt = req.prompt if req.prompt is not None else prompt
+        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+        height = req.height or self.default_sample_size * self.vae_scale_factor
+        width = req.width or self.default_sample_size * self.vae_scale_factor
+        sigmas = req.sigmas or sigmas
+        num_inference_steps = req.num_inference_steps or num_inference_steps
+        generator = req.generator or generator
+        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
+        if req_num_outputs and req_num_outputs > 0:
+            num_images_per_prompt = req_num_outputs
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+
+        negative_text_ids = None
+        if do_true_cfg:
+            (
+                negative_prompt_embeds,
+                negative_pooled_prompt_embeds,
+                negative_text_ids,
+            ) = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_2=negative_prompt_2,
+                prompt_embeds=negative_prompt_embeds,
+                pooled_prompt_embeds=negative_pooled_prompt_embeds,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        latents = self.diffuse(
+            prompt_embeds,
+            pooled_prompt_embeds,
+            negative_prompt_embeds,
+            negative_pooled_prompt_embeds,
+            latents,
+            latent_image_ids,
+            text_ids,
+            negative_text_ids,
+            timesteps,
+            do_true_cfg,
+            guidance,
+            true_cfg_scale,
+        )
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(output=image)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index d4429aff6..024352479 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -90,6 +90,11 @@ _DIFFUSION_MODELS = {
         "pipeline_flux2_klein",
         "Flux2KleinPipeline",
     ),
+    "FluxPipeline": (
+        "flux",
+        "pipeline_flux",
+        "FluxPipeline",
+    ),
 }
 
 
@@ -227,6 +232,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
     "Flux2KleinPipeline": "get_flux2_klein_post_process_func",
+    "FluxPipeline": "get_flux_post_process_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {

From 6fb405856691a1adcc52abf6811d51486894716b Mon Sep 17 00:00:00 2001
From: dengyunyang <584797741@qq.com>
Date: Wed, 28 Jan 2026 17:10:05 +0800
Subject: [PATCH 05/32] [BugFix] ignore mm data from stages to async omni
 (#954)

Signed-off-by: dengyunyang <584797741@qq.com>
---
 vllm_omni/entrypoints/omni_stage.py | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index 0a5ee55be..6bf1fc668 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -23,6 +23,7 @@ from vllm import PromptType, RequestOutput
 from vllm.inputs import TextPrompt
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
+from vllm.outputs import RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.usage.usage_lib import UsageContext
@@ -1423,7 +1424,7 @@ async def _stage_worker_async(
             batch_request_ids, batch_request_outputs, _gen_ms_list, batch_metrics
         ):
             try:
-                r_outputs = [output]
+                r_outputs = [output_strip(output, omni_stage)]
                 use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
                 if use_shm:
                     out_q.put(
@@ -1509,3 +1510,23 @@ def make_stage_stats(_agg_total_tokens: int, _agg_total_gen_time_ms: float):
     from vllm_omni.entrypoints.log_utils import StageStats
 
     return StageStats(total_token=_agg_total_tokens, total_gen_time=_agg_total_gen_time_ms)
+
+
+def output_strip(r_output: RequestOutput, omni_stage: OmniStage):
+    if omni_stage.final_output and omni_stage.final_output_type != "text":
+        return r_output
+
+    if getattr(r_output, "finished", False):
+        return r_output
+
+    mm_output = getattr(r_output, "multimodal_output", None)
+    if mm_output is not None:
+        r_output.multimodal_output = {}
+
+    outputs = getattr(r_output, "outputs", None)
+    if outputs is not None:
+        for out in outputs:
+            if getattr(out, "multimodal_output", None):
+                out.multimodal_output = {}
+
+    return r_output

From 91a3d465eeb890c88c5710c90b43a54f9469d365 Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Wed, 28 Jan 2026 17:41:28 +0800
Subject: [PATCH 06/32] Revert "[BugFix] ignore mm data from stages to async
 omni" (#1023)

---
 vllm_omni/entrypoints/omni_stage.py | 23 +----------------------
 1 file changed, 1 insertion(+), 22 deletions(-)

diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index 6bf1fc668..0a5ee55be 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -23,7 +23,6 @@ from vllm import PromptType, RequestOutput
 from vllm.inputs import TextPrompt
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
-from vllm.outputs import RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.usage.usage_lib import UsageContext
@@ -1424,7 +1423,7 @@ async def _stage_worker_async(
             batch_request_ids, batch_request_outputs, _gen_ms_list, batch_metrics
         ):
             try:
-                r_outputs = [output_strip(output, omni_stage)]
+                r_outputs = [output]
                 use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
                 if use_shm:
                     out_q.put(
@@ -1510,23 +1509,3 @@ def make_stage_stats(_agg_total_tokens: int, _agg_total_gen_time_ms: float):
     from vllm_omni.entrypoints.log_utils import StageStats
 
     return StageStats(total_token=_agg_total_tokens, total_gen_time=_agg_total_gen_time_ms)
-
-
-def output_strip(r_output: RequestOutput, omni_stage: OmniStage):
-    if omni_stage.final_output and omni_stage.final_output_type != "text":
-        return r_output
-
-    if getattr(r_output, "finished", False):
-        return r_output
-
-    mm_output = getattr(r_output, "multimodal_output", None)
-    if mm_output is not None:
-        r_output.multimodal_output = {}
-
-    outputs = getattr(r_output, "outputs", None)
-    if outputs is not None:
-        for out in outputs:
-            if getattr(out, "multimodal_output", None):
-                out.multimodal_output = {}
-
-    return r_output

From cb489682943153b5b6103d436474a8e7a9e28ca2 Mon Sep 17 00:00:00 2001
From: Canlin Guo <canlinguosdu@gmail.com>
Date: Wed, 28 Jan 2026 18:42:39 +0800
Subject: [PATCH 07/32] [Bugfix] Modify output to model_runner_output (#1026)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
---
 vllm_omni/worker/npu/npu_ar_model_runner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm_omni/worker/npu/npu_ar_model_runner.py b/vllm_omni/worker/npu/npu_ar_model_runner.py
index 0eedb61fb..1bb53b662 100644
--- a/vllm_omni/worker/npu/npu_ar_model_runner.py
+++ b/vllm_omni/worker/npu/npu_ar_model_runner.py
@@ -438,7 +438,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
             pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None),
             kv_connector_output=kv_connector_output,
         )
-        output.kv_extracted_req_ids = kv_extracted_req_ids
+        model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
         #  -------------------------------------- Omni-new -------------------------------------------------
 
         durations = ProfileExecuteDuration().pop_captured_sync()

From c0beebbfe305c94556e9e620809e4aa942d2870d Mon Sep 17 00:00:00 2001
From: Samit <285365963@qq.com>
Date: Wed, 28 Jan 2026 20:07:50 +0800
Subject: [PATCH 08/32] [Feature] Support cache-dit for Wan 2.2 inference
 (#1021)

Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>
---
 .../text_to_image/text_to_image.py            |   6 +
 .../text_to_video/text_to_video.py            |  47 +++++++-
 .../diffusion/cache/cache_dit_backend.py      | 113 ++++++++++--------
 vllm_omni/diffusion/data.py                   |   1 +
 .../worker/gpu_diffusion_model_runner.py      |  11 +-
 vllm_omni/diffusion/worker/npu/npu_worker.py  |   6 +-
 vllm_omni/entrypoints/async_omni.py           |   1 +
 vllm_omni/entrypoints/cli/serve.py            |   5 +
 8 files changed, 131 insertions(+), 59 deletions(-)

diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py
index 7fc6ec832..efc5010e0 100644
--- a/examples/offline_inference/text_to_image/text_to_image.py
+++ b/examples/offline_inference/text_to_image/text_to_image.py
@@ -73,6 +73,11 @@ def parse_args() -> argparse.Namespace:
             "Default: None (no cache acceleration)."
         ),
     )
+    parser.add_argument(
+        "--enable-cache-dit-summary",
+        action="store_true",
+        help="Enable cache-dit summary logging after diffusion forward passes.",
+    )
     parser.add_argument(
         "--ulysses_degree",
         type=int,
@@ -166,6 +171,7 @@ def main():
         vae_use_tiling=vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
+        enable_cache_dit_summary=args.enable_cache_dit_summary,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index 054db95b3..49522278f 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -36,8 +36,29 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
+    parser.add_argument(
+        "--cache_backend",
+        type=str,
+        default=None,
+        choices=["cache_dit"],
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
+            "Default: None (no cache acceleration)."
+        ),
+    )
+    parser.add_argument(
+        "--enable-cache-dit-summary",
+        action="store_true",
+        help="Enable cache-dit summary logging after diffusion forward passes.",
+    )
     parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
+    parser.add_argument(
+        "--enforce_eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -55,11 +76,6 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for ring sequence parallelism.",
     )
-    parser.add_argument(
-        "--enforce_eager",
-        action="store_true",
-        help="Disable torch.compile and force eager execution.",
-    )
     return parser.parse_args()
 
 
@@ -72,6 +88,24 @@ def main():
     vae_use_slicing = is_npu()
     vae_use_tiling = is_npu()
 
+    # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
+    cache_config = None
+    if args.cache_backend == "cache_dit":
+        cache_config = {
+            # DBCache parameters [cache-dit only]
+            "Fn_compute_blocks": 1,  # Optimized for single-transformer models
+            "Bn_compute_blocks": 0,  # Number of backward compute blocks
+            "max_warmup_steps": 4,  # Maximum warmup steps (works for few-step models)
+            "max_cached_steps": 20,
+            "residual_diff_threshold": 0.24,  # Higher threshold for more aggressive caching
+            "max_continuous_cached_steps": 3,  # Limit to prevent precision degradation
+            # TaylorSeer parameters [cache-dit only]
+            "enable_taylorseer": False,  # Disabled by default (not suitable for few-step models)
+            "taylorseer_order": 1,  # TaylorSeer polynomial order
+            # SCM (Step Computation Masking) parameters [cache-dit only]
+            "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
+            "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
+        }
     # Configure parallel settings (only SP is supported for Wan)
     # Note: cfg_parallel and tensor_parallel are not implemented for Wan models
     parallel_config = DiffusionParallelConfig(
@@ -88,6 +122,9 @@ def main():
         vae_use_tiling=vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
+        enable_cache_dit_summary=args.enable_cache_dit_summary,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index ff2c98002..c922121e9 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -28,6 +28,13 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig
 logger = init_logger(__name__)
 
 
+# Small helper to centralize cache-dit summaries.
+def cache_summary(pipeline: Any, details: bool = True) -> None:
+    cache_dit.summary(pipeline.transformer, details=details)
+    if hasattr(pipeline, "transformer_2"):
+        cache_dit.summary(pipeline.transformer_2, details=details)
+
+
 # Registry of custom cache-dit enablers for specific models
 # Maps pipeline names to their cache-dit enablement functions
 # Models in this registry require custom handling (e.g., dual-transformer architectures)
@@ -71,53 +78,46 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
         A refresh function that can be called to update cache context with new num_inference_steps.
     """
 
-    # Build DBCacheConfig for primary transformer
-    primary_cache_config = _build_db_cache_config(cache_config)
-
-    # FIXME: secondary cache shares the same config with primary cache for now,
-    # but we should support different config for secondary transformer in the future
-    # Build DBCacheConfig for secondary transformer (can use same or different config)
-    secondary_cache_config = _build_db_cache_config(cache_config)
-
-    # Build calibrator configs if TaylorSeer is enabled
-    primary_calibrator = None
-    secondary_calibrator = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        primary_calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        secondary_calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    # Build ParamsModifier for each transformer
-    primary_modifier = ParamsModifier(
-        cache_config=primary_cache_config,
-        calibrator_config=primary_calibrator,
-    )
-    secondary_modifier = ParamsModifier(
-        cache_config=secondary_cache_config,
-        calibrator_config=secondary_calibrator,
-    )
-
-    logger.info(
-        "Enabling cache-dit on Wan2.2 dual transformers with BlockAdapter: "
-        f"Fn={primary_cache_config.Fn_compute_blocks}, "
-        f"Bn={primary_cache_config.Bn_compute_blocks}, "
-        f"W={primary_cache_config.max_warmup_steps}, "
-    )
-
-    transformer = pipeline.transformer
-    transformer_2 = pipeline.transformer_2
-    transformer_blocks = transformer.blocks
-    transformer_2_blocks = transformer_2.blocks
-    # Enable cache-dit using BlockAdapter for both transformers simultaneously
     cache_dit.enable_cache(
         BlockAdapter(
-            transformer=[transformer, transformer_2],
-            blocks=[transformer_blocks, transformer_2_blocks],
-            forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
-            params_modifiers=[primary_modifier, secondary_modifier],
+            transformer=[
+                pipeline.transformer,
+                pipeline.transformer_2,
+            ],
+            blocks=[
+                pipeline.transformer.blocks,
+                pipeline.transformer_2.blocks,
+            ],
+            forward_pattern=[
+                ForwardPattern.Pattern_2,
+                ForwardPattern.Pattern_2,
+            ],
+            params_modifiers=[
+                # high-noise transformer only have 30% steps
+                ParamsModifier(
+                    cache_config=DBCacheConfig().reset(
+                        max_warmup_steps=cache_config.max_warmup_steps,
+                        max_cached_steps=cache_config.max_cached_steps,
+                    ),
+                ),
+                ParamsModifier(
+                    cache_config=DBCacheConfig().reset(
+                        max_warmup_steps=2,
+                        max_cached_steps=20,
+                    ),
+                ),
+            ],
             has_separate_cfg=True,
         ),
+        cache_config=DBCacheConfig(
+            Fn_compute_blocks=cache_config.Fn_compute_blocks,
+            Bn_compute_blocks=cache_config.Bn_compute_blocks,
+            max_warmup_steps=cache_config.max_warmup_steps,
+            max_cached_steps=cache_config.max_cached_steps,
+            max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+            residual_diff_threshold=cache_config.residual_diff_threshold,
+            num_inference_steps=None,
+        ),
     )
 
     # from https://github.com/vipshop/cache-dit/pull/542
@@ -133,8 +133,8 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
         Returns:
             A tuple of (num_high_noise_steps, num_low_noise_steps).
         """
-        if pipeline.config.boundary_ratio is not None:
-            boundary_timestep = pipeline.config.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
+        if pipeline.boundary_ratio is not None:
+            boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps
         else:
             boundary_timestep = None
 
@@ -158,10 +158,21 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
             pipeline: The Wan2.2 pipeline instance.
             num_inference_steps: New number of inference steps.
         """
+
         num_high_noise_steps, num_low_noise_steps = _split_inference_steps(num_inference_steps)
         # Refresh context for high-noise transformer
         if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose)
+            # cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose)
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                num_inference_steps=num_high_noise_steps,
+                verbose=verbose,
+            )
+            cache_dit.refresh_context(
+                pipeline.transformer_2,
+                num_inference_steps=num_low_noise_steps,
+                verbose=verbose,
+            )
         else:
             cache_dit.refresh_context(
                 pipeline.transformer,
@@ -175,10 +186,6 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
                 verbose=verbose,
             )
 
-        # Refresh context for low-noise transformer
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer_2, num_inference_steps=num_low_noise_steps, verbose=verbose)
-        else:
             cache_dit.refresh_context(
                 pipeline.transformer_2,
                 cache_config=DBCacheConfig().reset(
@@ -790,8 +797,10 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int],
 # Register custom cache-dit enablers after function definitions
 CUSTOM_DIT_ENABLERS.update(
     {
-        "WanPipeline": enable_cache_for_wan22,
-        # "FluxPipeline": enable_cache_for_flux,
+        "Wan22Pipeline": enable_cache_for_wan22,
+        "Wan22I2VPipeline": enable_cache_for_wan22,
+        "Wan22TI2VPipeline": enable_cache_for_wan22,
+        "FluxPipeline": enable_cache_for_flux,
         "LongCatImagePipeline": enable_cache_for_longcat_image,
         "LongCatImageEditPipeline": enable_cache_for_longcat_image,
         "StableDiffusion3Pipeline": enable_cache_for_sd3,
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 1a98462aa..286ab9153 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -258,6 +258,7 @@ class OmniDiffusionConfig:
     # Cache backend configuration (NEW)
     cache_backend: str = "none"  # "tea_cache", "deep_cache", etc.
     cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
+    enable_cache_dit_summary: bool = False
 
     # Distributed executor backend
     distributed_executor_backend: str = "mp"
diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py b/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
index bf56df155..0401655a8 100644
--- a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
@@ -20,6 +20,7 @@ from vllm.config import LoadConfig
 from vllm.logger import init_logger
 from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
 
+from vllm_omni.diffusion.cache.cache_dit_backend import cache_summary
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -284,11 +285,19 @@ class GPUDiffusionModelRunner:
             req.sampling_params.generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
 
         # Refresh cache context if needed
-        if self.cache_backend is not None and self.cache_backend.is_enabled():
+        if (
+            not getattr(req, "skip_cache_refresh", False)
+            and self.cache_backend is not None
+            and self.cache_backend.is_enabled()
+        ):
             self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps)
 
         with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
             with record_function("pipeline_forward"):
                 output = self.pipeline.forward(req)
 
+            # NOTE:
+            if self.od_config.cache_backend == "cache_dit" and self.od_config.enable_cache_dit_summary:
+                cache_summary(self.pipeline, details=True)
+
         return output
diff --git a/vllm_omni/diffusion/worker/npu/npu_worker.py b/vllm_omni/diffusion/worker/npu/npu_worker.py
index 446c29cae..c4d2c60e5 100644
--- a/vllm_omni/diffusion/worker/npu/npu_worker.py
+++ b/vllm_omni/diffusion/worker/npu/npu_worker.py
@@ -132,7 +132,11 @@ class NPUWorker:
             req.generator = torch.Generator(device=self.device).manual_seed(req.seed)
 
         # Refresh cache context if needed
-        if self.cache_backend is not None and self.cache_backend.is_enabled():
+        if (
+            not getattr(req, "skip_cache_refresh", False)
+            and self.cache_backend is not None
+            and self.cache_backend.is_enabled()
+        ):
             self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
         with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
             output = self.pipeline.forward(req)
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index a78710fa2..179f58ed2 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -162,6 +162,7 @@ class AsyncOmni(OmniBase):
                     "vae_use_tiling": kwargs.get("vae_use_tiling", False),
                     "cache_backend": cache_backend,
                     "cache_config": cache_config,
+                    "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
                     "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
                     "enforce_eager": kwargs.get("enforce_eager", False),
                 },
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index c3a37e3c8..bbd167111 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -175,6 +175,11 @@ class OmniServeCommand(CLISubcommand):
             default=None,
             help="JSON string of cache configuration (e.g., '{\"rel_l1_thresh\": 0.2}').",
         )
+        omni_config_group.add_argument(
+            "--enable-cache-dit-summary",
+            action="store_true",
+            help="Enable cache-dit summary logging after diffusion forward passes.",
+        )
 
         # VAE memory optimization parameters
         omni_config_group.add_argument(

From b11d436e199867050472be1ee7e72ef706aff184 Mon Sep 17 00:00:00 2001
From: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Date: Wed, 28 Jan 2026 20:18:02 +0800
Subject: [PATCH 09/32] [Doc]Format profiling doc (#993)

Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 docs/contributing/profiling.md | 35 +++++++++-------------------------
 1 file changed, 9 insertions(+), 26 deletions(-)

diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md
index 4d18d0fb4..a7df8c322 100644
--- a/docs/contributing/profiling.md
+++ b/docs/contributing/profiling.md
@@ -2,7 +2,7 @@
 
 > **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
 
-vLLM-Omni uses the PyTorch Profiler to analyze performance across both **Multi-Stage LLMs** and **Diffusion Models**.
+vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
 
 ### 1. Set the Output Directory
 Before running any script, set this environment variable. The system detects this and automatically saves traces here.
@@ -11,7 +11,7 @@ Before running any script, set this environment variable. The system detects thi
 export VLLM_TORCH_PROFILER_DIR=./profiles
 ```
 
-### 2. Start Profiling
+### 2. Profiling Omni-Modality Models
 
 It is best to limit profiling to one iteration to keep trace files manageable.
 
@@ -77,12 +77,12 @@ omni_llm.close()
 
 **Examples**:
 
-1. **Qwen-omni 2.5**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
+1. **Qwen2.5-Omni**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
 
-2. **Qwen-omni 3.0**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
+2. **Qwen3-Omni**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
 
 
-**For Diffusion Models as a single stage**
+### 3. Profiling diffusion models
 
 Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
 
@@ -131,27 +131,10 @@ python image_to_video.py \
 
 2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
 
+> **Note:**
+As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
 
-**Online Inference(Async)**
-
-For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
-
-```python
-from vllm_omni import AsyncOmni
-
-# Inside an async function:
-async_omni = AsyncOmni.from_engine_args(engine_args)
-
-await async_omni.start_profile()
-
-async for output in async_omni.generate(prompt, sampling_params, request_id):
-    # Process outputs...
-    pass
-
-await async_omni.stop_profile()
-```
-
-### 3. Analyzing Omni Traces
+### 4. Analyzing Omni Traces
 
 Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
 
@@ -163,4 +146,4 @@ Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
 - [Perfetto](https://ui.perfetto.dev/)(recommended)
 - ```chrome://tracing```(Chrome only)
 
-**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/latest/dev/profiling.html)
+**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)

From c9392dff01982eeec9a8bd24fcc75f6a1a7fd9b2 Mon Sep 17 00:00:00 2001
From: Canlin Guo <canlinguosdu@gmail.com>
Date: Wed, 28 Jan 2026 22:12:52 +0800
Subject: [PATCH 10/32] [Hardware] Support platforms and plugin system (#774)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
---
 .buildkite/pipeline.yml                       |  20 +-
 .buildkite/test-amd.yaml                      |   2 +-
 docs/configuration/stage_configs.md           |   6 +-
 .../stage_configs/qwen2_5_omni.yaml           |   6 +-
 .../image_to_image/image_edit.py              |  22 +-
 .../image_to_image/image_to_image.md          |   4 +
 .../image_to_video/README.md                  |   4 +
 .../image_to_video/image_to_video.py          |  23 +-
 .../text_to_audio/text_to_audio.py            |   5 +-
 .../offline_inference/text_to_image/README.md |   4 +
 .../text_to_image/gradio_demo.py              |   9 +-
 .../text_to_image/text_to_image.py            |  23 +-
 .../text_to_video/text_to_video.md            |   4 +
 .../text_to_video/text_to_video.py            |  23 +-
 .../qwen3_omni/qwen3_omni_moe_thinking.yaml   |   2 +-
 .../diffusion/attention/test_attention_sp.py  |  21 +-
 tests/diffusion/distributed/test_comm.py      |  28 +-
 ...ion_worker.py => test_diffusion_worker.py} |  62 ++-
 .../stage_configs/npu/qwen2_5_omni_ci.yaml    |   6 +-
 .../stage_configs/qwen2_5_omni_ci.yaml        |   6 +-
 .../stage_configs/qwen3_omni_ci.yaml          |   6 +-
 .../stage_configs/rocm/qwen2_5_omni_ci.yaml   |   6 +-
 .../stage_configs/rocm/qwen3_omni_ci.yaml     |   6 +-
 .../test_diffusion_cpu_offload.py             |   6 +-
 .../offline_inference/test_diffusion_lora.py  |   5 -
 .../offline_inference/test_qwen2_5_omni.py    |   6 +-
 .../e2e/offline_inference/test_qwen3_omni.py  |   4 +-
 .../test_sequence_parallel.py                 |  31 +-
 tests/e2e/offline_inference/test_t2i_model.py |   9 +-
 .../test_zimage_tensor_parallel.py            |   4 +-
 .../stage_configs/qwen3_omni_ci.yaml          |  10 +-
 .../stage_configs/qwen3_omni_thinker_ci.yaml  |   2 +-
 .../stage_configs/rocm/qwen3_omni_ci.yaml     |   6 +-
 .../test_images_generations_lora.py           |   4 -
 tests/e2e/online_serving/test_qwen3_omni.py   |  12 +-
 tests/e2e/stage_configs/qwen3_omni_ci.yaml    |   6 +-
 tests/entrypoints/test_stage_utils.py         |  55 +-
 .../diffusion/attention/backends/abstract.py  |  56 +-
 .../attention/backends/ascend_attn.py         |  92 ----
 .../attention/backends/flash_attn.py          |  54 +-
 .../diffusion/attention/backends/registry.py  | 152 ++++++
 .../diffusion/attention/backends/sage_attn.py |   2 +-
 .../diffusion/attention/backends/sdpa.py      |   4 +
 .../diffusion/attention/backends/utils/fa.py  |   4 +-
 vllm_omni/diffusion/attention/selector.py     | 127 ++---
 vllm_omni/diffusion/distributed/comm.py       |  10 +-
 .../distributed/group_coordinator.py          |  14 +-
 .../diffusion/distributed/parallel_state.py   |  18 +-
 vllm_omni/diffusion/distributed/utils.py      |   7 +-
 vllm_omni/diffusion/envs.py                   | 163 +-----
 .../diffusion/executor/multiproc_executor.py  |   7 +-
 .../diffusion/hooks/sequence_parallel.py      |   7 +-
 vllm_omni/diffusion/layers/adalayernorm.py    |   8 +
 vllm_omni/diffusion/layers/custom_op.py       |  14 +-
 vllm_omni/diffusion/layers/rope.py            |   8 +
 .../longcat_image_transformer.py              |   4 +-
 vllm_omni/diffusion/worker/__init__.py        |  10 +-
 ...el_runner.py => diffusion_model_runner.py} |   2 +-
 ...iffusion_worker.py => diffusion_worker.py} |  30 +-
 vllm_omni/diffusion/worker/npu/__init__.py    |   7 -
 vllm_omni/diffusion/worker/npu/npu_worker.py  | 255 ---------
 vllm_omni/engine/arg_utils.py                 |  13 +-
 vllm_omni/entrypoints/omni_stage.py           | 504 ++++++++++--------
 vllm_omni/entrypoints/stage_utils.py          |   6 +-
 vllm_omni/entrypoints/utils.py                |  16 +-
 .../models/qwen2_5_omni/qwen2_5_omni.py       |   4 +-
 .../qwen2_5_omni/qwen2_5_omni_token2wav.py    |   8 +-
 .../models/qwen3_omni/qwen3_omni.py           |   6 +-
 .../model_executor/stage_configs/bagel.yaml   |   2 +-
 .../stage_configs/bagel_multiconnector.yaml   |   2 +-
 .../stage_configs/qwen2_5_omni.yaml           |   6 +-
 .../qwen2_5_omni_multiconnector.yaml          |   6 +-
 .../stage_configs/qwen3_omni_moe.yaml         |   6 +-
 .../qwen3_omni_moe_async_chunk.yaml           |   6 +-
 .../qwen3_omni_moe_multiconnector.yaml        |   6 +-
 .../stage_configs/qwen3_tts.yaml              |   2 +-
 vllm_omni/platforms/__init__.py               | 191 +++++++
 vllm_omni/platforms/cuda/__init__.py          |   6 +
 vllm_omni/platforms/cuda/platform.py          | 101 ++++
 vllm_omni/platforms/interface.py              |  90 ++++
 vllm_omni/platforms/npu/__init__.py           |   6 +
 vllm_omni/platforms/npu/platform.py           |  84 +++
 .../npu/stage_configs}/qwen2_5_omni.yaml      |  12 +-
 .../npu/stage_configs}/qwen3_omni_moe.yaml    |   6 +-
 .../npu/stage_configs}/qwen3_tts.yaml         |   2 +-
 .../npu => platforms/npu/worker}/__init__.py  |   0
 .../npu/worker}/npu_ar_model_runner.py        |   2 +-
 .../npu/worker}/npu_ar_worker.py              |   5 +-
 .../worker}/npu_generation_model_runner.py    |   4 +-
 .../npu/worker}/npu_generation_worker.py      |   5 +-
 .../npu/worker}/npu_model_runner.py           |   0
 vllm_omni/platforms/rocm/__init__.py          |   6 +
 vllm_omni/platforms/rocm/platform.py          |  97 ++++
 .../rocm/stage_configs}/qwen2_5_omni.yaml     |   6 +-
 .../rocm/stage_configs}/qwen3_omni_moe.yaml   |   6 +-
 vllm_omni/platforms/xpu/__init__.py           |   6 +
 vllm_omni/platforms/xpu/platform.py           |  76 +++
 vllm_omni/plugins/__init__.py                 |  81 +++
 vllm_omni/utils/__init__.py                   |  13 -
 vllm_omni/utils/platform_utils.py             |  58 --
 vllm_omni/worker/gpu_ar_worker.py             |   3 +-
 vllm_omni/worker/gpu_generation_worker.py     |   3 +-
 vllm_omni/worker/mixins.py                    |  13 +
 103 files changed, 1770 insertions(+), 1218 deletions(-)
 rename tests/diffusion/{test_gpu_diffusion_worker.py => test_diffusion_worker.py} (81%)
 delete mode 100644 vllm_omni/diffusion/attention/backends/ascend_attn.py
 create mode 100644 vllm_omni/diffusion/attention/backends/registry.py
 rename vllm_omni/diffusion/worker/{gpu_diffusion_model_runner.py => diffusion_model_runner.py} (99%)
 rename vllm_omni/diffusion/worker/{gpu_diffusion_worker.py => diffusion_worker.py} (94%)
 delete mode 100644 vllm_omni/diffusion/worker/npu/__init__.py
 delete mode 100644 vllm_omni/diffusion/worker/npu/npu_worker.py
 create mode 100644 vllm_omni/platforms/__init__.py
 create mode 100644 vllm_omni/platforms/cuda/__init__.py
 create mode 100644 vllm_omni/platforms/cuda/platform.py
 create mode 100644 vllm_omni/platforms/interface.py
 create mode 100644 vllm_omni/platforms/npu/__init__.py
 create mode 100644 vllm_omni/platforms/npu/platform.py
 rename vllm_omni/{model_executor/stage_configs/npu => platforms/npu/stage_configs}/qwen2_5_omni.yaml (89%)
 rename vllm_omni/{model_executor/stage_configs/npu => platforms/npu/stage_configs}/qwen3_omni_moe.yaml (93%)
 rename vllm_omni/{model_executor/stage_configs/npu => platforms/npu/stage_configs}/qwen3_tts.yaml (89%)
 rename vllm_omni/{worker/npu => platforms/npu/worker}/__init__.py (100%)
 rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_ar_model_runner.py (99%)
 rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_ar_worker.py (74%)
 rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_generation_model_runner.py (99%)
 rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_generation_worker.py (72%)
 rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_model_runner.py (100%)
 create mode 100644 vllm_omni/platforms/rocm/__init__.py
 create mode 100644 vllm_omni/platforms/rocm/platform.py
 rename vllm_omni/{model_executor/stage_configs/rocm => platforms/rocm/stage_configs}/qwen2_5_omni.yaml (94%)
 rename vllm_omni/{model_executor/stage_configs/rocm => platforms/rocm/stage_configs}/qwen3_omni_moe.yaml (93%)
 create mode 100644 vllm_omni/platforms/xpu/__init__.py
 create mode 100644 vllm_omni/platforms/xpu/platform.py
 create mode 100644 vllm_omni/plugins/__init__.py
 delete mode 100644 vllm_omni/utils/platform_utils.py
 create mode 100644 vllm_omni/worker/mixins.py

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index d1e71aba2..78a400cd8 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -18,11 +18,23 @@ steps:
   #     queue: "ascend"
 
   - label: "Simple Unit Test"
-    depends_on: ~
+    depends_on: image-build
     commands:
-      - ".buildkite/scripts/simple_test.sh"
+      - pytest -v -s tests/entrypoints/
+      - pytest -v -s tests/diffusion/cache/
+      - pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
+      - pytest -v -s tests/worker/
     agents:
-      queue: "cpu_queue_premerge"
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
     timeout_in_minutes: 20
@@ -149,7 +161,7 @@ steps:
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
-      - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+      - pytest -s -v tests/diffusion/test_diffusion_worker.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index 86d65f15b..08a3f7a1c 100644
--- a/.buildkite/test-amd.yaml
+++ b/.buildkite/test-amd.yaml
@@ -54,7 +54,7 @@ steps:
   commands:
     - export MIOPEN_DEBUG_CONV_DIRECT=0
     - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+    - pytest -s -v tests/diffusion/test_diffusion_worker.py
 
 - label: "Omni Model Test Qwen2-5-Omni"
   timeout_in_minutes: 15
diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md
index 18972eb79..1f1a18902 100644
--- a/docs/configuration/stage_configs.md
+++ b/docs/configuration/stage_configs.md
@@ -40,7 +40,7 @@ stage_args:
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
+      worker_type: ar # The specific worker used
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
       gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
       enforce_eager: true  # Now we only support eager mode
@@ -66,7 +66,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -92,7 +92,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/docs/configuration/stage_configs/qwen2_5_omni.yaml b/docs/configuration/stage_configs/qwen2_5_omni.yaml
index e71f348de..e20e79c3e 100644
--- a/docs/configuration/stage_configs/qwen2_5_omni.yaml
+++ b/docs/configuration/stage_configs/qwen2_5_omni.yaml
@@ -8,7 +8,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -34,7 +34,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -60,7 +60,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py
index f2065e6b2..896c611b4 100644
--- a/examples/offline_inference/image_to_image/image_edit.py
+++ b/examples/offline_inference/image_to_image/image_edit.py
@@ -81,7 +81,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -280,6 +280,16 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -306,12 +316,8 @@ def main():
     else:
         input_image = input_images
 
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
@@ -344,8 +350,8 @@ def main():
     # Initialize Omni with appropriate pipeline
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
diff --git a/examples/offline_inference/image_to_image/image_to_image.md b/examples/offline_inference/image_to_image/image_to_image.md
index d624b4684..b9e25bf98 100644
--- a/examples/offline_inference/image_to_image/image_to_image.md
+++ b/examples/offline_inference/image_to_image/image_to_image.md
@@ -47,4 +47,8 @@ Key arguments:
 - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md
index afebb1963..52bb389e7 100644
--- a/examples/offline_inference/image_to_video/README.md
+++ b/examples/offline_inference/image_to_video/README.md
@@ -54,4 +54,8 @@ Key arguments:
 - `--num_inference_steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
+- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py
index 06c13fbee..4fca7825c 100644
--- a/examples/offline_inference/image_to_video/image_to_video.py
+++ b/examples/offline_inference/image_to_video/image_to_video.py
@@ -29,7 +29,7 @@ import torch
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -59,6 +59,16 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -80,8 +90,7 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Load input image
     image = PIL.Image.open(args.image).convert("RGB")
@@ -98,17 +107,13 @@ def main():
     # Resize image to target dimensions
     image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
 
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
         enable_cpu_offload=args.enable_cpu_offload,
diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py
index 1124dbdb4..0a9efcca5 100644
--- a/examples/offline_inference/text_to_audio/text_to_audio.py
+++ b/examples/offline_inference/text_to_audio/text_to_audio.py
@@ -22,7 +22,7 @@ import torch
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -118,8 +118,7 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     print(f"\n{'=' * 60}")
     print("Stable Audio Open - Text-to-Audio Generation")
diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md
index f4de19db8..ab28b115c 100644
--- a/examples/offline_inference/text_to_image/README.md
+++ b/examples/offline_inference/text_to_image/README.md
@@ -96,8 +96,12 @@ Key arguments:
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+
 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
 
 ## Web UI Demo
diff --git a/examples/offline_inference/text_to_image/gradio_demo.py b/examples/offline_inference/text_to_image/gradio_demo.py
index 30446f50e..614e40903 100644
--- a/examples/offline_inference/text_to_image/gradio_demo.py
+++ b/examples/offline_inference/text_to_image/gradio_demo.py
@@ -7,7 +7,7 @@ import torch
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
     "1:1": (1328, 1328),
@@ -62,8 +62,8 @@ def parse_args() -> argparse.Namespace:
 @lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
     # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    vae_use_slicing = current_omni_platform.is_npu()
+    vae_use_tiling = current_omni_platform.is_npu()
     return Omni(
         model=model_name,
         vae_use_slicing=vae_use_slicing,
@@ -72,7 +72,6 @@ def get_omni(model_name: str) -> Omni:
 
 
 def build_demo(args: argparse.Namespace) -> gr.Blocks:
-    device = detect_device_type()
     omni = get_omni(args.model)
 
     def run_inference(
@@ -99,7 +98,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
             raise gr.Error("Inference steps must be a positive integer.")
         if num_images not in {1, 2, 3, 4}:
             raise gr.Error("Number of images must be 1, 2, 3, or 4.")
-        generator = torch.Generator(device=device).manual_seed(seed)
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
         outputs = omni.generate(
             prompt.strip(),
             OmniDiffusionSamplingParams(
diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py
index efc5010e0..3e0628a15 100644
--- a/examples/offline_inference/text_to_image/text_to_image.py
+++ b/examples/offline_inference/text_to_image/text_to_image.py
@@ -12,7 +12,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -113,17 +113,22 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
-
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Configure cache based on backend type
     cache_config = None
@@ -167,8 +172,8 @@ def main():
 
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_cache_dit_summary=args.enable_cache_dit_summary,
diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md
index 7204ec58a..cd5bd2be4 100644
--- a/examples/offline_inference/text_to_video/text_to_video.md
+++ b/examples/offline_inference/text_to_video/text_to_video.md
@@ -29,4 +29,8 @@ Key arguments:
 - `--boundary_ratio`: Boundary split ratio for low/high DiT.
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index 49522278f..464ff9cca 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -13,7 +13,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -54,6 +54,16 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     parser.add_argument(
         "--enforce_eager",
         action="store_true",
@@ -81,12 +91,7 @@ def parse_args() -> argparse.Namespace:
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
-
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
     cache_config = None
@@ -118,8 +123,8 @@ def main():
 
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
         cache_backend=args.cache_backend,
diff --git a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml
index a69bd2fe1..34b63e26e 100644
--- a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml
+++ b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml
@@ -13,7 +13,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.9
       enforce_eager: true
diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py
index c0061daf9..447d1ea93 100644
--- a/tests/diffusion/attention/test_attention_sp.py
+++ b/tests/diffusion/attention/test_attention_sp.py
@@ -6,7 +6,6 @@ import tempfile
 
 import pytest
 import torch
-from vllm.platforms import current_platform
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import (
@@ -19,15 +18,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.utils.platform_utils import detect_device_type
-
-device_type = detect_device_type()
-if device_type == "cuda":
-    torch_device = torch.cuda
-elif device_type == "npu":
-    torch_device = torch.npu
-else:
-    raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
+from vllm_omni.platforms import current_omni_platform
 
 
 def update_environment_variables(envs_dict: dict[str, str]):
@@ -171,7 +162,7 @@ def test_sequence_parallel(
     sequence_parallel_size = ulysses_degree * ring_degree
 
     # Skip if not enough GPUs available
-    available_gpus = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
     if available_gpus < sequence_parallel_size:
         pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")
 
@@ -338,13 +329,13 @@ def ulysses_attention_on_test_model(
     """Run Ulysses attention test on a test model and save results for comparison."""
     # Use fixed seed for reproducibility across baseline and SP runs
     RANDOM_SEED = 42
-    current_platform.seed_everything(RANDOM_SEED)
+    current_omni_platform.seed_everything(RANDOM_SEED)
 
     mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
     print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
 
-    device = torch.device(f"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
 
@@ -423,7 +414,7 @@ def ulysses_attention_on_test_model(
             # Generate and save full input data with fixed seed
             # Reinitialize RNG to ensure reproducibility
             torch.manual_seed(42)
-            torch_device.manual_seed_all(42)
+            current_omni_platform.seed_everything(42)
             full_hidden_states = torch.randn(
                 (batch_size, seq_len, hidden_size),
                 dtype=dtype,
diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py
index 0e76f27b7..3e408c10a 100644
--- a/tests/diffusion/distributed/test_comm.py
+++ b/tests/diffusion/distributed/test_comm.py
@@ -14,15 +14,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
-from vllm_omni.utils.platform_utils import detect_device_type
-
-device_type = detect_device_type()
-if device_type == "cuda":
-    torch_device = torch.cuda
-elif device_type == "npu":
-    torch_device = torch.npu
-else:
-    raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
+from vllm_omni.platforms import current_omni_platform
 
 
 def update_environment_variables(envs_dict: dict[str, str]):
@@ -49,7 +41,7 @@ def test_4d_identity(
 ):
     """Test that two consecutive all-to-all operations return the original input."""
     # Skip if not enough GPUs available
-    available_gpus = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
     if available_gpus < world_size:
         pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
 
@@ -85,8 +77,8 @@ def _test_4d_identity_worker(
 ):
     """Worker function for test_4d_identity."""
     # Set device
-    device = torch.device(f"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
 
     # Set environment variables for distributed training
     update_environment_variables(
@@ -183,7 +175,7 @@ def test_5d_identity(
 ):
     """Test that two consecutive all-to-all operations return the original input."""
     # Skip if not enough GPUs available
-    available_gpus = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
     if available_gpus < world_size:
         pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
 
@@ -219,8 +211,8 @@ def _test_5d_identity_worker(
 ):
     """Worker function for test_5d_identity."""
     # Set device
-    device = torch.device(f"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
 
     # Set environment variables for distributed training
     update_environment_variables(
@@ -316,7 +308,7 @@ def test_ring_p2p(
 ):
     """Test Ring P2P communication (send_recv)."""
     # Skip if not enough GPUs available
-    available_gpus = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
     if available_gpus < world_size:
         pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
 
@@ -339,8 +331,8 @@ def _test_ring_p2p_worker(
     import sys
 
     # Set device
-    device = torch.device(f"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
 
     # Set env vars
     # Use a different port to avoid conflict with other tests if run in parallel
diff --git a/tests/diffusion/test_gpu_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py
similarity index 81%
rename from tests/diffusion/test_gpu_diffusion_worker.py
rename to tests/diffusion/test_diffusion_worker.py
index 7a43710c8..220f210a3 100644
--- a/tests/diffusion/test_gpu_diffusion_worker.py
+++ b/tests/diffusion/test_diffusion_worker.py
@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Unit tests for GPUDiffusionWorker class.
+Unit tests for DiffusionWorker class.
 
-This module tests the GPUDiffusionWorker implementation:
+This module tests the DiffusionWorker implementation:
 - load_weights: Loading model weights
 - sleep: Putting worker into sleep mode (levels 1 and 2)
 - wake_up: Waking worker from sleep mode
@@ -15,7 +15,7 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 
-from vllm_omni.diffusion.worker.gpu_diffusion_worker import GPUDiffusionWorker
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
 
 @pytest.fixture
@@ -33,18 +33,19 @@ def mock_od_config():
 
 @pytest.fixture
 def mock_gpu_worker(mock_od_config):
-    """Create a GPUDiffusionWorker with mocked initialization."""
-    with patch.object(GPUDiffusionWorker, "init_device"):
-        worker = GPUDiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
+    """Create a DiffusionWorker with mocked initialization."""
+    with patch.object(DiffusionWorker, "init_device"):
+        worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
         # Mock the model_runner with pipeline
         worker.model_runner = Mock()
         worker.model_runner.pipeline = Mock()
+        worker.device = torch.device("cuda", 0)
         worker._sleep_saved_buffers = {}
         return worker
 
 
-class TestGPUDiffusionWorkerLoadWeights:
-    """Test GPUDiffusionWorker.load_weights method."""
+class TestDiffusionWorkerLoadWeights:
+    """Test DiffusionWorker.load_weights method."""
 
     def test_load_weights_calls_pipeline(self, mock_gpu_worker):
         """Test that load_weights delegates to model_runner.load_weights."""
@@ -75,20 +76,21 @@ class TestGPUDiffusionWorkerLoadWeights:
         assert result == set()
 
 
-class TestGPUDiffusionWorkerSleep:
-    """Test GPUDiffusionWorker.sleep method."""
+class TestDiffusionWorkerSleep:
+    """Test DiffusionWorker.sleep method."""
 
-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_1(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
         # Setup memory info mocks
-        # Before sleep: 1GB free, 8GB total
-        # After sleep: 3GB free, 8GB total (freed 2GB)
-        mock_mem_info.side_effect = [
-            (1 * 1024**3, 8 * 1024**3),  # Before sleep
-            (3 * 1024**3, 8 * 1024**3),  # After sleep
+        # Before sleep: 1GB free
+        # After sleep: 3GB free (freed 2GB)
+        mock_platform.get_free_memory.side_effect = [
+            1 * 1024**3,  # Before sleep
+            3 * 1024**3,  # After sleep
         ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -104,15 +106,16 @@ class TestGPUDiffusionWorkerSleep:
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
 
-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_2(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
         # Setup memory info mocks
-        mock_mem_info.side_effect = [
-            (1 * 1024**3, 8 * 1024**3),  # Before sleep
-            (5 * 1024**3, 8 * 1024**3),  # After sleep (freed 4GB)
+        mock_platform.get_free_memory.side_effect = [
+            1 * 1024**3,  # Before sleep
+            5 * 1024**3,  # After sleep (freed 4GB)
         ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -141,15 +144,16 @@ class TestGPUDiffusionWorkerSleep:
         assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
         assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
 
-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
         # Simulate memory increase (should trigger assertion error)
-        mock_mem_info.side_effect = [
-            (3 * 1024**3, 8 * 1024**3),  # Before sleep: 3GB free
-            (1 * 1024**3, 8 * 1024**3),  # After sleep: 1GB free (negative freed!)
+        mock_platform.get_free_memory.side_effect = [
+            3 * 1024**3,  # Before sleep: 3GB free
+            1 * 1024**3,  # After sleep: 1GB free (negative freed!)
         ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         mock_allocator = Mock()
         mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
@@ -160,8 +164,8 @@ class TestGPUDiffusionWorkerSleep:
             mock_gpu_worker.sleep(level=1)
 
 
-class TestGPUDiffusionWorkerWakeUp:
-    """Test GPUDiffusionWorker.wake_up method."""
+class TestDiffusionWorkerWakeUp:
+    """Test DiffusionWorker.wake_up method."""
 
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
     def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
diff --git a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml
index d2aae1c0d..d6b04258a 100644
--- a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml
@@ -10,7 +10,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -40,7 +40,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -70,7 +70,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml
index eb9109adc..e093ec51b 100644
--- a/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml
@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
index 0780420fd..7ac435ea5 100644
--- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml
index 96e9d7fa7..474df5e79 100644
--- a/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml
@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
index 0780420fd..7ac435ea5 100644
--- a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py
index f6606dce6..35e106df8 100644
--- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py
@@ -7,7 +7,7 @@ from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
 from tests.utils import GPUMemoryMonitor
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.utils.platform_utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -20,7 +20,7 @@ models = ["riverclouds/qwen_image_random"]
 
 
 def inference(model_name: str, offload: bool = True):
-    torch.cuda.empty_cache()
+    current_omni_platform.empty_cache()
     device_index = torch.cuda.current_device()
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
@@ -45,7 +45,7 @@ def inference(model_name: str, offload: bool = True):
     return peak
 
 
-@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
+@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
 @pytest.mark.parametrize("model_name", models)
 def test_cpu_offload_diffusion_model(model_name: str):
     try:
diff --git a/tests/e2e/offline_inference/test_diffusion_lora.py b/tests/e2e/offline_inference/test_diffusion_lora.py
index ef734da61..2465bc607 100644
--- a/tests/e2e/offline_inference/test_diffusion_lora.py
+++ b/tests/e2e/offline_inference/test_diffusion_lora.py
@@ -9,7 +9,6 @@ from safetensors.torch import save_file
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import is_npu
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -25,10 +24,6 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 # model to reduce runtime and avoid extra downloads.
 models = ["Tongyi-MAI/Z-Image-Turbo"]
 
-# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly.
-if is_npu():
-    pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.", allow_module_level=True)
-
 
 @pytest.mark.parametrize("model_name", models)
 def test_diffusion_model(model_name: str, tmp_path: Path):
diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni.py b/tests/e2e/offline_inference/test_qwen2_5_omni.py
index 09b0306c1..3537ea127 100644
--- a/tests/e2e/offline_inference/test_qwen2_5_omni.py
+++ b/tests/e2e/offline_inference/test_qwen2_5_omni.py
@@ -13,7 +13,7 @@ from vllm.assets.video import VideoAsset
 from vllm.envs import VLLM_USE_MODELSCOPE
 from vllm.multimodal.image import convert_image_mode
 
-from vllm_omni.utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 from .conftest import OmniRunner
 from .utils import create_new_process_for_each_test
@@ -21,9 +21,9 @@ from .utils import create_new_process_for_each_test
 models = ["Qwen/Qwen2.5-Omni-3B"]
 
 # CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
-if is_npu():
+if current_omni_platform.is_npu():
     stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
-elif is_rocm():
+elif current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_config = str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
 else:
diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py
index edd5958a0..470cdb73b 100644
--- a/tests/e2e/offline_inference/test_qwen3_omni.py
+++ b/tests/e2e/offline_inference/test_qwen3_omni.py
@@ -14,14 +14,14 @@ from pathlib import Path
 import pytest
 from vllm.assets.video import VideoAsset
 
-from vllm_omni.utils import is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 from .conftest import OmniRunner
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if is_rocm():
+if current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py
index 367e2e59e..262c80243 100644
--- a/tests/e2e/offline_inference/test_sequence_parallel.py
+++ b/tests/e2e/offline_inference/test_sequence_parallel.py
@@ -30,8 +30,7 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.parallel_state import device_count
-from vllm_omni.diffusion.envs import get_device_name
+from vllm_omni.platforms import current_omni_platform
 
 models = ["riverclouds/qwen_image_random"]
 
@@ -63,9 +62,9 @@ def _cleanup_distributed():
 
     gc.collect()
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
 
     time.sleep(5)
 
@@ -88,7 +87,7 @@ def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height
                 width=width,
                 num_inference_steps=4,
                 guidance_scale=0.0,
-                generator=torch.Generator(get_device_name()).manual_seed(seed),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
                 num_outputs_per_prompt=1,
             ),
         )
@@ -125,7 +124,7 @@ def _run_sp(
                 width=width,
                 num_inference_steps=4,
                 guidance_scale=0.0,
-                generator=torch.Generator(get_device_name()).manual_seed(seed),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
                 num_outputs_per_prompt=1,
             ),
         )
@@ -157,8 +156,8 @@ def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
 @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 device_count() < 2:
-        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+    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
@@ -177,8 +176,8 @@ def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str
 @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 device_count() < 2:
-        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+    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
@@ -212,8 +211,8 @@ def test_sequence_parallel(
         )
 
     sp_size = ulysses_degree * ring_degree
-    if device_count() < sp_size:
-        pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available")
+    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
@@ -243,8 +242,10 @@ def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_ba
     ulysses_degree = 4
     ring_degree = 1
 
-    if device_count() < ulysses_degree * ring_degree:
-        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
+    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"
+        )
 
     height = 272
     width = 272
diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py
index ba6cac34c..e7351f157 100644
--- a/tests/e2e/offline_inference/test_t2i_model.py
+++ b/tests/e2e/offline_inference/test_t2i_model.py
@@ -7,7 +7,7 @@ import torch
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -21,12 +21,11 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
 models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
 
-# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly
 # Modelscope can't find riverclouds/qwen_image_random
 # TODO: When NPU support is ready, remove this branch.
-if is_npu():
-    models = ["Qwen/Qwen-Image"]
-elif is_rocm():
+if current_omni_platform.is_npu():
+    models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"]
+elif current_omni_platform.is_rocm():
     # TODO: When ROCm support is ready, remove this branch.
     # vLLM V0.11.0 has issues running riverclouds/qwen_image_random
     # on ROCm
diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py
index ea8f0acfe..0e3b97ec3 100644
--- a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py
+++ b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py
@@ -23,7 +23,7 @@ from tests.utils import GPUMemoryMonitor
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -129,7 +129,7 @@ def _run_zimage_generate(
 
 @pytest.mark.integration
 def test_zimage_tensor_parallel_tp2(tmp_path: Path):
-    if is_npu() or is_rocm():
+    if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
         pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
     if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
         pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")
diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
index c5987ebba..4db11e11f 100644
--- a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
@@ -6,13 +6,14 @@
 # The following config has been verified on 2x H100-80G GPUs.
 stage_args:
   - stage_id: 0
+    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "0,1"
       max_batch_size: 5
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -20,6 +21,7 @@ stage_args:
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
+      max_num_batched_tokens: 32768
       hf_config_name: thinker_config
       tensor_parallel_size: 2
       load_format: dummy
@@ -36,13 +38,14 @@ stage_args:
       repetition_penalty: 1.05
 
   - stage_id: 1
+    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
        devices: "1"
        max_batch_size: 5
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: false
@@ -67,13 +70,14 @@ stage_args:
       stop_token_ids: [2150]
 
   - stage_id: 2
+    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml
index 64c6b6ee3..89d03966d 100644
--- a/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml
+++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml
@@ -7,7 +7,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: true
diff --git a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
index ef0bae0fa..0488a06c9 100644
--- a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
+++ b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/tests/e2e/online_serving/test_images_generations_lora.py b/tests/e2e/online_serving/test_images_generations_lora.py
index 82044cd52..11165c6d2 100644
--- a/tests/e2e/online_serving/test_images_generations_lora.py
+++ b/tests/e2e/online_serving/test_images_generations_lora.py
@@ -27,8 +27,6 @@ from PIL import Image
 from safetensors.torch import save_file
 from vllm.utils.network_utils import get_open_port
 
-from vllm_omni.utils.platform_utils import is_npu
-
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 MODEL = "Tongyi-MAI/Z-Image-Turbo"
@@ -122,8 +120,6 @@ class OmniServer:
 
 @pytest.fixture(scope="module")
 def omni_server():
-    if is_npu():
-        pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.")
     with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
         yield server
 
diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py
index 8564104a0..eb134ab5e 100644
--- a/tests/e2e/online_serving/test_qwen3_omni.py
+++ b/tests/e2e/online_serving/test_qwen3_omni.py
@@ -27,7 +27,7 @@ from tests.conftest import (
     generate_synthetic_video,
     modify_stage_config,
 )
-from vllm_omni.utils import is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
@@ -57,7 +57,7 @@ def get_chunk_config():
 
 CHUNK_CONFIG_PATH = get_chunk_config()
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if is_rocm():
+if current_omni_platform.is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
@@ -144,7 +144,9 @@ def get_max_batch_size(size_type="few"):
     return batch_sizes.get(size_type, 5)
 
 
-@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues"
+)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
     """
@@ -224,7 +226,9 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     assert similarity > 0.9, "The audio content is not same as the text"
 
 
-@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues"
+)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
     """
diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml
index 60343dfdd..d7b07ddc9 100644
--- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml
@@ -12,7 +12,7 @@ stage_args:
   engine_args:
     model_stage: thinker
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+    worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
     gpu_memory_utilization: 0.6
     enforce_eager: false
@@ -43,7 +43,7 @@ stage_args:
   engine_args:
     model_stage: talker
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+    worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
     gpu_memory_utilization: 0.3
     enforce_eager: false
@@ -71,7 +71,7 @@ stage_args:
   engine_args:
     model_stage: code2wav
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+    worker_type: generation
     scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
     enforce_eager: true
     trust_remote_code: true
diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py
index 67a2550b7..ac503639b 100644
--- a/tests/entrypoints/test_stage_utils.py
+++ b/tests/entrypoints/test_stage_utils.py
@@ -1,5 +1,6 @@
 import os
 import sys
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -42,6 +43,14 @@ def _make_dummy_torch(call_log):
     return _Torch
 
 
+def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
+    """Create a mock platform for testing."""
+    mock_platform = MagicMock()
+    mock_platform.device_type = device_type
+    mock_platform.device_control_env_var = env_var
+    return mock_platform
+
+
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_respects_logical_ids(monkeypatch):
     # Preserve an existing logical mapping and ensure devices "0,1" map through it.
@@ -49,9 +58,51 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
     call_log: list[int] = []
     dummy_torch = _make_dummy_torch(call_log)
     monkeypatch.setitem(sys.modules, "torch", dummy_torch)
-    monkeypatch.setattr("vllm_omni.utils.detect_device_type", lambda: "cuda")
-    monkeypatch.setattr("vllm_omni.utils.get_device_control_env_var", lambda: "CUDA_VISIBLE_DEVICES")
+
+    # Mock the platform at the source module where it's defined
+    mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
+    monkeypatch.setattr(
+        "vllm_omni.platforms.current_omni_platform",
+        mock_platform,
+    )
 
     set_stage_devices(stage_id=0, devices="0,1")
 
     assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
+
+
+@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
+def test_set_stage_devices_npu_platform(monkeypatch):
+    """Test that set_stage_devices works correctly for NPU platform."""
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5")
+    call_log: list[int] = []
+
+    # Create NPU mock torch
+    class _Npu:
+        @staticmethod
+        def is_available():
+            return True
+
+        @staticmethod
+        def set_device(idx):
+            call_log.append(idx)
+
+        @staticmethod
+        def device_count():
+            return 2
+
+    class _NpuTorch:
+        npu = _Npu
+
+    monkeypatch.setitem(sys.modules, "torch", _NpuTorch)
+
+    # Mock NPU platform at the source module where it's defined
+    mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
+    monkeypatch.setattr(
+        "vllm_omni.platforms.current_omni_platform",
+        mock_platform,
+    )
+
+    set_stage_devices(stage_id=0, devices="0,1")
+
+    assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"
diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py
index 54906efc9..d0a62bcd9 100644
--- a/vllm_omni/diffusion/attention/backends/abstract.py
+++ b/vllm_omni/diffusion/attention/backends/abstract.py
@@ -7,12 +7,18 @@ from typing import Generic, TypeVar
 
 import torch
 
+from vllm_omni.platforms import current_omni_platform
+
 
 class AttentionBackend(ABC):
     """Abstract class for diffusion attention backends."""
 
     accept_output_buffer: bool = False
 
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return False
+
     @staticmethod
     @abstractmethod
     def get_name() -> str:
@@ -77,12 +83,58 @@ class AttentionImpl(ABC, Generic[T]):
     ) -> None:
         raise NotImplementedError
 
-    @abstractmethod
     def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: T,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        """Dispatch to platform-specific forward implementation."""
+        if current_omni_platform.is_rocm():
+            return self.forward_hip(query, key, value, attn_metadata)
+        elif current_omni_platform.is_cuda():
+            return self.forward_cuda(query, key, value, attn_metadata)
+        elif current_omni_platform.is_npu():
+            return self.forward_npu(query, key, value, attn_metadata)
+        elif current_omni_platform.is_xpu():
+            return self.forward_xpu(query, key, value, attn_metadata)
+        else:
+            raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
     ) -> torch.Tensor:
         raise NotImplementedError
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward_hip(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        # By default, HIP ops are compatible with CUDA ops.
+        return self.forward_cuda(query, key, value, attn_metadata)
diff --git a/vllm_omni/diffusion/attention/backends/ascend_attn.py b/vllm_omni/diffusion/attention/backends/ascend_attn.py
deleted file mode 100644
index 6563385c9..000000000
--- a/vllm_omni/diffusion/attention/backends/ascend_attn.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from importlib.util import find_spec
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
-
-logger = init_logger(__name__)
-
-
-class AscendAttentionBackend(AttentionBackend):
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [x for x in range(1024)]  # todo
-
-    @staticmethod
-    def get_name() -> str:
-        return "ASCEND"
-
-    @staticmethod
-    def get_impl_cls() -> type["AscendAttentionBackendImpl"]:
-        return AscendAttentionBackendImpl
-
-
-class AscendAttentionBackendImpl(AttentionImpl):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-
-    def forward_native(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
-        attention_mask = attn_metadata.attn_mask if attn_metadata else None
-
-        output = torch.nn.functional.scaled_dot_product_attention(
-            query,
-            key,
-            value,
-            attn_mask=attention_mask,
-            dropout_p=0.0,
-            is_causal=self.causal,
-            scale=self.softmax_scale,
-        )
-        out = output.permute(0, 2, 1, 3)
-        return out
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        if find_spec("mindiesd"):
-            from mindiesd import attention_forward
-
-            attention_mask = attn_metadata.attn_mask if attn_metadata else None
-            output = attention_forward(
-                query,
-                key,
-                value,
-                attn_mask=attention_mask,
-                opt_mode="manual",
-                op_type="fused_attn_score",
-                layout="BNSD",
-            )
-            return output
-        else:
-            return self.forward_native(query, key, value, attn_metadata)
diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py
index e9102b1d3..a9f2d925c 100644
--- a/vllm_omni/diffusion/attention/backends/flash_attn.py
+++ b/vllm_omni/diffusion/attention/backends/flash_attn.py
@@ -9,13 +9,6 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
-from vllm_omni.diffusion.attention.backends.utils.fa import (
-    _pad_input,
-    _unpad_input,
-    _upad_input,
-    flash_attn_func,
-    flash_attn_varlen_func,
-)
 
 logger = init_logger(__name__)
 
@@ -23,6 +16,10 @@ logger = init_logger(__name__)
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         return [64, 96, 128, 192, 256]
@@ -51,25 +48,22 @@ class FlashAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
 
-    def forward(
+    def forward_cuda(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
-        """
-        Flash attention implementation.
+        """CUDA/ROCm flash attention implementation."""
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+            flash_attn_func,
+            flash_attn_varlen_func,
+        )
 
-        Args:
-            query: (batch_size, seq_len, num_heads, head_dim)
-            key: (batch_size, seq_len, num_heads, head_dim)
-            value: (batch_size, seq_len, num_heads, head_dim)
-            attn_metadata: AttentionMetadata. Attention mask is supported as attn_metadata.attn_mask
-
-        Returns:
-            (batch_size, seq_len, num_heads, head_dim)
-        """
         query_length = query.size(1)
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         #  Contains at least one padding token in the sequence
@@ -106,3 +100,25 @@ class FlashAttentionImpl(AttentionImpl):
                 softmax_scale=self.softmax_scale,
             )
         return out
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        """NPU attention implementation using mindiesd."""
+        from mindiesd import attention_forward
+
+        attention_mask = attn_metadata.attn_mask if attn_metadata else None
+        output = attention_forward(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            opt_mode="manual",
+            op_type="fused_attn_score",
+            layout="BNSD",
+        )
+        return output
diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py
new file mode 100644
index 000000000..a77223d16
--- /dev/null
+++ b/vllm_omni/diffusion/attention/backends/registry.py
@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusion attention backend registry.
+
+This module provides an enum-based registry for diffusion attention backends,
+similar to vLLM's AttentionBackendEnum. Each backend registers its class path,
+and platforms can override or extend backends using register_backend().
+"""
+
+from collections.abc import Callable
+from enum import Enum, EnumMeta
+from typing import TYPE_CHECKING
+
+from vllm.logger import init_logger
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend
+
+logger = init_logger(__name__)
+
+
+class _DiffusionBackendEnumMeta(EnumMeta):
+    """Metaclass for DiffusionAttentionBackendEnum to provide better error messages."""
+
+    def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":
+        """Get backend by name with helpful error messages."""
+        try:
+            return super().__getitem__(name)  # type: ignore[return-value]
+        except KeyError:
+            members = list(cls.__members__.keys())
+            valid_backends = ", ".join(members)
+            raise ValueError(
+                f"Unknown diffusion attention backend: '{name}'. Valid options are: {valid_backends}"
+            ) from None
+
+
+class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
+    """Enumeration of all supported diffusion attention backends.
+
+    The enum value is the default class path, but this can be overridden
+    at runtime using register_backend().
+
+    To get the actual backend class (respecting overrides), use:
+        backend.get_class()
+
+    Example:
+        # Get backend class
+        backend = DiffusionAttentionBackendEnum.FLASH_ATTN
+        backend_cls = backend.get_class()
+
+        # Register custom backend
+        @register_diffusion_backend(DiffusionAttentionBackendEnum.CUSTOM)
+        class MyCustomBackend:
+            ...
+    """
+
+    # Common backends (available on most platforms)
+    FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
+    TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
+    SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
+
+    def get_path(self, include_classname: bool = True) -> str:
+        """Get the class path for this backend (respects overrides).
+
+        Returns:
+            The fully qualified class path string
+
+        Raises:
+            ValueError: If backend has empty path and is not registered
+        """
+        path = _DIFFUSION_ATTN_OVERRIDES.get(self, self.value)
+        if not path:
+            raise ValueError(
+                f"Backend {self.name} must be registered before use. "
+                f"Use register_diffusion_backend(DiffusionAttentionBackendEnum.{self.name}, "
+                f"'your.module.YourClass')"
+            )
+        if not include_classname:
+            path = path.rsplit(".", 1)[0]
+        return path
+
+    def get_class(self) -> "type[AttentionBackend]":
+        """Get the backend class (respects overrides).
+
+        Returns:
+            The backend class
+
+        Raises:
+            ImportError: If the backend class cannot be imported
+            ValueError: If backend has empty path and is not registered
+        """
+        return resolve_obj_by_qualname(self.get_path())
+
+    def is_overridden(self) -> bool:
+        """Check if this backend has been overridden.
+
+        Returns:
+            True if the backend has a registered override
+        """
+        return self in _DIFFUSION_ATTN_OVERRIDES
+
+    def clear_override(self) -> None:
+        """Clear any override for this backend, reverting to the default."""
+        _DIFFUSION_ATTN_OVERRIDES.pop(self, None)
+
+
+# Override registry
+_DIFFUSION_ATTN_OVERRIDES: dict[DiffusionAttentionBackendEnum, str] = {}
+
+
+def register_diffusion_backend(
+    backend: DiffusionAttentionBackendEnum,
+    class_path: str | None = None,
+) -> Callable[[type], type]:
+    """Register or override a diffusion backend implementation.
+
+    Args:
+        backend: The DiffusionAttentionBackendEnum member to register
+        class_path: Optional class path. If not provided and used as
+            decorator, will be auto-generated from the class.
+
+    Returns:
+        Decorator function if class_path is None, otherwise a no-op
+
+    Examples:
+        # Override an existing backend
+        @register_diffusion_backend(DiffusionAttentionBackendEnum.FLASH_ATTN)
+        class MyCustomFlashAttn:
+            ...
+
+        # Override an existing backend (e.g., ASCEND_ATTN)
+        @register_diffusion_backend(DiffusionAttentionBackendEnum.ASCEND_ATTN)
+        class CustomAscendAttentionBackend:
+            ...
+
+        # Direct registration
+        register_diffusion_backend(
+            DiffusionAttentionBackendEnum.CUSTOM,
+            "my.module.MyCustomBackend"
+        )
+    """
+
+    def decorator(cls: type) -> type:
+        _DIFFUSION_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
+        return cls
+
+    if class_path is not None:
+        _DIFFUSION_ATTN_OVERRIDES[backend] = class_path
+        return lambda x: x
+
+    return decorator
diff --git a/vllm_omni/diffusion/attention/backends/sage_attn.py b/vllm_omni/diffusion/attention/backends/sage_attn.py
index bfc06a839..af34a0633 100644
--- a/vllm_omni/diffusion/attention/backends/sage_attn.py
+++ b/vllm_omni/diffusion/attention/backends/sage_attn.py
@@ -54,7 +54,7 @@ class SageAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
 
-    def forward(
+    def forward_cuda(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py
index 89b3e84a7..16bc5fcf8 100644
--- a/vllm_omni/diffusion/attention/backends/sdpa.py
+++ b/vllm_omni/diffusion/attention/backends/sdpa.py
@@ -16,6 +16,10 @@ logger = init_logger(__name__)
 class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         return [x for x in range(1024)]  # todo
diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py
index 6a02435a7..ea12bc0f1 100644
--- a/vllm_omni/diffusion/attention/backends/utils/fa.py
+++ b/vllm_omni/diffusion/attention/backends/utils/fa.py
@@ -15,9 +15,9 @@
 import torch
 import torch.nn.functional as F
 
-from vllm_omni.utils.platform_utils import is_rocm
+from vllm_omni.platforms import current_omni_platform
 
-if is_rocm():
+if current_omni_platform.is_rocm():
     from vllm._aiter_ops import is_aiter_found_and_supported
 
     # Choose to enable this by default on ROCm
diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py
index ed5946098..920396e66 100644
--- a/vllm_omni/diffusion/attention/selector.py
+++ b/vllm_omni/diffusion/attention/selector.py
@@ -1,52 +1,54 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Diffusion attention backend selector.
+
+This module provides the interface for selecting diffusion attention backends.
+The actual backend selection logic is delegated to the platform layer
+(vllm_omni.platforms), similar to how vLLM handles attention backend selection.
+
+Usage:
+    from vllm_omni.diffusion.attention.selector import get_attn_backend
+
+    # Get the appropriate backend for current platform
+    backend_cls = get_attn_backend(head_size=64)
+
+    # Or override via environment variable
+    # export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+"""
 
 import importlib
 import os
 from functools import cache
 
-import torch
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
 )
-from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
-from vllm_omni.utils.platform_utils import detect_device_type, is_rocm
 
 logger = init_logger(__name__)
 
-# environment variable value -> backend module and class
-_BACKEND_CONFIG = {
-    "FLASH_ATTN": {
-        "module": "vllm_omni.diffusion.attention.backends.flash_attn",
-        "class": "FlashAttentionBackend",
-    },
-    "TORCH_SDPA": {
-        "module": "vllm_omni.diffusion.attention.backends.sdpa",
-        "class": "SDPABackend",
-    },
-    "SAGE_ATTN": {
-        "module": "vllm_omni.diffusion.attention.backends.sage_attn",
-        "class": "SageAttentionBackend",
-    },
-    "ASCEND": {"module": "vllm_omni.diffusion.attention.backends.ascend_attn", "class": "AscendAttentionBackend"},
-}
 
-_BACKENDS_SUPPORT_ATTENTION_MASK = ["SDPA", "ASCEND", "FLASH_ATTN"]
+def _load_backend_cls(cls_path: str) -> type[AttentionBackend]:
+    """Load a backend class from its fully qualified path.
 
+    Args:
+        cls_path: Fully qualified class path (e.g.,
+            "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend")
 
-def load_backend(backend_name: str) -> type[AttentionBackend]:
-    config = _BACKEND_CONFIG[backend_name]
-
+    Returns:
+        The loaded backend class
+    """
+    module_path, class_name = cls_path.rsplit(".", 1)
     try:
-        module = importlib.import_module(config["module"])
-        backend_class = getattr(module, config["class"])
+        module = importlib.import_module(module_path)
+        backend_class = getattr(module, class_name)
         return backend_class
     except ImportError as e:
-        raise ImportError(f"Failed to import module {config['module']}: {e}")
+        raise ImportError(f"Failed to import module {module_path}: {e}")
     except AttributeError as e:
-        raise AttributeError(f"Class {config['class']} not found in module: {e}")
+        raise AttributeError(f"Class {class_name} not found in module: {e}")
 
 
 @cache
@@ -54,65 +56,30 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
     """
     Get attention backend for diffusion models.
 
-    The backend is selected based on the following priority:
-    1. DIFFUSION_ATTENTION_BACKEND environment variable (if set, e.g. export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN)
-    2. Default backend (SDPA)
+    The backend selection is delegated to the current platform
+    (vllm_omni.platforms.current_omni_platform), which selects the
+    appropriate backend based on:
+    1. User override via DIFFUSION_ATTENTION_BACKEND environment variable
+    2. Platform-specific defaults and capabilities
+
+    This is similar to how vLLM's get_attn_backend_cls works, where the
+    platform layer decides which backend to use based on hardware capabilities.
 
     Args:
-        head_size: Head size (currently not used for selection, but kept for API compatibility)
+        head_size: Head size for attention computation (may affect backend selection)
 
     Returns:
         The selected attention backend class
     """
-    # Check environment variable
+    from vllm_omni.platforms import current_omni_platform
 
-    backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND", None)
+    # Check environment variable for user override
+    selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
 
-    if detect_device_type() == "cuda" and not is_rocm():
-        compute_capability = torch.cuda.get_device_capability()
-        major, minor = compute_capability
-        if 80 <= major * 10 + minor < 100:
-            if backend_name is None:
-                backend_name = "FLASH_ATTN"
-        else:
-            if backend_name == "FLASH_ATTN":
-                logger.warning(
-                    """Flash Attention requires GPU with compute capability >= 8.0 or < 10.0. "
-                               "Falling back to TORCH_SDPA backend."""
-                )
-                backend_name = "TORCH_SDPA"
-    elif detect_device_type() == "cuda" and is_rocm():
-        from vllm._aiter_ops import is_aiter_found_and_supported
+    # Delegate to platform for backend selection
+    backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
+        selected_backend=selected_backend,
+        head_size=head_size,
+    )
 
-        compute_capability = torch.cuda.get_device_capability()
-        major, minor = compute_capability
-
-        # Choose to enable this by default on ROCm
-        # Whenever possible as it is the fastest backend
-        # is_aiter_found_and_supported() checks if aiter library is found
-        # and is aiter supported on the current platform
-        # aiter currently only is supported on gfx942 and gfx950
-        # https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py
-        if is_aiter_found_and_supported() and 90 < major * 10 + minor < 100:
-            if backend_name is None:
-                backend_name = "FLASH_ATTN"
-        else:
-            if backend_name == "FLASH_ATTN":
-                logger.warning(
-                    "Flash Attention requires `aiter` library which is only supported "
-                    "on gfx942 and gfx950. Falling back to TORCH_SDPA backend."
-                )
-                backend_name = "TORCH_SDPA"
-
-    if backend_name is not None:
-        backend_name_upper = backend_name.upper()
-        if backend_name_upper not in _BACKEND_CONFIG:
-            valid_backends = list(_BACKEND_CONFIG.keys())
-            raise ValueError(
-                f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}"
-            )
-        logger.info(f"Using attention backend '{backend_name_upper}' for diffusion")
-        return load_backend(backend_name_upper)
-
-    # Default to SDPA
-    return SDPABackend
+    return _load_backend_cls(backend_cls_path)
diff --git a/vllm_omni/diffusion/distributed/comm.py b/vllm_omni/diffusion/distributed/comm.py
index ada03f4a1..f73b8442f 100644
--- a/vllm_omni/diffusion/distributed/comm.py
+++ b/vllm_omni/diffusion/distributed/comm.py
@@ -8,6 +8,8 @@ import torch
 import torch.distributed as dist
 from torch import Tensor
 
+from vllm_omni.platforms import current_omni_platform
+
 __all__ = ["all_to_all_4D", "all_to_all_5D", "SeqAllToAll4D", "SeqAllToAll5D", "RingComm"]
 
 
@@ -48,7 +50,7 @@ def all_to_all_4D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
         # if scattering the seq-dim, transpose the heads back to the original dimension
@@ -83,7 +85,7 @@ def all_to_all_4D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
 
@@ -155,7 +157,7 @@ def all_to_all_5D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
 
@@ -190,7 +192,7 @@ def all_to_all_5D(
         if seq_world_size > 1:
             dist.all_to_all_single(output, input_t, group=group)
             if use_sync:
-                torch.cuda.synchronize()
+                current_omni_platform.synchronize()
         else:
             output = input_t
 
diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py
index 8e33d6fb6..b722f61c0 100644
--- a/vllm_omni/diffusion/distributed/group_coordinator.py
+++ b/vllm_omni/diffusion/distributed/group_coordinator.py
@@ -9,18 +9,14 @@ from typing import Any
 
 import torch
 import torch.distributed
-from torch.cuda import synchronize
 from torch.distributed import Backend, ProcessGroup
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
-if envs._is_npu():
-    logger.info("torch.npu synchronize")
-    from torch.npu import synchronize
-
 
 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
@@ -124,7 +120,7 @@ class GroupCoordinator:
         assert self.cpu_group is not None
         assert self.device_group is not None
 
-        self.device = envs.get_device(local_rank)
+        self.device = current_omni_platform.get_torch_device(local_rank)
 
     @property
     def first_rank(self):
@@ -618,7 +614,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
         assert self.cpu_group is not None
         assert self.device_group is not None
 
-        self.device = envs.get_device(local_rank)
+        self.device = current_omni_platform.get_torch_device(local_rank)
 
         self.recv_buffer_set: bool = False
         self.recv_tasks_queue: list[tuple[str, int]] = []
@@ -768,7 +764,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
 
         # To protect against race condition when using batch_isend_irecv().
         # should take this out once the bug with batch_isend_irecv is resolved.
-        synchronize()
+        current_omni_platform.synchronize()
 
         ops = []
         recv_prev_shape_tensor = None
@@ -799,7 +795,7 @@ class PipelineGroupCoordinator(GroupCoordinator):
             for req in reqs:
                 req.wait()
 
-        synchronize()
+        current_omni_platform.synchronize()
 
         recv_prev_shape = [0, 0, 0]
         if recv_prev_shape_tensor is not None:
diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py
index b24951590..2f7c2f852 100644
--- a/vllm_omni/diffusion/distributed/parallel_state.py
+++ b/vllm_omni/diffusion/distributed/parallel_state.py
@@ -35,6 +35,7 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
+from vllm_omni.platforms import current_omni_platform
 
 from .group_coordinator import (
     GroupCoordinator,
@@ -42,14 +43,6 @@ from .group_coordinator import (
     SequenceParallelGroupCoordinator,
 )
 
-if envs._is_npu():
-    from torch.npu import device_count, set_device
-elif envs._is_musa():
-    from torch_musa.core.device import device_count, set_device
-else:
-    from torch.cuda import device_count, set_device
-
-
 env_info = envs.PACKAGES_CHECKER.get_packages_info()
 
 HAS_FLASH_ATTN = env_info["has_flash_attn"]
@@ -396,7 +389,7 @@ def init_distributed_environment(
     backend: str | None = None,
 ):
     if backend is None:
-        backend = envs.get_torch_distributed_backend()
+        backend = current_omni_platform.dist_backend
     logger.debug(
         "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",
         world_size,
@@ -416,7 +409,8 @@ def init_distributed_environment(
             world_size=world_size,
             rank=rank,
         )
-        set_device(torch.distributed.get_rank() % device_count())
+        device_id = torch.distributed.get_rank() % current_omni_platform.get_device_count()
+        current_omni_platform.set_device(current_omni_platform.get_torch_device(device_id))
     # set the local rank
     # local_rank is not available in torch ProcessGroup,
     # see https://github.com/pytorch/pytorch/issues/122816
@@ -572,7 +566,7 @@ def initialize_model_parallel(
     backend: str | None = None,
 ) -> None:
     if backend is None:
-        backend = envs.get_torch_distributed_backend()
+        backend = current_omni_platform.dist_backend
     """
     Initialize model parallel groups.
 
@@ -632,7 +626,7 @@ def initialize_model_parallel(
 
     # FIXME: Since the async p2p communication operation of NPU is not same as cuda in torch,
     # the pipefusion is not ready for npu yet
-    if envs._is_npu():
+    if current_omni_platform.is_npu():
         assert pipeline_parallel_size == 1, "Current pipefusion is not ready for NPU"
 
     dit_parallel_size = (
diff --git a/vllm_omni/diffusion/distributed/utils.py b/vllm_omni/diffusion/distributed/utils.py
index 2e2f0fa53..92d00ef4b 100644
--- a/vllm_omni/diffusion/distributed/utils.py
+++ b/vllm_omni/diffusion/distributed/utils.py
@@ -5,11 +5,10 @@ import os
 
 import torch
 
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.platforms import current_omni_platform
 
 
 def get_local_device() -> torch.device:
     """Return the torch device for the current rank based on detected device type."""
-    device_type = detect_device_type()
-    local_rank = os.environ.get("LOCAL_RANK", 0)
-    return torch.device(f"{device_type}:{local_rank}")
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    return current_omni_platform.get_torch_device(local_rank)
diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py
index f717f2ca0..be2b58ae8 100644
--- a/vllm_omni/diffusion/envs.py
+++ b/vllm_omni/diffusion/envs.py
@@ -5,20 +5,15 @@ import os
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
-import torch
-from packaging import version
 from vllm.logger import init_logger
 
-logger = init_logger(__name__)
+from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     MASTER_ADDR: str = ""
     MASTER_PORT: int | None = None
     CUDA_HOME: str | None = None
     LOCAL_RANK: int = 0
-    CUDA_VISIBLE_DEVICES: str | None = None
-    CUDA_VERSION: version.Version
-    TORCH_VERSION: version.Version
 
 environment_variables: dict[str, Callable[[], Any]] = {
     # ================== Runtime Env Vars ==================
@@ -32,127 +27,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # local rank of the process in the distributed setting, used to determine
     # the GPU device id
     "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
-    # used to control the visible devices in the distributed setting
-    "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),
 }
 
-
-def _is_hip():
-    has_rocm = torch.version.hip is not None
-    return has_rocm
-
-
-def _is_cuda():
-    has_cuda = torch.version.cuda is not None
-    return has_cuda
-
-
-def _is_musa():
-    try:
-        if hasattr(torch, "musa") and torch.musa.is_available():
-            return True
-    except ModuleNotFoundError:
-        return False
-
-
-def _is_mps():
-    return torch.backends.mps.is_available()
-
-
-def _is_npu():
-    try:
-        if hasattr(torch, "npu") and torch.npu.is_available():
-            return True
-    except ModuleNotFoundError:
-        return False
-
-
-def get_device(local_rank: int) -> torch.device:
-    if _is_cuda() or _is_hip():
-        return torch.device("cuda", local_rank)
-    elif _is_musa():
-        return torch.device("musa", local_rank)
-    elif _is_mps():
-        return torch.device("mps")
-    elif _is_npu():
-        return torch.device("npu", local_rank)
-    else:
-        return torch.device("cpu")
-
-
-def get_device_name() -> str:
-    if _is_cuda() or _is_hip():
-        return "cuda"
-    elif _is_musa():
-        return "musa"
-    elif _is_mps():
-        return "mps"
-    elif _is_npu():
-        return "npu"
-    else:
-        return "cpu"
-
-
-def get_device_version():
-    if _is_hip():
-        hip_version = torch.version.hip
-        hip_version = hip_version.split("-")[0]
-        return hip_version
-    elif _is_cuda():
-        return torch.version.cuda
-    elif _is_musa():
-        return torch.version.musa
-    elif _is_mps():
-        return None
-    elif _is_npu():
-        return None
-    else:
-        raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available")
-
-
-def get_torch_distributed_backend() -> str:
-    if _is_cuda() or _is_hip():
-        return "nccl"
-    elif _is_musa():
-        return "mccl"
-    elif _is_mps():
-        return "gloo"
-    elif _is_npu():
-        return "hccl"
-    else:
-        raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available")
-
-
-variables: dict[str, Callable[[], Any]] = {
-    # ================== Other Vars ==================
-    # used in version checking
-    "CUDA_VERSION": lambda: version.parse(get_device_version() or "0.0"),
-    "TORCH_VERSION": lambda: version.parse(version.parse(torch.__version__).base_version),
-}
-
-
-def _setup_musa(environment_variables, variables):
-    musa = getattr(torch, "musa", None)
-    if musa is None:
-        return
-    try:
-        if musa.is_available():
-            environment_variables["MUSA_HOME"] = lambda: os.environ.get("MUSA_HOME", None)
-            environment_variables["MUSA_VISIBLE_DEVICES"] = lambda: os.environ.get("MUSA_VISIBLE_DEVICES", None)
-            musa_ver = getattr(getattr(torch, "version", None), "musa", None)
-            if musa_ver:
-                variables["MUSA_VERSION"] = lambda: version.parse(musa_ver)
-    except Exception:
-        pass
-
-
-try:
-    _setup_musa(environment_variables, variables)
-except (AttributeError, ModuleNotFoundError):
-    pass
+logger = init_logger(__name__)
 
 
 class PackagesEnvChecker:
+    """Singleton class for checking package availability."""
+
     _instance = None
 
     def __new__(cls):
@@ -163,38 +45,39 @@ class PackagesEnvChecker:
 
     def initialize(self):
         packages_info = {}
-        packages_info["has_flash_attn"] = self.check_flash_attn(packages_info)
+        packages_info["has_flash_attn"] = self._check_flash_attn(packages_info)
         self.packages_info = packages_info
 
-    def check_flash_attn(self, packages_info):
-        if not torch.cuda.is_available():
+    def _check_flash_attn(self, packages_info) -> bool:
+        """Check if flash attention is available and compatible."""
+        platform = current_omni_platform
+
+        # Flash attention requires CUDA-like platforms (CUDA or ROCm)
+        if not platform.is_cuda_alike():
             return False
 
-        # Check if torch_npu is available
-        if _is_npu():
-            logger.info("`flash_attn` is not ready on torch_npu for now")
+        # Check if devices are available
+        if platform.get_device_count() == 0:
             return False
 
-        if _is_musa():
-            logger.info("Flash Attention library is not supported on MUSA for the moment.")
-            return False
         try:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            gpu_name = torch.cuda.get_device_name(device)
+            gpu_name = platform.get_device_name()
+            # Turing/Tesla/T4 GPUs don't support flash attention well
             if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
                 return False
-            else:
-                from flash_attn import __version__
 
-                if __version__ < "2.6.0":
-                    raise ImportError("install flash_attn >= 2.6.0")
-                return True
+            from flash_attn import __version__
+
+            if __version__ < "2.6.0":
+                raise ImportError("install flash_attn >= 2.6.0")
+            return True
         except ImportError:
             if not packages_info.get("has_aiter", False):
                 logger.warning('Flash Attention library "flash_attn" not found, using pytorch attention implementation')
             return False
 
-    def get_packages_info(self):
+    def get_packages_info(self) -> dict:
+        """Get the packages info dictionary."""
         return self.packages_info
 
 
@@ -205,8 +88,6 @@ def __getattr__(name):
     # lazy evaluation of environment variables
     if name in environment_variables:
         return environment_variables[name]()
-    if name in variables:
-        return variables[name]()
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py
index 21a94acd3..a421a01ed 100644
--- a/vllm_omni/diffusion/executor/multiproc_executor.py
+++ b/vllm_omni/diffusion/executor/multiproc_executor.py
@@ -10,7 +10,7 @@ from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import Scheduler
-from vllm_omni.utils.platform_utils import get_diffusion_worker_class
+from vllm_omni.diffusion.worker import WorkerProc
 
 logger = init_logger(__name__)
 
@@ -77,9 +77,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         mp.set_start_method("spawn", force=True)
         processes = []
 
-        # Get the appropriate worker class for current device
-        worker_proc = get_diffusion_worker_class()
-
         # Launch all worker processes
         scheduler_pipe_readers = []
         scheduler_pipe_writers = []
@@ -88,7 +85,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             reader, writer = mp.Pipe(duplex=False)
             scheduler_pipe_writers.append(writer)
             process = mp.Process(
-                target=worker_proc.worker_main,
+                target=WorkerProc.worker_main,
                 args=(
                     i,  # rank
                     od_config,
diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py
index 21f4a6fe1..8f2c2b755 100644
--- a/vllm_omni/diffusion/hooks/sequence_parallel.py
+++ b/vllm_omni/diffusion/hooks/sequence_parallel.py
@@ -335,10 +335,7 @@ class SequenceParallelSplitHook(ModelHook):
         2. Creates an attention mask indicating valid vs padding positions
         3. Stores the mask and padding info in ForwardContext
         """
-        from vllm_omni.diffusion.attention.selector import (
-            _BACKENDS_SUPPORT_ATTENTION_MASK,
-            get_attn_backend,
-        )
+        from vllm_omni.diffusion.attention.selector import get_attn_backend
         from vllm_omni.diffusion.distributed.parallel_state import (
             get_ring_parallel_world_size,
             get_sequence_parallel_rank,
@@ -359,7 +356,7 @@ class SequenceParallelSplitHook(ModelHook):
 
         # Check backend compatibility
         attn_backend = get_attn_backend(-1)
-        if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK:
+        if not attn_backend.supports_attention_mask:
             raise ValueError(
                 f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
                 f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "
diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py
index 835ee07c0..9a644e1d9 100644
--- a/vllm_omni/diffusion/layers/adalayernorm.py
+++ b/vllm_omni/diffusion/layers/adalayernorm.py
@@ -91,6 +91,14 @@ class AdaLayerNorm(CustomOp):
 
         return output, gate_result
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        mod_params: torch.Tensor,
+        index: torch.Tensor = None,
+    ) -> torch.Tensor:
+        return self.forward_native(x, mod_params, index)
+
     def forward_native(
         self,
         x: torch.Tensor,
diff --git a/vllm_omni/diffusion/layers/custom_op.py b/vllm_omni/diffusion/layers/custom_op.py
index 2f50073cc..321bcbf8a 100644
--- a/vllm_omni/diffusion/layers/custom_op.py
+++ b/vllm_omni/diffusion/layers/custom_op.py
@@ -3,7 +3,7 @@ from typing import Any
 
 import torch.nn as nn
 
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 
 class CustomOp(nn.Module):
@@ -14,16 +14,17 @@ class CustomOp(nn.Module):
 
     def __init__(self) -> None:
         super().__init__()
-        self.is_cuda = detect_device_type() == "cuda"
         self._forward_method = self.dispatch_forward()
 
     def dispatch_forward(self) -> Callable:
-        if is_rocm():
+        if current_omni_platform.is_rocm():
             return self.forward_hip
-        elif self.is_cuda:
+        elif current_omni_platform.is_cuda():
             return self.forward_cuda
-        elif is_npu():
+        elif current_omni_platform.is_npu():
             return self.forward_npu
+        elif current_omni_platform.is_xpu():
+            return self.forward_xpu
         else:
             return self.forward_native
 
@@ -44,6 +45,9 @@ class CustomOp(nn.Module):
     def forward_npu(self, *args, **kwargs):
         raise NotImplementedError
 
+    def forward_xpu(self, *args, **kwargs):
+        raise NotImplementedError
+
     def forward_hip(self, *args, **kwargs):
         # By default, we assume that HIP ops are compatible with CUDA ops.
         return self.forward_cuda(*args, **kwargs)
diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py
index 81e8c098f..461c25652 100644
--- a/vllm_omni/diffusion/layers/rope.py
+++ b/vllm_omni/diffusion/layers/rope.py
@@ -137,6 +137,14 @@ class RotaryEmbedding(CustomOp):
         else:
             return self.forward_native(x, cos, sin)
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x, cos, sin)
+
     def forward_native(
         self,
         x: torch.Tensor,
diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
index 6bb282a80..9d216363d 100644
--- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
+++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
@@ -24,7 +24,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_sp_group,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
-from vllm_omni.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -610,7 +610,7 @@ class LongCatImageTransformer2DModel(nn.Module):
 
         ids = torch.cat((txt_ids, img_ids), dim=0)
 
-        if is_npu():
+        if current_omni_platform.is_npu():
             freqs_cos, freqs_sin = self.pos_embed(ids.cpu())
             image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu())
         else:
diff --git a/vllm_omni/diffusion/worker/__init__.py b/vllm_omni/diffusion/worker/__init__.py
index dfec4596b..8af028385 100644
--- a/vllm_omni/diffusion/worker/__init__.py
+++ b/vllm_omni/diffusion/worker/__init__.py
@@ -2,14 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Worker classes for diffusion models."""
 
-from vllm_omni.diffusion.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
-from vllm_omni.diffusion.worker.gpu_diffusion_worker import (
-    GPUDiffusionWorker,
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import (
+    DiffusionWorker,
     WorkerProc,
 )
 
 __all__ = [
-    "GPUDiffusionModelRunner",
-    "GPUDiffusionWorker",
+    "DiffusionModelRunner",
+    "DiffusionWorker",
     "WorkerProc",
 ]
diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
similarity index 99%
rename from vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
rename to vllm_omni/diffusion/worker/diffusion_model_runner.py
index 0401655a8..7335162b0 100644
--- a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -34,7 +34,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 logger = init_logger(__name__)
 
 
-class GPUDiffusionModelRunner:
+class DiffusionModelRunner:
     """
     Model runner that handles model loading and execution for diffusion models.
 
diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
similarity index 94%
rename from vllm_omni/diffusion/worker/gpu_diffusion_worker.py
rename to vllm_omni/diffusion/worker/diffusion_worker.py
index d8b962a58..0fad90a5b 100644
--- a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -4,7 +4,7 @@
 Diffusion Worker for vLLM-Omni.
 
 Handles GPU infrastructure initialization and delegates model operations
-to GPUDiffusionModelRunner.
+to DiffusionModelRunner.
 """
 
 import multiprocessing as mp
@@ -31,13 +31,14 @@ from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.profiler import CurrentProfiler
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.lora.request import LoRARequest
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-class GPUDiffusionWorker:
+class DiffusionWorker:
     """
     A worker that manages GPU infrastructure and delegates to the model runner.
 
@@ -47,7 +48,7 @@ class GPUDiffusionWorker:
     - Memory management (sleep/wake)
 
     All model-related operations (loading, compilation, execution) are
-    delegated to GPUDiffusionModelRunner.
+    delegated to DiffusionModelRunner.
     """
 
     def __init__(
@@ -61,7 +62,7 @@ class GPUDiffusionWorker:
         self.od_config = od_config
         self.device: torch.device | None = None
         self.vllm_config: VllmConfig | None = None
-        self.model_runner: GPUDiffusionModelRunner | None = None
+        self.model_runner: DiffusionModelRunner | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
         self.init_device()
@@ -79,8 +80,8 @@ class GPUDiffusionWorker:
         os.environ["WORLD_SIZE"] = str(world_size)
 
         # Setup device
-        self.device = torch.device(f"cuda:{rank}")
-        torch.cuda.set_device(self.device)
+        self.device = current_omni_platform.get_torch_device(rank)
+        current_omni_platform.set_device(self.device)
 
         # Create vllm_config for parallel configuration
         vllm_config = VllmConfig()
@@ -105,7 +106,7 @@ class GPUDiffusionWorker:
             )
 
         # Create model runner and load model
-        self.model_runner = GPUDiffusionModelRunner(
+        self.model_runner = DiffusionModelRunner(
             vllm_config=self.vllm_config,
             od_config=self.od_config,
             device=self.device,
@@ -176,7 +177,7 @@ class GPUDiffusionWorker:
         """
         from vllm.device_allocator.cumem import CuMemAllocator
 
-        free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
+        free_bytes_before_sleep = current_omni_platform.get_free_memory()
 
         # Save the buffers before level 2 sleep
         if level == 2 and self.model_runner is not None:
@@ -185,7 +186,9 @@ class GPUDiffusionWorker:
 
         allocator = CuMemAllocator.get_instance()
         allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
-        free_bytes_after_sleep, total = torch.cuda.mem_get_info()
+        free_bytes_after_sleep = current_omni_platform.get_free_memory()
+        device_id = self.device.index if self.device.index is not None else 0
+        total = current_omni_platform.get_device_total_memory(device_id)
         freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
         used_bytes = total - free_bytes_after_sleep
         assert freed_bytes >= 0, "Memory usage increased after sleeping."
@@ -270,9 +273,9 @@ class WorkerProc:
         self.gpu_id = gpu_id
         self._running = True
 
-    def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> GPUDiffusionWorker:
+    def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> DiffusionWorker:
         """Create a worker instance. Override in subclasses for different worker types."""
-        return GPUDiffusionWorker(
+        return DiffusionWorker(
             local_rank=gpu_id,
             rank=gpu_id,
             od_config=od_config,
@@ -379,6 +382,9 @@ class WorkerProc:
         broadcast_handle,
     ) -> None:
         """Worker initialization and execution loops."""
+        from vllm_omni.plugins import load_omni_general_plugins
+
+        load_omni_general_plugins()
         worker_proc = WorkerProc(
             od_config,
             gpu_id=rank,
diff --git a/vllm_omni/diffusion/worker/npu/__init__.py b/vllm_omni/diffusion/worker/npu/__init__.py
deleted file mode 100644
index b2210b5e2..000000000
--- a/vllm_omni/diffusion/worker/npu/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""NPU worker classes for diffusion models."""
-
-from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorker, NPUWorkerProc
-
-__all__ = ["NPUWorker", "NPUWorkerProc"]
diff --git a/vllm_omni/diffusion/worker/npu/npu_worker.py b/vllm_omni/diffusion/worker/npu/npu_worker.py
deleted file mode 100644
index c4d2c60e5..000000000
--- a/vllm_omni/diffusion/worker/npu/npu_worker.py
+++ /dev/null
@@ -1,255 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import multiprocessing as mp
-import os
-import time
-from collections.abc import Iterable
-from contextlib import AbstractContextManager, nullcontext
-
-import torch
-from vllm.config import LoadConfig, VllmConfig
-from vllm.logger import init_logger
-from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
-
-from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.data import (
-    DiffusionOutput,
-    OmniDiffusionConfig,
-)
-from vllm_omni.diffusion.distributed.parallel_state import (
-    destroy_distributed_env,
-    init_distributed_environment,
-    initialize_model_parallel,
-)
-from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.gpu_diffusion_worker import WorkerProc
-
-logger = init_logger(__name__)
-
-
-class NPUWorker:
-    """
-    A worker that executes the model on a single NPU.
-    Inherits from GPUWorker and overrides device-specific initialization.
-    """
-
-    def __init__(
-        self,
-        local_rank: int,
-        rank: int,
-        od_config: OmniDiffusionConfig,
-    ):
-        self.local_rank = local_rank
-        self.rank = rank
-        self.od_config = od_config
-        self.pipeline = None
-        self.device = None
-        self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
-        self.init_device_and_model()
-
-    def init_device_and_model(self) -> None:
-        """Initialize the NPU device and load the model."""
-        world_size = self.od_config.num_gpus
-        rank = self.rank
-        # Set environment variables for distributed initialization
-        os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = str(self.od_config.master_port)
-        os.environ["LOCAL_RANK"] = str(self.local_rank)
-        os.environ["RANK"] = str(rank)
-        os.environ["WORLD_SIZE"] = str(world_size)
-
-        device = torch.device(f"npu:{rank}")
-        torch.npu.set_device(device)
-
-        # hack
-        vllm_config = VllmConfig()
-        vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
-        vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
-        self.vllm_config = vllm_config
-        with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config):
-            init_distributed_environment(world_size=world_size, rank=rank)
-            logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
-            parallel_config = self.od_config.parallel_config
-            initialize_model_parallel(
-                data_parallel_size=parallel_config.data_parallel_size,
-                cfg_parallel_size=parallel_config.cfg_parallel_size,
-                sequence_parallel_size=parallel_config.sequence_parallel_size,
-                ulysses_degree=parallel_config.ulysses_degree,
-                ring_degree=parallel_config.ring_degree,
-                tensor_parallel_size=parallel_config.tensor_parallel_size,
-                pipeline_parallel_size=parallel_config.pipeline_parallel_size,
-            )
-
-            load_config = LoadConfig()
-            model_loader = DiffusersPipelineLoader(load_config)
-            time_before_load = time.perf_counter()
-            with DeviceMemoryProfiler() as m:
-                self.pipeline = model_loader.load_model(
-                    od_config=self.od_config,
-                    load_device=f"npu:{rank}",
-                )
-            time_after_load = time.perf_counter()
-
-        logger.info(
-            "Model loading took %.4f GiB and %.6f seconds",
-            m.consumed_memory / GiB_bytes,
-            time_after_load - time_before_load,
-        )
-        logger.info(f"Worker {self.rank}: Model loaded successfully.")
-
-        # Setup cache backend based on type (both backends use enable()/reset() interface)
-        self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config)
-
-        if self.cache_backend is not None:
-            self.cache_backend.enable(self.pipeline)
-
-    def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
-        """
-        Generate output for the given requests.
-
-        Args:
-            requests: List of diffusion requests
-
-        Returns:
-            DiffusionOutput with generated results
-        """
-        return self.execute_model(requests, self.od_config)
-
-    @torch.inference_mode()
-    def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
-        """
-        Execute a forward pass.
-        """
-        assert self.pipeline is not None
-        if not reqs or len(reqs) == 0:
-            raise ValueError("Cannot execute model with empty request list")
-        # TODO: dealing with first req for now
-        req = reqs[0]
-
-        if req.generator is None and req.seed is not None:
-            req.generator = torch.Generator(device=self.device).manual_seed(req.seed)
-
-        # Refresh cache context if needed
-        if (
-            not getattr(req, "skip_cache_refresh", False)
-            and self.cache_backend is not None
-            and self.cache_backend.is_enabled()
-        ):
-            self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
-        with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
-            output = self.pipeline.forward(req)
-        return output
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        return self.pipeline.load_weights(weights)
-
-    def sleep(self, level: int = 1) -> bool:
-        """
-        Put the worker to sleep. The worker should not process any requests.
-        The caller should guarantee that no requests are being processed
-        during the sleep period, before `wake_up` is called.
-
-        Args:
-            level: The sleep level. Level 1 sleep will offload the model
-                weights and discard the kv cache.
-                Currently only support level 1.
-        """
-        from vllm.device_allocator.cumem import CuMemAllocator
-
-        free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
-
-        # Save the buffers before level 2 sleep
-        if level == 2:
-            model = self.pipeline
-            self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}
-
-        allocator = CuMemAllocator.get_instance()
-        allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
-        free_bytes_after_sleep, total = torch.cuda.mem_get_info()
-        freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
-        used_bytes = total - free_bytes_after_sleep
-        assert freed_bytes >= 0, "Memory usage increased after sleeping."
-        logger.info(
-            "Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.",
-            freed_bytes / GiB_bytes,
-            used_bytes / GiB_bytes,
-        )
-        return True
-
-    def wake_up(self, tags: list[str] | None = None) -> bool:
-        """
-        Wake up the worker from sleep mode. See the sleep function
-        method for more details.
-
-        Args:
-            tags: An optional list of tags to reallocate the worker memory
-                for specific memory allocations. Values must be in
-                `("weights")`. If None, all memory is reallocated.
-                wake_up should be called with all tags (or None) before the
-                worker is used again.
-        """
-        from vllm.device_allocator.cumem import CuMemAllocator
-
-        allocator = CuMemAllocator.get_instance()
-        allocator.wake_up(tags)
-
-        # Restore the buffers after level 2 sleep
-        if len(self._sleep_saved_buffers):
-            model = self.pipeline
-            for name, buffer in model.named_buffers():
-                if name in self._sleep_saved_buffers:
-                    buffer.data.copy_(self._sleep_saved_buffers[name].data)
-            self._sleep_saved_buffers = {}
-        return True
-
-    def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
-        if self.od_config.enable_sleep_mode:
-            from vllm.device_allocator.cumem import CuMemAllocator
-
-            allocator = CuMemAllocator.get_instance()
-            if tag == "weights":
-                assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
-            return allocator.use_memory_pool(tag=tag)
-        else:
-            return nullcontext()
-
-    def shutdown(self) -> None:
-        destroy_distributed_env()
-
-
-class NPUWorkerProc(WorkerProc):
-    """Wrapper that runs one NPUWorker in a separate process."""
-
-    def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> NPUWorker:
-        """Create an NPUWorker instead of GPUWorker."""
-        return NPUWorker(
-            local_rank=gpu_id,
-            rank=gpu_id,
-            od_config=od_config,
-        )
-
-    @staticmethod
-    def worker_main(
-        rank: int,
-        od_config: OmniDiffusionConfig,
-        pipe_writer: mp.connection.Connection,
-        broadcast_handle,
-    ) -> None:
-        """Worker initialization and execution loops."""
-
-        worker_proc = NPUWorkerProc(
-            od_config,
-            gpu_id=rank,
-            broadcast_handle=broadcast_handle,
-        )
-        logger.info(f"Worker {rank}: Scheduler loop started.")
-        pipe_writer.send(
-            {
-                "status": "ready",
-                "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
-            }
-        )
-        worker_proc.worker_busy_loop()
-        logger.info(f"Worker {rank}: Shutdown complete.")
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index cbe2001d7..21bb761d2 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -7,6 +7,7 @@ from vllm.transformers_utils.config import get_hf_text_config
 from vllm.v1.engine.async_llm import AsyncEngineArgs
 
 from vllm_omni.config import OmniModelConfig
+from vllm_omni.plugins import load_omni_general_plugins
 
 logger = init_logger(__name__)
 
@@ -89,7 +90,11 @@ class OmniEngineArgs(EngineArgs):
                 f"Config attribute '{hf_config_name}' not found in hf_config, "
                 "falling back to default get_hf_text_config"
             )
-            return get_hf_text_config(hf_config)
+        return get_hf_text_config(hf_config)
+
+    def __post_init__(self) -> None:
+        load_omni_general_plugins()
+        super().__post_init__()
 
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):
@@ -187,7 +192,11 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
                 f"Config attribute '{hf_config_name}' not found in hf_config, "
                 "falling back to default get_hf_text_config"
             )
-            return get_hf_text_config(hf_config)
+        return get_hf_text_config(hf_config)
+
+    def __post_init__(self) -> None:
+        load_omni_general_plugins()
+        super().__post_init__()
 
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):
diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index 0a5ee55be..2adb1c342 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -51,11 +51,27 @@ from vllm_omni.entrypoints.stage_utils import (
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils import detect_device_type
 
 logger = init_logger(__name__)
 
 
+def _resolve_worker_cls(engine_args: dict[str, Any]) -> None:
+    worker_type = engine_args.pop("worker_type", None)
+    if not worker_type:
+        return
+    if engine_args.get("worker_cls"):
+        return
+    from vllm_omni.platforms import current_omni_platform
+
+    worker_type = str(worker_type).lower()
+    if worker_type == "ar":
+        engine_args["worker_cls"] = current_omni_platform.get_omni_ar_worker_cls()
+    elif worker_type == "generation":
+        engine_args["worker_cls"] = current_omni_platform.get_omni_generation_worker_cls()
+    else:
+        raise ValueError(f"Unknown worker_type: {worker_type}")
+
+
 def _build_od_config(engine_args: dict[str, Any], model: str) -> dict[str, Any]:
     """Build OmniDiffusionConfig kwargs from engine args."""
     od_config = engine_args.get("od_config", {})
@@ -483,6 +499,9 @@ def _stage_worker(
     import os as _os
     import time as _time
 
+    from vllm_omni.plugins import load_omni_general_plugins
+
+    load_omni_general_plugins()
     # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker /
     # GPUARModelRunner) are spawned with a fork-safe method.
     # Mooncake / gRPC / RDMA and CUDA/NCCL can deadlock under fork-with-threads.
@@ -503,6 +522,9 @@ def _stage_worker(
     connectors_config = stage_payload.get("connectors_config", {})
     stage_type: Literal["llm", "diffusion"] = stage_payload.get("stage_type", "llm")
 
+    if stage_type != "diffusion":
+        _resolve_worker_cls(engine_args)
+
     # Aggregates for running average
     _agg_total_tokens = 0
     _agg_total_gen_time_ms = 0.0
@@ -512,7 +534,9 @@ def _stage_worker(
     # Device mapping
     device_type = None
     try:
-        device_type = detect_device_type()
+        from vllm_omni.platforms import current_omni_platform
+
+        device_type = current_omni_platform.device_type
         set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
     except Exception as e:
         logger.warning("Device setup failed: %s", e)
@@ -521,139 +545,137 @@ def _stage_worker(
     # when multiple instances start simultaneously
     # For TP/PP/DP/SP, we need to lock ALL devices that will be used by this stage
     lock_files = []
-    if device_type == "cuda":
-        try:
-            import torch
+    try:
+        # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
+        if "parallel_config" in engine_args:
+            parallel_config = engine_args["parallel_config"]
+            tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
+            data_parallel_size = parallel_config.get("data_parallel_size", 1)
+            prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
+            cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
+        else:
+            tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
+            data_parallel_size = engine_args.get("data_parallel_size", 1)
+            prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = 1  # not use in omni model
+            cfg_parallel_size = 1  # not used in omni model
 
-            if torch.cuda.is_available():
-                # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
-                if "parallel_config" in engine_args:
-                    parallel_config = engine_args["parallel_config"]
-                    tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
-                    pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
-                    data_parallel_size = parallel_config.get("data_parallel_size", 1)
-                    prefill_context_parallel_size = 1  # not used for diffusion
-                    sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
-                    cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
-                else:
-                    tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
-                    pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
-                    data_parallel_size = engine_args.get("data_parallel_size", 1)
-                    prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
-                    sequence_parallel_size = 1  # not use in omni model
-                    cfg_parallel_size = 1  # not used in omni model
+        # Calculate total number of devices needed for this stage
+        # For a single stage worker:
+        # - TP: splits model across devices (always needed)
+        # - PP: splits layers across pipeline stages, but each stage uses TP devices
+        # - DP: replicates model, but each replica uses TP devices
+        # - PCP: context parallelism, typically uses TP devices
+        # - SP: sequence parallelism, typically uses TP devices
+        # - CFG: Classifier-Free Guidance parallelism for diffusion models
+        # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
+        # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
+        num_devices_per_stage = (
+            tensor_parallel_size
+            * pipeline_parallel_size
+            * data_parallel_size
+            * prefill_context_parallel_size
+            * sequence_parallel_size
+            * cfg_parallel_size
+        )
 
-                # Calculate total number of devices needed for this stage
-                # For a single stage worker:
-                # - TP: splits model across GPUs (always needed)
-                # - PP: splits layers across pipelinestages, but each stage uses TP devices
-                # - DP: replicates model, but each replica uses TP devices
-                # - PCP: context parallelism, typically uses TP devices
-                # - SP: sequence parallelism, typically uses TP devices
-                # - CFG: Classifier-Free Guidance parallelism for diffusion models
-                # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
-                # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
-                num_devices_per_stage = (
-                    tensor_parallel_size
-                    * pipeline_parallel_size
-                    * data_parallel_size
-                    * prefill_context_parallel_size
-                    * sequence_parallel_size
-                    * cfg_parallel_size
-                )
+        # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES)
+        # After set_stage_devices, this env var is set to physical device(s)
+        device_control_env = current_omni_platform.device_control_env_var
+        visible_devices_str = _os.environ.get(device_control_env)
+        physical_devices = []
 
-                # Get physical device IDs from CUDA_VISIBLE_DEVICES
-                # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s)
-                cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES")
-                physical_devices = []
+        if visible_devices_str:
+            try:
+                physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
+            except (ValueError, IndexError):
+                pass
 
-                if cuda_visible_devices:
+        if not physical_devices:
+            # Fallback: use logical device count if device control env var not set
+            num_devices = current_omni_platform.get_device_count()
+            physical_devices = list(range(num_devices))
+
+        # Determine which devices will be used (min of devices per stage and available devices)
+        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+        devices_to_lock = physical_devices[:num_devices_to_lock]
+
+        # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
+        devices_to_lock = sorted(devices_to_lock)
+
+        logger.debug(
+            "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
+            tensor_parallel_size,
+            pipeline_parallel_size,
+            data_parallel_size,
+            prefill_context_parallel_size,
+            sequence_parallel_size,
+            cfg_parallel_size,
+            num_devices_to_lock,
+            devices_to_lock,
+        )
+
+        # Acquire exclusive locks for all devices using fcntl.flock
+        # Locks are automatically released when process dies
+        wait_start = _time.time()
+        acquired_lock_fds = []  # Store file descriptors to keep locks alive
+
+        for device_id in devices_to_lock:
+            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+            lock_acquired = False
+
+            while not lock_acquired:
+                try:
+                    # Open or create the lock file
+                    lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
+
+                    # Try to acquire exclusive lock (non-blocking first)
                     try:
-                        physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()]
-                    except (ValueError, IndexError):
-                        pass
+                        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                        # Successfully acquired lock - write PID
+                        _os.ftruncate(lock_fd, 0)  # Clear file
+                        _os.write(lock_fd, f"{_os.getpid()}\n".encode())
+                        _os.fsync(lock_fd)  # Ensure written to disk
+                        lock_acquired = True
+                        acquired_lock_fds.append(lock_fd)
+                        logger.debug("Acquired exclusive lock for device %s", device_id)
+                    except BlockingIOError:
+                        # Lock is held by another process
+                        _os.close(lock_fd)
 
-                if not physical_devices:
-                    # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set
-                    num_devices = torch.cuda.device_count()
-                    physical_devices = list(range(num_devices))
-
-                # Determine which devices will be used (min of devices per stage and available devices)
-                num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
-                devices_to_lock = physical_devices[:num_devices_to_lock]
-
-                # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
-                devices_to_lock = sorted(devices_to_lock)
-
-                logger.debug(
-                    "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d; will lock %d devices: %s",
-                    tensor_parallel_size,
-                    pipeline_parallel_size,
-                    data_parallel_size,
-                    prefill_context_parallel_size,
-                    sequence_parallel_size,
-                    num_devices_to_lock,
-                    devices_to_lock,
-                )
-
-                # Acquire exclusive locks for all devices using fcntl.flock
-                # Locks are automatically released when process dies
-                wait_start = _time.time()
-                acquired_lock_fds = []  # Store file descriptors to keep locks alive
-
-                for device_id in devices_to_lock:
-                    lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
-                    lock_acquired = False
-
-                    while not lock_acquired:
-                        try:
-                            # Open or create the lock file
-                            lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
-
-                            # Try to acquire exclusive lock (non-blocking first)
-                            try:
-                                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                                # Successfully acquired lock - write PID
-                                _os.ftruncate(lock_fd, 0)  # Clear file
-                                _os.write(lock_fd, f"{_os.getpid()}\n".encode())
-                                _os.fsync(lock_fd)  # Ensure written to disk
-                                lock_acquired = True
-                                acquired_lock_fds.append(lock_fd)
-                                logger.debug("Acquired exclusive lock for device %s", device_id)
-                            except BlockingIOError:
-                                # Lock is held by another process
-                                _os.close(lock_fd)
-
-                                # Check if we've been waiting too long
-                                if _time.time() - wait_start > stage_init_timeout:
-                                    logger.warning(
-                                        "Timeout waiting for device %s initialization lock, proceeding anyway",
-                                        device_id,
-                                    )
-                                    break
-
-                                # Wait a bit before retrying
-                                _time.sleep(0.1)
-                        except OSError as e:
-                            # Other error - log and continue without lock
-                            logger.debug(
-                                "Failed to acquire lock for device %s: %s, continuing anyway",
+                        # Check if we've been waiting too long
+                        if _time.time() - wait_start > stage_init_timeout:
+                            logger.warning(
+                                "Timeout waiting for device %s initialization lock, proceeding anyway",
                                 device_id,
-                                e,
                             )
-                            try:
-                                _os.close(lock_fd)
-                            except (OSError, NameError):
-                                pass
                             break
 
-                lock_files = acquired_lock_fds
-        except Exception as e:
-            logger.debug(
-                "[Stage-%s] Failed to set up sequential initialization lock: %s",
-                stage_id,
-                e,
-            )
+                        # Wait a bit before retrying
+                        _time.sleep(0.1)
+                except OSError as e:
+                    # Other error - log and continue without lock
+                    logger.debug(
+                        "Failed to acquire lock for device %s: %s, continuing anyway",
+                        device_id,
+                        e,
+                    )
+                    try:
+                        _os.close(lock_fd)
+                    except (OSError, NameError):
+                        pass
+                    break
+
+        lock_files = acquired_lock_fds
+    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):
@@ -1001,6 +1023,9 @@ async def _stage_worker_async(
     import os as _os
     import time as _time
 
+    from vllm_omni.plugins import load_omni_general_plugins
+
+    load_omni_general_plugins()
     # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker /
     # GPUARModelRunner) are spawned with a fork-safe method.
     if _os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn":
@@ -1018,6 +1043,9 @@ async def _stage_worker_async(
     connectors_config = stage_payload.get("connectors_config", {})
     stage_type = stage_payload.get("stage_type", "llm")
 
+    if stage_type != "diffusion":
+        _resolve_worker_cls(engine_args)
+
     in_q = omni_stage._in_q
     out_q = omni_stage._out_q
 
@@ -1031,9 +1059,9 @@ async def _stage_worker_async(
     # Device mapping
     device_type = None
     try:
-        from vllm_omni.utils import detect_device_type
+        from vllm_omni.platforms import current_omni_platform
 
-        device_type = detect_device_type()
+        device_type = current_omni_platform.device_type
         set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
     except Exception as e:
         logger.warning("Device setup failed: %s", e)
@@ -1051,120 +1079,136 @@ async def _stage_worker_async(
 
     # Sequential initialization on the same device to avoid memory calculation errors
     # when multiple instances start simultaneously
-    # For TP, we need to lock ALL devices that will be used by this stage
+    # For TP/PP/DP/PCP, we need to lock ALL devices that will be used by this stage
     lock_files = []
-    if device_type == "cuda":
-        try:
-            import torch
+    try:
+        # Get all parallel sizes from engine_args or parallel_config (defaults to 1)
+        if "parallel_config" in engine_args:
+            parallel_config = engine_args["parallel_config"]
+            tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1)
+            data_parallel_size = parallel_config.get("data_parallel_size", 1)
+            prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1)
+            cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1)
+        else:
+            tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
+            data_parallel_size = engine_args.get("data_parallel_size", 1)
+            prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = 1  # not use in omni model
+            cfg_parallel_size = 1  # not used in omni model
 
-            if torch.cuda.is_available():
-                # Get all parallel sizes from engine_args (defaults to 1)
-                tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
-                pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
-                data_parallel_size = engine_args.get("data_parallel_size", 1)
-                prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+        # Calculate total number of devices needed for this stage
+        # For a single stage worker:
+        # - TP: splits model across devices (always needed)
+        # - PP: splits layers across pipeline stages, but each stage uses TP devices
+        # - DP: replicates model, but each replica uses TP devices
+        # - PCP: context parallelism, typically uses TP devices
+        # - SP: sequence parallelism, typically uses TP devices
+        # - CFG: Classifier-Free Guidance parallelism for diffusion models
+        # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size
+        # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
+        num_devices_per_stage = (
+            tensor_parallel_size
+            * pipeline_parallel_size
+            * data_parallel_size
+            * prefill_context_parallel_size
+            * sequence_parallel_size
+            * cfg_parallel_size
+        )
 
-                # Calculate total number of devices needed for this stage
-                # For a single stage worker in omni:
-                # - TP: splits model across GPUs (always needed)
-                # - PP: splits layers across stages, but each stage uses TP devices
-                # - DP: replicates model, but each replica uses TP devices
-                # - PCP: context parallelism, typically uses TP devices
-                # The number of devices per stage is determined by TP * PP * DP * PCP size
-                # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
-                num_devices_per_stage = (
-                    tensor_parallel_size * pipeline_parallel_size * data_parallel_size * prefill_context_parallel_size
-                )
+        # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES)
+        # After set_stage_devices, this env var is set to physical device(s)
+        device_control_env = current_omni_platform.device_control_env_var
+        visible_devices_str = _os.environ.get(device_control_env)
+        physical_devices = []
 
-                # Get physical device IDs from CUDA_VISIBLE_DEVICES
-                # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s)
-                cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES")
-                physical_devices = []
+        if visible_devices_str:
+            try:
+                physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
+            except (ValueError, IndexError):
+                pass
 
-                if cuda_visible_devices:
+        if not physical_devices:
+            # Fallback: use logical device count if device control env var not set
+            num_devices = current_omni_platform.get_device_count()
+            physical_devices = list(range(num_devices))
+
+        # Determine which devices will be used (min of devices per stage and available devices)
+        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+        devices_to_lock = physical_devices[:num_devices_to_lock]
+
+        # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
+        devices_to_lock = sorted(devices_to_lock)
+
+        logger.debug(
+            "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
+            tensor_parallel_size,
+            pipeline_parallel_size,
+            data_parallel_size,
+            prefill_context_parallel_size,
+            sequence_parallel_size,
+            cfg_parallel_size,
+            num_devices_to_lock,
+            devices_to_lock,
+        )
+
+        # Acquire exclusive locks for all devices using fcntl.flock
+        # Locks are automatically released when process dies
+        wait_start = _time.time()
+        acquired_lock_fds = []  # Store file descriptors to keep locks alive
+
+        for device_id in devices_to_lock:
+            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+            lock_acquired = False
+
+            while not lock_acquired:
+                try:
+                    # Open or create the lock file
+                    lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
+
+                    # Try to acquire exclusive lock (non-blocking first)
                     try:
-                        physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()]
-                    except (ValueError, IndexError):
-                        pass
+                        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                        # Successfully acquired lock - write PID
+                        _os.ftruncate(lock_fd, 0)  # Clear file
+                        _os.write(lock_fd, f"{_os.getpid()}\n".encode())
+                        _os.fsync(lock_fd)  # Ensure written to disk
+                        lock_acquired = True
+                        acquired_lock_fds.append(lock_fd)
+                        logger.debug("Acquired exclusive lock for device %s", device_id)
+                    except BlockingIOError:
+                        # Lock is held by another process
+                        _os.close(lock_fd)
 
-                if not physical_devices:
-                    # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set
-                    num_devices = torch.cuda.device_count()
-                    physical_devices = list(range(num_devices))
-
-                # Determine which devices will be used (min of devices per stage and available devices)
-                num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
-                devices_to_lock = physical_devices[:num_devices_to_lock]
-
-                # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
-                devices_to_lock = sorted(devices_to_lock)
-
-                logger.debug(
-                    "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
-                    tensor_parallel_size,
-                    pipeline_parallel_size,
-                    data_parallel_size,
-                    prefill_context_parallel_size,
-                    num_devices_to_lock,
-                    devices_to_lock,
-                )
-
-                # Acquire exclusive locks for all devices using fcntl.flock
-                # Locks are automatically released when process dies
-                wait_start = _time.time()
-                acquired_lock_fds = []  # Store file descriptors to keep locks alive
-
-                for device_id in devices_to_lock:
-                    lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
-                    lock_acquired = False
-
-                    while not lock_acquired:
-                        try:
-                            # Open or create the lock file
-                            lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
-
-                            # Try to acquire exclusive lock (non-blocking first)
-                            try:
-                                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                                # Successfully acquired lock - write PID
-                                _os.ftruncate(lock_fd, 0)  # Clear file
-                                _os.write(lock_fd, f"{_os.getpid()}\n".encode())
-                                _os.fsync(lock_fd)  # Ensure written to disk
-                                lock_acquired = True
-                                acquired_lock_fds.append(lock_fd)
-                                logger.debug("Acquired exclusive lock for device %s", device_id)
-                            except BlockingIOError:
-                                # Lock is held by another process
-                                _os.close(lock_fd)
-
-                                # Check if we've been waiting too long
-                                if _time.time() - wait_start > stage_init_timeout:
-                                    logger.warning(
-                                        "Timeout waiting for device %s initialization lock, "
-                                        "proceeding anyway with timeout %s",
-                                        device_id,
-                                        stage_init_timeout,
-                                    )
-                                    break
-
-                                # Wait a bit before retrying
-                                _time.sleep(0.1)
-                        except OSError as e:
-                            # Other error - log and continue without lock
-                            logger.debug(
-                                "Failed to acquire lock for device %s: %s, continuing anyway",
+                        # Check if we've been waiting too long
+                        if _time.time() - wait_start > stage_init_timeout:
+                            logger.warning(
+                                "Timeout waiting for device %s initialization lock, proceeding anyway with timeout %s",
                                 device_id,
-                                e,
+                                stage_init_timeout,
                             )
-                            try:
-                                _os.close(lock_fd)
-                            except (OSError, NameError):
-                                pass
                             break
 
-                lock_files = acquired_lock_fds
-        except Exception as e:
-            logger.debug("Failed to set up sequential initialization lock: %s", e)
+                        # Wait a bit before retrying
+                        _time.sleep(0.1)
+                except OSError as e:
+                    # Other error - log and continue without lock
+                    logger.debug(
+                        "Failed to acquire lock for device %s: %s, continuing anyway",
+                        device_id,
+                        e,
+                    )
+                    try:
+                        _os.close(lock_fd)
+                    except (OSError, NameError):
+                        pass
+                    break
+
+        lock_files = acquired_lock_fds
+    except Exception as e:
+        logger.debug("Failed to set up sequential initialization lock: %s", e)
 
     # Init engine based on stage_type
     logger.debug(
diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py
index da837ce9a..74ad42f04 100644
--- a/vllm_omni/entrypoints/stage_utils.py
+++ b/vllm_omni/entrypoints/stage_utils.py
@@ -57,12 +57,12 @@ def set_stage_devices(
         - CUDA: Sets CUDA_VISIBLE_DEVICES and calls torch.cuda.set_device()
         - NPU: Sets ASCEND_RT_VISIBLE_DEVICES and calls torch.npu.set_device()
     """
-    from vllm_omni.utils import detect_device_type, get_device_control_env_var
+    from vllm_omni.platforms import current_omni_platform
 
     if device_type is None:
-        device_type = detect_device_type()
+        device_type = current_omni_platform.device_type
 
-    env_var = get_device_control_env_var()
+    env_var = current_omni_platform.device_control_env_var
 
     try:
         selected_physical: int | None = None
diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py
index f1ed4d843..4ebff3290 100644
--- a/vllm_omni/entrypoints/utils.py
+++ b/vllm_omni/entrypoints/utils.py
@@ -10,7 +10,7 @@ from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
 from vllm.transformers_utils.repo_utils import file_or_path_exists
 
 from vllm_omni.entrypoints.stage_utils import _to_dict
-from vllm_omni.utils import detect_device_type, is_rocm
+from vllm_omni.platforms import current_omni_platform
 
 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
@@ -164,16 +164,12 @@ def resolve_model_config_path(model: str) -> str:
                 f"Model is not in standard transformers format and does not have model_index.json. "
                 f"Please ensure the model has proper configuration files with 'model_type' field"
             )
-    device_type = detect_device_type()
 
-    # Try device-specific config first
-    if device_type != "cuda" or is_rocm():
-        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
-        if is_rocm():
-            device_config_file = f"vllm_omni/model_executor/stage_configs/rocm/{model_type}.yaml"
-        device_config_path = PROJECT_ROOT / device_config_file
-        if os.path.exists(device_config_path):
-            return str(device_config_path)
+    default_config_path = current_omni_platform.get_default_stage_config_path()
+    model_type_str = f"{model_type}.yaml"
+    complete_config_path = PROJECT_ROOT / default_config_path / model_type_str
+    if os.path.exists(complete_config_path):
+        return str(complete_config_path)
 
     # Fall back to default config
     stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index 45a1447b3..f4afe675a 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -36,7 +36,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import download_weights_
 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.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -247,7 +247,7 @@ class Qwen2_5OmniForConditionalGeneration(
             if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
                 inputs_embeds = inputs_embeds.to(thinker_dev)
 
-            if is_npu():
+            if current_omni_platform.is_npu():
                 # TODO: remove this hack when NPU supports batched inputs properly
                 thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids
                 # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
index c60015c0e..9d176ac6c 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
@@ -35,7 +35,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames
-from vllm_omni.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
@@ -732,7 +732,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
         beta = 0.0
 
     # TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument
-    if is_npu():
+    if current_omni_platform.is_npu():
         kaiser_window = torch.kaiser_window(
             kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
         ).to("npu")
@@ -796,7 +796,7 @@ class UpSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        if is_npu():
+        if current_omni_platform.is_npu():
             # TODO: When F.pad supports replicate mode on NPU, remove this branch
             input_dtype = hidden_states.dtype
             # F.pad in NPU doesn't support BF16 when mode is replicate.
@@ -843,7 +843,7 @@ class DownSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        if is_npu():
+        if current_omni_platform.is_npu():
             input_dtype = hidden_states.dtype
             # F.pad in NPU doesn't support BF16 when mode is replicate.
             # To ensure the accuracy, manually pad the input tensor.
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
index a42bd8e76..37753ee47 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
@@ -35,7 +35,7 @@ from vllm.v1.sample.sampler import Sampler
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape
-from vllm_omni.utils.platform_utils import is_npu
+from vllm_omni.platforms import current_omni_platform
 
 # Special token IDs for Qwen3 Omni MoE
 # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
@@ -269,7 +269,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         # ========== Stage 1: Thinker ==========
         if self.model_stage == "thinker":
             thinker_dev = self._module_device(self.thinker)
-            if is_npu():
+            if current_omni_platform.is_npu():
                 # Normalize to batched inputs if needed
                 _added_batch_dim = False
                 if input_ids is not None and input_ids.ndim == 1:
@@ -308,7 +308,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     "capture_layer_indices": [0, int(accept_layer)],
                     "return_hidden_states": True,
                 }
-            if is_npu():
+            if current_omni_platform.is_npu():
                 # TODO: remove this hack when NPU supports batched inputs properly
                 thinker_input_ids = input_ids[0] if input_ids is not None and _added_batch_dim else input_ids
                 thinker_inputs_embeds = (
diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml
index 96bdc01b6..93d2b611a 100644
--- a/vllm_omni/model_executor/stage_configs/bagel.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel.yaml
@@ -9,7 +9,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: BagelForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.4
       enforce_eager: true
diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
index 678d5775f..86c0330a6 100644
--- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
@@ -9,7 +9,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: BagelForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.4
       enforce_eager: true
diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
index e3d740ad5..e6ed97660 100644
--- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -40,7 +40,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -69,7 +69,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml
index aae194d37..68f5817a7 100644
--- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml
@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.3
       enforce_eager: true
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
index 9c422d76e..d0380a39d 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
@@ -14,7 +14,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: true
@@ -45,7 +45,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.3
       enforce_eager: true
@@ -76,7 +76,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
index 55d455f56..42748960e 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
@@ -14,7 +14,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -46,7 +46,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.3
       enforce_eager: false
@@ -77,7 +77,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml
index 4abd3835a..69102df6f 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml
@@ -13,7 +13,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: true
@@ -46,7 +46,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
        enforce_eager: true
@@ -82,7 +82,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls:  vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml
index 40799fe7f..d408dbab9 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml
@@ -7,7 +7,7 @@ stage_args:
     engine_args:
       model_stage: qwen3_tts
       model_arch: Qwen3TTSForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/vllm_omni/platforms/__init__.py b/vllm_omni/platforms/__init__.py
new file mode 100644
index 000000000..b76c59a67
--- /dev/null
+++ b/vllm_omni/platforms/__init__.py
@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+import traceback
+from itertools import chain
+from typing import TYPE_CHECKING
+
+from vllm.utils.import_utils import resolve_obj_by_qualname
+from vllm.utils.torch_utils import supports_xccl
+
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+from vllm_omni.plugins import (
+    OMNI_PLATFORM_PLUGINS_GROUP,
+    load_omni_plugins_by_group,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def cuda_omni_platform_plugin() -> str | None:
+    """Check if CUDA OmniPlatform should be activated."""
+    is_cuda = False
+    logger.debug("Checking if CUDA OmniPlatform is available.")
+    try:
+        from vllm.utils.import_utils import import_pynvml
+
+        pynvml = import_pynvml()
+        pynvml.nvmlInit()
+        try:
+            if pynvml.nvmlDeviceGetCount() > 0:
+                is_cuda = True
+                logger.debug("Confirmed CUDA OmniPlatform is available.")
+            else:
+                logger.debug("CUDA OmniPlatform is not available because no GPU is found.")
+        finally:
+            pynvml.nvmlShutdown()
+    except Exception as e:
+        logger.debug("CUDA OmniPlatform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.cuda.platform.CudaOmniPlatform" if is_cuda else None
+
+
+def rocm_omni_platform_plugin() -> str | None:
+    """Check if ROCm OmniPlatform should be activated."""
+    is_rocm = False
+    logger.debug("Checking if ROCm OmniPlatform is available.")
+    try:
+        import amdsmi
+
+        amdsmi.amdsmi_init()
+        try:
+            if len(amdsmi.amdsmi_get_processor_handles()) > 0:
+                is_rocm = True
+                logger.debug("Confirmed ROCm OmniPlatform is available.")
+            else:
+                logger.debug("ROCm OmniPlatform is not available because no GPU is found.")
+        finally:
+            amdsmi.amdsmi_shut_down()
+    except Exception as e:
+        logger.debug("ROCm OmniPlatform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.rocm.platform.RocmOmniPlatform" if is_rocm else None
+
+
+def npu_omni_platform_plugin() -> str | None:
+    """Check if NPU OmniPlatform should be activated."""
+    is_npu = False
+    logger.debug("Checking if NPU OmniPlatform is available.")
+    try:
+        import torch
+
+        if hasattr(torch, "npu") and torch.npu.is_available():
+            is_npu = True
+            logger.debug("Confirmed NPU OmniPlatform is available.")
+    except Exception as e:
+        logger.debug("NPU OmniPlatform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.npu.platform.NPUOmniPlatform" if is_npu else None
+
+
+def xpu_omni_platform_plugin() -> str | None:
+    """Check if XPU OmniPlatform should be activated."""
+    is_xpu = False
+    logger.debug("Checking if XPU OmniPlatform is available.")
+    try:
+        # installed IPEX if the machine has XPUs.
+        import intel_extension_for_pytorch  # noqa: F401
+        import torch
+
+        if supports_xccl():
+            dist_backend = "xccl"
+        else:
+            dist_backend = "ccl"
+            import oneccl_bindings_for_pytorch  # noqa: F401
+
+        if hasattr(torch, "xpu") and torch.xpu.is_available():
+            is_xpu = True
+            from vllm_omni.platforms.xpu import XPUOmniPlatform
+
+            XPUOmniPlatform.dist_backend = dist_backend
+            logger.debug("Confirmed %s backend is available.", XPUOmniPlatform.dist_backend)
+            logger.debug("Confirmed XPU platform is available.")
+    except Exception as e:
+        logger.debug("XPU omni platform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.xpu.platform.XPUOmniPlatform" if is_xpu else None
+
+
+builtin_omni_platform_plugins = {
+    "cuda": cuda_omni_platform_plugin,
+    "rocm": rocm_omni_platform_plugin,
+    "npu": npu_omni_platform_plugin,
+    "xpu": xpu_omni_platform_plugin,
+}
+
+
+def resolve_current_omni_platform_cls_qualname() -> str:
+    """Resolve the current OmniPlatform class qualified name."""
+    platform_plugins = load_omni_plugins_by_group(OMNI_PLATFORM_PLUGINS_GROUP)
+
+    activated_plugins = []
+
+    for name, func in chain(builtin_omni_platform_plugins.items(), platform_plugins.items()):
+        try:
+            assert callable(func)
+            platform_cls_qualname = func()
+            if platform_cls_qualname is not None:
+                activated_plugins.append(name)
+        except Exception:
+            pass
+
+    activated_builtin_plugins = list(set(activated_plugins) & set(builtin_omni_platform_plugins.keys()))
+    activated_oot_plugins = list(set(activated_plugins) & set(platform_plugins.keys()))
+
+    if len(activated_oot_plugins) >= 2:
+        raise RuntimeError(f"Only one OmniPlatform plugin can be activated, but got: {activated_oot_plugins}")
+    elif len(activated_oot_plugins) == 1:
+        platform_cls_qualname = platform_plugins[activated_oot_plugins[0]]()
+        logger.info("OmniPlatform plugin %s is activated", activated_oot_plugins[0])
+    elif len(activated_builtin_plugins) >= 2:
+        raise RuntimeError(f"Only one OmniPlatform plugin can be activated, but got: {activated_builtin_plugins}")
+    elif len(activated_builtin_plugins) == 1:
+        platform_cls_qualname = builtin_omni_platform_plugins[activated_builtin_plugins[0]]()
+        logger.debug("Automatically detected OmniPlatform %s.", activated_builtin_plugins[0])
+    else:
+        platform_cls_qualname = "vllm_omni.platforms.interface.UnspecifiedOmniPlatform"
+        logger.debug("No platform detected, vLLM-Omni is running on UnspecifiedOmniPlatform")
+
+    return platform_cls_qualname
+
+
+_current_omni_platform = None
+_init_trace: str = ""
+
+if TYPE_CHECKING:
+    current_omni_platform: OmniPlatform
+
+
+def __getattr__(name: str):
+    if name == "current_omni_platform":
+        # Lazy init current_omni_platform
+        global _current_omni_platform
+        if _current_omni_platform is None:
+            platform_cls_qualname = resolve_current_omni_platform_cls_qualname()
+            _current_omni_platform = resolve_obj_by_qualname(platform_cls_qualname)()
+            global _init_trace
+            _init_trace = "".join(traceback.format_stack())
+        return _current_omni_platform
+    elif name in globals():
+        return globals()[name]
+    else:
+        raise AttributeError(f"No attribute named '{name}' exists in {__name__}.")
+
+
+def __setattr__(name: str, value):  # noqa: N807
+    if name == "current_omni_platform":
+        global _current_omni_platform
+        _current_omni_platform = value
+    elif name in globals():
+        globals()[name] = value
+    else:
+        raise AttributeError(f"No attribute named '{name}' exists in {__name__}.")
+
+
+__all__ = [
+    "OmniPlatform",
+    "OmniPlatformEnum",
+    "current_omni_platform",
+    "_init_trace",
+]
diff --git a/vllm_omni/platforms/cuda/__init__.py b/vllm_omni/platforms/cuda/__init__.py
new file mode 100644
index 000000000..91ef864be
--- /dev/null
+++ b/vllm_omni/platforms/cuda/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+__all__ = ["CudaOmniPlatform"]
diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py
new file mode 100644
index 000000000..9c65877be
--- /dev/null
+++ b/vllm_omni/platforms/cuda/platform.py
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms.cuda import CudaPlatformBase
+from vllm.platforms.interface import DeviceCapability
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
+    """CUDA/GPU implementation of OmniPlatform (default).
+
+    Inherits all CUDA-specific implementations from vLLM's CudaPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.CUDA
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/model_executor/stage_configs"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        # Check compute capability for Flash Attention support
+        # Flash Attention requires compute capability >= 8.0 and < 10.0
+        compute_capability = cls.get_device_capability()
+        flash_attn_supported = False
+        if compute_capability is not None:
+            major, minor = compute_capability
+            capability = major * 10 + minor
+            flash_attn_supported = 80 <= capability < 100
+
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
+                logger.warning(
+                    "Flash Attention requires GPU with compute capability >= 8.0 "
+                    "and < 10.0. Falling back to TORCH_SDPA backend."
+                )
+                logger.info("Defaulting to diffusion attention backend SDPA")
+                return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        if flash_attn_supported:
+            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+        logger.info("Defaulting to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("cuda")
+        return torch.device("cuda", local_rank)
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        major, minor = torch.cuda.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.cuda.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        return torch.version.cuda
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.cuda.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        free, _ = torch.cuda.mem_get_info(device)
+        return free
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(device_id)
diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py
new file mode 100644
index 000000000..2aac4b7a1
--- /dev/null
+++ b/vllm_omni/platforms/interface.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+
+import torch
+from vllm.platforms import Platform
+
+
+class OmniPlatformEnum(Enum):
+    """Enum for supported Omni platforms."""
+
+    CUDA = "cuda"
+    ROCM = "rocm"
+    NPU = "npu"
+    XPU = "xpu"
+    UNSPECIFIED = "unspecified"
+
+
+class OmniPlatform(Platform):
+    """
+    Abstract base class for vllm-omni Platform.
+
+    Inherits from vLLM's Platform and adds Omni-specific interfaces.
+    This gives OmniPlatform all vLLM Platform capabilities plus
+    Omni-specific methods.
+    """
+
+    _omni_enum: OmniPlatformEnum
+
+    def is_npu(self) -> bool:
+        return self._omni_enum == OmniPlatformEnum.NPU
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        """Get the diffusion attention backend class path for this platform.
+
+        This method selects the appropriate attention backend for diffusion
+        models based on platform capabilities and user preferences.
+
+        Args:
+            selected_backend: User-selected backend name (e.g., "FLASH_ATTN",
+                "TORCH_SDPA", "SAGE_ATTN"). If None, uses platform default.
+            head_size: Attention head size.
+
+        Returns:
+            Fully qualified class path of the selected backend.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        raise NotImplementedError
+
+    @classmethod
+    def synchronize(cls) -> None:
+        raise NotImplementedError
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        raise NotImplementedError
+
+
+class UnspecifiedOmniPlatform(OmniPlatform):
+    _omni_enum = OmniPlatformEnum.UNSPECIFIED
+    device_type = ""
diff --git a/vllm_omni/platforms/npu/__init__.py b/vllm_omni/platforms/npu/__init__.py
new file mode 100644
index 000000000..fb95f7b29
--- /dev/null
+++ b/vllm_omni/platforms/npu/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+__all__ = ["NPUOmniPlatform"]
diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py
new file mode 100644
index 000000000..429297bbd
--- /dev/null
+++ b/vllm_omni/platforms/npu/platform.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm_ascend.platform import NPUPlatform
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class NPUOmniPlatform(OmniPlatform, NPUPlatform):
+    """NPU/Ascend implementation of OmniPlatform.
+
+    Inherits all NPU-specific implementations from vllm-ascend's NPUPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.NPU
+    dist_backend: str = "hccl"
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker"
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/platforms/npu/stage_configs"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        from importlib.util import find_spec
+
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        # Try FLASH_ATTN if mindiesd is available, otherwise fall back to SDPA
+        if find_spec("mindiesd"):
+            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+        logger.info("Falling back to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("npu")
+        return torch.device("npu", local_rank)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.npu.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        return None
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.npu.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        free, _ = torch.npu.mem_get_info(device)
+        return free
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.npu.get_device_properties(device_id)
+        return device_props.total_memory
diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml
similarity index 89%
rename from vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml
rename to vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml
index 49085276f..a02d84ed6 100644
--- a/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml
+++ b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml
@@ -9,10 +9,10 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
-      enforce_eager: true  # Now we only support eager mode
+      enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent
       enable_prefix_caching: false
@@ -36,10 +36,10 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
-      enforce_eager: true
+      enforce_eager: false
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: latent
@@ -58,12 +58,12 @@ stage_args:
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       process: true
-      devices: "0"            # Example: use a different NPU than the previous stage; use "0" if single NPU
+      devices: "2"            # Example: use a different NPU than the previous stage; use "0" if single NPU
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
similarity index 93%
rename from vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml
rename to vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
index eadfea570..fdf817239 100644
--- a/vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml
+++ b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: false
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+       worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.2
        enforce_eager: false
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen3_tts.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml
similarity index 89%
rename from vllm_omni/model_executor/stage_configs/npu/qwen3_tts.yaml
rename to vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml
index 86b5ca60e..d408dbab9 100644
--- a/vllm_omni/model_executor/stage_configs/npu/qwen3_tts.yaml
+++ b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml
@@ -7,7 +7,7 @@ stage_args:
     engine_args:
       model_stage: qwen3_tts
       model_arch: Qwen3TTSForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/vllm_omni/worker/npu/__init__.py b/vllm_omni/platforms/npu/worker/__init__.py
similarity index 100%
rename from vllm_omni/worker/npu/__init__.py
rename to vllm_omni/platforms/npu/worker/__init__.py
diff --git a/vllm_omni/worker/npu/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
similarity index 99%
rename from vllm_omni/worker/npu/npu_ar_model_runner.py
rename to vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
index 1bb53b662..0ba57e038 100644
--- a/vllm_omni/worker/npu/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -37,7 +37,7 @@ from vllm_ascend.utils import ProfileExecuteDuration
 
 from vllm_omni.core.sched.omni_ar_scheduler import KVCacheTransferData
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner
+from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 
 
 class ExecuteModelState(NamedTuple):
diff --git a/vllm_omni/worker/npu/npu_ar_worker.py b/vllm_omni/platforms/npu/worker/npu_ar_worker.py
similarity index 74%
rename from vllm_omni/worker/npu/npu_ar_worker.py
rename to vllm_omni/platforms/npu/worker/npu_ar_worker.py
index c4895ee1b..1fa3e59ed 100644
--- a/vllm_omni/worker/npu/npu_ar_worker.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_worker.py
@@ -4,10 +4,11 @@
 from vllm.v1.worker.workspace import init_workspace_manager
 from vllm_ascend.worker.worker import NPUWorker
 
-from vllm_omni.worker.npu.npu_ar_model_runner import NPUARModelRunner
+from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
 
 
-class NPUARWorker(NPUWorker):
+class NPUARWorker(OmniWorkerMixin, NPUWorker):
     """NPU AR worker for thinker/talker stages in Omni model."""
 
     def init_device(self):
diff --git a/vllm_omni/worker/npu/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
similarity index 99%
rename from vllm_omni/worker/npu/npu_generation_model_runner.py
rename to vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
index e1beef982..e8559bb46 100644
--- a/vllm_omni/worker/npu/npu_generation_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
@@ -27,8 +27,8 @@ from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable
 
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.npu.npu_ar_model_runner import ExecuteModelState
-from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner
+from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState
+from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 
 
 class NPUGenerationModelRunner(OmniNPUModelRunner):
diff --git a/vllm_omni/worker/npu/npu_generation_worker.py b/vllm_omni/platforms/npu/worker/npu_generation_worker.py
similarity index 72%
rename from vllm_omni/worker/npu/npu_generation_worker.py
rename to vllm_omni/platforms/npu/worker/npu_generation_worker.py
index 588d21fc1..edbaa9f64 100644
--- a/vllm_omni/worker/npu/npu_generation_worker.py
+++ b/vllm_omni/platforms/npu/worker/npu_generation_worker.py
@@ -4,10 +4,11 @@
 from vllm.v1.worker.workspace import init_workspace_manager
 from vllm_ascend.worker.worker import NPUWorker
 
-from vllm_omni.worker.npu.npu_generation_model_runner import NPUGenerationModelRunner
+from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
 
 
-class NPUGenerationWorker(NPUWorker):
+class NPUGenerationWorker(OmniWorkerMixin, NPUWorker):
     """NPU generation worker for code2wav stage in Omni model."""
 
     def init_device(self):
diff --git a/vllm_omni/worker/npu/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py
similarity index 100%
rename from vllm_omni/worker/npu/npu_model_runner.py
rename to vllm_omni/platforms/npu/worker/npu_model_runner.py
diff --git a/vllm_omni/platforms/rocm/__init__.py b/vllm_omni/platforms/rocm/__init__.py
new file mode 100644
index 000000000..999dc5a08
--- /dev/null
+++ b/vllm_omni/platforms/rocm/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.rocm.platform import RocmOmniPlatform
+
+__all__ = ["RocmOmniPlatform"]
diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py
new file mode 100644
index 000000000..f2c2d3a85
--- /dev/null
+++ b/vllm_omni/platforms/rocm/platform.py
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms.rocm import RocmPlatform
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class RocmOmniPlatform(OmniPlatform, RocmPlatform):
+    """ROCm/AMD GPU implementation of OmniPlatform.
+
+    Inherits all ROCm-specific implementations from vLLM's RocmPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.ROCM
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        from vllm._aiter_ops import is_aiter_found_and_supported
+
+        # Check if aiter is available for Flash Attention support
+        # aiter currently only is supported on gfx942 and gfx950
+        # https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py
+        compute_capability = torch.cuda.get_device_capability()
+        major, minor = compute_capability
+        capability = major * 10 + minor
+        aiter_supported = is_aiter_found_and_supported() and 90 < capability < 100
+
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            if backend_upper == "FLASH_ATTN" and not aiter_supported:
+                logger.warning(
+                    "Flash Attention requires `aiter` library which is only supported "
+                    "on gfx942 and gfx950. Falling back to TORCH_SDPA backend."
+                )
+                logger.info("Defaulting to diffusion attention backend SDPA")
+                return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        # Choose to enable Flash Attention by default on ROCm
+        # whenever possible as it is the fastest backend
+        if aiter_supported:
+            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+        logger.info("Defaulting to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/platforms/rocm/stage_configs"
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("cuda")
+        return torch.device("cuda", local_rank)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.cuda.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        if torch.version.hip is not None:
+            hip_version = torch.version.hip
+            return hip_version.split("-")[0]
+        return None
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.cuda.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        free, _ = torch.cuda.mem_get_info(device)
+        return free
diff --git a/vllm_omni/model_executor/stage_configs/rocm/qwen2_5_omni.yaml b/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml
similarity index 94%
rename from vllm_omni/model_executor/stage_configs/rocm/qwen2_5_omni.yaml
rename to vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml
index c646aa76a..7887cd2bb 100644
--- a/vllm_omni/model_executor/stage_configs/rocm/qwen2_5_omni.yaml
+++ b/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml
@@ -10,7 +10,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -38,7 +38,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -66,7 +66,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
diff --git a/vllm_omni/model_executor/stage_configs/rocm/qwen3_omni_moe.yaml b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml
similarity index 93%
rename from vllm_omni/model_executor/stage_configs/rocm/qwen3_omni_moe.yaml
rename to vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml
index 73f65ecb5..93e05ea9c 100644
--- a/vllm_omni/model_executor/stage_configs/rocm/qwen3_omni_moe.yaml
+++ b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml
@@ -12,7 +12,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: true
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.3
       enforce_eager: true
@@ -73,7 +73,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
diff --git a/vllm_omni/platforms/xpu/__init__.py b/vllm_omni/platforms/xpu/__init__.py
new file mode 100644
index 000000000..8ec0986c2
--- /dev/null
+++ b/vllm_omni/platforms/xpu/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.xpu.platform import XPUOmniPlatform
+
+__all__ = ["XPUOmniPlatform"]
diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py
new file mode 100644
index 000000000..5c126c6d9
--- /dev/null
+++ b/vllm_omni/platforms/xpu/platform.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms.xpu import XPUPlatform
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class XPUOmniPlatform(OmniPlatform, XPUPlatform):
+    """XPU/Intel GPU implementation of OmniPlatform.
+
+    Inherits all XPU-specific implementations from vLLM's XPUPlatform,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.XPU
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker"
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        logger.info("Defaulting to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/platforms/xpu/stage_configs"
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        if local_rank is None:
+            return torch.device("xpu")
+        return torch.device("xpu", local_rank)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return torch.xpu.device_count()
+
+    @classmethod
+    def get_device_version(cls) -> str | None:
+        # XPU does not have a version string like CUDA
+        return None
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.xpu.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        if device is None:
+            device_id = 0
+        else:
+            device_id = device.index if device.index is not None else 0
+        props = torch.xpu.get_device_properties(device_id)
+        return props.total_memory
diff --git a/vllm_omni/plugins/__init__.py b/vllm_omni/plugins/__init__.py
new file mode 100644
index 000000000..41bf398ca
--- /dev/null
+++ b/vllm_omni/plugins/__init__.py
@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+import vllm.envs as envs
+
+logger = logging.getLogger(__name__)
+
+# Default plugins group will be loaded in all processes (process0, engine core
+# process and worker processes).
+OMNI_DEFAULT_PLUGINS_GROUP = "vllm_omni.general_plugins"
+# Platform plugins group will be loaded in all processes when
+# `vllm_omni.platforms.current_omni_platform` is called and the value not
+# initialized.
+OMNI_PLATFORM_PLUGINS_GROUP = "vllm_omni.platform_plugins"
+
+# Make sure one process only loads plugins once.
+omni_plugins_loaded = False
+
+
+def load_omni_plugins_by_group(group: str) -> dict[str, Callable[[], Any]]:
+    from importlib.metadata import entry_points
+
+    allowed_plugins = envs.VLLM_PLUGINS
+
+    discovered_plugins = entry_points(group=group)
+    if len(discovered_plugins) == 0:
+        logger.debug("No plugins for group %s found.", group)
+        return {}
+
+    # Check if the only discovered plugin is the default one.
+    is_default_group = group == OMNI_DEFAULT_PLUGINS_GROUP
+    # Use INFO for non-default groups and DEBUG for the default group.
+    log_level = logger.debug if is_default_group else logger.info
+
+    log_level("Available plugins for group %s:", group)
+    for plugin in discovered_plugins:
+        log_level("- %s -> %s", plugin.name, plugin.value)
+
+    if allowed_plugins is None:
+        log_level("All plugins in this group will be loaded. Set `VLLM_PLUGINS` to control which plugins to load.")
+
+    plugins: dict[str, Callable[[], Any]] = {}
+    for plugin in discovered_plugins:
+        if allowed_plugins is None or plugin.name in allowed_plugins:
+            if allowed_plugins is not None:
+                log_level("Loading plugin %s", plugin.name)
+
+            try:
+                func = plugin.load()
+                plugins[plugin.name] = func
+            except Exception:
+                logger.exception("Failed to load plugin %s", plugin.name)
+
+    return plugins
+
+
+def load_omni_general_plugins() -> None:
+    """WARNING: plugins can be loaded for multiple times in different
+    processes. They should be designed in a way that they can be loaded
+    multiple times without causing issues.
+    """
+    global omni_plugins_loaded
+    if omni_plugins_loaded:
+        return
+    omni_plugins_loaded = True
+
+    plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP)
+    # General plugins: we only need to execute the loaded functions.
+    for func in plugins.values():
+        func()
+
+
+__all__ = [
+    "load_omni_general_plugins",
+    "OMNI_DEFAULT_PLUGINS_GROUP",
+    "OMNI_PLATFORM_PLUGINS_GROUP",
+]
diff --git a/vllm_omni/utils/__init__.py b/vllm_omni/utils/__init__.py
index 34b2545db..e69de29bb 100644
--- a/vllm_omni/utils/__init__.py
+++ b/vllm_omni/utils/__init__.py
@@ -1,13 +0,0 @@
-from vllm_omni.utils.platform_utils import (
-    detect_device_type,
-    get_device_control_env_var,
-    is_npu,
-    is_rocm,
-)
-
-__all__ = [
-    "detect_device_type",
-    "get_device_control_env_var",
-    "is_npu",
-    "is_rocm",
-]
diff --git a/vllm_omni/utils/platform_utils.py b/vllm_omni/utils/platform_utils.py
deleted file mode 100644
index fb47018e7..000000000
--- a/vllm_omni/utils/platform_utils.py
+++ /dev/null
@@ -1,58 +0,0 @@
-from __future__ import annotations
-
-import torch
-from vllm.platforms import current_platform
-
-
-def detect_device_type() -> str:
-    device_type = getattr(current_platform, "device_type", None)
-    if isinstance(device_type, str) and device_type:
-        return device_type.lower()
-    if torch.cuda.is_available():
-        return "cuda"
-    if hasattr(torch, "npu") and torch.npu.is_available():  # type: ignore[attr-defined]
-        return "npu"
-    return "cpu"
-
-
-def is_npu() -> bool:
-    return detect_device_type() == "npu"
-
-
-def is_rocm() -> bool:
-    return current_platform.is_rocm()
-
-
-def get_device_control_env_var() -> str:
-    """Return the environment variable name for device visibility control."""
-    if hasattr(current_platform, "device_control_env_var"):
-        env_var = getattr(current_platform, "device_control_env_var", None)
-        if isinstance(env_var, str) and env_var:
-            return env_var
-
-    device_type = detect_device_type()
-    if device_type == "npu":
-        return "ASCEND_RT_VISIBLE_DEVICES"
-    return "CUDA_VISIBLE_DEVICES"  # fallback
-
-
-def get_diffusion_worker_class() -> type:
-    """Get the appropriate diffusion WorkerProc class based on current device type.
-
-    Returns:
-        The WorkerProc class for the detected device type (either NPUWorkerProc or WorkerProc).
-
-    Raises:
-        ImportError: If the worker module for the detected device is not available.
-    """
-    device_type = detect_device_type()
-
-    if device_type == "npu":
-        from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorkerProc
-
-        return NPUWorkerProc
-    else:
-        # Default to GPU worker for cuda and other devices
-        from vllm_omni.diffusion.worker.gpu_diffusion_worker import WorkerProc
-
-        return WorkerProc
diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py
index d2dafce68..a1d6b4665 100644
--- a/vllm_omni/worker/gpu_ar_worker.py
+++ b/vllm_omni/worker/gpu_ar_worker.py
@@ -13,11 +13,12 @@ from vllm.v1.worker.utils import request_memory
 from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
 
 logger = init_logger(__name__)
 
 
-class GPUARWorker(GPUWorker):
+class GPUARWorker(OmniWorkerMixin, GPUWorker):
     """GPU worker for autoregressive omni model stages.
 
     Extends the base GPUWorker to initialize and manage autoregressive
diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py
index 19f8ab84b..0aa13ed76 100644
--- a/vllm_omni/worker/gpu_generation_worker.py
+++ b/vllm_omni/worker/gpu_generation_worker.py
@@ -13,11 +13,12 @@ from vllm.v1.worker.utils import request_memory
 from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
 
 logger = init_logger(__name__)
 
 
-class GPUGenerationWorker(GPUWorker):
+class GPUGenerationWorker(OmniWorkerMixin, GPUWorker):
     """GPU Worker for Generation model (non-autoregressive waveform generation).
 
     Usage in stage config:
diff --git a/vllm_omni/worker/mixins.py b/vllm_omni/worker/mixins.py
new file mode 100644
index 000000000..5b25b8362
--- /dev/null
+++ b/vllm_omni/worker/mixins.py
@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class OmniWorkerMixin:
+    """Mixin to ensure Omni plugins are loaded in worker processes."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        from vllm_omni.plugins import load_omni_general_plugins
+
+        load_omni_general_plugins()

From 741f7e282605a1e32a63a8631d4289605a7730df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= <wangzhipeng628@gmail.com>
Date: Wed, 28 Jan 2026 23:12:45 +0800
Subject: [PATCH 11/32] [Core]: KV Cache Transfer Encapsulation (#979)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
---
 .../omni_connectors/test_kv_flow.py           | 423 +++++++++-------
 .../worker/diffusion_model_runner.py          | 133 +----
 .../omni_connectors/kv_transfer_manager.py    | 459 ++++++++++++++++++
 .../stage_configs/bagel_multiconnector.yaml   |   6 -
 vllm_omni/worker/gpu_ar_model_runner.py       | 197 +-------
 5 files changed, 714 insertions(+), 504 deletions(-)
 create mode 100644 vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py
index 2bb06d4e0..8c7ff79ca 100644
--- a/tests/distributed/omni_connectors/test_kv_flow.py
+++ b/tests/distributed/omni_connectors/test_kv_flow.py
@@ -1,194 +1,251 @@
-import unittest
-from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
-
+import pytest
 import torch
 
-from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
+from tests.utils import hardware_test
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+)
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 
 
-class MockGPUARModelRunner(GPUARModelRunner):
-    """Subclass to bypass heavy initialization."""
-
-    def __init__(self, kv_caches, input_batch):
-        self.kv_caches = kv_caches
-        self.input_batch = input_batch
-        self.device = "cpu"
-        self.cache_config = MagicMock()
-        self.cache_config.block_size = 16
-        self.cache_config.cache_dtype = "auto"
-        self.logger = MagicMock()
-
-
-class MockBagelPipeline(BagelPipeline):
-    """Subclass to bypass heavy initialization."""
-
+class MockConnector:
     def __init__(self):
-        self.device = "cpu"
-        self.od_config = MagicMock()
-        self.bagel = MagicMock()
-        self.tokenizer = MagicMock()
-        self.language_model = MagicMock()
-        self.new_token_ids = {}
+        self.store = {}
+
+    def put(self, from_stage, to_stage, put_key, data):
+        # The manager now passes full key as put_key
+        key = f"{from_stage}->{to_stage}:{put_key}"
+        self.store[key] = data
+        return True, len(str(data)), None  # (success, size, metadata)
+
+    def get(self, from_stage, to_stage, get_key, metadata=None):
+        # The manager now passes full key as get_key
+        key = f"{from_stage}->{to_stage}:{get_key}"
+        if key in self.store:
+            return self.store[key], len(str(self.store[key]))
+        return None
 
 
-class TestKVFlow(unittest.TestCase):
-    def setUp(self):
-        # Common constants
-        self.num_layers = 2
-        self.num_heads = 4
-        self.head_dim = 16
-        self.block_size = 8
-        self.x = 8  # PagedAttention factor
-        self.seq_len = 20
-        self.req_id = "req_test_1"
-
-    def test_sender_extraction_logic(self):
-        """Test extraction logic in GPUARModelRunner."""
-        if GPUARModelRunner is object:
-            self.skipTest("vLLM not installed")
-
-        # 1. Setup KV Cache (List of tuples (K, V))
-        # Shape: [num_blocks, block_size, num_heads, head_dim] matching 4D expectation
-        num_blocks = 10
-        kv_caches = []
-        for _ in range(self.num_layers):
-            k_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
-            v_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
-            # Stack K and V to create [2, num_blocks, block_size, n_heads, head_dim]
-            layer_cache = torch.stack([k_cache, v_cache], dim=0)
-            kv_caches.append(layer_cache)
-
-        # 2. Setup Input Batch Mock
-        block_ids = [1, 3, 5]
-        mock_input_batch = MagicMock()
-        mock_input_batch.req_id_to_index = {self.req_id: 0}
-        mock_input_batch.block_table.get_row.return_value = block_ids
-
-        # 3. Instantiate Runner
-        runner = MockGPUARModelRunner(kv_caches, mock_input_batch)
-
-        # 4. Run Extraction
-        # calling _extract_kv_cache directly: (req_id, block_ids, seq_len)
-        result = runner._extract_kv_cache(self.req_id, block_ids, self.seq_len)
-
-        # 5. Verify Result
-        self.assertIsNotNone(result)
-        self.assertEqual(result.request_id, self.req_id)
-
-        # Check keys "key_cache" and "value_cache" exist (length 2)
-        self.assertEqual(len(result.layer_blocks["key_cache"]), 2)
-        self.assertEqual(len(result.layer_blocks["value_cache"]), 2)
-
-        # Check Tensor Shape: [seq_len, num_heads, head_dim]
-        # result.layer_blocks["key_cache"] is a list of tensors
-        expected_shape = (self.seq_len, self.num_heads, self.head_dim)
-        self.assertEqual(result.layer_blocks["key_cache"][0].shape, expected_shape)
-
-        return result  # Return for use in next test
-
-    def test_receiver_injection_logic(self):
-        """Test injection logic in BagelPipeline."""
-        if BagelPipeline is object:
-            self.skipTest("vLLM not installed")
-
-        # 1. Get Data from Sender Test (simulate transfer)
-        # Re-create data manually to be independent
-        key_cache = []
-        value_cache = []
-        expected_shape = (self.seq_len, self.num_heads, self.head_dim)
-        for i in range(self.num_layers):
-            key_cache.append(torch.randn(expected_shape))
-            value_cache.append(torch.randn(expected_shape))
-
-        layer_blocks = {"key_cache": key_cache, "value_cache": value_cache}
-
-        transfer_data = MagicMock()  # Mock KVCacheTransferData
-        transfer_data.layer_blocks = layer_blocks
-        transfer_data.metadata = {"kv_lens": [self.seq_len], "ropes": [0]}
-
-        # 2. Setup Request with Injected Data
-        sp = OmniDiffusionSamplingParams(
-            past_key_values=SimpleNamespace(**layer_blocks),
-            kv_metadata=transfer_data.metadata,
-        )
-
-        req = OmniDiffusionRequest(["test"], sp)
-
-        # 3. Setup Pipeline
-        pipeline = MockBagelPipeline()
-
-        # Mock Bagel's NaiveCache
-        class RealNaiveCache:  # Minimal impl
-            def __init__(self, n):
-                self.key_cache = {i: None for i in range(n)}
-                self.value_cache = {i: None for i in range(n)}
-
-        pipeline.bagel.config.llm_config.num_hidden_layers = self.num_layers
-
-        captured_context = {}
-
-        def mock_prepare_prompts(curr_kvlens, curr_rope, **kwargs):
-            # Capture the state passed to prepare_prompts
-            captured_context["kv_lens"] = curr_kvlens
-            captured_context["ropes"] = curr_rope
-            return {}, [0], [0]
-
-        pipeline.bagel.prepare_prompts = MagicMock(side_effect=mock_prepare_prompts)
-        with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.NaiveCache") as MockNaiveCacheCls:
-            # Setup the instance returned by the constructor
-            mock_cache_instance = RealNaiveCache(self.num_layers)
-            MockNaiveCacheCls.return_value = mock_cache_instance
-
-        # Verification of Injection Logic (Simulation)
-        current_cache = RealNaiveCache(self.num_layers)
-
-        # --- Logic from Source Code ---
-        injected_kv = req.sampling_params.past_key_values
-        if isinstance(current_cache, RealNaiveCache) and hasattr(injected_kv, "key_cache"):
-            # Assuming injected_kv is SimpleNamespace or object with list attrs
-            for layer_idx in range(len(injected_kv.key_cache)):
-                if injected_kv.key_cache[layer_idx] is not None:
-                    k_tensor = injected_kv.key_cache[layer_idx]
-                    v_tensor = injected_kv.value_cache[layer_idx]
-
-                    if k_tensor.device != pipeline.device:
-                        k_tensor = k_tensor.to(pipeline.device)
-                    if v_tensor.device != pipeline.device:
-                        v_tensor = v_tensor.to(pipeline.device)
-
-                    current_cache.key_cache[layer_idx] = k_tensor
-                    current_cache.value_cache[layer_idx] = v_tensor
-
-        self.assertTrue(torch.allclose(current_cache.key_cache[0], layer_blocks["key_cache"][0]))
-        self.assertTrue(torch.allclose(current_cache.value_cache[1], layer_blocks["value_cache"][1]))
-
-    def test_integration(self):
-        """Simulate the flow from Sender -> Connector (Dict) -> Receiver."""
-        if GPUARModelRunner is object or BagelPipeline is object:
-            self.skipTest("vLLM not installed")
-
-        # 1. Sender (Extraction)
-        runner_test_result = self.test_sender_extraction_logic()  # Get KVCacheTransferData
-
-        # 2. Connector (Serialize/Deserialize Simulation)
-        # KVCacheTransferData has to_dict
-        data_dict = runner_test_result.to_dict()
-
-        # 3. Receiver (Request Setup)
-        sp = OmniDiffusionSamplingParams(
-            past_key_values=data_dict["layer_blocks"],
-            kv_metadata=data_dict["metadata"],
-        )
-        req = OmniDiffusionRequest(["integration_test"], sp)  # noqa: F841
-
-        # 4. Receiver (Injection Simulation)
-        # Use the logic verification again
-        pass
+@pytest.fixture
+def mock_connector():
+    return MockConnector()
 
 
-if __name__ == "__main__":
-    unittest.main()
+@pytest.fixture
+def kv_config():
+    return OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage="stage1",
+        to_stage="stage2",
+        stage_id="stage2",  # Acting as receiver for some tests
+        need_recv_cache=True,
+        need_send_cache=True,
+        recv_timeout=1.0,  # Short timeout for tests
+    )
+
+
+@pytest.fixture
+def common_constants():
+    return {
+        "num_layers": 2,
+        "num_heads": 4,
+        "head_dim": 16,
+        "block_size": 8,
+        "seq_len": 20,
+        "req_id": "req_test_1",
+    }
+
+
+@pytest.mark.cache
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_manager_extraction(kv_config, mock_connector, common_constants):
+    """Test extraction and sending logic in OmniKVTransferManager."""
+    num_layers = common_constants["num_layers"]
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    seq_len = common_constants["seq_len"]
+    req_id = common_constants["req_id"]
+
+    num_blocks = 10
+    kv_caches = []
+    for _ in range(num_layers):
+        k_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
+        v_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
+        # Stack K and V to create [2, num_blocks, block_size, n_heads, head_dim]
+        layer_cache = torch.stack([k_cache, v_cache], dim=0)
+        kv_caches.append(layer_cache)
+
+    block_ids = [1, 3, 5]
+    finished_reqs = {req_id: {"block_ids": block_ids, "seq_len": seq_len}}
+
+    manager = OmniKVTransferManager(kv_config)
+    # Mock the connector factory or injection
+    manager._connector = mock_connector
+
+    processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
+
+    assert req_id in processed
+
+    # Check if data was put into connector
+    # Manager builds full key: omni_{from}_to_{to}_kv_cache_{req_id}
+    full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
+    expected_key = f"stage1->stage2:{full_request_id}"
+    assert expected_key in mock_connector.store
+
+    data = mock_connector.store[expected_key]
+    assert data["request_id"] == req_id
+    assert "layer_blocks" in data
+    assert len(data["layer_blocks"]["key_cache"]) == num_layers
+
+    # Verify shape of extracted tensor: [seq_len, heads, dim]
+    # Note: Manager detaches and moves to CPU
+    expected_shape = (seq_len, num_heads, head_dim)
+    assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
+
+
+@pytest.mark.cache
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_manager_reception(kv_config, mock_connector, common_constants):
+    """Test reception and injection logic in OmniKVTransferManager."""
+    num_layers = common_constants["num_layers"]
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    seq_len = common_constants["seq_len"]
+    req_id = common_constants["req_id"]
+
+    expected_shape = (seq_len, num_heads, head_dim)
+    key_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+    value_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+
+    layer_blocks = {"key_cache": key_cache, "value_cache": value_cache}
+    metadata = {
+        "block_size": block_size,
+        "num_layers": num_layers,
+        "dtype": "float32",
+        "seq_len": seq_len,
+    }
+
+    data_to_receive = {
+        "request_id": req_id,
+        "layer_blocks": layer_blocks,
+        "metadata": metadata,
+        "block_ids": [],
+    }
+
+    # In setUp, from_stage="stage1", stage_id="stage2". recv_stages=("stage1", "stage2")
+
+    manager = OmniKVTransferManager(kv_config)
+    manager._connector = mock_connector
+
+    # Pre-populate connector with data
+    # Manager builds full key: omni_{from}_to_{to}_kv_cache_{req_id}
+    full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
+    store_key = f"stage1->stage2:{full_request_id}"
+    mock_connector.store[store_key] = data_to_receive
+
+    req = OmniDiffusionRequest(
+        prompts=["test_recv"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=[req_id],
+    )
+    # req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
+    # But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
+    success = manager.receive_kv_cache(req, target_device=torch.device("cpu"))
+
+    assert success
+    assert hasattr(req, "past_key_values")
+    assert hasattr(req, "kv_metadata")
+
+    assert len(req.past_key_values.key_cache) == num_layers
+    assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0])
+    assert req.kv_metadata["seq_len"] == seq_len
+
+
+@pytest.mark.cache
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_integration_flow(common_constants):
+    """Simulate extraction -> connector -> reception."""
+    num_layers = common_constants["num_layers"]
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    req_id = common_constants["req_id"]
+
+    sender_config = OmniKVCacheConfig(
+        connector_config={"type": "mock"}, from_stage="sender", to_stage="receiver", need_send_cache=True
+    )
+    sender_manager = OmniKVTransferManager(sender_config)
+    connector = MockConnector()
+    sender_manager._connector = connector  # Shared connector
+
+    # Create Data
+    num_blocks = 5
+    kv_caches = []
+    for _ in range(num_layers):
+        layer = torch.randn(2, num_blocks, block_size, num_heads, head_dim)
+        kv_caches.append(layer)
+
+    finished_reqs = {req_id: {"block_ids": [0, 1], "seq_len": 10}}
+
+    # Send
+    sender_manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
+
+    receiver_config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage="sender",
+        stage_id="receiver",
+        need_recv_cache=True,
+        recv_timeout=1.0,
+    )
+    receiver_manager = OmniKVTransferManager(receiver_config)
+    receiver_manager._connector = connector  # Share the same mock connector instance
+
+    req = OmniDiffusionRequest(
+        prompts=["test_integ"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=[req_id],
+    )
+
+    # Receive
+    success = receiver_manager.receive_kv_cache(req)
+
+    # Verify
+    assert success
+    assert req.past_key_values is not None
+    assert req.kv_metadata["seq_len"] == 10
+
+
+@pytest.mark.cache
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_manager_extraction_no_connector(kv_config, common_constants):
+    """Test extraction when connector is unavailable (should still return IDs)."""
+    block_size = common_constants["block_size"]
+    req_id = common_constants["req_id"]
+
+    manager = OmniKVTransferManager(kv_config)
+    # Force connector to be None
+    manager._connector = None
+    manager.config.connector_config = None
+    finished_reqs = {req_id: {"block_ids": [1, 2], "seq_len": 10}}
+
+    processed = manager.handle_finished_requests_kv_transfer(
+        finished_reqs, kv_caches=[], block_size=block_size, cache_dtype="float32"
+    )
+
+    assert req_id in processed
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index 7335162b0..b60d2c2f5 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -28,8 +28,7 @@ from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.offload import apply_offload_hooks
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
-from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 
 logger = init_logger(__name__)
 
@@ -62,10 +61,9 @@ class DiffusionModelRunner:
         self.device = device
         self.pipeline = None
         self.cache_backend = None
-        self.connector = None
 
-        # Initialize OmniConnector after vllm_config is available (via init_device_and_model)
-        self._init_omni_connector()
+        # Initialize KV cache manager for connector management
+        self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
 
     def load_model(
         self,
@@ -138,126 +136,6 @@ class DiffusionModelRunner:
 
         logger.info("Model runner: Initialization complete.")
 
-    def _init_omni_connector(self) -> None:
-        # TODO(wzliu)! get real connector from yaml file instead of hardcode
-        """Initialize OmniConnector for KV cache transfer."""
-        try:
-            connector_config = None
-
-            # 1. Try to get from omni_kv_config (injected from YAML)
-            # Use self.od_config because self.vllm_config is a dummy VllmConfig without model_config
-            if self.od_config.omni_kv_config:
-                connector_config = self.od_config.omni_kv_config.get("connector_config")
-
-            if not connector_config:
-                logger.warning("No OmniConnector config found, skipping initialization")
-                return
-
-            logger.info(f"Initializing OmniConnector with config: {connector_config}")
-
-            c_type = connector_config.get("type")
-            if not c_type:
-                logger.error("Connector config missing 'type'")
-                return
-
-            c_extra = {k: v for k, v in connector_config.items() if k != "type"}
-            connector_spec = ConnectorSpec(name=c_type, extra=c_extra)
-
-            self.connector = OmniConnectorFactory.create_connector(connector_spec)
-
-        except Exception as e:
-            logger.error(f"Failed to initialize OmniConnector: {e}")
-            import traceback
-
-            traceback.print_exc()
-
-    def _receive_kv_cache_for_request(self, req: OmniDiffusionRequest) -> None:
-        """Receive KV cache for a request via OmniConnector."""
-        # TODO(wzliu)! must get control info from stage queue instead of hardcode
-        if not req.request_ids:
-            logger.warning("Request has no ID, cannot receive KV cache")
-            return
-        request_id = req.request_ids[0]
-
-        try:
-            logger.info(f"Attempting to receive KV cache for request {request_id}")
-
-            # TODO: Key used for transfer (must match sender side)
-            # key = f"kv_cache_{req.request_id}"
-
-            # Get data from connector
-            # Determine from_stage and to_stage dynamically
-            omni_kv_config = self.od_config.omni_kv_config
-            stage_id = omni_kv_config.get("stage_id")
-            engine_input_source = omni_kv_config.get("engine_input_source", [])
-
-            to_stage = stage_id
-            # Default to stage_id - 1 if input source is not explicit
-            if engine_input_source:
-                from_stage = engine_input_source[0]
-            elif isinstance(stage_id, int):
-                from_stage = stage_id - 1
-            else:
-                raise ValueError("Invalid stage id")
-            logger.info(f"Wait for KV cache for request {request_id} from stage {from_stage} to {to_stage}...")
-
-            # Check if we should receive KV cache based on config
-            need_recv_cache = omni_kv_config.get("need_recv_cache", False)
-            if need_recv_cache:
-                # Default timeout 30 seconds to prevent infinite hanging
-                timeout = omni_kv_config.get("recv_timeout", 30.0)
-                start_time = time.time()
-
-                while True:
-                    get_key = f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"
-                    result = self.connector.get(
-                        from_stage=from_stage,
-                        to_stage=to_stage,
-                        get_key=get_key,
-                    )
-                    if result:
-                        break
-
-                    if time.time() - start_time > timeout:
-                        logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s")
-                        result = None
-                        break
-
-                    time.sleep(0.5)
-            else:
-                logger.info(f"Skip receiving KV cache for {request_id} (need_recv_cache=False)")
-                result = None
-
-            if result:
-                data, size = result
-                logger.info(f"Successfully received KV cache for {request_id}")
-
-                # Assume data structure matches KVCacheTransferData.to_dict()
-                if isinstance(data, dict) and "layer_blocks" in data:
-                    # Get layer blocks and ensure they are on the correct device
-                    layer_blocks = data["layer_blocks"]
-
-                    # Move tensors to GPU if needed (OmniSerializer should handle tensor reconstruction)
-                    for cache_list in [layer_blocks["key_cache"], layer_blocks["value_cache"]]:
-                        for i, tensor in enumerate(cache_list):
-                            if isinstance(tensor, torch.Tensor) and tensor.device != self.pipeline.device:
-                                cache_list[i] = tensor.to(self.pipeline.device).contiguous()
-                    from types import SimpleNamespace
-
-                    req.sampling_params.past_key_values = SimpleNamespace(**layer_blocks)
-
-                if "metadata" in data:
-                    req.sampling_params.kv_metadata = data["metadata"]
-
-            else:
-                logger.warning(f"No KV cache received for {request_id} (timeout or empty)")
-
-        except Exception as e:
-            logger.error(f"Error receiving KV cache for {request_id}: {e}")
-            import traceback
-
-            traceback.print_exc()
-
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights into the pipeline."""
         return self.pipeline.load_weights(weights)
@@ -277,9 +155,8 @@ class DiffusionModelRunner:
         if len(req.prompts) == 0:
             raise ValueError("Cannot execute model with empty request list")
 
-        # [Omni] KV Cache Receiving Logic
-        if req.sampling_params.need_kv_receive and self.connector is not None:
-            self._receive_kv_cache_for_request(req)
+        # The manager handles the check for need_recv_cache internally
+        self.kv_transfer_manager.receive_kv_cache(req, target_device=getattr(self.pipeline, "device", None))
 
         if req.sampling_params.generator is None and req.sampling_params.seed is not None:
             req.sampling_params.generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
new file mode 100644
index 000000000..82c06fdaf
--- /dev/null
+++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
@@ -0,0 +1,459 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unified OmniConnector and KV cache transfer management."""
+
+import time
+from collections.abc import Callable
+from dataclasses import asdict, dataclass
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from .factory import OmniConnectorFactory
+from .utils.config import ConnectorSpec
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class OmniKVCacheConfig:
+    """Configuration for OmniKVTransferManager."""
+
+    connector_config: dict[str, Any] | None = None
+    from_stage: str | None = None
+    to_stage: str | None = None
+    stage_id: str | int | None = None
+    engine_input_source: list[str | int] | None = None
+    need_recv_cache: bool = False
+    need_send_cache: bool = False
+    recv_timeout: float = 30.0
+
+
+@dataclass
+class KVCacheTransferData:
+    """Container for KV cache transfer data."""
+
+    request_id: str
+    layer_blocks: dict[str, Any]
+    block_ids: list[int]
+    metadata: dict[str, Any]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return asdict(self)
+
+
+class OmniKVTransferManager:
+    """Unified management for OmniConnector and KV cache transfer.
+
+    This class encapsulates all KV cache related operations:
+    - Connector initialization and lazy creation
+    - KV cache extraction from GPU blocks
+    - KV cache transfer with retry logic
+    - KV cache receiving with timeout
+    """
+
+    def __init__(self, config: OmniKVCacheConfig):
+        self.config = config
+        self._connector = None
+
+        # Pre-calculate send stages (from_stage, to_stage)
+        self.send_stages = (
+            (str(config.from_stage), str(config.to_stage)) if config.from_stage and config.to_stage else (None, None)
+        )
+
+        # Pre-calculate receive stages (from_stage, to_stage)
+        recv_from = config.from_stage
+        if config.engine_input_source:
+            recv_from = config.engine_input_source[0]
+        elif isinstance(config.stage_id, int):
+            recv_from = config.stage_id - 1
+
+        self.recv_stages = (
+            (str(recv_from), str(config.stage_id))
+            if recv_from is not None and config.stage_id is not None
+            else (None, None)
+        )
+
+    @classmethod
+    def _create(cls, cfg: dict | None) -> "OmniKVTransferManager":
+        """Create manager from raw config dict."""
+        if not cfg or not isinstance(cfg, dict):
+            return cls(OmniKVCacheConfig())
+        return cls(
+            OmniKVCacheConfig(
+                connector_config=cfg.get("connector_config"),
+                from_stage=cfg.get("omni_from_stage"),
+                to_stage=cfg.get("omni_to_stage"),
+                stage_id=cfg.get("stage_id"),
+                engine_input_source=cfg.get("engine_input_source", []),
+                need_recv_cache=cfg.get("need_recv_cache", False),
+                need_send_cache=cfg.get("need_send_cache", False),
+                recv_timeout=cfg.get("recv_timeout", 30.0),
+            )
+        )
+
+    @classmethod
+    def from_model_config(cls, config: Any) -> "OmniKVTransferManager":
+        """Create from model config (for AR model runner)."""
+        return cls._create(getattr(config, "omni_kv_config", None))
+
+    @classmethod
+    def from_od_config(cls, config: Any) -> "OmniKVTransferManager":
+        """Create from OmniDiffusion config (for diffusion runner)."""
+        return cls._create(getattr(config, "omni_kv_config", None))
+
+    @classmethod
+    def from_vllm_config(cls, vllm_config: Any, model_config: Any) -> "OmniKVTransferManager":
+        """Create from vllm config with fallback to kv_transfer_config."""
+        # Primary: omni_kv_config from model_config
+        omni_kv = getattr(model_config, "omni_kv_config", None)
+        if isinstance(omni_kv, dict):
+            return cls._create(omni_kv)
+
+        # Fallback: check kv_transfer_config
+        kv_cfg = getattr(vllm_config, "kv_transfer_config", None)
+        if kv_cfg:
+            direct = getattr(kv_cfg, "omni_connector_config", None)
+            if isinstance(direct, dict) and direct:
+                return cls._create({"connector_config": direct})
+            extra = getattr(kv_cfg, "kv_connector_extra_config", None)
+            if isinstance(extra, dict):
+                omni = extra.get("omni_connector_config")
+                if isinstance(omni, dict) and omni:
+                    return cls._create({"connector_config": omni})
+
+        return cls(OmniKVCacheConfig())
+
+    @property
+    def connector(self):
+        """Lazy initialization of connector."""
+        # If a previous initialization attempt failed, don't retry on every access.
+        if self._connector is False:
+            return None
+
+        if self._connector is None:
+            cfg = self.config.connector_config
+            if cfg and (c_type := cfg.get("type")):
+                try:
+                    logger.info(f"Initializing OmniConnector with config: {cfg}")
+                    c_extra = {k: v for k, v in cfg.items() if k != "type"}
+                    self._connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra))
+                except Exception as e:
+                    logger.error(f"Failed to initialize OmniConnector: {e}")
+                    import traceback
+
+                    traceback.print_exc()
+                    # Cache failure sentinel to avoid repeated initialization attempts in hot paths.
+                    self._connector = False
+
+        return self._connector if self._connector else None
+
+    def get_connector(self):
+        """Get connector (compatibility wrapper for existing code)."""
+        return self.connector
+
+    def handle_finished_requests_kv_transfer(
+        self,
+        finished_reqs: dict[str, dict[str, Any]],
+        kv_caches: list[torch.Tensor],
+        block_size: int,
+        cache_dtype: str,
+        request_id_resolver: Callable[[str], str] | None = None,
+    ) -> list[str]:
+        """Handle KV cache transfer for finished requests.
+
+        This method extracts KV cache from GPU blocks and transfers them
+        to the downstream stage via the connector.
+
+        Args:
+            finished_reqs: Dict mapping request_id to {block_ids, seq_len}
+            kv_caches: List of KV cache tensors per layer
+            block_size: Size of each cache block
+            cache_dtype: Data type of the cache
+            request_id_resolver: Optional function to resolve global request ID
+
+        Returns:
+            List of request IDs that were processed
+        """
+        if not finished_reqs:
+            return []
+
+        if not self.config.need_send_cache:
+            return list(finished_reqs.keys())
+
+        if not self.connector:
+            logger.warning("No connector available, skipping KV transfer but freeing resources")
+            return list(finished_reqs.keys())
+
+        logger.debug(f"Processing KV transfer for {len(finished_reqs)} requests")
+
+        extracted_ids = []
+        for req_id, data in finished_reqs.items():
+            try:
+                seq_len = data.get("seq_len", 0)
+                block_ids = data.get("block_ids", [])
+                if not block_ids:
+                    logger.warning(f"Request {req_id} has no block IDs, skipping")
+                    continue
+
+                # Extract KV cache from GPU blocks -> CPU tensors
+                kv_data = self._extract_kv_cache(req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype)
+                if kv_data:
+                    # Resolve global request ID if available
+                    transfer_req_id = request_id_resolver(req_id) if request_id_resolver else req_id
+
+                    # Transfer to downstream stage via connector
+                    self._transfer_kv_cache(kv_data, transfer_req_id)
+
+            except Exception as e:
+                logger.error(f"Failed KV transfer for {req_id}: {e}")
+            finally:
+                extracted_ids.append(req_id)
+
+        return extracted_ids
+
+    def _extract_kv_cache(
+        self,
+        req_id: str,
+        block_ids: list[int],
+        seq_len: int,
+        kv_caches: list[torch.Tensor],
+        block_size: int,
+        cache_dtype: str,
+    ) -> KVCacheTransferData | None:
+        """Extract KV cache from GPU blocks for a single request.
+
+        Args:
+            req_id: Request identifier
+            block_ids: List of block IDs to extract
+            seq_len: Sequence length
+            kv_caches: List of KV cache tensors per layer
+            block_size: Size of each cache block
+            cache_dtype: Data type of the cache
+
+        Returns:
+            KVCacheTransferData if extraction successful, None otherwise
+        """
+        num_layers = len(kv_caches)
+        key_cache: list[torch.Tensor | None] = [None] * num_layers
+        value_cache: list[torch.Tensor | None] = [None] * num_layers
+
+        for layer_idx, kv_tensor in enumerate(kv_caches):
+            # Validate block IDs - shape: [2, num_blocks, block_size, n_heads, head_dim]
+            max_block = kv_tensor.shape[1] - 1
+            valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block]
+            if not valid_ids:
+                continue
+
+            # Extract and reshape: [2, n_blocks, block_size, n_heads, head_dim]
+            # -> [2, seq_len, n_heads, head_dim]
+            selected = kv_tensor[:, valid_ids]  # [2, n_valid, block_size, n_heads, head_dim]
+            n_kv, n_blks, blk_sz, n_heads, d_head = selected.shape
+            flat = selected.reshape(n_kv, n_blks * blk_sz, n_heads, d_head)
+            if seq_len < flat.shape[1]:
+                flat = flat[:, :seq_len]
+
+            # Move to CPU
+            flat_cpu = flat.detach().cpu().contiguous()
+            key_cache[layer_idx] = flat_cpu[0]
+            value_cache[layer_idx] = flat_cpu[1]
+
+        if not any(k is not None for k in key_cache):
+            return None
+
+        return KVCacheTransferData(
+            request_id=req_id,
+            layer_blocks={"key_cache": key_cache, "value_cache": value_cache},
+            block_ids=block_ids,
+            metadata={
+                "block_size": block_size,
+                "num_layers": num_layers,
+                "dtype": str(cache_dtype),
+                "seq_len": seq_len,
+            },
+        )
+
+    def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None:
+        """Transfer KV cache data to downstream stage via OmniConnector.
+
+        Args:
+            kv_data: The extracted KV cache data
+            transfer_req_id: The request ID to use for transfer
+        """
+        from_stage, to_stage = self.send_stages
+        if not from_stage or not to_stage:
+            raise ValueError("Transfer stages (omni_from_stage, omni_to_stage) not configured")
+
+        # Prepare data and transfer with retry
+        data_dict = kv_data.to_dict()
+        data_dict["request_id"] = transfer_req_id
+
+        success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict)
+
+        if success:
+            logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes")
+        else:
+            logger.error(f"KV transfer FAILED: {transfer_req_id}")
+
+    def _transfer_with_retry(
+        self,
+        from_stage: str,
+        to_stage: str,
+        request_id: str,
+        data: dict[str, Any],
+        max_retries: int = 3,
+    ) -> tuple[bool, int, dict[str, Any] | None]:
+        """Transfer data with retry and exponential backoff.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Target stage identifier
+            request_id: Request identifier for the key
+            data: Data to transfer
+            max_retries: Maximum number of retry attempts
+
+        Returns:
+            Tuple of (success, size, metadata)
+        """
+        for attempt in range(max_retries):
+            try:
+                # Build the full key for connector
+                full_request_id = f"omni_{from_stage}_to_{to_stage}_{request_id}"
+                success, size, metadata = self.connector.put(
+                    from_stage=from_stage, to_stage=to_stage, put_key=full_request_id, data=data
+                )
+                if success:
+                    return success, size, metadata
+                logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}")
+            except Exception as e:
+                logger.warning(f"Transfer attempt {attempt + 1} exception: {e}")
+
+            if attempt < max_retries - 1:
+                time.sleep(0.1 * (2**attempt))
+
+        return False, 0, None
+
+    @torch.inference_mode()
+    def receive_kv_cache_for_request(
+        self,
+        request_id: str,
+        target_device: torch.device | None = None,
+    ) -> tuple[dict[str, Any] | None, int]:
+        """Receive KV cache for a specific request.
+
+        This implements the receiving logic from gpu_diffusion_model_runner.py.
+
+        Args:
+            request_id: The request ID to receive KV cache for
+            target_device: Optional device to move tensors to
+
+        Returns:
+            Tuple of (data dict, size) if successful, (None, 0) otherwise
+        """
+        if not self.connector:
+            logger.warning("No connector available for receiving KV cache")
+            return None, 0
+
+        from_stage, to_stage = self.recv_stages
+        if not from_stage or not to_stage:
+            logger.warning("Receive stages not configured")
+            return None, 0
+
+        # Check if we should receive KV cache based on config
+        if not self.config.need_recv_cache:
+            logger.info(f"Skip receiving KV cache for {request_id} (need_recv_cache=False)")
+            return None, 0
+
+        timeout = self.config.recv_timeout
+        start_time = time.time()
+
+        logger.info(f"Wait for KV cache for request {request_id} from stage {from_stage} to {to_stage}...")
+
+        try:
+            while True:
+                # Build the full key for connector
+                full_request_id = f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"
+                result = self.connector.get(
+                    from_stage=from_stage,
+                    to_stage=to_stage,
+                    get_key=full_request_id,
+                )
+                if result:
+                    data, size = result
+                    logger.info(f"Successfully received KV cache for {request_id}, {size} bytes")
+
+                    # Move tensors to target device if specified
+                    if target_device is not None and isinstance(data, dict) and "layer_blocks" in data:
+                        layer_blocks = data["layer_blocks"]
+                        for cache_list in [
+                            layer_blocks.get("key_cache", []),
+                            layer_blocks.get("value_cache", []),
+                        ]:
+                            for i, tensor in enumerate(cache_list):
+                                if isinstance(tensor, torch.Tensor) and tensor.device != target_device:
+                                    cache_list[i] = tensor.to(target_device).contiguous()
+
+                    return data, size
+
+                if time.time() - start_time > timeout:
+                    logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s")
+                    return None, 0
+
+                time.sleep(0.5)
+
+        except Exception as e:
+            logger.error(f"Error receiving KV cache for {request_id}: {e}")
+            import traceback
+
+            traceback.print_exc()
+            return None, 0
+
+    def apply_kv_cache_to_request(self, req: Any, data: dict[str, Any]) -> None:
+        """Apply received KV cache data to a request object.
+
+        Args:
+            req: The request object to apply KV cache to
+            data: The received KV cache data dictionary
+        """
+        if isinstance(data, dict) and "layer_blocks" in data:
+            layer_blocks = data["layer_blocks"]
+            from types import SimpleNamespace
+
+            kv_obj = SimpleNamespace(**layer_blocks)
+            req.past_key_values = kv_obj
+
+            # [Omni] Also attach to sampling_params for BagelPipeline compatibility
+            # BagelPipeline checks req.sampling_params.past_key_values
+            if hasattr(req, "sampling_params") and req.sampling_params is not None:
+                req.sampling_params.past_key_values = kv_obj
+
+        if "metadata" in data:
+            req.kv_metadata = data["metadata"]
+
+    # Legacy compatibility method
+    def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:
+        """Receive KV cache and populate request object (legacy interface).
+
+        Args:
+            req: Request object with request_id attribute
+            target_device: Optional device to move tensors to
+
+        Returns:
+            True if successful, False otherwise
+        """
+        request_id = getattr(req, "request_id", None)
+        if not request_id and hasattr(req, "request_ids") and req.request_ids:
+            # Adaptation for new OmniDiffusionRequest which has list of prompts/ids
+            request_id = req.request_ids[0]
+
+        if not request_id:
+            logger.warning("Request has no ID, cannot receive KV cache")
+            return False
+
+        data, size = self.receive_kv_cache_for_request(request_id, target_device)
+        if data:
+            self.apply_kv_cache_to_request(req, data)
+            return True
+        return False
diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
index 86c0330a6..e378fe675 100644
--- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
@@ -61,13 +61,7 @@ stage_args:
     final_output_type: image
     is_comprehension: false
     default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 2048
       seed: 52
-      detokenize: True
-      repetition_penalty: 1.0
     input_connectors:
       from_stage_0: mooncake_connector
 
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index ecc839fd3..8440cb5b8 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -35,7 +35,7 @@ from vllm.v1.worker.gpu_model_runner import (
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
-from vllm_omni.core.sched.omni_ar_scheduler import KVCacheTransferData
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
@@ -70,7 +70,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
         # each model stage has their own hidden size
         self.hidden_size = self.model_config.hf_text_config.hidden_size
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
-        self.omni_connector = None
+        # Initialize KV cache manager (preserve vllm_config fallback behavior)
+        self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
 
     def _make_buffer(self, *size, dtype, numpy=True):
         # Prevent ray from pinning the buffer due to large size
@@ -95,7 +96,13 @@ class GPUARModelRunner(OmniGPUModelRunner):
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
         # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
-        self.kv_extracted_req_ids = self._handle_finished_requests_kv_transfer(scheduler_output)
+        self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
+            finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
+            kv_caches=self.kv_caches,
+            block_size=self.cache_config.block_size,
+            cache_dtype=str(self.cache_config.cache_dtype),
+            request_id_resolver=self._resolve_global_request_id,
+        )
 
         if self.vllm_config.model_config.enable_return_routed_experts:
             capturer = RoutedExpertsCapturer.get_instance()
@@ -577,161 +584,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         return async_output
 
-    def _handle_finished_requests_kv_transfer(self, scheduler_output: SchedulerOutput) -> list[str]:
-        """Handle KV cache transfer for finished requests.
-
-        Returns list of request IDs that were processed (for Scheduler to free blocks).
-        """
-        finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
-        if not finished_reqs:
-            return []
-
-        logger.debug(f"Processing KV transfer for {len(finished_reqs)} requests")
-
-        extracted_ids = []
-        for req_id, data in finished_reqs.items():
-            try:
-                seq_len = data.get("seq_len", 0)
-                block_ids = data.get("block_ids", [])
-                if not block_ids:
-                    logger.warning(f"Request {req_id} has no block IDs, skipping")
-                    continue
-
-                # Extract KV cache from GPU blocks -> CPU tensors
-                kv_data = self._extract_kv_cache(req_id, block_ids, seq_len)
-                if kv_data:
-                    # Transfer to downstream stage via connector
-                    self._transfer_kv_cache(kv_data)
-
-            except Exception as e:
-                logger.error(f"Failed KV transfer for {req_id}: {e}")
-            finally:
-                extracted_ids.append(req_id)
-
-        return extracted_ids
-
-    def _extract_kv_cache(self, req_id: str, block_ids: list[int], seq_len: int) -> KVCacheTransferData | None:
-        """Extract KV cache from GPU blocks for a single request."""
-        num_layers = len(self.kv_caches)
-        key_cache = [None] * num_layers
-        value_cache = [None] * num_layers
-
-        for layer_idx, kv_tensor in enumerate(self.kv_caches):
-            # Validate block IDs
-            max_block = kv_tensor.shape[1] - 1
-            valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block]
-            if not valid_ids:
-                continue
-
-            # Extract and reshape: [2, n_blocks, block_size, n_heads, head_dim]
-            # -> [2, seq_len, n_heads, head_dim]
-            selected = kv_tensor[:, valid_ids]  # [2, n_valid, block_size, n_heads, head_dim]
-            n_kv, n_blks, blk_sz, n_heads, d_head = selected.shape
-            flat = selected.reshape(n_kv, n_blks * blk_sz, n_heads, d_head)
-            if seq_len < flat.shape[1]:
-                flat = flat[:, :seq_len]
-
-            # Move to CPU
-            flat_cpu = flat.detach().cpu().contiguous()
-            key_cache[layer_idx] = flat_cpu[0]
-            value_cache[layer_idx] = flat_cpu[1]
-
-        if not any(k is not None for k in key_cache):
-            return None
-
-        return KVCacheTransferData(
-            request_id=req_id,
-            layer_blocks={"key_cache": key_cache, "value_cache": value_cache},
-            block_ids=block_ids,
-            metadata={
-                "block_size": self.cache_config.block_size,
-                "num_layers": num_layers,
-                "dtype": str(self.cache_config.cache_dtype),
-                "seq_len": seq_len,
-            },
-        )
-
-    def _transfer_kv_cache(self, kv_data: KVCacheTransferData) -> None:
-        """Transfer KV cache data to downstream stage via OmniConnector."""
-        connector = self._get_or_create_connector()
-        if not connector:
-            return
-
-        # Resolve global request ID if available
-        transfer_req_id = self._resolve_global_request_id(kv_data.request_id)
-        from_stage, to_stage = self._detect_transfer_stages()
-
-        # Prepare data and transfer with retry
-        data_dict = kv_data.to_dict()
-        data_dict["request_id"] = transfer_req_id
-
-        success, size, _ = self._transfer_with_retry(
-            connector, from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict
-        )
-
-        if success:
-            logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes")
-        else:
-            logger.error(f"KV transfer FAILED: {transfer_req_id}")
-
-    def _get_or_create_connector(self) -> Any | None:
-        """Get existing connector or create one from config."""
-        if self.omni_connector:
-            return self.omni_connector
-
-        from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
-        from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
-
-        config = self._get_omni_connector_config()
-        if not config or not isinstance(config, dict):
-            logger.warning("No valid OmniConnector config found")
-            return None
-
-        c_type = config.get("type")
-        if not c_type:
-            logger.error("OmniConnector config missing 'type' field")
-            return None
-
-        c_extra = {k: v for k, v in config.items() if k != "type"}
-        self.omni_connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra))
-        return self.omni_connector
-
-    def _get_omni_connector_config(self) -> dict[str, Any] | None:
-        """Get OmniConnector configuration from model config."""
-        # Primary: omni_kv_config from YAML
-        omni_kv = getattr(self.model_config, "omni_kv_config", None)
-        if isinstance(omni_kv, dict):
-            cfg = omni_kv.get("connector_config")
-            if isinstance(cfg, dict) and cfg:
-                return cfg
-
-        # Fallback: kv_transfer_config
-        kv_cfg = getattr(self.vllm_config, "kv_transfer_config", None)
-        if kv_cfg:
-            direct = getattr(kv_cfg, "omni_connector_config", None)
-            if isinstance(direct, dict) and direct:
-                return direct
-            extra = getattr(kv_cfg, "kv_connector_extra_config", None)
-            if isinstance(extra, dict):
-                omni = extra.get("omni_connector_config")
-                if isinstance(omni, dict) and omni:
-                    return omni
-
-        return None
-
-    def _detect_transfer_stages(self) -> tuple[str, str]:
-        """Detect source and target stages for KV transfer."""
-        omni_kv = getattr(self.model_config, "omni_kv_config", None)
-        if isinstance(omni_kv, dict):
-            from_s = omni_kv.get("omni_from_stage")
-            to_s = omni_kv.get("omni_to_stage")
-            if from_s and to_s:
-                return str(from_s), str(to_s)
-
-        raise ValueError(
-            "KV transfer stages not configured. Please set 'omni_from_stage' and 'omni_to_stage' in omni_kv_config."
-        )
-
     def _resolve_global_request_id(self, req_id: str) -> str:
         """Resolve global request ID from request state."""
         req_state = self.requests.get(req_id)
@@ -747,32 +599,3 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 return global_id.decode("utf-8")
             return str(global_id)
         return req_id
-
-    def _transfer_with_retry(
-        self,
-        connector: Any,
-        from_stage: str,
-        to_stage: str,
-        request_id: str,
-        data: dict[str, Any],
-        max_retries: int = 3,
-    ) -> tuple[bool, int, dict[str, Any] | None]:
-        """Transfer data with retry and exponential backoff."""
-        import time
-
-        for attempt in range(max_retries):
-            try:
-                put_key = f"omni_{from_stage}_to_{to_stage}_{request_id}"
-                success, size, metadata = connector.put(
-                    from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data
-                )
-                if success:
-                    return success, size, metadata
-                logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}")
-            except Exception as e:
-                logger.warning(f"Transfer attempt {attempt + 1} exception: {e}")
-
-            if attempt < max_retries - 1:
-                time.sleep(0.1 * (2**attempt))
-
-        return False, 0, None

From 20e7310eb184e57f912c7ba4a997f7a7a8d55839 Mon Sep 17 00:00:00 2001
From: wangyu <53896905+yenuo26@users.noreply.github.com>
Date: Thu, 29 Jan 2026 11:21:24 +0800
Subject: [PATCH 12/32] [Test]Delete skip mark for amd ci test and fix CI
 failure (#927)

Signed-off-by: wangyu31577 <wangyu31577@hundsun.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: wangyu31577 <wangyu31577@hundsun.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
---
 tests/conftest.py                             | 39 +++-------
 tests/e2e/offline_inference/conftest.py       |  7 +-
 .../stage_configs/rocm/qwen3_omni_ci.yaml     |  8 +-
 tests/e2e/online_serving/test_qwen3_omni.py   |  8 +-
 tests/e2e/stage_configs/qwen3_omni_ci.yaml    |  2 +
 tests/utils.py                                | 76 ++++++++++++++-----
 6 files changed, 78 insertions(+), 62 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index ace891db7..e9b15aeee 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -10,6 +10,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 if "VLLM_TARGET_DEVICE" not in os.environ:
     os.environ["VLLM_TARGET_DEVICE"] = "cpu"
 
+import gc
 import socket
 import subprocess
 import sys
@@ -59,13 +60,12 @@ def clean_gpu_memory_between_tests():
     _run_post_test_cleanup()
 
 
-def _run_pre_test_cleanup():
-    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
+def _run_pre_test_cleanup(enable_force=False):
+    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
         print("GPU cleanup disabled")
         return
 
     print("Pre-test GPU status:")
-    _print_simple_gpu_status()
 
     num_gpus = torch.cuda.device_count()
     if num_gpus > 0:
@@ -74,44 +74,25 @@ def _run_pre_test_cleanup():
 
             wait_for_gpu_memory_to_clear(
                 devices=list(range(num_gpus)),
-                threshold_ratio=0.1,
+                threshold_ratio=0.05,
             )
         except Exception as e:
             print(f"Pre-test cleanup note: {e}")
 
 
-def _run_post_test_cleanup():
-    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
+def _run_post_test_cleanup(enable_force=False):
+    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
+        print("GPU cleanup disabled")
         return
 
-    import gc
-
     if torch.cuda.is_available():
         gc.collect()
         torch.cuda.empty_cache()
 
         print("Post-test GPU status:")
-        _print_simple_gpu_status()
         _print_gpu_processes()
 
 
-def _print_simple_gpu_status():
-    """Print simple GPU memory status"""
-    if not torch.cuda.is_available():
-        print("  CUDA not available")
-        return
-
-    num_devices = torch.cuda.device_count()
-    for device_id in range(num_devices):
-        try:
-            torch.cuda.set_device(device_id)
-            allocated = torch.cuda.memory_allocated(device_id) / (1024**2)
-            reserved = torch.cuda.memory_reserved(device_id) / (1024**2)
-            print(f"  GPU {device_id}: Allocated: {allocated:.1f}MB, Reserved: {reserved:.1f}MB")
-        except Exception:
-            print(f"  GPU {device_id}: Error reading status")
-
-
 def _print_gpu_processes():
     """Print GPU information including nvidia-smi and system processes"""
 
@@ -871,6 +852,9 @@ class OmniServer:
         *,
         env_dict: dict[str, str] | None = None,
     ) -> None:
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
+        cleanup_dist_env_and_memory()
         self.model = model
         self.serve_args = serve_args
         self.env_dict = env_dict
@@ -986,5 +970,6 @@ class OmniServer:
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self.proc:
             self._kill_process_tree(self.proc.pid)
-        _run_post_test_cleanup()
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
         cleanup_dist_env_and_memory()
diff --git a/tests/e2e/offline_inference/conftest.py b/tests/e2e/offline_inference/conftest.py
index 9ee0e4b74..891709837 100644
--- a/tests/e2e/offline_inference/conftest.py
+++ b/tests/e2e/offline_inference/conftest.py
@@ -10,7 +10,7 @@ import pytest
 from vllm import TextPrompt
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import _run_post_test_cleanup
+from tests.conftest import _run_post_test_cleanup, _run_pre_test_cleanup
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -51,6 +51,9 @@ class OmniRunner:
             stage_configs_path: Optional path to YAML stage config file
             **kwargs: Additional arguments passed to Omni
         """
+        cleanup_dist_env_and_memory()
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
         self.model_name = model_name
         self.seed = seed
 
@@ -337,7 +340,7 @@ class OmniRunner:
         self.close()
         del self.omni
         cleanup_dist_env_and_memory()
-        _run_post_test_cleanup()
+        _run_post_test_cleanup(enable_force=True)
 
     def close(self):
         """Close and cleanup the Omni instance."""
diff --git a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
index 0488a06c9..3297347c4 100644
--- a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
+++ b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
@@ -2,7 +2,6 @@
 # Stage 0: Thinker (multimodal understanding + text generation)
 # Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
 # Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
 # The following config has been verified on 2x H100-80G GPUs.
 stage_args:
   - stage_id: 0
@@ -22,7 +21,6 @@ stage_args:
       enable_prefix_caching: false
       hf_config_name: thinker_config
       tensor_parallel_size: 2
-      load_format: dummy
     final_output: true
     final_output_type: text
     is_comprehension: true
@@ -52,7 +50,6 @@ stage_args:
        enable_prefix_caching: false
        distributed_executor_backend: "mp"
        hf_config_name: talker_config
-       load_format: dummy
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
     # final_output: true
@@ -60,7 +57,7 @@ stage_args:
     default_sampling_params:
       temperature: 0.9
       top_k: 50
-      max_tokens: 100
+      max_tokens: 1000
       seed: 42
       detokenize: False
       repetition_penalty: 1.05
@@ -83,7 +80,6 @@ stage_args:
       distributed_executor_backend: "mp"
       max_num_batched_tokens: 1000000
       hf_config_name: thinker_config
-      load_format: dummy
       async_scheduling: false
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
@@ -93,7 +89,7 @@ stage_args:
       temperature: 0.0
       top_p: 1.0
       top_k: -1
-      max_tokens: 200
+      max_tokens: 2000
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py
index eb134ab5e..54c36633e 100644
--- a/tests/e2e/online_serving/test_qwen3_omni.py
+++ b/tests/e2e/online_serving/test_qwen3_omni.py
@@ -133,7 +133,7 @@ def dummy_messages_from_video_data(
 
 def get_prompt(prompt_type="text_only"):
     prompts = {
-        "text_only": "What is the capital of China?",
+        "text_only": "What is the capital of China? Answer in 20 words.",
         "mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
     }
     return prompts.get(prompt_type, prompts["text_only"])
@@ -144,9 +144,6 @@ def get_max_batch_size(size_type="few"):
     return batch_sizes.get(size_type, 5)
 
 
-@pytest.mark.skipif(
-    current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues"
-)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
     """
@@ -226,9 +223,6 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     assert similarity > 0.9, "The audio content is not same as the text"
 
 
-@pytest.mark.skipif(
-    current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues"
-)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
     """
diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml
index d7b07ddc9..148a949d8 100644
--- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml
@@ -20,6 +20,7 @@ stage_args:
     engine_output_type: latent  # Output hidden states for talker
     distributed_executor_backend: "mp"
     max_num_batched_tokens: 32768
+    max_model_len: 32768
     enable_prefix_caching: false
     hf_config_name: thinker_config
     tensor_parallel_size: 2
@@ -51,6 +52,7 @@ stage_args:
     engine_output_type: latent  # Output codec codes for code2wav
     enable_prefix_caching: false
     max_num_batched_tokens: 32768
+    max_model_len: 32768
     distributed_executor_backend: "mp"
     hf_config_name: talker_config
   engine_input_source: [0]
diff --git a/tests/utils.py b/tests/utils.py
index ead286760..f5c513a4d 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -77,11 +77,36 @@ def wait_for_gpu_memory_to_clear(
     threshold_ratio: float | None = None,
     timeout_s: float = 120,
 ) -> None:
+    import gc
+
     assert threshold_bytes is not None or threshold_ratio is not None
     # Use nvml instead of pytorch to reduce measurement error from torch cuda
     # context.
     devices = get_physical_device_indices(devices)
     start_time = time.time()
+
+    # Print waiting start information
+    device_list = ", ".join(str(d) for d in devices)
+    if threshold_bytes is not None:
+        threshold_str = f"{threshold_bytes / 2**30:.2f} GiB"
+        condition_str = f"Memory usage ≤ {threshold_str}"
+    else:
+        threshold_percent = threshold_ratio * 100
+        threshold_str = f"{threshold_percent:.1f}%"
+        condition_str = f"Memory usage ratio ≤ {threshold_str}"
+
+    print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}")
+
+    # Define the is_free function based on threshold type
+    if threshold_bytes is not None:
+
+        def is_free(used, total):
+            return used <= threshold_bytes / 2**30
+    else:
+
+        def is_free(used, total):
+            return used / total <= threshold_ratio
+
     while True:
         output: dict[int, str] = {}
         output_raw: dict[int, tuple[float, float]] = {}
@@ -97,33 +122,44 @@ def wait_for_gpu_memory_to_clear(
                 gb_used = mem_info.used / 2**30
                 gb_total = mem_info.total / 2**30
             output_raw[device] = (gb_used, gb_total)
-            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
+            # Format to more readable form
+            usage_percent = (gb_used / gb_total) * 100 if gb_total > 0 else 0
+            output[device] = f"{gb_used:.1f}GiB/{gb_total:.1f}GiB ({usage_percent:.1f}%)"
 
-        print("gpu memory used/total (GiB): ", end="")
-        for k, v in output.items():
-            print(f"{k}={v}; ", end="")
-        print("")
-
-        if threshold_bytes is not None:
-
-            def is_free(used, total):
-                return used <= threshold_bytes / 2**30  # noqa E731
-
-            threshold = f"{threshold_bytes / 2**30} GiB"
-        else:
-
-            def is_free(used, total):
-                return used / total <= threshold_ratio  # noqa E731
-
-            threshold = f"{threshold_ratio:.2f}"
+        # Optimized GPU memory status print
+        print("[GPU Memory Status] Current usage:")
+        for device_id, mem_info in output.items():
+            print(f"  GPU {device_id}: {mem_info}")
 
+        # Calculate waiting duration
         dur_s = time.time() - start_time
+        elapsed_minutes = dur_s / 60
+
+        # Check if all devices meet the condition
         if all(is_free(used, total) for used, total in output_raw.values()):
-            print(f"Done waiting for free GPU memory on devices {devices=} ({threshold=}) {dur_s=:.02f}")
+            # Optimized completion message
+            print(f"[GPU Memory Freed] Devices {device_list} meet memory condition")
+            print(f"   Condition: {condition_str}")
+            print(f"   Wait time: {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)")
+            print("   Final status:")
+            for device_id, mem_info in output.items():
+                print(f"     GPU {device_id}: {mem_info}")
             break
 
+        # Check timeout
         if dur_s >= timeout_s:
-            raise ValueError(f"Memory of devices {devices=} not free after {dur_s=:.02f} ({threshold=})")
+            raise ValueError(
+                f"[GPU Memory Timeout] Devices {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n"
+                f"Condition: {condition_str}\n"
+                f"Current status:\n" + "\n".join(f"  GPU {device}: {output[device]}" for device in devices)
+            )
+
+        # Add waiting hint (optional)
+        if dur_s > 10 and int(dur_s) % 10 == 0:  # Show hint every 10 seconds
+            print(f"Waiting... Already waited {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)")
+
+        gc.collect()
+        torch.cuda.empty_cache()
 
         time.sleep(5)
 

From 144fd6e1e0efd4376fbb46d9b6e26e681001c14f Mon Sep 17 00:00:00 2001
From: kYLe <yellowsea@gmail.com>
Date: Wed, 28 Jan 2026 21:46:48 -0600
Subject: [PATCH 13/32] [Bugfix][Doc]Specify Qwen3-TTS model name for each task
 type (#1036)

Signed-off-by: Kyle Huang <yellowsea@gmail.com>
---
 examples/online_serving/qwen3_tts/README.md               | 2 ++
 examples/online_serving/qwen3_tts/openai_speech_client.py | 1 +
 2 files changed, 3 insertions(+)

diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md
index 2cddd0fbf..d4626e7c0 100644
--- a/examples/online_serving/qwen3_tts/README.md
+++ b/examples/online_serving/qwen3_tts/README.md
@@ -41,12 +41,14 @@ python openai_speech_client.py \
 
 # VoiceDesign: Describe the voice style
 python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
     --task-type VoiceDesign \
     --text "哥哥，你回来啦" \
     --instructions "体现撒娇稚嫩的萝莉女声，音调偏高"
 
 # Base: Voice cloning
 python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
     --task-type Base \
     --text "Hello, this is a cloned voice" \
     --ref-audio /path/to/reference.wav \
diff --git a/examples/online_serving/qwen3_tts/openai_speech_client.py b/examples/online_serving/qwen3_tts/openai_speech_client.py
index 272e84265..71ba5b70a 100644
--- a/examples/online_serving/qwen3_tts/openai_speech_client.py
+++ b/examples/online_serving/qwen3_tts/openai_speech_client.py
@@ -90,6 +90,7 @@ def run_tts_generation(args) -> None:
     if args.x_vector_only:
         payload["x_vector_only_mode"] = True
 
+    print(f"Model: {args.model}")
     print(f"Task type: {args.task_type or 'CustomVoice'}")
     print(f"Text: {args.text}")
     print(f"Voice: {args.voice}")

From d4190115f7ddfe890ec0ead52e8936d1c2fa34e6 Mon Sep 17 00:00:00 2001
From: Jiangyun Zhu <riverclouds.zhu@qq.com>
Date: Thu, 29 Jan 2026 13:11:33 +0800
Subject: [PATCH 14/32] [Misc] pin version of fa3-fwd (#1051)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 809d1f53c..3e1149099 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ dependencies = [
     "cache-dit==1.2.0",
     "tqdm>=4.66.0",
     "torchsde>=0.2.6",  # Required for Stable Audio scheduler
-    "fa3-fwd", # flash attention 3, maintained by @ZJY0516
+    "fa3-fwd==0.0.1", # flash attention 3, maintained by @ZJY0516
     "openai-whisper>=20250625",
     "imageio[ffmpeg]>=2.37.2",
     "onnxruntime>=1.19.0",

From 2ac141026b5237b5b1c70c53867872c919cd0536 Mon Sep 17 00:00:00 2001
From: TJian <tunjian.tan@embeddedllm.com>
Date: Thu, 29 Jan 2026 14:32:25 +0800
Subject: [PATCH 15/32] [CI] [ROCm] Add more AMD CI tests (#1039)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
---
 .../scripts/hardware_ci/run-amd-test.sh       |  6 ++
 .buildkite/test-amd.yaml                      | 76 +++++++++++--------
 2 files changed, 49 insertions(+), 33 deletions(-)

diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh
index 5f7440a80..f86b4b5d9 100755
--- a/.buildkite/scripts/hardware_ci/run-amd-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh
@@ -116,6 +116,9 @@ if [[ $commands == *"--shard-id="* ]]; then
         --shm-size=16gb \
         --group-add "$render_gid" \
         --rm \
+        -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+        -e MIOPEN_DEBUG_CONV_GEMM=0 \
+        -e VLLM_ROCM_USE_AITER=1 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
         -e AWS_ACCESS_KEY_ID \
@@ -148,6 +151,9 @@ else
           --shm-size=16gb \
           --group-add "$render_gid" \
           --rm \
+          -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+          -e MIOPEN_DEBUG_CONV_GEMM=0 \
+          -e VLLM_ROCM_USE_AITER=1 \
           -e HF_TOKEN \
           -e AWS_ACCESS_KEY_ID \
           -e AWS_SECRET_ACCESS_KEY \
diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index 08a3f7a1c..930ef7a40 100644
--- a/.buildkite/test-amd.yaml
+++ b/.buildkite/test-amd.yaml
@@ -8,14 +8,32 @@ steps:
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
+- label: "Diffusion Images API LoRA E2E"
+  timeout_in_minutes: 20
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+
+- label: "Diffusion Model CPU offloading Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+
 - label: "Diffusion Cache Backend Test"
   timeout_in_minutes: 15
   agent_pool: mi325_1
@@ -26,25 +44,30 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
 
-- label: "Diffusion Parallelism Test"
-  timeout_in_minutes: 15
+- label: "Diffusion Sequence Parallelism Test"
+  timeout_in_minutes: 20
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
 
+- label: "Diffusion Tensor Parallelism Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
+
 - label: "Diffusion GPU Worker Test"
   timeout_in_minutes: 20
   agent_pool: mi325_2
@@ -52,8 +75,6 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
     - pytest -s -v tests/diffusion/test_diffusion_worker.py
 
 - label: "Omni Model Test Qwen2-5-Omni"
@@ -66,12 +87,6 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
 
 - label: "Omni Model Test Qwen3-Omni"
@@ -83,9 +98,10 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
+    - 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
 
 - label: "Diffusion Image Edit Test"
   timeout_in_minutes: 15
@@ -97,10 +113,4 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py

From 43c6f520ca24cab5e33a2bcb795f0cacc3eca472 Mon Sep 17 00:00:00 2001
From: Jiangyun Zhu <riverclouds.zhu@qq.com>
Date: Thu, 29 Jan 2026 15:18:19 +0800
Subject: [PATCH 16/32] [Bugfix] fix qwen image layerd in dummy run (#1027)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
---
 examples/offline_inference/image_to_image/image_edit.py  | 3 +--
 vllm_omni/diffusion/diffusion_engine.py                  | 9 +++++++--
 vllm_omni/diffusion/models/interface.py                  | 1 +
 .../models/qwen_image/pipeline_qwen_image_layered.py     | 2 ++
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py
index 896c611b4..5ffab87f0 100644
--- a/examples/offline_inference/image_to_image/image_edit.py
+++ b/examples/offline_inference/image_to_image/image_edit.py
@@ -77,7 +77,7 @@ from pathlib import Path
 import torch
 from PIL import Image
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -428,7 +428,6 @@ def main():
 
     if not outputs:
         raise ValueError("No output generated from omni.generate()")
-    logger.info("Outputs: %s", outputs)
 
     # Extract images from OmniRequestOutput
     # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index b0b9b89db..a44596983 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -30,6 +30,11 @@ def supports_image_input(model_class_name: str) -> bool:
     return bool(getattr(model_cls, "support_image_input", False))
 
 
+def image_color_format(model_class_name: str) -> str:
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    return getattr(model_cls, "color_format", "RGB")
+
+
 def supports_audio_output(model_class_name: str) -> bool:
     model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
     if model_cls is None:
@@ -315,8 +320,8 @@ class DiffusionEngine:
         width = 1024
         if supports_image_input(self.od_config.model_class_name):
             # Provide a dummy image input if the model supports it
-
-            dummy_image = PIL.Image.new("RGB", (width, height), color=(0, 0, 0))
+            color_format = image_color_format(self.od_config.model_class_name)
+            dummy_image = PIL.Image.new(color_format, (width, height))
         else:
             dummy_image = None
         prompt: OmniTextPrompt = {"prompt": "dummy run", "multi_modal_data": {"image": dummy_image}}
diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py
index 8e8ecbb77..abdfe1e50 100644
--- a/vllm_omni/diffusion/models/interface.py
+++ b/vllm_omni/diffusion/models/interface.py
@@ -10,6 +10,7 @@ from typing import (
 @runtime_checkable
 class SupportImageInput(Protocol):
     support_image_input: ClassVar[bool] = True
+    color_format: ClassVar[str] = "RGB"  # Default color format
 
 
 @runtime_checkable
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
index ae27ba708..1001e8a14 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
@@ -192,6 +192,8 @@ def retrieve_latents(
 
 
 class QwenImageLayeredPipeline(nn.Module, SupportImageInput):
+    color_format = "RGBA"
+
     def __init__(
         self,
         *,

From 8306c91ea9c8da550925133d016fd1b797ec2f2a Mon Sep 17 00:00:00 2001
From: NATURE <wzliu@connect.hku.hk>
Date: Thu, 29 Jan 2026 15:39:26 +0800
Subject: [PATCH 17/32] [BugFix] Fix noisy output without setting a seed in
 Qwen Image (#1043)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
---
 vllm_omni/entrypoints/openai/api_server.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index f80fe95fa..6c987f4d2 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -5,6 +5,7 @@ import multiprocessing.forkserver as forkserver
 import os
 
 # Image generation API imports
+import random
 import time
 from argparse import Namespace
 from collections.abc import AsyncIterator
@@ -984,6 +985,13 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             gen_params.true_cfg_scale = request.true_cfg_scale
         if request.seed is not None:
             gen_params.seed = request.seed
+        else:
+            # If seed is not provided, generate a random one to ensure
+            # a proper generator is initialized in the backend.
+            # This fixes issues where using the default global generator
+            # might produce blurry images in some environments.
+            gen_params.seed = random.randint(0, 2**32 - 1)
+
         request_id = f"img_gen_{int(time.time())}"
 
         logger.info(f"Generating {request.n} image(s) {size_str}")

From cac6504adf05902e9377a78aa38c1ce7c72b5da6 Mon Sep 17 00:00:00 2001
From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Date: Thu, 29 Jan 2026 03:02:59 -0500
Subject: [PATCH 18/32] [bugfix] remove vllm speech route (#1060)

Signed-off-by: linyueqian <linyueqian@outlook.com>
---
 vllm_omni/entrypoints/openai/api_server.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 6c987f4d2..aedaceee3 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -736,6 +736,9 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
     return StreamingResponse(content=generator, media_type="text/event-stream")
 
 
+_remove_route_from_router(router, "/v1/audio/speech", {"POST"})
+
+
 @router.post(
     "/v1/audio/speech",
     dependencies=[Depends(validate_json_request)],

From def39561aee26284aca299b6be38bcbc2e3280cc Mon Sep 17 00:00:00 2001
From: Zhou Taichang <tzhouam@connect.ust.hk>
Date: Thu, 29 Jan 2026 16:03:27 +0800
Subject: [PATCH 19/32] [Debug] Update GLM-Image Pipeline (#1049)

Co-authored-by: root <root@hk01dgx028.cm.cluster>
---
 vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
index d222342b5..f5c167c32 100644
--- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
+++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py
@@ -491,7 +491,7 @@ class GlmImagePipeline(nn.Module):
             condition_grid = image_grid_thw[:-1]
             prior_token_image_embed = self.vision_language_encoder.get_image_features(
                 inputs["pixel_values"], condition_grid
-            )
+            ).pooler_output
             prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0)
             flat_prior_token_image_ids = self.vision_language_encoder.get_image_tokens(
                 prior_token_image_embed, condition_grid
@@ -859,7 +859,7 @@ class GlmImagePipeline(nn.Module):
         preprocessed_images = (
             None
             if isinstance(first_prompt, str)
-            else first_prompt.get("additional_information", {}).get("preprocessed_image")
+            else [first_prompt.get("additional_information", {}).get("preprocessed_image")]
         )
         condition_images = (
             None

From 5b8ade8446cec56ac6919aa90b8c91ffe030b7f2 Mon Sep 17 00:00:00 2001
From: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Date: Thu, 29 Jan 2026 16:44:58 +0800
Subject: [PATCH 20/32] [Diffusion][Bugfix] Fix the flash_attn backends
 selection logic (#983)

Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
---
 .../attention/backends/flash_attn.py          | 31 ++++--
 .../attention/backends/ring/ring_globals.py   | 55 ++++++++---
 .../attention/backends/ring/ring_kernels.py   | 98 ++++++-------------
 .../attention/backends/ring/ring_selector.py  | 10 --
 .../diffusion/attention/backends/utils/fa.py  | 54 ++++++++--
 .../diffusion/attention/parallel/ring.py      | 19 ++--
 vllm_omni/diffusion/envs.py                   | 22 ++++-
 vllm_omni/platforms/cuda/platform.py          | 24 +++--
 8 files changed, 185 insertions(+), 128 deletions(-)

diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py
index a9f2d925c..dc54acf04 100644
--- a/vllm_omni/diffusion/attention/backends/flash_attn.py
+++ b/vllm_omni/diffusion/attention/backends/flash_attn.py
@@ -10,8 +10,26 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
 )
 
+# Import flash attention functions with fallback chain from utils/fa.py
+# FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
+from vllm_omni.diffusion.attention.backends.utils.fa import (
+    HAS_FLASH_ATTN,
+    _pad_input,
+    _unpad_input,
+    _upad_input,
+    flash_attn_func,
+    flash_attn_varlen_func,
+)
+
 logger = init_logger(__name__)
 
+if not HAS_FLASH_ATTN:
+    raise ImportError(
+        "FlashAttentionBackend requires Flash Attention. "
+        "Please install one of: fa3-fwd, flash-attention, or flash-attn. "
+        "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+    )
+
 
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
@@ -56,14 +74,6 @@ class FlashAttentionImpl(AttentionImpl):
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
         """CUDA/ROCm flash attention implementation."""
-        from vllm_omni.diffusion.attention.backends.utils.fa import (
-            _pad_input,
-            _unpad_input,
-            _upad_input,
-            flash_attn_func,
-            flash_attn_varlen_func,
-        )
-
         query_length = query.size(1)
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         #  Contains at least one padding token in the sequence
@@ -92,13 +102,16 @@ class FlashAttentionImpl(AttentionImpl):
             out = _pad_input(out_unpad, indices_q, query.size(0), query_length)
 
         else:
-            out: torch.Tensor = flash_attn_func(
+            out = flash_attn_func(
                 query,
                 key,
                 value,
                 causal=self.causal,
                 softmax_scale=self.softmax_scale,
             )
+            # FA3 may return (out, lse) tuple, FA2 returns just out
+            if isinstance(out, tuple):
+                out = out[0]
         return out
 
     def forward_npu(
diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py
index 7ec092785..d80d81fb4 100644
--- a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py
+++ b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py
@@ -3,30 +3,53 @@
 # Copyright (c) 2024, Jiarui Fang.
 # Adapted from https://github.com/feifeibear/long-context-attention
 
-
-# test if flash_attn is available
+# test if flash_attn (FA2) is available
 try:
     import flash_attn  # noqa: F401
-    from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward  # noqa: F401
+    from flash_attn.flash_attn_interface import _flash_attn_forward  # noqa: F401
 
     HAS_FLASH_ATTN = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_FLASH_ATTN = False
 
-try:
-    from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward  # noqa: F401
-    from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper  # noqa: F401
-    from flash_attn_interface import flash_attn_func as flash3_attn_func  # noqa: F401
+# FA3 detection: try multiple sources (forward only, no backward needed for inference)
+# Source 1: flash_attn_interface (from flash-attention source build)
+# Source 2: fa3_fwd_interface (from fa3-fwd PyPI package, supports Ampere/Ada/Hopper)
+# Note: FA3 high-level API may or may not return softmax_lse depending on version.
+#       For Ring Attention which requires LSE, we fall back to low-level API if needed.
+HAS_FA3 = False
+fa3_fwd_func = None  # Low-level forward function (_flash_attn_forward)
+fa3_attn_func = None  # High-level attention function (flash_attn_func)
 
-    HAS_FLASH_ATTN_HOPPER = True
-except ImportError:
-    HAS_FLASH_ATTN_HOPPER = False
+# Try flash_attn_interface first (from flash-attention source build)
+try:
+    from flash_attn_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401
+    from flash_attn_interface import flash_attn_func as fa3_attn_func  # noqa: F401
+
+    HAS_FA3 = True
+except (ImportError, ModuleNotFoundError):
+    pass
+
+# Fallback: try fa3_fwd_interface (PyPI package, supports Ampere/Ada/Hopper)
+if not HAS_FA3:
+    try:
+        from fa3_fwd_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401
+        from fa3_fwd_interface import flash_attn_func as fa3_attn_func  # noqa: F401
+
+        HAS_FA3 = True
+    except (ImportError, ModuleNotFoundError):
+        pass
+
+# Legacy aliases for backward compatibility
+HAS_FLASH_ATTN_HOPPER = HAS_FA3
+flash_attn_forward_hopper = fa3_fwd_func
+flash3_attn_func = fa3_attn_func
 
 try:
     from flashinfer.prefill import single_prefill_with_kv_cache  # noqa: F401
 
     HAS_FLASHINFER = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_FLASHINFER = False
 
 try:
@@ -34,26 +57,26 @@ try:
     from aiter import flash_attn_func as flash_attn_func_aiter  # noqa: F401
 
     HAS_AITER = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_AITER = False
 
 try:
     import sageattention  # noqa: F401
 
     HAS_SAGE_ATTENTION = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_SAGE_ATTENTION = False
 
 try:
     import spas_sage_attn  # noqa: F401
 
     HAS_SPARSE_SAGE_ATTENTION = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_SPARSE_SAGE_ATTENTION = False
 
 try:
     import torch_npu  # noqa: F401
 
     HAS_NPU = True
-except ImportError:
+except (ImportError, ModuleNotFoundError):
     HAS_NPU = False
diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py
index 2444dd894..ad61f9a20 100644
--- a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py
+++ b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py
@@ -6,7 +6,13 @@ import math
 
 import torch
 
-from .ring_globals import HAS_AITER, HAS_FLASH_ATTN, HAS_FLASH_ATTN_HOPPER, HAS_FLASHINFER, HAS_NPU
+from .ring_globals import (
+    HAS_AITER,
+    HAS_FA3,
+    HAS_FLASH_ATTN,
+    HAS_FLASHINFER,
+    fa3_fwd_func,
+)
 
 _scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention
 _scaled_dot_product_efficient_attention = torch.ops.aten._scaled_dot_product_efficient_attention
@@ -26,23 +32,11 @@ if HAS_FLASH_ATTN:
     import flash_attn
     from flash_attn.flash_attn_interface import _flash_attn_forward
 
-if HAS_FLASH_ATTN_HOPPER:
-    from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward
-    from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper
-    from flash_attn_interface import flash_attn_func as flash3_attn_func
-else:
-    flash_attn_forward_hopper = None
-    flash_attn_func_hopper_backward = None
-    flash3_attn_func = None
-
 if HAS_FLASHINFER:
     from flashinfer.prefill import single_prefill_with_kv_cache
 
     _LOG2_E = math.log2(math.e)
 
-if HAS_NPU:
-    import torch_npu
-
 
 def pytorch_attn_forward(
     q: torch.Tensor,
@@ -146,52 +140,35 @@ def flash_attn_forward(
     return block_out, block_lse
 
 
-def flash_attn3_func_forward(
-    q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax
-):
-    assert HAS_FLASH_ATTN_HOPPER
-    # current signature of flash_attn_forward_hopper:
-    # (q, k, v, softmax_scale, causal, window_size, descale_q=None, descale_k=None, descale_v=None, gqa_parallel=False)
+def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax):
+    """FA3 forward pass for inference.
 
-    out, softmax_lse, *unused = flash_attn_forward_hopper(
-        q=q,
-        k=k,
-        v=v,
-        k_new=None,
-        v_new=None,
-        qv=None,
-        out=None,
-        cu_seqlens_q=None,
-        cu_seqlens_k=None,
-        cu_seqlens_k_new=None,
-        seqused_q=None,
-        seqused_k=None,
-        max_seqlen_q=None,
-        max_seqlen_k=None,
-        page_table=None,
-        kv_batch_idx=None,
-        leftpad_k=None,
-        rotary_cos=None,
-        rotary_sin=None,
-        seqlens_rotary=None,
-        q_descale=None,
-        k_descale=None,
-        v_descale=None,
+    FA3 supports Ampere, Ada, and Hopper GPUs. Dropout is ignored since FA3 is inference-only.
+    Uses low-level API (_flash_attn_forward) which always returns softmax_lse,
+    required for Ring Attention's correct accumulation.
+    """
+    assert HAS_FA3, "FA3 is not available"
+    assert fa3_fwd_func is not None, "FA3 low-level API (fa3_fwd_func) not available"
+
+    # Low-level API always returns (out, softmax_lse, S_dmask, rng_state)
+    out, softmax_lse, *_ = fa3_fwd_func(
+        q,
+        k,
+        v,
         softmax_scale=softmax_scale,
-        causal=False,
-        window_size=(-1, -1),
-        attention_chunk=0,
-        softcap=0.0,
-        rotary_interleaved=True,
-        scheduler_metadata=None,
-        num_splits=0,
-        pack_gqa=None,
-        sm_margin=0,
+        causal=causal,
+        window_size_left=window_size[0] if window_size else -1,
+        window_size_right=window_size[1] if window_size else -1,
+        softcap=softcap if softcap else 0.0,
     )
 
     return out, softmax_lse
 
 
+# Legacy alias for backward compatibility
+flash_attn3_func_forward = fa3_forward
+
+
 def flash_attn_forward_aiter(
     q,
     k,
@@ -264,20 +241,3 @@ def flashinfer_attn_forward(
         raise ValueError(f"Invalid input shape: {q.shape}")
     lse = lse / _LOG2_E
     return out, lse
-
-
-def npu_attn_forward(q, k, v, softmax_scale=None, layout="BSND"):
-    assert HAS_NPU, "torch_npu is not available"
-    softmax_scale = q.shape[-1] ** (-0.5)
-    block_out, block_lse = torch_npu.npu_fused_infer_attention_score(
-        q,
-        k,
-        v,
-        num_heads=q.shape[-2],
-        input_layout=layout,
-        scale=softmax_scale,
-        softmax_lse_flag=True,
-        pre_tokens=65535,
-        next_tokens=65535,
-    )
-    return block_out, block_lse.squeeze(dim=-1)
diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_selector.py b/vllm_omni/diffusion/attention/backends/ring/ring_selector.py
index 56e4a2e6b..77189c61b 100644
--- a/vllm_omni/diffusion/attention/backends/ring/ring_selector.py
+++ b/vllm_omni/diffusion/attention/backends/ring/ring_selector.py
@@ -10,7 +10,6 @@ from functools import partial
 import torch
 
 from .ring_globals import (
-    HAS_NPU,
     HAS_SAGE_ATTENTION,
     HAS_SPARSE_SAGE_ATTENTION,
 )
@@ -28,9 +27,6 @@ if HAS_SAGE_ATTENTION:
 if HAS_SPARSE_SAGE_ATTENTION:
     from spas_sage_attn.autotune import SparseAttentionMeansim
 
-if HAS_NPU:
-    from torch_npu import npu_fused_infer_attention_score
-
 
 class AttnType(Enum):
     AITER = "aiter"
@@ -44,7 +40,6 @@ class AttnType(Enum):
     SAGE_FP8 = "sage_fp8"
     SAGE_FP8_SM90 = "sage_fp8_sm90"
     SPARSE_SAGE = "sparse_sage"
-    NPU = "npu"
 
     @classmethod
     def from_string(cls, s: str):
@@ -157,11 +152,6 @@ def select_flash_attn_impl(
 
         return fn
 
-    elif impl_type == AttnType.NPU:
-        if not HAS_NPU:
-            raise ImportError("torch_npu is not available!")
-        return npu_fused_infer_attention_score
-
     elif attn_processor is not None:
         return attn_processor
 
diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py
index ea12bc0f1..3344c9963 100644
--- a/vllm_omni/diffusion/attention/backends/utils/fa.py
+++ b/vllm_omni/diffusion/attention/backends/utils/fa.py
@@ -17,17 +17,53 @@ import torch.nn.functional as F
 
 from vllm_omni.platforms import current_omni_platform
 
-if current_omni_platform.is_rocm():
-    from vllm._aiter_ops import is_aiter_found_and_supported
+# Flash Attention function detection with fallback chain
+flash_attn_func = None
+flash_attn_varlen_func = None
 
-    # Choose to enable this by default on ROCm
-    # Whenever possible as it is the fastest backend
-    if is_aiter_found_and_supported():
-        from aiter import flash_attn_func, flash_attn_varlen_func  # noqa: F401
-    else:
-        raise ImportError("Aiter is not found and supported on currentROCm device.")
+if current_omni_platform.is_rocm():
+    # ROCm: try Aiter first
+    try:
+        from vllm._aiter_ops import is_aiter_found_and_supported
+
+        if is_aiter_found_and_supported():
+            from aiter import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+    except (ImportError, ModuleNotFoundError):
+        pass
 else:
-    from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+    # CUDA: try FA3 -> FA2 fallback chain
+    # Try FA3 from fa3-fwd PyPI package
+    try:
+        from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+    except (ImportError, ModuleNotFoundError):
+        pass
+
+    # Fallback: Try FA3 from flash-attention source build
+    if flash_attn_func is None:
+        try:
+            from flash_attn_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+    # Fallback: Try FA2 from flash-attn package (try multiple import paths)
+    if flash_attn_func is None:
+        try:
+            from flash_attn import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+    if flash_attn_func is None:
+        try:
+            from flash_attn.flash_attn_interface import (  # noqa: F401
+                flash_attn_func,
+                flash_attn_varlen_func,
+            )
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+# If no FA backend available, SDPA backend will be selected at the platform level
+# flash_attn_func and flash_attn_varlen_func will be None
+HAS_FLASH_ATTN = flash_attn_func is not None
 
 
 def _index_first_axis(tensor, indices):
diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py
index 2d63f40c0..090d2ef61 100644
--- a/vllm_omni/diffusion/attention/parallel/ring.py
+++ b/vllm_omni/diffusion/attention/parallel/ring.py
@@ -10,7 +10,7 @@ import torch
 from vllm.logger import init_logger
 
 # import torch.distributed as dist # Not used directly here, but good practice if needed
-from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FLASH_ATTN
+from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FA3, HAS_FLASH_ATTN
 from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
 from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
@@ -116,11 +116,14 @@ class RingParallelAttention:
             except Exception:
                 backend_pref = None
 
-        # Fallback for FP32 or if Flash Attention is not available
-        if query.dtype == torch.float32 or not HAS_FLASH_ATTN:
-            if not HAS_FLASH_ATTN and backend_pref != "sdpa":
+        # Determine attention type with fallback chain: FA3 -> FA2 -> SDPA
+        # FP32 is not supported by Flash Attention, force SDPA
+        if query.dtype == torch.float32:
+            backend_pref = "sdpa"
+        elif not HAS_FA3 and not HAS_FLASH_ATTN:
+            if backend_pref != "sdpa":
                 logger = init_logger(__name__)
-                logger.warning_once("Flash Attention is not available! Force enabling SDPA.")
+                logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.")
             backend_pref = "sdpa"
 
         # Extract joint tensors
@@ -150,6 +153,9 @@ class RingParallelAttention:
 
         from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
 
+        # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper)
+        attn_type = AttnType.FA3 if HAS_FA3 else AttnType.FA
+
         return ring_flash_attn_func(
             query,
             key,
@@ -161,9 +167,8 @@ class RingParallelAttention:
             softcap=0.0,
             alibi_slopes=None,
             deterministic=False,
-            # return_attn_probs=False, # Removed as it might not be supported in signature
             group=self._sp_group.ring_group,
-            attn_type=AttnType.FA,
+            attn_type=attn_type,
             joint_tensor_key=joint_key,
             joint_tensor_value=joint_value,
             joint_strategy=joint_strategy,
diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py
index be2b58ae8..dd566b059 100644
--- a/vllm_omni/diffusion/envs.py
+++ b/vllm_omni/diffusion/envs.py
@@ -66,14 +66,32 @@ class PackagesEnvChecker:
             if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
                 return False
 
+            # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn)
+            # Try FA3 from fa3-fwd PyPI package
+            try:
+                import fa3_fwd_interface  # noqa: F401
+
+                return True
+            except (ImportError, ModuleNotFoundError):
+                pass
+
+            # Try FA3 from flash-attention source build
+            try:
+                import flash_attn_interface  # noqa: F401
+
+                return True
+            except (ImportError, ModuleNotFoundError):
+                pass
+
+            # Try FA2 from flash-attn package
             from flash_attn import __version__
 
             if __version__ < "2.6.0":
                 raise ImportError("install flash_attn >= 2.6.0")
             return True
-        except ImportError:
+        except (ImportError, ModuleNotFoundError):
             if not packages_info.get("has_aiter", False):
-                logger.warning('Flash Attention library "flash_attn" not found, using pytorch attention implementation')
+                logger.warning("No Flash Attention backend found, using pytorch SDPA implementation")
             return False
 
     def get_packages_info(self) -> dict:
diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py
index 9c65877be..89e403f99 100644
--- a/vllm_omni/platforms/cuda/platform.py
+++ b/vllm_omni/platforms/cuda/platform.py
@@ -39,22 +39,34 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         selected_backend: str | None,
         head_size: int,
     ) -> str:
+        from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+
         # Check compute capability for Flash Attention support
         # Flash Attention requires compute capability >= 8.0 and < 10.0
         compute_capability = cls.get_device_capability()
-        flash_attn_supported = False
+        compute_supported = False
         if compute_capability is not None:
             major, minor = compute_capability
             capability = major * 10 + minor
-            flash_attn_supported = 80 <= capability < 100
+            compute_supported = 80 <= capability < 100
+
+        # Check if FA packages are available
+        packages_info = PACKAGES_CHECKER.get_packages_info()
+        packages_available = packages_info.get("has_flash_attn", False)
+
+        # Both compute capability and packages must be available for FA
+        flash_attn_supported = compute_supported and packages_available
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
-                logger.warning(
-                    "Flash Attention requires GPU with compute capability >= 8.0 "
-                    "and < 10.0. Falling back to TORCH_SDPA backend."
-                )
+                if not compute_supported:
+                    logger.warning(
+                        "Flash Attention requires GPU with compute capability >= 8.0 "
+                        "and < 10.0. Falling back to TORCH_SDPA backend."
+                    )
+                elif not packages_available:
+                    logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.")
                 logger.info("Defaulting to diffusion attention backend SDPA")
                 return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             backend = DiffusionAttentionBackendEnum[backend_upper]

From dd518973de11e610b427d50c0ebb851ac92a61f6 Mon Sep 17 00:00:00 2001
From: amy-why-3459 <wuhaiyan17@huawei.com>
Date: Thu, 29 Jan 2026 19:30:11 +0800
Subject: [PATCH 21/32] [BugFix] Fix the accuracy issue of multimodal input.
 (#1020)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Rein Yang <ruiruyang2@gmail.com>
---
 tests/conftest.py                             | 30 +++++++
 tests/e2e/online_serving/test_qwen3_omni.py   | 16 +---
 .../distributed/omni_connectors/adapter.py    |  9 +-
 vllm_omni/engine/output_processor.py          |  2 +-
 vllm_omni/entrypoints/async_omni.py           | 88 +++++++++----------
 vllm_omni/entrypoints/openai/serving_chat.py  |  8 +-
 .../models/qwen3_omni/qwen3_omni_code2wav.py  |  2 +-
 7 files changed, 87 insertions(+), 68 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index e9b15aeee..7737a7dc2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -604,6 +604,36 @@ def convert_audio_to_text(audio_data):
         return ""
 
 
+def merge_base64_and_convert_to_text(base64_list):
+    """
+    Merge a list of base64 encoded audio data and convert to text.
+    """
+    import whisper
+    from pydub import AudioSegment
+
+    merged_audio = None
+    for base64_data in base64_list:
+        audio_data = base64.b64decode(base64_data.split(",", 1)[-1])
+        seg = AudioSegment.from_file(io.BytesIO(audio_data))
+        if merged_audio is None:
+            merged_audio = seg
+        else:
+            merged_audio += seg
+    output_path = f"./test_{int(time.time())}"
+    merged_audio.export(output_path, format="wav")
+    model = whisper.load_model("base")
+    text = model.transcribe(
+        output_path,
+        temperature=0.0,
+        word_timestamps=True,
+        condition_on_previous_text=False,
+    )["text"]
+    if text:
+        return text
+    else:
+        return ""
+
+
 def modify_stage_config(
     yaml_path: str,
     updates: dict[str, Any],
diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py
index 54c36633e..073419fb8 100644
--- a/tests/e2e/online_serving/test_qwen3_omni.py
+++ b/tests/e2e/online_serving/test_qwen3_omni.py
@@ -25,6 +25,7 @@ from tests.conftest import (
     generate_synthetic_audio,
     generate_synthetic_image,
     generate_synthetic_video,
+    merge_base64_and_convert_to_text,
     modify_stage_config,
 )
 from vllm_omni.platforms import current_omni_platform
@@ -154,11 +155,6 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     Input Setting: stream=True
     Datasets: single request
     """
-    # TODO：This skip will be removed when the chunk scenario supports multimodal input.
-    param = request.node.callspec.params.get("omni_server")
-
-    if param[1] == CHUNK_CONFIG_PATH:
-        pytest.skip("The current chunk scenario does not support multimodal.")
 
     # Test single completion
     e2e_list = list()
@@ -178,7 +174,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True)
 
     text_content = ""
-    audio_data = None
+    audio_data = []
     for chunk in chat_completion:
         for choice in chunk.choices:
             if hasattr(choice, "delta"):
@@ -189,11 +185,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
             modality = getattr(chunk, "modality", None)
 
             if modality == "audio" and content:
-                # Audio chunk - content
-                if audio_data is None:
-                    audio_data = content
-                else:
-                    audio_data += content
+                audio_data.append(content)
             elif modality == "text" and content:
                 # Text chunk - accumulate text content
                 text_content += content if content else ""
@@ -215,7 +207,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     ), "The output does not contain any of the keywords."
 
     # Verify text output same as audio output
-    audio_content = convert_audio_to_text(audio_data)
+    audio_content = merge_base64_and_convert_to_text(audio_data)
     print(f"text content is: {text_content}")
     print(f"audio content is: {audio_content}")
     similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py
index a634595ac..b03202525 100644
--- a/vllm_omni/distributed/omni_connectors/adapter.py
+++ b/vllm_omni/distributed/omni_connectors/adapter.py
@@ -247,7 +247,7 @@ def get_through_connector(connector, target_stage_id, stage_id, req_id, connecto
                 connector.get_requests[req_id] += 1
                 logger.debug(f"[Stage-{stage_id}] Received one chunk for request {connector_get_key}")
                 break
-        time.sleep(0.1)
+        time.sleep(0.01)
     return payload_data
 
 
@@ -330,10 +330,11 @@ def put_chunk(
 
         if stage_id == 0 and chunk_id == 0:
             if connector.request_payload.get(request_id) is None:
-                connector.request_payload[request_id] = payload_data
-                return
+                if not payload_data.get("finished"):
+                    connector.request_payload[request_id] = payload_data
+                    return
             else:
-                save_payload = connector.request_payload.get(request_id)
+                save_payload = connector.request_payload.pop(request_id)
                 payload_data["thinker_embeddings"] = torch.cat(
                     (save_payload.get("thinker_embeddings"), payload_data.get("thinker_embeddings")), dim=0
                 )
diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py
index ab5cf878c..65a69e0d4 100644
--- a/vllm_omni/engine/output_processor.py
+++ b/vllm_omni/engine/output_processor.py
@@ -114,7 +114,7 @@ class OmniRequestState(RequestState):
                         if k == "audio":
                             # When the audio tensor shape is inconsistent, torch.cat will fail.
                             # We need to use torch.cat in -1 dimension.
-                            self.mm_accumulated[k] = torch.cat(v, dim=-1)
+                            continue
                         else:
                             self.mm_accumulated[k] = torch.cat(v, dim=0)
                     except Exception:
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 179f58ed2..ec2a18ac4 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -1,12 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
+import copy
 import time
 import weakref
 from collections.abc import AsyncGenerator, Iterable, Sequence
 from dataclasses import asdict
 from pprint import pformat
-from typing import Any, cast
+from typing import Any
 
 from vllm.config import VllmConfig
 from vllm.inputs.preprocess import InputPreprocessor
@@ -32,7 +33,7 @@ from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
 )
-from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams, OmniTokensPrompt
+from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams
 
 # Internal imports (our code)
 from vllm_omni.lora.request import LoRARequest
@@ -304,38 +305,25 @@ class AsyncOmni(OmniBase):
             req_state = ClientRequestState(request_id)
             req_state.metrics = metrics
             self.request_states[request_id] = req_state
-
+            sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
+            task = {
+                "request_id": request_id,
+                "engine_inputs": prompt,
+                "sampling_params": sp0,
+            }
+            self.stage_list[0].submit(task)
+            metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
+            _req_start_ts[request_id] = time.time()
+            logger.info(
+                f"[{self._name}] Entering scheduling loop: stages={num_stages}, final_stage={final_stage_id_for_e2e}"
+            )
             if self.async_chunk:
                 stage_queues = {stage_id: asyncio.Queue() for stage_id in range(num_stages)}
                 req_state.stage_queues = stage_queues
-                for i in range(num_stages):
-                    sp: SamplingParams = cast(SamplingParams, sampling_params_list[i])
-                    engine_inputs = cast(OmniTokensPrompt, prompt)
-                    if i != 0:
-                        prompt_token_ids = engine_inputs["prompt_token_ids"]
-                        prompt_1 = engine_inputs.copy()
-                        prompt_1["prompt_token_ids"] = [0] * compute_talker_prompt_ids_length(prompt_token_ids)
-                        prompt_1["multi_modal_data"] = prompt_1["mm_processor_kwargs"] = None
-                        engine_inputs = prompt_1
-
-                    task = {
-                        "request_id": request_id,
-                        "engine_inputs": engine_inputs,
-                        "sampling_params": sp,
-                    }
-                    self.stage_list[i].submit(task)
-                    metrics.stage_first_ts[i] = metrics.stage_first_ts[0] or time.time()
-
-                    logger.info(f"[{self._name}] Enqueued request {request_id} to stage-{str(i)}")
-
-                _req_start_ts[request_id] = time.time()
-
-                logger.info(
-                    f"[{self._name}] Entering scheduling loop: "
-                    f"stages={num_stages}, final_stage={final_stage_id_for_e2e}"
-                )
                 async for output in self._process_async_results(
                     request_id,
+                    prompt,
+                    sampling_params_list,
                     req_state,
                     metrics,
                     final_stage_id_for_e2e,
@@ -344,22 +332,6 @@ class AsyncOmni(OmniBase):
                 ):
                     yield output
             else:
-                sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
-                task = {
-                    "request_id": request_id,
-                    "engine_inputs": prompt,
-                    "sampling_params": sp0,
-                }
-                self.stage_list[0].submit(task)
-
-                _req_start_ts[request_id] = time.time()
-                # Mark first input time for stage-0
-                metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
-                logger.info(
-                    f"[{self._name}] Entering scheduling loop: "
-                    f"stages={num_stages}, final_stage={final_stage_id_for_e2e}"
-                )
-
                 async for output in self._process_sequential_results(
                     request_id,
                     req_state,
@@ -390,6 +362,8 @@ class AsyncOmni(OmniBase):
     async def _process_async_results(
         self,
         request_id: str,
+        prompt: Any,
+        sampling_params_list: list[SamplingParams],
         req_state: ClientRequestState,
         metrics: OrchestratorMetrics,
         final_stage_id_for_e2e: int,
@@ -397,16 +371,34 @@ class AsyncOmni(OmniBase):
         wall_start_ts: float,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)}
+        submit_flag = True
         while not all(all_stages_finished.values()):
             for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
                 if all_stages_finished[stage_id]:
                     continue
-                result = await req_state.stage_queues[stage_id].get()
-                logger.info(f"[{self._name}] Received result from stage-{stage_id}: {result}")
+                try:
+                    result = req_state.stage_queues[stage_id].get_nowait()
+                except asyncio.QueueEmpty:
+                    await asyncio.sleep(0.001)
+                    continue
+
                 engine_outputs, finished, output_to_yield = self._process_single_result(
                     result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
                 )
-
+                if submit_flag and stage_id == 0:
+                    submit_flag = False
+                    prompt_token_ids = engine_outputs.prompt_token_ids
+                    engine_input = copy.deepcopy(prompt)
+                    engine_input["prompt_token_ids"] = [0] * compute_talker_prompt_ids_length(prompt_token_ids)
+                    engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None
+                    for i in range(1, len(self.stage_list)):
+                        task = {
+                            "request_id": request_id,
+                            "engine_inputs": engine_input,
+                            "sampling_params": sampling_params_list[i],
+                        }
+                        self.stage_list[i].submit(task)
+                        metrics.stage_first_ts[i] = time.time()
                 all_stages_finished[stage_id] = finished
 
                 if output_to_yield:
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index db8c972a1..61b00d803 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -9,6 +9,7 @@ from io import BytesIO
 from typing import TYPE_CHECKING, Any, Final, Optional, cast
 
 import jinja2
+import torch
 from fastapi import Request
 from PIL import Image
 from pydantic import TypeAdapter
@@ -1656,10 +1657,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     ):
         choices: list[ChatCompletionResponseChoice] = []
         final_res = omni_outputs.request_output
+        audio_data = final_res.multimodal_output.get("audio")
         if stream:
-            audio_tensor = final_res.multimodal_output["audio"][-1].float().detach().cpu().numpy()
+            audio_tensor = audio_data[-1].float().detach().cpu().numpy()
         else:
-            audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy()
+            if isinstance(audio_data, list):
+                audio_data = torch.cat(audio_data, dim=-1)
+            audio_tensor = audio_data.float().detach().cpu().numpy()
 
         # Ensure audio is 1D (flatten if needed)
         if audio_tensor.ndim > 1:
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
index bc2242f70..7adb6f96f 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
@@ -219,7 +219,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
         end_index = codes.shape[-1]
         # TODO: need to optimize algorithms, current only support
         # chunk_size = left_context_size = 25
-        if end_index == chunk_size:
+        if end_index <= chunk_size:
             context_size = 0
         else:
             context_size = left_context_size

From 0c8b5ac2d587f97db5709346457d597702735fbc Mon Sep 17 00:00:00 2001
From: Canlin Guo <canlinguosdu@gmail.com>
Date: Thu, 29 Jan 2026 20:02:39 +0800
Subject: [PATCH 22/32] [Bugfix] Set VaeImageProcessor `do_convert_rgb` True
 (#1032)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
---
 .../diffusion/models/qwen_image/pipeline_qwen_image_edit.py | 6 +++---
 .../models/qwen_image/pipeline_qwen_image_edit_plus.py      | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
index ba65bf07e..4241370ba 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
@@ -61,7 +61,7 @@ def get_qwen_image_edit_pre_process_func(
         vae_config = json.load(f)
         vae_scale_factor = 2 ** len(vae_config["temporal_downsample"]) if "temporal_downsample" in vae_config else 8
 
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
     latent_channels = vae_config.get("z_dim", 16)
 
     def pre_process_func(
@@ -136,7 +136,7 @@ def get_qwen_image_edit_post_process_func(
         vae_config = json.load(f)
         vae_scale_factor = 2 ** len(vae_config["temporal_downsample"]) if "temporal_downsample" in vae_config else 8
 
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
 
     def post_process_func(
         images: torch.Tensor,
@@ -261,7 +261,7 @@ class QwenImageEditPipeline(
 
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
         self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_convert_rgb=True)
         self.tokenizer_max_length = 1024
         # Edit prompt template - different from generation template
         self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"  # noqa: E501
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
index 52f1d64ec..9402c1f7c 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
@@ -67,7 +67,7 @@ def get_qwen_image_edit_plus_pre_process_func(
         vae_config = json.load(f)
         vae_scale_factor = 2 ** len(vae_config["temporal_downsample"]) if "temporal_downsample" in vae_config else 8
 
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
     latent_channels = vae_config.get("z_dim", 16)
 
     def pre_process_func(
@@ -158,7 +158,7 @@ def get_qwen_image_edit_plus_post_process_func(
         vae_config = json.load(f)
         vae_scale_factor = 2 ** len(vae_config["temporal_downsample"]) if "temporal_downsample" in vae_config else 8
 
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
 
     def post_process_func(
         images: torch.Tensor,
@@ -214,7 +214,7 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput):
 
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
         self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_convert_rgb=True)
         self.tokenizer_max_length = 1024
         # Edit prompt template - different from generation template, supports multiple images
         self.prompt_template_encode = (

From ee0f7a7797bb0dd9ac0befb9594eb18990e38fc1 Mon Sep 17 00:00:00 2001
From: JohnJan <wuzhongjian_yewu@cmss.chinamobile.com>
Date: Thu, 29 Jan 2026 20:19:31 +0800
Subject: [PATCH 23/32] [feat]: adapt batch request for flux (#1028)

Signed-off-by: wuzhongjian wuzhongjian_yewu@cmss.chinamobile.com
---
 .../diffusion/models/flux/pipeline_flux.py    | 39 ++++++++++++-------
 1 file changed, 25 insertions(+), 14 deletions(-)

diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py
index 01e0bfba9..a91601b7e 100644
--- a/vllm_omni/diffusion/models/flux/pipeline_flux.py
+++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py
@@ -557,10 +557,10 @@ class FluxPipeline(
     def forward(
         self,
         req: OmniDiffusionRequest,
-        prompt: str | list[str] = None,
-        prompt_2: str | list[str] = None,
-        negative_prompt: str | list[str] = None,
-        negative_prompt_2: str | list[str] = None,
+        prompt: str | list[str] | None = None,
+        prompt_2: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        negative_prompt_2: str | list[str] | None = None,
         true_cfg_scale: float = 1.0,
         height: int | None = None,
         width: int | None = None,
@@ -581,16 +581,27 @@ class FluxPipeline(
         max_sequence_length: int = 512,
     ):
         """Forward pass for flux."""
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
-        height = req.height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or self.default_sample_size * self.vae_scale_factor
-        sigmas = req.sigmas or sigmas
-        num_inference_steps = req.num_inference_steps or num_inference_steps
-        generator = req.generator or generator
-        req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
-        if req_num_outputs and req_num_outputs > 0:
-            num_images_per_prompt = req_num_outputs
+        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
+        # TODO: May be some data formatting operations on the API side. Hack for now.
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
+            negative_prompt = None
+        elif req.prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
+        )
+        generator = req.sampling_params.generator or generator
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(

From c48a9f33adc950aa0af7d2c71ce83aa31f0eaf37 Mon Sep 17 00:00:00 2001
From: Ziming Huang <hzm414167@alibaba-inc.com>
Date: Thu, 29 Jan 2026 21:10:09 +0800
Subject: [PATCH 24/32] [CI] Change Qwen3 Omni stage placement strategy 
 (#1072)

Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
---
 .../offline_inference/stage_configs/qwen3_omni_ci.yaml | 10 +++++-----
 .../stage_configs/rocm/qwen3_omni_ci.yaml              | 10 +++++-----
 .../online_serving/stage_configs/qwen3_omni_ci.yaml    | 10 +++++-----
 .../stage_configs/rocm/qwen3_omni_ci.yaml              | 10 +++++-----
 tests/e2e/stage_configs/qwen3_omni_ci.yaml             | 10 +++++-----
 .../model_executor/stage_configs/qwen3_omni_moe.yaml   | 10 +++++-----
 .../stage_configs/qwen3_omni_moe_async_chunk.yaml      | 10 +++++-----
 .../stage_configs/qwen3_omni_moe_multiconnector.yaml   | 10 +++++-----
 .../platforms/rocm/stage_configs/qwen3_omni_moe.yaml   | 10 +++++-----
 9 files changed, 45 insertions(+), 45 deletions(-)

diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
index 7ac435ea5..477e6e59f 100644
--- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
@@ -7,21 +7,21 @@
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
       load_format: dummy
     final_output: true
     final_output_type: text
@@ -44,7 +44,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.6
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -68,7 +68,7 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
diff --git a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
index 7ac435ea5..477e6e59f 100644
--- a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
@@ -7,21 +7,21 @@
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
       load_format: dummy
     final_output: true
     final_output_type: text
@@ -44,7 +44,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.6
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -68,7 +68,7 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
index 4db11e11f..8f0161edd 100644
--- a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
@@ -8,14 +8,14 @@ stage_args:
   - stage_id: 0
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 5
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
@@ -23,7 +23,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
       load_format: dummy
     final_output: true
     final_output_type: text
@@ -47,7 +47,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.6
        enforce_eager: false
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -72,7 +72,7 @@ stage_args:
   - stage_id: 2
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
diff --git a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
index 3297347c4..59642a77b 100644
--- a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
+++ b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
@@ -6,21 +6,21 @@
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 5
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
     final_output: true
     final_output_type: text
     is_comprehension: true
@@ -42,7 +42,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.6
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -65,7 +65,7 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml
index 148a949d8..479e4d6e9 100644
--- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml
@@ -7,14 +7,14 @@
 stage_args:
 - stage_id: 0
   runtime:
-    devices: "0,1"
+    devices: "0"
     max_batch_size: 5
   engine_args:
     model_stage: thinker
     model_arch: Qwen3OmniMoeForConditionalGeneration
     worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.6
+    gpu_memory_utilization: 0.9
     enforce_eager: false
     trust_remote_code: true
     engine_output_type: latent  # Output hidden states for talker
@@ -23,7 +23,7 @@ stage_args:
     max_model_len: 32768
     enable_prefix_caching: false
     hf_config_name: thinker_config
-    tensor_parallel_size: 2
+    tensor_parallel_size: 1
   final_output: true
   final_output_type: text
   is_comprehension: true
@@ -46,7 +46,7 @@ stage_args:
     model_arch: Qwen3OmniMoeForConditionalGeneration
     worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.3
+    gpu_memory_utilization: 0.6
     enforce_eager: false
     trust_remote_code: true
     engine_output_type: latent  # Output codec codes for code2wav
@@ -68,7 +68,7 @@ stage_args:
 
 - stage_id: 2
   runtime:
-    devices: "0"
+    devices: "1"
     max_batch_size: 1
   engine_args:
     model_stage: code2wav
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
index d0380a39d..e3dcf940f 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
@@ -9,14 +9,14 @@ stage_args:
   - stage_id: 0
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 64
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
@@ -24,7 +24,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
     final_output: true
     final_output_type: text
     is_comprehension: true
@@ -47,7 +47,7 @@ stage_args:
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.3
+      gpu_memory_utilization: 0.6
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output codec codes for code2wav
@@ -71,7 +71,7 @@ stage_args:
   - stage_id: 2
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
index 42748960e..1a8b60aa1 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
@@ -9,14 +9,14 @@ stage_args:
   - stage_id: 0
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 64
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
@@ -24,7 +24,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
       custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
     final_output: true
     final_output_type: text
@@ -41,14 +41,14 @@ stage_args:
   - stage_id: 1
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 64
     engine_args:
       model_stage: talker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.3
+      gpu_memory_utilization: 0.6
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output codec codes for code2wav
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml
index 69102df6f..0bc6e48e5 100644
--- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml
@@ -8,21 +8,21 @@ stage_args:
   - stage_id: 0
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
     final_output: true
     final_output_type: text
     is_comprehension: true
@@ -48,7 +48,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.6
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -77,7 +77,7 @@ stage_args:
   - stage_id: 2
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
diff --git a/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml
index 93e05ea9c..31312673a 100644
--- a/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml
+++ b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml
@@ -7,14 +7,14 @@
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0,1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
@@ -22,7 +22,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      tensor_parallel_size: 1
     final_output: true
     final_output_type: text
     is_comprehension: true
@@ -44,7 +44,7 @@ stage_args:
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.3
+      gpu_memory_utilization: 0.6
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output codec codes for code2wav
@@ -68,7 +68,7 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "0"
+      devices: "1"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav

From f8dd566abda9cfa482e0a024aecee896123fe36e Mon Sep 17 00:00:00 2001
From: dsinghvi <divyanshsinghvi@gmail.com>
Date: Thu, 29 Jan 2026 18:50:39 +0530
Subject: [PATCH 25/32] [BugFix] Fix to use correct attn backend (#1038)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
---
 vllm_omni/diffusion/attention/layer.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py
index 0d3bf8b7d..83cfe84d7 100644
--- a/vllm_omni/diffusion/attention/layer.py
+++ b/vllm_omni/diffusion/attention/layer.py
@@ -12,6 +12,7 @@ import torch.nn as nn
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend
@@ -45,6 +46,14 @@ class Attention(nn.Module):
             causal=causal,
             num_kv_heads=num_kv_heads,
         )
+        # Instantiate fallback backend for float32 support
+        self.sdpa_fallback = SDPABackend.get_impl_cls()(
+            num_heads=num_heads,
+            head_size=head_size,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            num_kv_heads=num_kv_heads,
+        )
         self.backend_pref = None
 
         self.softmax_scale = softmax_scale
@@ -104,12 +113,12 @@ class Attention(nn.Module):
         return out
 
     def _run_local_attention(self, query, key, value, attn_metadata):
-        if self.backend_pref == "flash_attn" and query.dtype == torch.float32:
-            logger.warning(
-                "Flash Attention does not support float32. Overriding user config "
+        if query.dtype == torch.float32:
+            logger.warning_once(
+                f"Only SDPA supports float32. Overriding user config {type(self.attention)} "
                 f"attention_backend='{self.backend_pref}' to 'sdpa' for dtype={query.dtype}."
             )
-            self.backend_pref = "sdpa"
+            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
 
         # Fallback to standard attention
         return self.attention.forward(query, key, value, attn_metadata)

From a909d364153ab98da3863b39f1201822469a4751 Mon Sep 17 00:00:00 2001
From: Ziming Huang <hzm414167@alibaba-inc.com>
Date: Thu, 29 Jan 2026 21:23:19 +0800
Subject: [PATCH 26/32] [Perf] Qwen3 Omni talker mtp optimization (#1005)

Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
---
 .../qwen3_omni_moe_code_predictor_mtp.py      | 323 ++++++++----------
 .../qwen3_omni/qwen3_omni_moe_talker.py       |  62 +---
 vllm_omni/worker/gpu_model_runner.py          |  12 +-
 3 files changed, 157 insertions(+), 240 deletions(-)

diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py
index bcefe13e8..6c135926d 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py
@@ -13,9 +13,15 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import Cache, PretrainedConfig
-from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeRotaryEmbedding
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 from vllm.compilation.decorators import support_torch_compile
-from vllm.config import CacheConfig, ModelConfig, VllmConfig
+from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
+from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -24,81 +30,12 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.utils.torch_utils import direct_register_custom_op
 
 logger = init_logger(__name__)
 
-
-# ============================================================================
-# Rotary Embeddings and Helper Functions
-# ============================================================================
-
-
-class Qwen3OmniCodePredictorRotaryEmbedding(nn.Module):
-    """Rotary positional embeddings for the code predictor."""
-
-    def __init__(self, dim, max_position_embeddings=2048, base=10000):
-        super().__init__()
-        self.dim = dim
-        self.max_position_embeddings = max_position_embeddings
-        self.base = base
-
-        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-
-    def forward(self, x: torch.Tensor, position_ids: torch.LongTensor) -> tuple[torch.Tensor, torch.Tensor]:
-        inv_freq = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
-        position = position_ids[:, None, :].float()
-        freqs = (inv_freq @ position).transpose(1, 2)
-        emb = torch.cat((freqs, freqs), dim=-1)
-        cos = emb.cos().to(x.dtype)
-        sin = emb.sin().to(x.dtype)
-        return cos, sin
-
-
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """Repeat key/value heads for grouped query attention."""
-    if n_rep == 1:
-        return hidden_states
-    bsz, num_kv_heads, seq_len, head_dim = hidden_states.shape
-    hidden_states = hidden_states[:, :, None, :, :].expand(bsz, num_kv_heads, n_rep, seq_len, head_dim)
-    return hidden_states.reshape(bsz, num_kv_heads * n_rep, seq_len, head_dim)
-
-
 # ============================================================================
 # Code Predictor Attention Layer
 # ============================================================================
@@ -149,6 +86,12 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
             prefix=f"{prefix}.o_proj",
             disable_tp=True,
         )
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=config.code_predictor_config.max_position_embeddings,
+            rope_parameters=None,
+            dual_chunk_attention_config=None,
+        )
 
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_key_value_heads * self.head_dim
@@ -159,12 +102,19 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
         self.is_causal = True
         self.config = config
 
+        self.attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
+        cudagraph_mode = get_current_vllm_config().compilation_config.cudagraph_mode
+        if "flash_attention_2" in ALL_ATTENTION_FUNCTIONS and cudagraph_mode.has_full_cudagraphs():
+            logger.warning(
+                f"CUDAGraphMode.{cudagraph_mode.name} is currently not supported "
+                f"with flash attention for Qwen3-Omni talker MTP."
+                f"removing flash attention from attention_backends"
+            )
+            self.attention_backends.remove("flash_attention_2")
+
     def forward(
         self,
         hidden_states: torch.Tensor,
-        causal_mask: torch.Tensor,
-        cos: torch.Tensor,
-        sin: torch.Tensor,
         past_key_values: Cache | None = None,
         cache_position: torch.LongTensor | None = None,
         use_cache: bool = False,
@@ -176,34 +126,39 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
 
         # Reshape for attention
-        q = q.view(bsz, seq_len, self.num_heads, self.head_dim)
-        k = k.view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
-        v = v.view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+        q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+        v = v.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
 
         # Apply normalization
-        q = self.q_norm(q)
-        k = self.k_norm(k)
+        q = self.q_norm(q).contiguous()
+        k = self.k_norm(k).contiguous()
+        q = q.reshape(-1, self.q_size)
+        k = k.reshape(-1, self.kv_size)
 
         # Apply RoPE
+        q, k = self.rotary_emb(position_ids, q, k)
+
+        # Reshape for attention
+        q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+
+        v_heads = v.transpose(1, 2).contiguous()
         q_heads = q.transpose(1, 2).contiguous()
         k_heads = k.transpose(1, 2).contiguous()
-        q_heads, k_heads = apply_rotary_pos_emb(q_heads, k_heads, cos, sin)
-        v_heads = v.transpose(1, 2).contiguous()
 
         if past_key_values is not None:
+            sin, cos = self.rotary_emb.get_cos_sin(seq_len)
             # sin and cos are specific to RoPE models; cache_position needed for the static cache
             cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
             k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs)
 
-        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
-
         # Try attention backends in order of preference, with runtime error handling
         # This handles cases where the backend is registered but not actually available
-        attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
         attn_output = None
         last_error = None
 
-        for backend_name in attention_backends:
+        for backend_name in self.attention_backends:
             if backend_name not in ALL_ATTENTION_FUNCTIONS:
                 continue
 
@@ -219,15 +174,10 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
                     scaling=self.head_dim**-0.5,
                     sliding_window=None,
                     use_cache=use_cache,
-                    position_ids=position_ids,
+                    position_ids=position_ids[:seq_len].unsqueeze(0),
                     output_hidden_states=True,
                     output_attentions=False,
                 )
-                # Success - log fallback if not using flash_attention_2
-                if backend_name != "flash_attention_2":
-                    logger.warning_once(
-                        f"Using {backend_name} attention backend (flash_attention_2 not available or failed)"
-                    )
                 break
             except (ValueError, ImportError, RuntimeError, AttributeError) as e:
                 # Store error and try next backend
@@ -335,9 +285,6 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module):
     def mtp_block(
         self,
         hidden_states: torch.Tensor,
-        causal_mask: torch.Tensor,
-        cos: torch.Tensor,
-        sin: torch.Tensor,
         past_key_values: Cache | None = None,
         cache_position: torch.LongTensor | None = None,
         use_cache: bool = False,
@@ -346,9 +293,7 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module):
         # Self-attention with residual
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(
-            hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
-        )
+        hidden_states = self.self_attn(hidden_states, past_key_values, cache_position, use_cache, position_ids)
         hidden_states = residual + hidden_states
 
         # MLP with residual
@@ -359,47 +304,6 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module):
 
         return hidden_states
 
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        previous_hidden_states: torch.Tensor,
-        inputs_embeds: torch.Tensor | None = None,
-        spec_step_index: int = 0,
-    ) -> torch.Tensor:
-        assert inputs_embeds is not None, "inputs_embeds required for MTP"
-
-        # Mask position 0 (not needed for MTP)
-        inputs_embeds[positions == 0] = 0
-
-        hidden_states = torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
-
-        # Get position info for RoPE
-        batch_size, seq_len, _ = hidden_states.shape
-        device = hidden_states.device
-        position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(batch_size, -1)
-
-        # Get RoPE embeddings
-        head_dim = self.self_attn.head_dim
-        rotary_emb = Qwen3OmniCodePredictorRotaryEmbedding(
-            head_dim, max_position_embeddings=self.config.code_predictor_config.max_position_embeddings
-        )
-        rotary_emb = Qwen3OmniMoeRotaryEmbedding(self.config)
-        cos, sin = rotary_emb(hidden_states, position_ids)
-
-        # Create causal mask
-        causal_mask = (
-            torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
-            .unsqueeze(0)
-            .unsqueeze(0)
-        )
-        causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
-
-        # Forward through MTP block
-        hidden_states = self.mtp_block(hidden_states, causal_mask, cos, sin)
-
-        return hidden_states
-
 
 class Qwen3OmniCodePredictorBaseModel(nn.Module):
     """
@@ -444,9 +348,6 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
 
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        # RoPE
-        self.rotary_emb = Qwen3OmniMoeRotaryEmbedding(config=config)
-
     def forward(
         self,
         inputs_embeds: torch.Tensor,
@@ -462,7 +363,6 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
 
         Args:
             inputs_embeds: [batch, seq_len, hidden_size]
-            attention_mask: Optional attention mask tensor
             position_ids: Optional position IDs tensor
             past_key_values: Optional cached key-value pairs
             use_cache: Whether to use cache
@@ -473,37 +373,11 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
             Named tuple with .last_hidden_state and .past_key_values attributes
         """
         batch_size, seq_len, _ = inputs_embeds.shape
-
-        # Create positions tensor if not provided
-        # positions must be [num_tokens] or [batch_size, seq_len]
-        if position_ids is None:
-            if cache_position is not None:
-                position_ids = cache_position  # [num_tokens]
-            else:
-                position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0)  # [1, seq_len]
-        else:
-            position_ids = position_ids.flatten()  # Ensure [num_tokens]
-
-        # Extract cos/sin from rotary_emb cache
-        # The cos_sin_cache is [max_pos, rotary_dim * 2]
-        cos, sin = self.rotary_emb(inputs_embeds, position_ids)
-
-        # Create causal mask
-        device = inputs_embeds.device
-        causal_mask = (
-            torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
-            .unsqueeze(0)
-            .unsqueeze(0)
-        )
-        causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
-
         # Forward through decoder layers
         hidden_states = inputs_embeds
 
         for layer in self.layers:
-            hidden_states = layer.mtp_block(
-                hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
-            )
+            hidden_states = layer.mtp_block(hidden_states, past_key_values, cache_position, use_cache, position_ids)
 
         # Final norm
         hidden_states = self.norm(hidden_states)
@@ -517,6 +391,32 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
         return self.codec_embedding
 
 
+def code_predictor_sample(
+    logits: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    forward_context = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    logits = self.logits_processors(None, logits[:, -1])
+    probs = F.softmax(logits, dim=-1)
+    code = torch.multinomial(probs.squeeze(1), num_samples=1)  # [batch, 1]
+    return code
+
+
+def code_predictor_sample_fake(
+    logits: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    return torch.empty((logits.shape[0], 1), dtype=torch.int64, device=logits.device)
+
+
+direct_register_custom_op(
+    op_name="qwen3_omni_code_predictor_sample",
+    op_func=code_predictor_sample,
+    fake_impl=code_predictor_sample_fake,
+)
+
+
 @support_torch_compile
 class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
     """
@@ -552,30 +452,89 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
                 for _ in range(self.num_code_groups - 1)
             ]
         )
+        self.logits_processors = LogitsProcessorList(
+            [
+                TopKLogitsWarper(top_k=50),
+                TopPLogitsWarper(top_p=0.8),
+            ]
+        )
+        max_batch_size = vllm_config.scheduler_config.max_num_seqs
+        self.position_ids_buffer = []
+        for layer_idx in range(self.num_code_groups - 1):
+            self.position_ids_buffer.append(torch.arange(layer_idx + 2, dtype=torch.int64).repeat(max_batch_size))
+
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+        self.layer_name = prefix
 
     def forward(
         self,
-        inputs_embeds: torch.Tensor,  # [batch, seq_len, hidden_size]
-        layer_idx: int,  # Which layer to predict (0-num_layers-2 for layers 1-num_layers-1)
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Forward pass for code predictor.
 
         Args:
-            inputs_embeds: Input embeddings [batch_size, seq_len, hidden_size]
-            layer_idx: Which residual layer to predict (0-num_layers-2 for layers 1-num_layers-1)
+            layer0_code:
+                Code index for code-group (layer) 0.
+                Shape: [batch_size, 1], dtype typically int64.
+
+            last_talker_hidden:
+
+                Shape: [batch_size, hidden_size].
 
         Returns:
-            logits: Predicted logits [batch_size, seq_len, vocab_size]
-            hidden_states: Output hidden states [batch_size, seq_len, hidden_size]
+            pos_all_layers:
+                Predicted codes for all code groups, including `layer0_code`.
+                Shape: [batch_size, num_code_groups, 1].
+
+            current_input:
+                The final input embedding sequence after appending embeddings of all
+                predicted codes (one token per predicted layer).
+                Shape: [batch_size, num_code_groups + 2, hidden_size].
         """
-        # Pass through base model
-        hidden_states = self.model(inputs_embeds)
+        pos_codes = [layer0_code]  # Start with layer 0: [batch, 1]
+        try:
+            current_input = torch.cat([last_talker_hidden, layer0_embed], dim=1)  # [batch, 2, hidden_size]
+        except Exception as e:
+            print(f"Error in current_input: {e}")
+            print(f"last_talker_hidden shape: {last_talker_hidden.shape}")
+            print(f"prev_embed shape: {layer0_embed.shape}")
+            raise e
+        batch_size = current_input.shape[0]
 
-        # Get logits from corresponding head
-        logits = self.lm_head[layer_idx](hidden_states)
+        # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
+        begin_pos = 0
+        for layer_idx in range(self.num_code_groups - 1):
+            # Input for this layer: [last_talker_hidden, prev_layer_embed]
 
-        return logits, hidden_states
+            # Forward through code_predictor model
+            position_ids = self.position_ids_buffer[layer_idx][: (layer_idx + 2) * batch_size]  # [1, seq_len]
+            begin_pos += layer_idx + 2
+            outputs = self.model(
+                inputs_embeds=current_input,
+                attention_mask=None,
+                position_ids=position_ids,
+                past_key_values=None,
+                use_cache=False,
+                cache_position=None,
+            )
+            hidden_state = outputs.last_hidden_state  # [batch, 2, hidden_size]
+
+            # Use the corresponding lm_head for this layer
+            logits = self.lm_head[layer_idx](hidden_state[:, -1:, :])
+            code = torch.ops.vllm.qwen3_omni_code_predictor_sample(logits, self.layer_name)
+            pos_codes.append(code)
+            # Update prev_embed for next layer (if not last layer)
+            # layer_idx=0 predicts layer 1, embed with codec_embedding[1]
+            new_embed = self.model.codec_embedding[layer_idx](code)  # [batch, 1, hidden_size]
+            current_input = torch.cat([current_input, new_embed], dim=1)  # [batch, 3~n, hidden_size]
+        pos_all_layers = torch.stack(pos_codes, dim=1)  # [batch, num_code_groups, 1]
+        return pos_all_layers, current_input
 
     def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights with mapping for fused QKV and gate_up projections.
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
index 5b032d86e..65cb43d0b 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
@@ -4,11 +4,6 @@ from typing import Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers.generation.logits_process import (
-    LogitsProcessorList,
-    TopKLogitsWarper,
-    TopPLogitsWarper,
-)
 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeTalkerConfig,
 )
@@ -140,7 +135,10 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         self.code_predictor = Qwen3OmniMoeTalkerCodePredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor")
         )
-        self.empty_code = torch.empty((1, 0), dtype=torch.long)
+        self.layer0_embed_buffer = torch.zeros(
+            (vllm_config.scheduler_config.max_num_seqs, 1, self.config.text_config.hidden_size),
+            dtype=vllm_config.model_config.dtype,
+        )
 
     def code_predictor_forward(
         self,
@@ -192,62 +190,18 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         all_codes_per_position = []
         middle_hidden_states = []  # Collect hidden states for each position
 
-        logits_processors = LogitsProcessorList(
-            [
-                TopKLogitsWarper(top_k=top_k),
-                TopPLogitsWarper(top_p=top_p),
-            ]
-        )
         # Generate residual layers for each position
         for pos in range(seq_len):
             layer0_code = input_ids[:, pos : pos + 1]  # [batch, 1]
 
-            # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
-            pos_codes = [layer0_code]  # Start with layer 0: [batch, 1]
-
             # Initial input: [last_talker_hidden, layer0_embed]
             layer0_embed = self.embed_input_ids(layer0_code)
-            prev_embed = layer0_embed  # Track previous layer embedding
-            try:
-                current_input = torch.cat([last_talker_hidden, prev_embed], dim=1)  # [batch, 2, hidden_size]
-            except Exception as e:
-                print(f"Error in current_input: {e}")
-                print(f"last_talker_hidden shape: {last_talker_hidden.shape}")
-                print(f"prev_embed shape: {prev_embed.shape}")
-                raise e
-
-            for layer_idx in range(self.num_code_groups - 1):
-                # Input for this layer: [last_talker_hidden, prev_layer_embed]
-
-                # Forward through code_predictor model
-                outputs = self.code_predictor.model(
-                    inputs_embeds=current_input,
-                    attention_mask=None,
-                    position_ids=None,
-                    past_key_values=None,
-                    use_cache=False,
-                    cache_position=None,
-                )
-
-                hidden_state = outputs.last_hidden_state  # [batch, 2, hidden_size]
-
-                # Use the corresponding lm_head for this layer
-                logits = self.code_predictor.lm_head[layer_idx](hidden_state[:, -1:, :])  # [batch, 1, vocab_size]
-
-                logits = logits_processors(None, logits[:, -1])
-
-                # Sample from the filtered distribution
-                probs = F.softmax(logits, dim=-1)
-                code = torch.multinomial(probs.squeeze(1), num_samples=1)  # [batch, 1]
-                pos_codes.append(code)
-
-                # Update prev_embed for next layer (if not last layer)
-                # layer_idx=0 predicts layer 1, embed with codec_embedding[1]
-                new_embed = self.code_predictor.model.codec_embedding[layer_idx](code)  # [batch, 1, hidden_size]
-                current_input = torch.cat([current_input, new_embed], dim=1)  # [batch, 3~n, hidden_size]
+            self.layer0_embed_buffer[:batch_size].copy_(layer0_embed)
+            pos_all_layers, current_input = self.code_predictor(
+                layer0_code, self.layer0_embed_buffer[:batch_size], last_talker_hidden
+            )
 
             # Stack all layers for this position: [batch, num_code_groups, 1]
-            pos_all_layers = torch.stack(pos_codes, dim=1)  # [batch, num_code_groups, 1]
             all_codes_per_position.append(pos_all_layers)
             middle_hidden_states.append(current_input[:, 2:-1, :])
 
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index a486212aa..1e2beb046 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -643,12 +643,16 @@ class OmniGPUModelRunner(GPUModelRunner):
                 ),
             ):
                 if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
+                    num_tokens_padded_talker_mtp = num_tokens_padded
+                    if num_tokens_padded_talker_mtp == self.max_num_tokens:
+                        num_tokens_padded_talker_mtp = min(num_tokens_padded, self.max_num_reqs)
                     outputs = self.talker_mtp(
-                        self.talker_mtp_input_ids.gpu[:num_tokens_padded],
-                        self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded],
-                        self.last_talker_hidden.gpu[:num_tokens_padded],
-                        self.text_step.gpu[:num_tokens_padded],
+                        self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
+                        self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
+                        self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp],
+                        self.text_step.gpu[:num_tokens_padded_talker_mtp],
                     )
+                    self.compilation_config.cache_dir = None
                 outputs = self.model(
                     input_ids=input_ids,
                     positions=positions,

From c7f89ef0852a44ce378cc5f3ddafe782f3ffda66 Mon Sep 17 00:00:00 2001
From: Fanli Lin <fanli.lin@intel.com>
Date: Thu, 29 Jan 2026 22:02:24 +0800
Subject: [PATCH 27/32] [Wan2.2] Optimize memory usage with conditional
 transformer loading (#980)

Signed-off-by: Lin, Fanli <fanli.lin@intel.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
---
 .../text_to_video/text_to_video.md            |  2 +-
 .../text_to_video/text_to_video.py            |  7 +-
 .../models/wan2_2/pipeline_wan2_2.py          | 96 ++++++++++++++-----
 3 files changed, 80 insertions(+), 25 deletions(-)

diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md
index cd5bd2be4..b5d0b2adc 100644
--- a/examples/offline_inference/text_to_video/text_to_video.md
+++ b/examples/offline_inference/text_to_video/text_to_video.md
@@ -26,7 +26,7 @@ Key arguments:
 - `--num_frames`: Number of frames (Wan default is 81).
 - `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
 - `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
-- `--boundary_ratio`: Boundary split ratio for low/high DiT.
+- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
 - `--vae_use_slicing`: enable VAE slicing for memory optimization.
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index 464ff9cca..b5c003dfb 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -32,7 +32,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--width", type=int, default=1280, help="Video width.")
     parser.add_argument("--num_frames", type=int, default=81, help="Number of frames (Wan default is 81).")
     parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
-    parser.add_argument("--boundary_ratio", type=float, default=0.875, help="Boundary split ratio for low/high DiT.")
+    parser.add_argument(
+        "--boundary_ratio",
+        type=float,
+        default=0.875,
+        help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
+    )
     parser.add_argument(
         "--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
index 712c05400..562d8eec5 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
@@ -227,17 +227,30 @@ class Wan22Pipeline(nn.Module):
             except Exception:
                 pass
 
+        self.boundary_ratio = od_config.boundary_ratio
+
+        # Determine which transformers to load based on boundary_ratio
+        # boundary_ratio=1.0: only load transformer_2 (low-noise stage only)
+        # boundary_ratio=0.0: only load transformer (high-noise stage only)
+        # otherwise: load both transformers
+        load_transformer = self.boundary_ratio != 1.0 if self.boundary_ratio is not None else True
+        load_transformer_2 = self.has_transformer_2 and (
+            self.boundary_ratio != 0.0 if self.boundary_ratio is not None else True
+        )
+
         # Set up weights sources for transformer(s)
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder="transformer",
-                revision=None,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            ),
-        ]
-        if self.has_transformer_2:
+        self.weights_sources = []
+        if load_transformer:
+            self.weights_sources.append(
+                DiffusersPipelineLoader.ComponentSource(
+                    model_or_path=od_config.model,
+                    subfolder="transformer",
+                    revision=None,
+                    prefix="transformer.",
+                    fall_back_to_pt=True,
+                )
+            )
+        if load_transformer_2:
             self.weights_sources.append(
                 DiffusersPipelineLoader.ComponentSource(
                     model_or_path=od_config.model,
@@ -257,14 +270,26 @@ class Wan22Pipeline(nn.Module):
         ).to(self.device)
 
         # Initialize transformers with correct config (weights loaded via load_weights)
-        transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
-        if self.has_transformer_2:
+        if load_transformer:
+            transformer_config = load_transformer_config(model, "transformer", local_files_only)
+            self.transformer = create_transformer_from_config(transformer_config)
+        else:
+            self.transformer = None
+
+        if load_transformer_2:
             transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
             self.transformer_2 = create_transformer_from_config(transformer_2_config)
         else:
             self.transformer_2 = None
 
+        # Store the active transformer config
+        if load_transformer:
+            self.transformer_config = self.transformer.config
+        elif load_transformer_2:
+            self.transformer_config = self.transformer_2.config
+        else:
+            raise RuntimeError("No transformer loaded")
+
         # Initialize UniPC scheduler
         flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0  # default for 720p
         self.scheduler = FlowUniPCMultistepScheduler(
@@ -275,7 +300,6 @@ class Wan22Pipeline(nn.Module):
 
         self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4
         self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8
-        self.boundary_ratio = od_config.boundary_ratio
 
         self._guidance_scale = None
         self._guidance_scale_2 = None
@@ -333,7 +357,7 @@ class Wan22Pipeline(nn.Module):
 
         # Ensure dimensions are compatible with VAE and patch size
         # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size)
-        patch_size = self.transformer.config.patch_size
+        patch_size = self.transformer_config.patch_size
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
@@ -374,7 +398,14 @@ class Wan22Pipeline(nn.Module):
         num_frames = max(num_frames, 1)
 
         device = self.device
-        dtype = self.transformer.dtype
+        # Get dtype from whichever transformer is loaded
+        if self.transformer is not None:
+            dtype = self.transformer.dtype
+        elif self.transformer_2 is not None:
+            dtype = self.transformer_2.dtype
+        else:
+            # Fallback to text_encoder dtype if no transformer loaded
+            dtype = self.text_encoder.dtype
 
         # Seed / generator
         if generator is None:
@@ -444,7 +475,7 @@ class Wan22Pipeline(nn.Module):
                 image_tensor = image
 
             # Use out_channels for noise latents (not in_channels which includes condition)
-            num_channels_latents = self.transformer.config.out_channels
+            num_channels_latents = self.transformer_config.out_channels
             batch_size = prompt_embeds.shape[0]
 
             # Prepare noise latents
@@ -489,7 +520,7 @@ class Wan22Pipeline(nn.Module):
             first_frame_mask[:, :, 0] = 0
         else:
             # T2V mode: standard latent preparation
-            num_channels_latents = self.transformer.config.in_channels
+            num_channels_latents = self.transformer_config.in_channels
             latents = self.prepare_latents(
                 batch_size=prompt_embeds.shape[0],
                 num_channels_latents=num_channels_latents,
@@ -508,11 +539,30 @@ class Wan22Pipeline(nn.Module):
         # Denoising
         for t in timesteps:
             self._current_timestep = t
-            current_model = self.transformer
-            current_guidance_scale = guidance_low
-            if boundary_timestep is not None and t < boundary_timestep and self.transformer_2 is not None:
-                current_model = self.transformer_2
+
+            # Select model based on timestep and boundary_ratio
+            # High noise stage (t >= boundary_timestep): use transformer
+            # Low noise stage (t < boundary_timestep): use transformer_2
+            if boundary_timestep is not None and t < boundary_timestep:
+                # Low noise stage - always use guidance_high for this stage
                 current_guidance_scale = guidance_high
+                if self.transformer_2 is not None:
+                    current_model = self.transformer_2
+                elif self.transformer is not None:
+                    # Fallback to transformer if transformer_2 not loaded
+                    current_model = self.transformer
+                else:
+                    raise RuntimeError("No transformer available for low-noise stage")
+            else:
+                # High noise stage - always use guidance_low for this stage
+                current_guidance_scale = guidance_low
+                if self.transformer is not None:
+                    current_model = self.transformer
+                elif self.transformer_2 is not None:
+                    # Fallback to transformer_2 if transformer not loaded
+                    current_model = self.transformer_2
+                else:
+                    raise RuntimeError("No transformer available for high-noise stage")
 
             if self.expand_timesteps and latent_condition is not None:
                 # I2V mode: blend condition with latents using mask
@@ -520,7 +570,7 @@ class Wan22Pipeline(nn.Module):
                 latent_model_input = latent_model_input.to(dtype)
 
                 # Expand timesteps per patch - use floor division to match patch embedding
-                patch_size = self.transformer.config.patch_size
+                patch_size = self.transformer_config.patch_size
                 num_latent_frames = latents.shape[2]
                 patch_height = latents.shape[3] // patch_size[1]
                 patch_width = latents.shape[4] // patch_size[2]

From bfbf3e5b69492b12e8349fa7f6b62145a57c58df Mon Sep 17 00:00:00 2001
From: Fanli Lin <fanli.lin@intel.com>
Date: Thu, 29 Jan 2026 23:33:05 +0800
Subject: [PATCH 28/32] [Feat] Support XPU Backend in vLLM-Omni (#191)

Signed-off-by: Fanli Lin <fanli.lin@intel.com>
Signed-off-by: Fanli Lin <fanli0116@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 vllm_omni/diffusion/offload.py                |  26 +++--
 .../qwen2_5_omni/qwen2_5_omni_token2wav.py    |   4 +
 vllm_omni/platforms/interface.py              |   9 ++
 .../xpu/stage_configs/qwen2_5_omni.yaml       | 101 ++++++++++++++++++
 .../xpu/stage_configs/qwen3_omni_moe.yaml     |  99 +++++++++++++++++
 vllm_omni/platforms/xpu/utils.py              |  16 +++
 vllm_omni/platforms/xpu/worker/__init__.py    |   0
 .../xpu/worker/xpu_ar_model_runner.py         |  19 ++++
 .../platforms/xpu/worker/xpu_ar_worker.py     |  15 +++
 .../xpu/worker/xpu_generation_model_runner.py |  19 ++++
 .../xpu/worker/xpu_generation_worker.py       |  15 +++
 11 files changed, 315 insertions(+), 8 deletions(-)
 create mode 100644 vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml
 create mode 100644 vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml
 create mode 100644 vllm_omni/platforms/xpu/utils.py
 create mode 100644 vllm_omni/platforms/xpu/worker/__init__.py
 create mode 100644 vllm_omni/platforms/xpu/worker/xpu_ar_model_runner.py
 create mode 100644 vllm_omni/platforms/xpu/worker/xpu_ar_worker.py
 create mode 100644 vllm_omni/platforms/xpu/worker/xpu_generation_model_runner.py
 create mode 100644 vllm_omni/platforms/xpu/worker/xpu_generation_worker.py

diff --git a/vllm_omni/diffusion/offload.py b/vllm_omni/diffusion/offload.py
index 4220ee9ff..174c64342 100644
--- a/vllm_omni/diffusion/offload.py
+++ b/vllm_omni/diffusion/offload.py
@@ -18,6 +18,8 @@ import torch
 from torch import nn
 from vllm.logger import init_logger
 
+from vllm_omni.platforms import current_omni_platform
+
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -63,8 +65,8 @@ class SequentialOffloader:
         module.to("cpu", non_blocking=True)
 
         # Release allocator blocks when tensors leave the GPU.
-        if previous_device.type == "cuda" and torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        if previous_device.type != "cpu":
+            current_omni_platform.empty_cache()
 
         if self.pin_memory:
             for p in module.parameters():
@@ -87,7 +89,9 @@ class SequentialOffloader:
         for enc in self.encoders:
             self._to_cpu(enc)
         self._to_gpu(module)
-        torch.cuda.synchronize()
+
+        current_omni_platform.synchronize()
+
         logger.debug("Swapped: encoders -> CPU, DiT -> GPU")
 
     def _encoder_pre_hook(self, module: nn.Module, args: tuple) -> None:
@@ -95,7 +99,9 @@ class SequentialOffloader:
         for dit_mod in self.dits:
             self._to_cpu(dit_mod)
         self._to_gpu(module)
-        torch.cuda.synchronize()
+
+        current_omni_platform.synchronize()
+
         logger.debug("Swapped: DiT -> CPU, encoder -> GPU")
 
     def register(self) -> None:
@@ -166,7 +172,10 @@ def apply_offload_hooks(
         try:
             device = next(dit_modules[0].parameters()).device
         except StopIteration:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            try:
+                device = current_omni_platform.get_torch_device()
+            except (NotImplementedError, AttributeError):
+                device = torch.device("cpu")
 
     # Collect all encoders
     encoders: list[nn.Module] = []
@@ -184,9 +193,10 @@ def apply_offload_hooks(
     pin = getattr(od_config, "pin_cpu_memory", True)
     for dit_mod in dit_modules:
         dit_mod.to("cpu")
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    if pin and torch.cuda.is_available():
+
+    current_omni_platform.empty_cache()
+
+    if pin:
         for dit_mod in dit_modules:
             for p in dit_mod.parameters():
                 if p.data.device.type == "cpu" and not p.data.is_pinned():
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
index 9d176ac6c..84ff8b55f 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
@@ -736,6 +736,10 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
         kaiser_window = torch.kaiser_window(
             kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
         ).to("npu")
+    elif current_omni_platform.is_xpu():
+        kaiser_window = torch.kaiser_window(
+            kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
+        ).to("xpu")
     else:
         kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
 
diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py
index 2aac4b7a1..db541ebce 100644
--- a/vllm_omni/platforms/interface.py
+++ b/vllm_omni/platforms/interface.py
@@ -31,6 +31,15 @@ class OmniPlatform(Platform):
     def is_npu(self) -> bool:
         return self._omni_enum == OmniPlatformEnum.NPU
 
+    def is_xpu(self) -> bool:
+        return self._omni_enum == OmniPlatformEnum.XPU
+
+    def is_cuda(self) -> bool:
+        return self._omni_enum == OmniPlatformEnum.CUDA
+
+    def is_rocm(self) -> bool:
+        return self._omni_enum == OmniPlatformEnum.ROCM
+
     @classmethod
     def get_omni_ar_worker_cls(cls) -> str:
         raise NotImplementedError
diff --git a/vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml b/vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml
new file mode 100644
index 000000000..3fe1b79e4
--- /dev/null
+++ b/vllm_omni/platforms/xpu/stage_configs/qwen2_5_omni.yaml
@@ -0,0 +1,101 @@
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
+
+# The following config has been verified on 2x 1550-64G XPUs.
+stage_args:
+  - stage_id: 0
+    stage_type: llm # Use llm stage type to launch OmniLLM
+    runtime:
+      process: true # Run this stage in a separate process
+      devices: "0" # Visible devices for this stage
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: false
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+  - stage_id: 1
+    stage_type: llm # Use llm stage type to launch OmniLLM
+    runtime:
+      process: true
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: talker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: false
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
+
+  - stage_id: 2
+    stage_type: llm # Use llm stage type to launch OmniLLM
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      gpu_memory_utilization: 0.15
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio
+    engine_input_source: [1]
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1 # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1 # Simplified: process serially within each stage
+
+  edges:
+    - from: 0 # thinker → talker: trigger only after receiving full input (-1)
+      to: 1
+      window_size: -1
+    - from: 1 # talker → code2wav: trigger only after receiving full input (-1)
+      to: 2
+      window_size: -1
diff --git a/vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml
new file mode 100644
index 000000000..0803d7353
--- /dev/null
+++ b/vllm_omni/platforms/xpu/stage_configs/qwen3_omni_moe.yaml
@@ -0,0 +1,99 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+# The following config has been verified on 4x 1550-64G XPUs.
+stage_args:
+  - stage_id: 0
+    stage_type: llm # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "2"
+      max_batch_size: 1
+    engine_args:
+      model_stage: talker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.3
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent # Output codec codes for code2wav
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      distributed_executor_backend: "mp"
+      hf_config_name: talker_config
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    stage_type: llm # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "3"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio # Final output: audio waveform
+      gpu_memory_utilization: 0.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
diff --git a/vllm_omni/platforms/xpu/utils.py b/vllm_omni/platforms/xpu/utils.py
new file mode 100644
index 000000000..18437cae3
--- /dev/null
+++ b/vllm_omni/platforms/xpu/utils.py
@@ -0,0 +1,16 @@
+from contextlib import contextmanager
+
+import torch
+
+
+@contextmanager
+def torch_cuda_wrapper():
+    try:
+        # replace cuda APIs with xpu APIs, this should work by default
+        torch.cuda.Stream = torch.xpu.Stream
+        torch.cuda.default_stream = torch.xpu.current_stream
+        torch.cuda.current_stream = torch.xpu.current_stream
+        torch.cuda.stream = torch.xpu.stream
+        yield
+    finally:
+        pass
diff --git a/vllm_omni/platforms/xpu/worker/__init__.py b/vllm_omni/platforms/xpu/worker/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_omni/platforms/xpu/worker/xpu_ar_model_runner.py b/vllm_omni/platforms/xpu/worker/xpu_ar_model_runner.py
new file mode 100644
index 000000000..2235d3548
--- /dev/null
+++ b/vllm_omni/platforms/xpu/worker/xpu_ar_model_runner.py
@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm_omni.platforms.xpu.utils import torch_cuda_wrapper
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+
+class XPUARModelRunner(GPUARModelRunner):
+    def __init__(self, *args, **kwargs):
+        with torch_cuda_wrapper():
+            super().__init__(*args, **kwargs)
+
+    def _init_device_properties(self):
+        self.num_sms = None
+
+    def _sync_device(self) -> None:
+        torch.xpu.synchronize()
diff --git a/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py b/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py
new file mode 100644
index 000000000..ea6349264
--- /dev/null
+++ b/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py
@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.v1.worker.xpu_worker import XPUWorker
+
+from vllm_omni.platforms.xpu.worker.xpu_ar_model_runner import XPUARModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
+
+
+class XPUARWorker(OmniWorkerMixin, XPUWorker):
+    """XPU AR worker for thinker/talker stages in Omni model."""
+
+    def init_device(self):
+        super().init_device()
+        self.model_runner: XPUARModelRunner = XPUARModelRunner(self.vllm_config, self.device)
diff --git a/vllm_omni/platforms/xpu/worker/xpu_generation_model_runner.py b/vllm_omni/platforms/xpu/worker/xpu_generation_model_runner.py
new file mode 100644
index 000000000..e865ac5e0
--- /dev/null
+++ b/vllm_omni/platforms/xpu/worker/xpu_generation_model_runner.py
@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm_omni.platforms.xpu.utils import torch_cuda_wrapper
+from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+
+
+class XPUGenerationModelRunner(GPUGenerationModelRunner):
+    def __init__(self, *args, **kwargs):
+        with torch_cuda_wrapper():
+            super().__init__(*args, **kwargs)
+
+    def _init_device_properties(self):
+        self.num_sms = None
+
+    def _sync_device(self) -> None:
+        torch.xpu.synchronize()
diff --git a/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py b/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py
new file mode 100644
index 000000000..9c6d530eb
--- /dev/null
+++ b/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py
@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.v1.worker.xpu_worker import XPUWorker
+
+from vllm_omni.platforms.xpu.worker.xpu_generation_model_runner import XPUGenerationModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
+
+
+class XPUGenerationWorker(OmniWorkerMixin, XPUWorker):
+    """XPU generation worker for the code2wav (non-AR waveform generation) stage in the Omni model."""
+
+    def init_device(self):
+        super().init_device()
+        self.model_runner: XPUGenerationModelRunner = XPUGenerationModelRunner(self.vllm_config, self.device)

From 3eecd2057b018b3984f607141916552f9be54932 Mon Sep 17 00:00:00 2001
From: dongbo910220 <32610838+dongbo910220@users.noreply.github.com>
Date: Fri, 30 Jan 2026 00:28:29 +0800
Subject: [PATCH 29/32] [Fix] stabilize diffusion images LoRA E2E across CI
 drift (#1075)

Signed-off-by: dongbo910220 <1275604947@qq.com>
---
 .../test_images_generations_lora.py           | 128 +++++-------------
 1 file changed, 33 insertions(+), 95 deletions(-)

diff --git a/tests/e2e/online_serving/test_images_generations_lora.py b/tests/e2e/online_serving/test_images_generations_lora.py
index 11165c6d2..e912c420d 100644
--- a/tests/e2e/online_serving/test_images_generations_lora.py
+++ b/tests/e2e/online_serving/test_images_generations_lora.py
@@ -12,10 +12,6 @@ This validates:
 import base64
 import json
 import os
-import signal
-import subprocess
-import sys
-import time
 from io import BytesIO
 from pathlib import Path
 
@@ -25,7 +21,8 @@ import requests
 import torch
 from PIL import Image
 from safetensors.torch import save_file
-from vllm.utils.network_utils import get_open_port
+
+from tests.conftest import OmniServer
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -37,87 +34,6 @@ SIZE = "256x256"
 SEED = 42
 
 
-class OmniServer:
-    """Omniserver for vLLM-Omni tests."""
-
-    def __init__(
-        self,
-        model: str,
-        serve_args: list[str],
-        *,
-        env_dict: dict[str, str] | None = None,
-    ) -> None:
-        self.model = model
-        self.serve_args = serve_args
-        self.env_dict = env_dict
-        self.proc: subprocess.Popen | None = None
-        self.host = "127.0.0.1"
-        self.port = get_open_port()
-
-    def _start_server(self) -> None:
-        env = os.environ.copy()
-        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-        if self.env_dict is not None:
-            env.update(self.env_dict)
-
-        cmd = [
-            sys.executable,
-            "-m",
-            "vllm_omni.entrypoints.cli.main",
-            "serve",
-            self.model,
-            "--omni",
-            "--host",
-            self.host,
-            "--port",
-            str(self.port),
-        ] + self.serve_args
-
-        print(f"Launching OmniServer with: {' '.join(cmd)}")
-        self.proc = subprocess.Popen(
-            cmd,
-            env=env,
-            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),  # vllm-omni root
-            start_new_session=True,
-        )
-
-        # Wait for server to be ready.
-        max_wait = 1200
-        url = f"http://{self.host}:{self.port}/v1/models"
-        start_time = time.time()
-        while time.time() - start_time < max_wait:
-            try:
-                resp = requests.get(url, headers={"Authorization": "Bearer EMPTY"}, timeout=10)
-                if resp.status_code == 200:
-                    print(f"Server ready on {self.host}:{self.port}")
-                    return
-            except Exception:
-                pass
-            time.sleep(2)
-
-        raise RuntimeError(f"Server failed to become ready within {max_wait} seconds")
-
-    def __enter__(self):
-        self._start_server()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.proc is None:
-            return
-        try:
-            os.killpg(self.proc.pid, signal.SIGTERM)
-        except ProcessLookupError:
-            pass
-        try:
-            self.proc.wait(timeout=30)
-        except subprocess.TimeoutExpired:
-            try:
-                os.killpg(self.proc.pid, signal.SIGKILL)
-            except ProcessLookupError:
-                pass
-            self.proc.wait()
-
-
 @pytest.fixture(scope="module")
 def omni_server():
     with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
@@ -183,17 +99,30 @@ def _image_blue_tail_slice(img: Image.Image) -> np.ndarray:
     return tail
 
 
-def _assert_slice_close(actual: np.ndarray, expected: np.ndarray, *, label: str) -> None:
+def _slice_diff_stats(actual: np.ndarray, expected: np.ndarray) -> tuple[float, float]:
+    diff = np.abs(actual - expected)
+    return float(diff.max()), float(diff.mean())
+
+
+def _assert_slice_close(
+    actual: np.ndarray,
+    expected: np.ndarray,
+    *,
+    label: str,
+    base_max: float,
+    base_mean: float,
+) -> None:
     assert actual.shape == (3, 3)
     assert expected.shape == (3, 3)
-    diff = np.abs(actual - expected)
-    max_diff = float(diff.max())
-    mean_diff = float(diff.mean())
+    max_diff, mean_diff = _slice_diff_stats(actual, expected)
     # NOTE: Different attention backends / torch.compile can introduce small
     # floating-point drift that shows up as a few LSBs in uint8 pixels. Keep
     # the reset check tolerant but bounded to avoid flaky CI.
-    assert max_diff <= 5.0 and mean_diff <= 3.0, (
-        f"{label} slice mismatch (max={max_diff:.1f}, mean={mean_diff:.1f}): {actual.tolist()}"
+    max_thresh = max(10.0, base_max + 4.0)
+    mean_thresh = max(6.0, base_mean + 4.0)
+    assert max_diff <= max_thresh and mean_diff <= mean_thresh, (
+        f"{label} slice mismatch (max={max_diff:.1f} > {max_thresh:.1f} or "
+        f"mean={mean_diff:.1f} > {mean_thresh:.1f}): {actual.tolist()}"
     )
 
 
@@ -219,6 +148,9 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer,
     # Base generation.
     base_img = _post_images(omni_server, _basic_payload())
     base_slice = _image_blue_tail_slice(base_img)
+    base_ref_img = _post_images(omni_server, _basic_payload())
+    base_ref_slice = _image_blue_tail_slice(base_ref_img)
+    base_ref_max, base_ref_mean = _slice_diff_stats(base_ref_slice, base_slice)
 
     # Adapter A: apply delta to V slice only.
     lora_a_dir = tmp_path / "zimage_lora_a"
@@ -245,11 +177,17 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer,
     # Ensure switching back to no-LoRA restores the base output.
     base_img_2 = _post_images(omni_server, _basic_payload())
     base_slice_2 = _image_blue_tail_slice(base_img_2)
-    _assert_slice_close(base_slice_2, base_slice, label="base_after_reset")
-    base_reset = float(np.abs(base_slice_2 - base_slice).mean())
+    _, base_reset_mean = _slice_diff_stats(base_slice_2, base_slice)
+    _assert_slice_close(
+        base_slice_2,
+        base_slice,
+        label="base_after_reset",
+        base_max=base_ref_max,
+        base_mean=base_ref_mean,
+    )
 
     # Ensure LoRA effects are clearly above the baseline drift.
-    min_delta = base_reset + 0.5
+    min_delta = max(base_reset_mean + 1.0, 1.5)
     assert a_vs_base > min_delta, f"lora_a_vs_base drift too small: {a_vs_base} <= {min_delta}"
     assert b_vs_base > min_delta, f"lora_b_vs_base drift too small: {b_vs_base} <= {min_delta}"
     assert b_vs_a > min_delta, f"lora_b_vs_lora_a drift too small: {b_vs_a} <= {min_delta}"

From 2c4974c6eabf5616c99eb7ab2e52da48ef0fa9d4 Mon Sep 17 00:00:00 2001
From: Canlin Guo <canlinguosdu@gmail.com>
Date: Fri, 30 Jan 2026 00:55:53 +0800
Subject: [PATCH 30/32] [Bugfix][Test] Re-enable the log simple tests (#1065)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
---
 .buildkite/pipeline.yml           |  2 ++
 .buildkite/scripts/simple_test.sh | 58 -------------------------------
 2 files changed, 2 insertions(+), 58 deletions(-)
 delete mode 100755 .buildkite/scripts/simple_test.sh

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 78a400cd8..6d731f127 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -22,8 +22,10 @@ steps:
     commands:
       - pytest -v -s tests/entrypoints/
       - pytest -v -s tests/diffusion/cache/
+      - pytest -v -s tests/diffusion/lora/
       - pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
       - pytest -v -s tests/worker/
+      - pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py
     agents:
       queue: "gpu_1_queue"
     plugins:
diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh
deleted file mode 100755
index e76a5c70d..000000000
--- a/.buildkite/scripts/simple_test.sh
+++ /dev/null
@@ -1,58 +0,0 @@
-#!/usr/bin/env bash
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
-UV_BIN="uv"
-PYTHON_BOOTSTRAP="${PYTHON:-python3}"
-VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
-TARGET_PY_VERSION="3.12"
-
-if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
-  PYTHON_BOOTSTRAP="python"
-fi
-
-cd "${REPO_ROOT}"
-
-# Ensure pip-installed scripts (including uv) are on PATH
-PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
-USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
-PATH="${PYTHON_SCRIPTS}:${PATH}"
-if [[ -n "${USER_BASE}" ]]; then
-  PATH="${USER_BASE}/bin:${PATH}"
-fi
-
-if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
-  "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
-  "${PYTHON_BOOTSTRAP}" -m pip install uv
-  hash -r
-fi
-
-TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  "${UV_BIN}" python install "${TARGET_PY_VERSION}"
-  TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-fi
-
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
-  exit 1
-fi
-
-if [[ ! -d "${VENV_DIR}" ]]; then
-  "${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
-fi
-
-VENV_PYTHON="${VENV_DIR}/bin/python"
-[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
-
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.14.0
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
-
-"${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
-"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
-"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/lora/
-"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
-"${VENV_PYTHON}" -m pytest -v -s tests/worker/
-"${VENV_PYTHON}" -m pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py

From 47a0969976ae8dac9d2c567c2663720370a64df9 Mon Sep 17 00:00:00 2001
From: dengyunyang <584797741@qq.com>
Date: Fri, 30 Jan 2026 01:09:34 +0800
Subject: [PATCH 31/32] [Bugfix] pr conflict fix, bugfix ignore mm data from
 stages to async omni (#1025)

Signed-off-by: dengyunyang <584797741@qq.com>
---
 vllm_omni/entrypoints/omni_stage.py | 31 ++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index 2adb1c342..fc01c147c 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -1467,7 +1467,7 @@ async def _stage_worker_async(
             batch_request_ids, batch_request_outputs, _gen_ms_list, batch_metrics
         ):
             try:
-                r_outputs = [output]
+                r_outputs = [output_strip(output, omni_stage)]
                 use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
                 if use_shm:
                     out_q.put(
@@ -1553,3 +1553,32 @@ def make_stage_stats(_agg_total_tokens: int, _agg_total_gen_time_ms: float):
     from vllm_omni.entrypoints.log_utils import StageStats
 
     return StageStats(total_token=_agg_total_tokens, total_gen_time=_agg_total_gen_time_ms)
+
+
+def output_strip(r_output: RequestOutput | OmniRequestOutput, omni_stage: OmniStage):
+    """
+    Strip unnecessary multimodal outputs from stages results,
+    in order to:
+    - reduce memory usage
+    - reduce transfer & serialization overhead
+    """
+
+    # check multimodal data is required by stage output config.
+    if omni_stage.final_output and omni_stage.final_output_type != "text":
+        return r_output
+
+    # If the request has already finished, should not be altered.
+    if getattr(r_output, "finished", False):
+        return r_output
+
+    mm_output = getattr(r_output, "multimodal_output", None)
+    if mm_output is not None:
+        r_output.multimodal_output = {}
+
+    outputs = getattr(r_output, "outputs", None)
+    if outputs is not None:
+        for out in outputs:
+            if getattr(out, "multimodal_output", None):
+                out.multimodal_output = {}
+
+    return r_output

From 0e07eb6ff920937e003399cd9a92b6055c37796d Mon Sep 17 00:00:00 2001
From: Ding Zuhao <e1583181@u.nus.edu>
Date: Fri, 30 Jan 2026 07:56:10 +0800
Subject: [PATCH 32/32] [Doc][Bagel] Add BAGEL-7B-MoT documentation and edit
 the default stage configuration (#987)

Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Signed-off-by: jzz <e1583181@u.nus.edu>
---
 .../examples/offline_inference/bagel.md       | 179 ++++++++++++++
 .../examples/online_serving/bagel.md          | 232 ++++++++++++++++++
 examples/offline_inference/bagel/README.md    | 177 +++++++++++++
 .../bagel/run_single_prompt.sh                |   6 -
 examples/online_serving/bagel/README.md       | 230 +++++++++++++++++
 .../model_executor/stage_configs/bagel.yaml   |   4 +-
 .../stage_configs/bagel_multiconnector.yaml   |   2 +-
 7 files changed, 821 insertions(+), 9 deletions(-)
 create mode 100644 docs/user_guide/examples/offline_inference/bagel.md
 create mode 100644 docs/user_guide/examples/online_serving/bagel.md
 create mode 100644 examples/offline_inference/bagel/README.md
 delete mode 100644 examples/offline_inference/bagel/run_single_prompt.sh
 create mode 100644 examples/online_serving/bagel/README.md

diff --git a/docs/user_guide/examples/offline_inference/bagel.md b/docs/user_guide/examples/offline_inference/bagel.md
new file mode 100644
index 000000000..b2ee5d241
--- /dev/null
+++ b/docs/user_guide/examples/offline_inference/bagel.md
@@ -0,0 +1,179 @@
+# BAGEL-7B-MoT
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
+
+## Set up
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Run examples
+
+**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
+
+Get into the bagel folder
+
+```bash
+cd examples/offline_inference/bagel
+```
+
+### Modality Control
+
+BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
+
+#### Text to Image (text2img)
+
+- **Pipeline**: Text → Thinker  → DiT → VAE Decode → Image
+- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
+- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
+
+Generate images from text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat"
+```
+
+#### Image to Image (img2img)
+
+- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
+- **Stages Used**: Stage 1 (DiT) only
+- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
+
+Transform images based on text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2img \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Let the woman wear a blue dress"
+```
+
+#### Image to Text (img2text)
+
+- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
+
+Generate text descriptions from images:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image in detail"
+```
+
+#### Text to Text (text2text)
+
+- **Pipeline**: Text → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: No visual components involved, operates as pure language model
+
+Pure text generation:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?"
+
+# You can load prompts from a text file (one prompt per line):  
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --txt-prompts /path/to/prompts.txt
+```
+
+### Inference Steps
+
+Control the number of inference steps for image generation:
+
+```bash
+# You can adjust steps to 100 to improve image quality
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --steps 50 \
+                  --prompts "A cute cat"
+```
+
+### Key arguments
+
+BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
+
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+
+#### 📌 Command Line Arguments (end2end.py)
+
+| Argument               | Type   | Default                       | Description                                                  |
+| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
+| `--model`              | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name                                           |
+| `--modality`           | choice | `text2img`                    | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
+| `--prompts`            | list   | `None`                        | Input text prompts directly                                  |
+| `--txt-prompts`        | string | `None`                        | Path to txt file with one prompt per line                    |
+| `--image-path`         | string | `None`                        | Input image path (for `img2img`/`img2text`)                  |
+| `--steps`              | int    | `50`                          | Number of inference steps                                    |
+| `--stage-configs-path` | string | `None`                        | Custom stage config file path                                |
+| `--worker-backend`     | choice | `process`                     | Worker backend: `process` or `ray`                           |
+| `--ray-address`        | string | `None`                        | Ray cluster address                                          |
+| `--enable-stats`       | flag   | `False`                       | Enable statistics logging                                    |
+| `--init-sleep-seconds` | int    | `20`                          | Initialization sleep time                                    |
+| `--batch-timeout`      | int    | `5`                           | Batch timeout                                                |
+| `--init-timeout`       | int    | `300`                         | Initialization timeout                                       |
+
+------
+
+#### ⚙️ Stage Configuration Parameters (bagel.yaml)
+
+ **Stage 0 - Thinker (LLM Stage)**
+
+| Parameter                        | Value                           | Description              |
+| :------------------------------- | :------------------------------ | :----------------------- |
+| `stage_type`                     | `llm`                           | Stage type               |
+| `devices`                        | `"0"`                           | GPU device ID            |
+| `max_batch_size`                 | `1`                             | Maximum batch size       |
+| `model_stage`                    | `thinker`                       | Model stage identifier   |
+| `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
+| `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
+| `tensor_parallel_size`           | `1`                             | Tensor parallel size     |
+| `max_num_batched_tokens`         | `32768`                         | Maximum batched tokens   |
+| `omni_kv_config.need_send_cache` | `true`                          | Whether to send KV cache |
+
+------
+
+**Stage 1 - DiT (Diffusion Stage)**
+
+| Parameter                        | Value       | Description                 |
+| :------------------------------- | :---------- | :-------------------------- |
+| `stage_type`                     | `diffusion` | Stage type                  |
+| `devices`                        | `"0"`       | GPU device ID               |
+| `max_batch_size`                 | `1`         | Maximum batch size          |
+| `model_stage`                    | `dit`       | Model stage identifier      |
+| `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
+| `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |
+| `engine_input_source`            | `[0]`       | Input source from Stage 0   |
+
+------
+
+#### 🔗 Runtime Configuration
+
+| Parameter             | Value   | Description                      |
+| :-------------------- | :------ | :------------------------------- |
+| `window_size`         | `-1`    | Window size (-1 means unlimited) |
+| `max_inflight`        | `1`     | Maximum inflight requests        |
+| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
+
+## FAQ
+
+- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
+
+```bash
+sudo apt update
+sudo apt install ffmpeg
+```
+
+- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
+
+| Stage               | VRAM                         |
+| :------------------ | :--------------------------- |
+| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
+| Stage-1 (DiT)       | **26.50 GiB**                |
+| Total               | **~42 GiB + KV Cache**       |
diff --git a/docs/user_guide/examples/online_serving/bagel.md b/docs/user_guide/examples/online_serving/bagel.md
new file mode 100644
index 000000000..107e42449
--- /dev/null
+++ b/docs/user_guide/examples/online_serving/bagel.md
@@ -0,0 +1,232 @@
+# BAGEL-7B-MoT
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (BAGEL-7B-MoT)
+
+**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
+
+### Launch the Server
+
+```bash
+# Use default configuration
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
+```
+
+Or use the convenience script:
+
+```bash
+cd /workspace/vllm-omni/examples/online_serving/bagel
+bash run_server.sh
+```
+
+If you have a custom stage configs file, launch the server with the command below:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+
+### Send Multi-modal Request
+
+Get into the bagel folder:
+
+```bash
+cd examples/online_serving/bagel
+```
+
+Send request via Python
+
+```bash
+python openai_chat_client.py --prompt "A cute cat" --modality text2img
+```
+
+The Python client supports the following command-line arguments:
+
+- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
+- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
+- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
+- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
+- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
+- `--height`: Image height in pixels (default: 512)
+- `--width`: Image width in pixels (default: 512)
+- `--steps`: Number of inference steps (default: 25)
+- `--seed`: Random seed (default: 42)
+- `--negative`: Negative prompt for image generation
+
+Example with custom parameters:
+
+```bash
+python openai_chat_client.py \
+    --prompt "A futuristic city" \
+    --modality text2img \
+    --height 768 \
+    --width 768 \
+    --steps 50 \
+    --seed 42 \
+    --negative "blurry, low quality"
+```
+
+## Modality Control
+
+BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
+
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+
+| Modality    | Input        | Output | Description                            |
+| ----------- | ------------ | ------ | -------------------------------------- |
+| `text2img`  | Text         | Image  | Generate images from text prompts      |
+| `img2img`   | Image + Text | Image  | Transform images using text guidance   |
+| `img2text`  | Image + Text | Text   | Generate text descriptions from images |
+| `text2text` | Text         | Text   | Pure text generation                   |
+
+### Text to Image (text2img)
+
+Generate images from text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "A beautiful sunset over mountains" \
+    --modality text2img \
+    --output sunset.png \
+    --steps 50
+```
+
+**Using curl**
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
+    "modalities": ["image"],
+    "height": 512,
+    "width": 512,
+    "num_inference_steps": 50,
+    "seed": 42
+  }'
+```
+
+
+### Image to Image (img2img)
+
+Transform images based on text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Make the cat stand up" \
+    --modality img2img \
+    --image-url /path/to/input.jpg \
+    --output transformed.png
+```
+
+**Using curl**
+
+```bash
+IMAGE_BASE64=$(base64 -w 0 cat.jpg)
+
+cat <<EOF > payload.json
+{
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
+        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
+      ]
+    }],
+    "modalities": ["image"],
+    "height": 512,
+    "width": 512,
+    "num_inference_steps": 50,
+    "seed": 42
+}
+EOF
+
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @payload.json
+
+```
+
+### Image to Text (img2text)
+
+Generate text descriptions from images:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Describe this image in detail" \
+    --modality img2text \
+    --image-url /path/to/image.jpg
+```
+
+**Using curl**
+
+```bash
+IMAGE_BASE64=$(base64 -w 0 cat.jpg)
+
+cat <<EOF > payload.json
+{
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
+      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
+    ]
+  }],
+  "modalities": ["text"]
+}
+EOF
+
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @payload.json
+```
+
+### Text to Text (text2text)
+
+Pure text generation:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "What is the capital of France?" \
+    --modality text2text
+```
+
+**Using curl**
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
+    "modalities": ["text"]
+  }'
+```
+
+## FAQ
+
+- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
+
+```bash
+sudo apt update
+sudo apt install ffmpeg
+```
+
+- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
+
+| Stage               | VRAM                         |
+| :------------------ | :--------------------------- |
+| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
+| Stage-1 (DiT)       | **26.50 GiB**                |
+| Total               | **~42 GiB + KV Cache**       |
diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md
new file mode 100644
index 000000000..7bcfb42fa
--- /dev/null
+++ b/examples/offline_inference/bagel/README.md
@@ -0,0 +1,177 @@
+# BAGEL-7B-MoT
+
+## Set up
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Run examples
+
+**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
+
+Get into the bagel folder
+
+```bash
+cd examples/offline_inference/bagel
+```
+
+### Modality Control
+
+BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
+
+#### Text to Image (text2img)
+
+- **Pipeline**: Text → Thinker  → DiT → VAE Decode → Image
+- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
+- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
+
+Generate images from text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat"
+```
+
+#### Image to Image (img2img)
+
+- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
+- **Stages Used**: Stage 1 (DiT) only
+- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
+
+Transform images based on text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2img \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Let the woman wear a blue dress"
+```
+
+#### Image to Text (img2text)
+
+- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
+
+Generate text descriptions from images:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image in detail"
+```
+
+#### Text to Text (text2text)
+
+- **Pipeline**: Text → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: No visual components involved, operates as pure language model
+
+Pure text generation:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?"
+
+# You can load prompts from a text file (one prompt per line):  
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --txt-prompts /path/to/prompts.txt
+```
+
+### Inference Steps
+
+Control the number of inference steps for image generation:
+
+```bash
+# You can adjust steps to 100 to improve image quality
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --steps 50 \
+                  --prompts "A cute cat"
+```
+
+### Key arguments
+
+BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
+
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+
+#### 📌 Command Line Arguments (end2end.py)
+
+| Argument               | Type   | Default                       | Description                                                  |
+| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
+| `--model`              | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name                                           |
+| `--modality`           | choice | `text2img`                    | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
+| `--prompts`            | list   | `None`                        | Input text prompts directly                                  |
+| `--txt-prompts`        | string | `None`                        | Path to txt file with one prompt per line                    |
+| `--image-path`         | string | `None`                        | Input image path (for `img2img`/`img2text`)                  |
+| `--steps`              | int    | `50`                          | Number of inference steps                                    |
+| `--stage-configs-path` | string | `None`                        | Custom stage config file path                                |
+| `--worker-backend`     | choice | `process`                     | Worker backend: `process` or `ray`                           |
+| `--ray-address`        | string | `None`                        | Ray cluster address                                          |
+| `--enable-stats`       | flag   | `False`                       | Enable statistics logging                                    |
+| `--init-sleep-seconds` | int    | `20`                          | Initialization sleep time                                    |
+| `--batch-timeout`      | int    | `5`                           | Batch timeout                                                |
+| `--init-timeout`       | int    | `300`                         | Initialization timeout                                       |
+
+------
+
+#### ⚙️ Stage Configuration Parameters (bagel.yaml)
+
+ **Stage 0 - Thinker (LLM Stage)**
+
+| Parameter                        | Value                           | Description              |
+| :------------------------------- | :------------------------------ | :----------------------- |
+| `stage_type`                     | `llm`                           | Stage type               |
+| `devices`                        | `"0"`                           | GPU device ID            |
+| `max_batch_size`                 | `1`                             | Maximum batch size       |
+| `model_stage`                    | `thinker`                       | Model stage identifier   |
+| `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
+| `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
+| `tensor_parallel_size`           | `1`                             | Tensor parallel size     |
+| `max_num_batched_tokens`         | `32768`                         | Maximum batched tokens   |
+| `omni_kv_config.need_send_cache` | `true`                          | Whether to send KV cache |
+
+------
+
+**Stage 1 - DiT (Diffusion Stage)**
+
+| Parameter                        | Value       | Description                 |
+| :------------------------------- | :---------- | :-------------------------- |
+| `stage_type`                     | `diffusion` | Stage type                  |
+| `devices`                        | `"0"`       | GPU device ID               |
+| `max_batch_size`                 | `1`         | Maximum batch size          |
+| `model_stage`                    | `dit`       | Model stage identifier      |
+| `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
+| `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |
+| `engine_input_source`            | `[0]`       | Input source from Stage 0   |
+
+------
+
+#### 🔗 Runtime Configuration
+
+| Parameter             | Value   | Description                      |
+| :-------------------- | :------ | :------------------------------- |
+| `window_size`         | `-1`    | Window size (-1 means unlimited) |
+| `max_inflight`        | `1`     | Maximum inflight requests        |
+| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
+
+## FAQ
+
+- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
+
+```bash
+sudo apt update
+sudo apt install ffmpeg
+```
+
+- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
+
+| Stage               | VRAM                         |
+| :------------------ | :--------------------------- |
+| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
+| Stage-1 (DiT)       | **26.50 GiB**                |
+| Total               | **~42 GiB + KV Cache**       |
diff --git a/examples/offline_inference/bagel/run_single_prompt.sh b/examples/offline_inference/bagel/run_single_prompt.sh
deleted file mode 100644
index 2099620e3..000000000
--- a/examples/offline_inference/bagel/run_single_prompt.sh
+++ /dev/null
@@ -1,6 +0,0 @@
-prompt="<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
-
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                                 --prompt_type text \
-                                 --init-sleep-seconds 0 \
-                                 --prompts ${prompt}
diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md
new file mode 100644
index 000000000..3fbea0550
--- /dev/null
+++ b/examples/online_serving/bagel/README.md
@@ -0,0 +1,230 @@
+# BAGEL-7B-MoT
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (BAGEL-7B-MoT)
+
+**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
+
+### Launch the Server
+
+```bash
+# Use default configuration
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
+```
+
+Or use the convenience script:
+
+```bash
+cd /workspace/vllm-omni/examples/online_serving/bagel
+bash run_server.sh
+```
+
+If you have a custom stage configs file, launch the server with the command below:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+
+### Send Multi-modal Request
+
+Get into the bagel folder:
+
+```bash
+cd examples/online_serving/bagel
+```
+
+Send request via Python
+
+```bash
+python openai_chat_client.py --prompt "A cute cat" --modality text2img
+```
+
+The Python client supports the following command-line arguments:
+
+- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
+- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
+- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
+- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
+- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
+- `--height`: Image height in pixels (default: 512)
+- `--width`: Image width in pixels (default: 512)
+- `--steps`: Number of inference steps (default: 25)
+- `--seed`: Random seed (default: 42)
+- `--negative`: Negative prompt for image generation
+
+Example with custom parameters:
+
+```bash
+python openai_chat_client.py \
+    --prompt "A futuristic city" \
+    --modality text2img \
+    --height 768 \
+    --width 768 \
+    --steps 50 \
+    --seed 42 \
+    --negative "blurry, low quality"
+```
+
+## Modality Control
+
+BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
+
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+
+| Modality    | Input        | Output | Description                            |
+| ----------- | ------------ | ------ | -------------------------------------- |
+| `text2img`  | Text         | Image  | Generate images from text prompts      |
+| `img2img`   | Image + Text | Image  | Transform images using text guidance   |
+| `img2text`  | Image + Text | Text   | Generate text descriptions from images |
+| `text2text` | Text         | Text   | Pure text generation                   |
+
+### Text to Image (text2img)
+
+Generate images from text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "A beautiful sunset over mountains" \
+    --modality text2img \
+    --output sunset.png \
+    --steps 50
+```
+
+**Using curl**
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
+    "modalities": ["image"],
+    "height": 512,
+    "width": 512,
+    "num_inference_steps": 50,
+    "seed": 42
+  }'
+```
+
+
+### Image to Image (img2img)
+
+Transform images based on text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Make the cat stand up" \
+    --modality img2img \
+    --image-url /path/to/input.jpg \
+    --output transformed.png
+```
+
+**Using curl**
+
+```bash
+IMAGE_BASE64=$(base64 -w 0 cat.jpg)
+
+cat <<EOF > payload.json
+{
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
+        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
+      ]
+    }],
+    "modalities": ["image"],
+    "height": 512,
+    "width": 512,
+    "num_inference_steps": 50,
+    "seed": 42
+}
+EOF
+
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @payload.json
+
+```
+
+### Image to Text (img2text)
+
+Generate text descriptions from images:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Describe this image in detail" \
+    --modality img2text \
+    --image-url /path/to/image.jpg
+```
+
+**Using curl**
+
+```bash
+IMAGE_BASE64=$(base64 -w 0 cat.jpg)
+
+cat <<EOF > payload.json
+{
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
+      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
+    ]
+  }],
+  "modalities": ["text"]
+}
+EOF
+
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @payload.json
+```
+
+### Text to Text (text2text)
+
+Pure text generation:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "What is the capital of France?" \
+    --modality text2text
+```
+
+**Using curl**
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
+    "modalities": ["text"]
+  }'
+```
+
+## FAQ
+
+- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
+
+```bash
+sudo apt update
+sudo apt install ffmpeg
+```
+
+- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
+
+| Stage               | VRAM                         |
+| :------------------ | :--------------------------- |
+| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
+| Stage-1 (DiT)       | **26.50 GiB**                |
+| Total               | **~42 GiB + KV Cache**       |
diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml
index 93d2b611a..ee2b17729 100644
--- a/vllm_omni/model_executor/stage_configs/bagel.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel.yaml
@@ -4,7 +4,7 @@ stage_args:
   - stage_id: 0
     stage_type: llm
     runtime:
-      devices: "0"  # Adjust as needed
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
@@ -38,7 +38,7 @@ stage_args:
   - stage_id: 1
     stage_type: diffusion
     runtime:
-      devices: "1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: dit
diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
index e378fe675..eddc844d5 100644
--- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
@@ -4,7 +4,7 @@ stage_args:
   - stage_id: 0
     stage_type: llm
     runtime:
-      devices: "0"  # Adjust as needed
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: thinker