Revert "[BugFix] ignore mm data from stages to async omni (#954)"

This reverts commit 6fb4058566.

.buildkite/pipeline.yml

@@ -29,7 +29,23 @@ steps:
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
-      - pytest -s -v tests/e2e/offline_inference/test_ovis_image.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Images API LoRA E2E"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -93,11 +109,28 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion Parallelism Test"
-    timeout_in_minutes: 25
+  - label: "Diffusion Sequence Parallelism Test"
+    timeout_in_minutes: 20
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Tensor Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
       - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -154,7 +187,10 @@ steps:
     depends_on: image-build
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_async_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
     agents:
       queue: "mithril-h100-pool"
     plugins:

.buildkite/scripts/simple_test.sh

@@ -47,9 +47,12 @@ 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.12.0
+"${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

.buildkite/test-template-amd-omni.j2

@@ -14,7 +14,7 @@
         soft_fail: false
         commands:
           - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
-          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --progress plain ."
+          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"
         env:

.pre-commit-config.yaml

@@ -29,7 +29,10 @@ repos:
         # only for staged files
 
   - repo: https://github.com/rhysd/actionlint
-    rev: v1.7.9
+    # v1.7.8+ sets `go 1.24.0` in go.mod, which older Go toolchains (and most
+    # current CI images) cannot parse. Pin to v1.7.7 until actionlint fixes the
+    # go.mod directive.
+    rev: v1.7.7
     hooks:
       - id: actionlint
         files: ^\.github/workflows/.*\.ya?ml$

README.md

@@ -15,7 +15,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
-
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -973,10 +973,10 @@ if __name__ == "__main__":
         "--base-url",
         type=str,
         default=None,
-        help="Base URL of the server (e.g., http://localhost:30000). Overrides host/port.",
+        help="Base URL of the server (e.g., http://localhost:8091). Overrides host/port.",
     )
     parser.add_argument("--host", type=str, default="localhost", help="Server host.")
-    parser.add_argument("--port", type=int, default=30000, help="Server port.")
+    parser.add_argument("--port", type=int, default=8091, help="Server port.")
     parser.add_argument("--model", type=str, default="default", help="Model name.")
     parser.add_argument(
         "--dataset",

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.12.0
+ARG VLLM_BASE_TAG=v0.14.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.rocm

@@ -1,9 +1,7 @@
-ARG BASE_IMAGE=rocm/vllm-dev:nightly_main_20251205
-FROM ${BASE_IMAGE}
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0
+FROM ${BASE_IMAGE} AS final
 
 ARG COMMON_WORKDIR=/app
-ARG VLLM_VERSION=v0.12.0
-ARG PYTORCH_ROCM_ARCH="gfx942;gfx950"
 
 WORKDIR ${COMMON_WORKDIR}
 
@@ -13,30 +11,22 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Step 2: Reinstall vllm from source
-RUN python3 -m pip uninstall -y vllm && rm -rf vllm &&\
-    git clone https://github.com/vllm-project/vllm.git && \
-    cd vllm && \
-    git checkout ${VLLM_VERSION} && \
-    python3 -m pip install -r requirements/rocm.txt && \
-    python3 setup.py clean --all && \
-    PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
-    cd ../ && \
-    rm -rf vllm/.git
-
 RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
 
-# Step 3: Copy vllm-omni code and install without uv
+# Step 2: Copy vllm-omni code and install without uv
 COPY . ${COMMON_WORKDIR}/vllm-omni
-RUN cd ${COMMON_WORKDIR}/vllm-omni && python3 -m pip install --no-cache-dir ".[dev]"
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+
+# When we are installing onnxruntime-rocm, we need to uninstall the system-installed onnxruntime first.
+# These are the dependencies of Qwen3-TTS.
+RUN uv pip uninstall onnxruntime --system && uv pip install --no-cache-dir onnxruntime-rocm sox --system
 
-# Create python symlink
-# `GPU_ARCHS` is an environment variable that is used to set the GPU archs for the AITER.
-# This is needed to prevent the AITER automatic GPU arch detection from failing on MI325X.
-# The AITER version used in this dockerfile has issues with handling
-# the GPU archs of MI325X (CI machine) correctly. So we manually set the GPU archs here.
-# We reuse AITER_ROCM_ARCH from the base image to avoid duplication.
-ENV GPU_ARCHS=${AITER_ROCM_ARCH}
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 
+CMD ["/bin/bash"]
+
 ENTRYPOINT []
+
+#Set entrypoint for vllm-openai official images
+FROM final AS vllm-openai
+ENTRYPOINT ["vllm", "serve", "--omni"]

docs/.nav.yml

@@ -14,6 +14,7 @@ nav:
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
+      - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
     - Online Serving:

docs/api/README.md

@@ -36,7 +36,6 @@ Input data structures for multi-modal inputs.
 
 Engine classes for offline and online inference.
 
-- [vllm_omni.diffusion.diffusion_engine.BackgroundResources][]
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
@@ -54,41 +53,16 @@ Engine classes for offline and online inference.
 
 Core scheduling and caching components.
 
+- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][]
 - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
+- [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
-
-## Model Executor
-
-Model execution components.
-
-- [vllm_omni.model_executor.custom_process_mixin.CustomProcessMixin][]
-- [vllm_omni.model_executor.models.output_templates.OmniOutput][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniConditionalGenerationMixin][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavBigVGANModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavDiTModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavForConditionalGenerationVLLM][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerSharedExpertWrapper][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.VectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.preprocess][]
 
 ## Configuration
 
@@ -98,6 +72,17 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2Config][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2DecoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1Config][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderBigVGANConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
 
 ## Workers
 

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.12.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 

docs/contributing/model/adding_diffusion_model.md

@@ -44,7 +44,7 @@ vllm_omni/
 
 # Step-by-step-implementation
 ## Step 1: Model Implementation
-The diffusion pipeline’s implementation follows **HuggingFace Diffusers**, and components that do not need modification can be imported directly.
+The diffusion pipeline’s implementation follows **HuggingFace Diffusers**.
 ### 1.1 Define the Pipeline Class
 Define the pipeline class, e.g., `QwenImageEditPipeline`, and initialize all required submodules, either from HuggingFace `diffusers` or custom implementations. In `QwenImageEditPipeline`, only `QwenImageTransformer2DModel` is re-implemented to support optimizations such as Ulysses-SP. When adding new models in the future, you can either reuse this re-implemented `QwenImageTransformer2DModel` or extend it as needed.
 
@@ -73,10 +73,75 @@ def get_qwen_image_edit_post_process_func(
 The forward function of `QwenImageEditPipeline` follows the HuggingFace `diffusers` design for the most part. The key differences are:
 + As described in the overview, arguments are passed through `OnniDiffusionRequest`, so we need to get user parameters from it accordingly.
 ```python
-prompt = req.prompt if req.prompt is not None else prompt
+prompt = req.prompt
 ```
 + pre/post-processing are handled by the framework elsewhere, so skip them.
 
+### 1.4 Replace some ops or layers in DiT component
+
+vLLM-Omni provides a set of optimized operators with better performance and built-in support for parallelism, including attention, rotary embeddings (RoPE), and linear layers.
+
+Below is an example showing how to replace standard Transformer attention and FFN layers with vLLM-Omni implementations:
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+class MyAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = Attention()
+        self.to_qkv = QKVParallelLinear()
+        self.to_out = RowParallelLinear()
+        self.rope = RotaryEmbedding(is_neox_style=False)
+
+    def forward(self, hidden_states):
+        qkv, _ = self.to_qkv(hidden_states)
+        q, k, v = qkv.split(...)
+        q, k = self.rope(...)
+        attn_output = self.attn(q, k, v)
+        output = self.to_out(attn_output)
+
+class MyFFN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = ColumnParallelLinear()
+        self.fc2 = RowParallelLinear()
+        self.act = F.gelu
+
+    def forward(self, hidden_states):
+        hidden, _ = self.fc1(hidden_states)
+        hidden = self.act(hidden)
+        output = self.fc2(hidden)
+        return output
+```
+
+In this example:
+
++ Attention uses vLLM-Omni’s optimized attention kernel together with parallel QKV projection and RoPE.
+
++ Linear layers are replaced with column- and row-parallel variants to enable tensor parallelism.
+
++ The FFN follows a standard two-layer structure and can be further optimized (e.g., using fused or merged projections) if needed.
+
+
+### 1.5 Provide a `_repeated_blocks` in DiT model
+`_repeated_blocks` is the small and frequently-repeated block(s) of a model -- typically a transformer layer.
+
+It's used for torch compile optimizations.
+```python
+_repeated_blocks = ["QwenImageTransformerBlock"]
+```
+
+
+### 1.6 (Optional) implement sequence parallelism
+vLLM-Omni has a non-intrusive `_sp_plan` that enable sequence parallel without modifying `forward()` logic.
+You can refer to [How to parallelize a new model](../../user_guide/diffusion/parallelism_acceleration.md)
+
+
+### 1.7 (Optional) integrate with Cache-Dit
+vLLM-Omni supports acceleration via [Cache-Dit](../../user_guide/diffusion/cache_dit_acceleration.md). Most models compatible with Diffusers can use Cache-Dit seamlessly. For new models, you can extend support by modifying`cache_dit_backend.py`
+
 ## Step 2: Extend OmniDiffusionRequest Fields
 User-provided inputs are ultimately passed to the model’s forward method through OmniDiffusionRequest, so we add the required fields here to support the new model.
 ```python
@@ -139,6 +204,21 @@ Key point for writing the example:
 
 + Save or display the generated results so users can validate the integration.
 
+## Step 5: Open a Pull Request
+
+When submitting a pull request to add support for a new model, please include the following information in the PR description:
+
++ Output verification: provide generation outputs to verify correctness and model behavior.
+
++ Inference speed: provide a comparison with the corresponding implementation in Diffusers.
+
++ Parallelism support: specify the supported parallel sizes and any relevant limitations.
+
++ Cache acceleration: check whether the model can be accelerated using Cache-Dit or not.
+
+
+Providing these details helps reviewers evaluate correctness, performance improvements, and parallel scalability of the new model integration.
+
 # Testing
 For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
 

docs/contributing/profiling.md

@@ -1,48 +1,30 @@
 # Profiling vLLM-Omni
 
-Performance Profiling Guidelines Profiling capabilities in vLLM-Omni are reserved for development and maintenance tasks aimed at temporal analysis of the codebase. Production use is **strongly discouraged**; enabling the profiler incurs a substantial overhead that negatively impacts inference latency.
+> **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
 
-**Mechanism**: vLLM-Omni implements cross-stage profiling via the PyTorch Profiler. To accommodate the architecture—where stages operate as distinct engine instances in separate processes—the profiling interface supports both holistic capturing (all stages) and targeted capturing (specific stages).
+vLLM-Omni uses the PyTorch Profiler to analyze performance across both **Multi-Stage LLMs** and **Diffusion Models**.
 
-**1. Enabling the Profiler**
+### 1. Set the Output Directory
+Before running any script, set this environment variable. The system detects this and automatically saves traces here.
 
-Before running your script, you must set the ```VLLM_TORCH_PROFILER_DIR``` environment
-variable.
-
-```Bash
-export VLLM_TORCH_PROFILER_DIR=/path/to/save/traces
+```bash
+export VLLM_TORCH_PROFILER_DIR=./profiles
 ```
 
-**Highly Recommended: Limit Profiling to a Single Iteration**  
-For most use cases (especially when profiling audio stages), you should limit the profiler to just **one iteration** to keep trace files small and readable.
+### 2. Start Profiling
 
+It is best to limit profiling to one iteration to keep trace files manageable.
 
 ```bash
 export VLLM_PROFILER_MAX_ITERS=1
 ```
 
-**2. Offline Inference**
-
-For offlinie processing using ```OmniLLM```, you can wrap your ```generate``` calls with ```start_profile``` and ```stop_profile()```.
-
-Basic Usage(All Stages)
-```Python
-from vllm_omni import OmniLLM
-
-omni_llm = OmniLLM.from_engine_args(engine_args)
-
-# Start profiling all active stages
-omni_llm.start_profile()
-
-outputs = omni_llm.generate(prompts, sampling_params)
-
-# Stop profiling and save traces
-omni_llm.stop_profile()
-```
-
 **Selective Stage Profiling**
 The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
 ```python
+# Profile all stages
+omni_llm.start_profile()
+
 # Only profile Stage 1
 omni_llm.start_profile(stages=[1])
 ```
@@ -52,13 +34,105 @@ omni_llm.start_profile(stages=[1])
 omni_llm.start_profile(stages=[0, 2])
 ```
 
+**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
+
+```python
+from vllm_omni import omni_llm
+
+profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
+# 1. Start profiling if enabled
+if profiler_enabled:
+    omni_llm.start_profile(stages=[0])
+
+# Initialize generator
+omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+
+total_requests = len(prompts)
+processed_count = 0
+
+# Main Processing Loop
+for stage_outputs in omni_generator:
+
+    # ... [Output processing logic for text/audio would go here] ...
+
+    # Update count to track when to stop profiling
+    processed_count += len(stage_outputs.request_output)
+
+    # 2. Check if all requests are done to stop the profiler safely
+    if profiler_enabled and processed_count >= total_requests:
+        print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
+
+        # Stop the profiler while workers are still active
+        omni_llm.stop_profile()
+
+        # Wait for traces to flush to disk
+        print("[Info] Waiting 30s for workers to write trace files to disk...")
+        time.sleep(30)
+        print("[Info] Trace export wait time finished.")
+
+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)
 
 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)
 
-**3. Online Inference(Async)**
+
+**For Diffusion Models as a single stage**
+
+Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
+
+**CLI Usage:**
+```python
+
+python image_to_video.py \
+    --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+    --image qwen-bear.png \
+    --prompt "A cat playing with yarn, smooth motion" \
+    \
+    # Minimize Spatial Dimensions (Optional but helpful):
+    #    Drastically reduces memory usage so the profiler doesn't
+    #    crash due to overhead, though for accurate performance
+    #    tuning you often want target resolutions.
+    --height 48 \
+    --width 64 \
+    \
+    # Minimize Temporal Dimension (Frames):
+    #    Video models process 3D tensors (Time, Height, Width).
+    #    Reducing frames to the absolute minimum (2) keeps the
+    #    tensor size small, ensuring the trace file doesn't become
+    #    multi-gigabytes in size.
+    --num_frames 2 \
+    \
+    # Minimize Iteration Loop (Steps):
+    #    This is the most critical setting for profiling.
+    #    Diffusion models run the same loop X times.
+    #    Profiling 2 steps gives you the exact same performance
+    #    data as 50 steps, but saves minutes of runtime and
+    #    prevents the trace viewer from freezing.
+    --num_inference_steps 2 \
+    \
+    --guidance_scale 5.0 \
+    --guidance_scale_high 6.0 \
+    --boundary_ratio 0.875 \
+    --flow_shift 12.0 \
+    --fps 16 \
+    --output i2v_output.mp4
+
+```
+
+**Examples**:
+
+1. **Qwen image edit**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.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)
+
+
+**Online Inference(Async)**
 
 For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
 
@@ -77,23 +151,16 @@ async for output in async_omni.generate(prompt, sampling_params, request_id):
 await async_omni.stop_profile()
 ```
 
-**4. Analyzing Omni Traces**
+### 3. Analyzing Omni Traces
 
-After ``stop_profile()`` completes (and the file write wait time has passed), the directory specified in ```VLLM_TORCH_PROFILER_DIR``` will contain the trace files.
+Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
 
-```
-Output/
-│── ...rank-0.pt.trace.json.gz   # GPU 0 trace (TP=2 Example)
-│── ...rank-1.pt.trace.json.gz   # GPU 1 trace (TP=2 Example)
-│       # Load these into Perfetto to visualize synchronization
-│
-│── profiler_out_.txt            # Summary tables (CPU/CUDA time %)
-```
+**Output**
+**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
 
 **Viewing Tools:**
-     - [Perfetto](https://ui.perfetto.dev/): (Recommended): Best for handling large audio trace files.
-     - ```chrome://tracing```: Good for smaller text-only traces.
 
+- [Perfetto](https://ui.perfetto.dev/)(recommended)
+- ```chrome://tracing```(Chrome only)
 
-**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM.  
-For more advanced configuration options (memory profiling, custom activities, etc.), 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/latest/dev/profiling.html)

docs/design/architecture_overview.md

@@ -192,4 +192,4 @@ curl -sS -X POST http://localhost:8091/v1/chat/completions \
 }
 ```
 
-For more usages, please refer to [examples](../examples/README.md).
+For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).

docs/design/feature/disaggregated_inference.md

@@ -1,27 +1,35 @@
 # Disaggregated Inference for Omni-Modality Models
 
-This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines.
+This guide explains how to configure and use distributed connectors
+(`vllm_omni/distributed/omni_connectors`) in vllm-omni for multi-stage pipelines.
 
-## 1. Overview
+Backend-specific setup lives in separate docs:
+
+- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
+- [MooncakeConnector](omni_connectors/mooncake_connector.md)
+- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
+
+## Overview
 
 Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker).
-Currently supported connectors operate in **D2H2D (Device-to-Host-to-Device)** mode:
-1. **SharedMemoryConnector**: Uses system shared memory.
-2. **MooncakeConnector**: Uses [Mooncake](https://github.com/kvcache-ai/Mooncake).
+Current connectors operate in D2H2D (device to host to device) mode.
 
-*   **SharedMemoryConnector (Default)**: Zero-copy (on host), lowest latency. Best for **single-node** deployments. Auto-configured if no connectors are specified.
-*   **MooncakeConnector**: TCP/RDMA based. Best for **multi-node** distributed deployments. Requires a Mooncake Master service.
+## Connector Choices
 
-## 2. API Design
+| Use Case | Recommended Connector | Notes |
+| :--- | :--- | :--- |
+| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
+| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
+| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
 
-The connector system is built around the `OmniConnectorBase` abstraction, which decouples data transport from stage logic.
+## Core API
 
-### Core Interface
+The connector system is built around `OmniConnectorBase`.
 
 ```python
 class OmniConnectorBase(ABC):
     @abstractmethod
-    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]:
+    def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, Optional[dict]]:
         """
         Store data.
         Returns: (success, serialized_size, metadata)
@@ -29,145 +37,23 @@ class OmniConnectorBase(ABC):
         pass
 
     @abstractmethod
-    def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
+    def get(self, from_stage: str, to_stage: str, get_key: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
         """
         Retrieve data.
-        Args: metadata - Transport-specific handles returned by put() (e.g., SHM name).
+        Args: metadata - transport-specific handles returned by put() (e.g., SHM name).
         Returns: (object, serialized_size)
         """
         pass
 ```
 
-### Key Concept: Metadata Passing
-Unlike a pure key-value store, some connectors (like `SharedMemoryConnector`) generate transient resources (e.g., a shared memory block name) during `put()`. This `metadata` **must be passed via the control plane** (e.g., HTTP headers, queue messages) from the producer stage to the consumer stage so `get()` can locate the data.
+### Metadata Passing
 
-## 3. Backends & Use Cases
+Some connectors (e.g., SharedMemoryConnector) generate transient resources during `put()`.
+This `metadata` must be passed through the control plane so `get()` can locate the data.
 
-### 3.1 SharedMemoryConnector
-**Best for:** Single-node, high-performance IPC.
+## Configuration Model
 
-*   **Mechanism:**
-    *   **Small Payloads (< Threshold)**: Data is serialized and passed directly "inline" within the `metadata` dictionary. This avoids the overhead of creating SHM blocks for tiny messages.
-    *   **Large Payloads (>= Threshold)**: Data is written to a named System Shared Memory block. The block name is returned in `metadata`.
-*   **Configuration:**
-    *   `shm_threshold_bytes`: Size in bytes to switch from inline to SHM (default: 64KB).
-
-### 3.2 MooncakeConnector
-**Best for:** Multi-node distributed inference.
-
-*   **Mechanism:** Uses Mooncake's distributed KVCache store.
-    *   **Data Plane**: TCP or RDMA for high-bandwidth transfer.
-    *   **Control Plane**: Uses a centralized Mooncake Master and Metadata Server.
-    *   **Keying**: Deterministic keys based on `request_id/from_stage_to_stage`.
-*   **Requirements**: Requires a running Mooncake Master service.
-
-## 4. Relationship with vLLM
-
-vLLM provides specialized distributed mechanisms for specific artifacts:
-*   **KV Transfer** (`vllm.distributed.kv_transfer`): Optimized for transferring KV caches between prefill and decode instances (using NCCL, Mooncake, etc.).
-*   **EC Transfer** (`vllm.distributed.ec_transfer`): Optimized for sharing encoder embeddings.
-*   **Device Communicators** (`vllm.distributed.device_communicators`): Low-level primitives (NCCL, SHM) for Tensor/Pipeline Parallelism.
-
-`vllm-omni` complements this by introducing a **Generalized Connector Abstraction** (`OmniConnector`) for multimodal pipelines. While vLLM's connectors are artifact-specific, `vllm-omni`:
-
-1.  **Unifies Transport**: Provides a single API (`put`/`get`) to transport *any* stage artifact (Input Embeddings, Hidden States, Audio/Image Tensors, KV Cache, Final Output) between arbitrary pipeline stages (e.g., AudioEncoder -> LLM -> AudioGenerator).
-2.  **Extends Connectivity**: Enables flexible construction of complex DAGs (Directed Acyclic Graphs) where stages can run in the same process, same node, or across nodes, using the most appropriate backend (SHM, Mooncake, etc.) for each edge.
-3.  **Wraps & Adapts**: Can internally utilize vLLM's specialized `kv_transfer` for KV paths while using generic transports (SHM/Mooncake) for other data types, presenting a consistent interface to the application layer.
-
-## 5. Installation (Mooncake)
-
-If using `MooncakeConnector`, install the library first:
-
-```bash
-# For CUDA-enabled systems (Recommended)
-pip install mooncake-transfer-engine
-
-# For non-CUDA systems
-pip install mooncake-transfer-engine-non-cuda
-```
-
-## 6. Using MooncakeConnector
-
-### 6.1 Start Mooncake Master
-
-Start the master service on your primary node:
-
-```bash
-# if you use mooncake SSD storage
-mkdir -p ./mc_storage
-
-mooncake_master \
-  --rpc_port=50051 \
-  --enable_http_metadata_server=true \
-  --http_metadata_server_host=0.0.0.0 \
-  --http_metadata_server_port=8080 \
-  --metrics_port=9003 \
-  --root_fs_dir=./mc_storage/ \
-  --cluster_id=mc-local-1 &
-```
-
-### 6.2 Configuration (YAML)
-
-Edit your stage config (e.g., `qwen2_5_omni.yaml`).
-
-**Step 1: Define Connector in Global Runtime**
-
-```yaml
-runtime:
-  connectors:
-    connector_of_mooncake:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"           # Local Worker IP
-        metadata_server: "http://<MASTER_IP>:8080/metadata"
-        master: "<MASTER_IP>:50051"
-        segment: 512000000          # 512MB segment
-        localbuf: 64000000          # 64MB buffer
-        proto: "tcp"                # "tcp" or "rdma"
-    ```
-
-**Mooncake Configuration Parameters:**
-
-*   **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it.
-*   **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA).
-*   **master**: The address of the Mooncake Master Server (e.g., `<MASTER_IP>:50051`). This is used for global state management and control plane operations.
-*   **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer.
-*   **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations.
-*   **proto**: The transport protocol to use. Options:
-    *   `tcp`: Standard TCP/IP (easier setup, universal compatibility).
-    *   `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware).
-
-For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake).
-
-    **Step 2: Reference in Stages**
-
-Explicitly link stages using `input_connectors` and `output_connectors`:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    # ...
-    output_connectors:
-      to_stage_1: connector_of_mooncake
-
-  - stage_id: 1
-    # ...
-    input_connectors:
-      from_stage_0: connector_of_mooncake
-```
-
-## 7. Using SharedMemoryConnector (Auto-Mode)
-
-**Best for single-node.**
-
-The system will automatically create `SharedMemoryConnector`s for any pipeline edge that does not have an explicit connector defined. This is inferred from:
-1.  `runtime.edges` list in the config.
-2.  `engine_input_source` dependencies defined in `stage_args`.
-
-### Threshold Configuration
-By default, payloads larger than **64KB** (default threshold) are transferred via shared memory, while smaller ones use the control queue (inline).
-
-To adjust this threshold (e.g., to 1GB), add the following to your `runtime.connectors`:
+Define connectors in runtime:
 
 ```yaml
 runtime:
@@ -175,26 +61,46 @@ runtime:
     connector_of_shared_memory:
       name: SharedMemoryConnector
       extra:
-        shm_threshold_bytes: 1024 # 1KB threshold
+        shm_threshold_bytes: 65536
 ```
 
-## 8. Summary
+Wire stages to connectors:
 
-| Use Case | Recommended Connector | Configuration |
-| :--- | :--- | :--- |
-| **Single Node** | `SharedMemoryConnector` | **None** (Automatic) or Custom Threshold |
-| **Multi Node** | `MooncakeConnector` | Explicit YAML + Mooncake Master |
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
 
-## 9. Operational Notes (important)
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+```
 
-- **Fail-fast config validation**: the loader raises if any expected edge is missing a connector. Define `input_connectors`/`output_connectors` or rely on auto-SHM filling; otherwise startup aborts.
-- **Missing payloads halt the stage**: workers expect connector payloads; if metadata or connector config is missing, the stage raises and stops. Verify connector wiring and metadata propagation before production.
+If a pipeline edge has no explicit connector, the system auto-creates a
+SharedMemoryConnector for that edge.
 
-## 10. Future Roadmap: Device-to-Device (D2D) Transport
+## Relationship with vLLM
 
-The current implementations (`SharedMemoryConnector`, `MooncakeConnector`) utilize a **D2H2D (Device-to-Host-to-Device)** data path. Tensors are moved to CPU memory (Host) for transport, which incurs PCIe overhead.
+vLLM provides specialized distributed mechanisms for specific artifacts:
 
-As outlined in the design RFC, future versions will introduce **D2D (Device-to-Device)** connectors:
+- KV Transfer (`vllm.distributed.kv_transfer`): optimized for KV caches.
+- EC Transfer (`vllm.distributed.ec_transfer`): optimized for encoder embeddings.
+- Device Communicators (`vllm.distributed.device_communicators`): low-level primitives (NCCL, SHM).
 
-*   **Goal**: Direct GPU-to-GPU transfer (via NCCL, UCX, or IPC) to minimize latency for large tensor payloads.
-*   **Mechanism**: The `OmniConnector` API allows `put()` to initiate a transfer and return a lightweight handle (metadata) via the control plane, while the heavy payload flows directly between devices.
+vllm-omni complements this with a generalized connector abstraction:
+
+1. Unifies transport via a single `put`/`get` API for any stage artifact.
+2. Enables DAG-style pipelines across processes or nodes with per-edge transports.
+3. Can wrap vLLM-specific transfers for KV paths while keeping a consistent interface.
+
+## Operational Notes
+
+- Fail-fast config validation: missing expected edges cause startup failures.
+- Missing payloads halt stages: verify connector wiring and metadata propagation.
+
+## Future Roadmap: D2D Transport
+
+Current connectors use D2H2D paths. Future versions will introduce direct
+device-to-device connectors (NCCL, UCX, IPC) to reduce latency for large
+tensor payloads.

docs/design/feature/omni_connectors/mooncake_connector.md

@@ -0,0 +1,74 @@
+# MooncakeConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Mooncake.
+
+## Installation
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+## Start Mooncake Master
+
+```bash
+# If you use Mooncake SSD storage
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_mooncake:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://<MASTER_IP>:8080/metadata"
+        master: "<MASTER_IP>:50051"
+        segment: 512000000
+        localbuf: 64000000
+        proto: "tcp"
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_mooncake
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_mooncake
+```
+
+Parameters:
+
+- host: local worker IP registered in the metadata server.
+- metadata_server: metadata server URL for discovery and setup.
+- master: Mooncake Master address.
+- segment: global memory segment size in bytes.
+- localbuf: local buffer size in bytes.
+- proto: transport protocol ("tcp" or "rdma").
+
+For more details, refer to the
+[Mooncake repository](https://github.com/kvcache-ai/Mooncake).

docs/design/feature/omni_connectors/shared_memory_connector.md

@@ -0,0 +1,27 @@
+# SharedMemoryConnector
+
+## When to Use
+
+Best for single-node deployments where stages run on the same host. It is
+auto-configured when no explicit connector is specified for an edge.
+
+## How It Works
+
+- Small payloads (< threshold): serialized and passed inline in metadata.
+- Large payloads (>= threshold): stored in shared memory; the SHM name is
+  returned in metadata.
+
+## Configuration
+
+```yaml
+runtime:
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+```
+
+## Notes
+
+- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.

docs/design/feature/omni_connectors/yuanrong_connector.md

@@ -0,0 +1,101 @@
+# YuanrongConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Yuanrong Datasystem.
+
+## Mechanism
+
+Uses Yuanrong Datasystem's distributed KV store (`datasystem.kv_client`).
+
+- Data Plane: TCP or RDMA for high-bandwidth transfer.
+- Control Plane: Yuanrong Datasystem workers and etcd.
+- Keying: deterministic keys based on `put_key` (often composed as `request_id:fromStage_toStage`).
+
+## Installation
+
+```bash
+pip install openyuanrong-datasystem
+```
+
+## Start etcd
+
+```bash
+# Download and install etcd (v3.5.12 or higher)
+ETCD_VERSION="v3.5.12"
+ETCD_ARCH="linux-arm64"
+wget https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+tar -xvf etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+cd etcd-${ETCD_VERSION}-${ETCD_ARCH}
+sudo cp etcd etcdctl /usr/local/bin/
+
+# Start etcd
+etcd \
+  --name etcd-single \
+  --data-dir /tmp/etcd-data \
+  --listen-client-urls http://0.0.0.0:2379 \
+  --advertise-client-urls http://0.0.0.0:2379 \
+  --listen-peer-urls http://0.0.0.0:2380 \
+  --initial-advertise-peer-urls http://0.0.0.0:2380 \
+  --initial-cluster etcd-single=http://0.0.0.0:2380 &
+
+# Verify etcd is running
+etcdctl --endpoints "127.0.0.1:2379" put key "value"
+etcdctl --endpoints "127.0.0.1:2379" get key
+```
+
+For production environments, refer to the
+[official etcd clustering documentation](https://etcd.io/docs/current/op-guide/clustering/).
+
+## Start Datasystem Worker
+
+```bash
+# Replace ${ETCD_IP} with etcd node IP, ${WORKER_IP} with local node IP
+dscli start -w \
+  --worker_address "${WORKER_IP}:31501" \
+  --etcd_address "${ETCD_IP}:2379" \
+  --shared_memory_size_mb 20480
+```
+
+To stop the worker:
+
+```bash
+dscli stop --worker_address "${WORKER_IP}:31501"
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_yuanrong:
+      name: YuanrongConnector
+      extra:
+        host: "127.0.0.1"
+        port: 31501
+        get_sub_timeout_ms: 1000
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_yuanrong
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_yuanrong
+```
+
+Parameters:
+
+- host: datasystem worker host.
+- port: datasystem worker port.
+- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
+
+For more details, refer to the
+[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).

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

@@ -17,12 +17,12 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]
 
 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1 and 0.12.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.12.0 --torch-backend=auto
+uv pip install vllm==0.14.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -36,13 +36,13 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.12.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
 
 ```bash
-uv pip install vllm==0.12.0 --torch-backend=auto
+uv pip install vllm==0.14.0 --torch-backend=auto
 ```
 
-The release 0.12.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,11 +58,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.12.0
+git checkout v0.14.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.12.0/vllm-0.12.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.14.0/vllm-0.14.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.12.0rc1 \
+    vllm/vllm-omni:v0.14.0rc1 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -23,19 +23,30 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm-omni-rocm .
 ```
 
-If you want to specify which GPU Arch to build for to cutdown build time:
-
-```bash
-DOCKER_BUILDKIT=1 docker build \
-  -f docker/Dockerfile.rocm \
-  --build-arg PYTORCH_ROCM_ARCH="gfx942;gfx950" \
-  -t vllm-omni-rocm .
-```
-
 #### Launch the docker image
 
+##### Launch with OpenAI API Server
+
 ```
-docker run -it \
+docker run --rm \
+--group-add=video \
+--ipc=host \
+--cap-add=SYS_PTRACE \
+--security-opt seccomp=unconfined \
+--device /dev/kfd \
+--device /dev/dri \
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--env "HF_TOKEN=$HF_TOKEN" \
+-p 8091:8091 \
+--ipc=host \
+vllm-omni-rocm \
+--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+```
+
+##### Launch with interactive session for development
+
+```
+docker run --rm -it \
 --network=host \
 --group-add=video \
 --ipc=host \
@@ -44,8 +55,9 @@ docker run -it \
 --device /dev/kfd \
 --device /dev/dri \
 -v <path/to/model>:/app/model \
-vllm-omni-rocm \
-bash
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--entrypoint bash \
+vllm-omni-rocm
 ```
 
 # --8<-- [end:build-docker]
@@ -57,8 +69,7 @@ vLLM-Omni offers an official docker image for deployment. These images are built
 #### Launch vLLM-Omni Server
 Here's an example deployment command that has been verified on 2 x MI300's:
 ```bash
-docker run -it \
-  --network=host \
+docker run --rm \
   --group-add=video \
   --ipc=host \
   --cap-add=SYS_PTRACE \
@@ -68,14 +79,15 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  vllm/vllm-omni-rocm:v0.12.0rc1 \
-  vllm serve --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+  -p 8091:8091 \
+  vllm/vllm-omni-rocm:v0.14.0rc1 \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
 #### Launch an interactive terminal with prebuilt docker image.
 If you want to run in dev environment you can launch the docker image as follows:
 ```bash
-docker run -it \
+docker run --rm -it \
   --network=host \
   --group-add=video \
   --ipc=host \
@@ -86,8 +98,8 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  vllm/vllm-omni-rocm:v0.12.0rc1 \
-  bash
+  --entrypoint bash \
+  vllm/vllm-omni-rocm:v0.14.0rc1
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/installation/npu/npu.inc.md

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -42,7 +42,7 @@ source ~/.bashrc
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.12.0rc1 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn

docs/getting_started/quickstart.md

@@ -17,7 +17,13 @@ For installation on GPU from source:
 ```bash
 uv venv --python 3.12 --seed
 source .venv/bin/activate
-uv pip install vllm==0.12.0 --torch-backend=auto
+
+# On CUDA
+uv pip install vllm==0.14.0 --torch-backend=auto
+
+# On ROCm
+uv pip install vllm==0.14.0 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .
@@ -40,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

docs/mkdocs/hooks/generate_api_readme.py

@@ -34,10 +34,10 @@ CATEGORIES = {
         "name": "Core",
         "description": "Core scheduling and caching components.",
     },
-    "model_executor": {
-        "name": "Model Executor",
-        "description": "Model execution components.",
-    },
+    # "model_executor": {
+    #     "name": "Model Executor",
+    #     "description": "Model execution components.",
+    # },
     "config": {
         "name": "Configuration",
         "description": "Configuration classes.",

docs/models/supported_models.md

@@ -33,7 +33,11 @@ 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` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
 
 
 ## List of Supported Models for NPU

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),
 )
 ```
 
@@ -69,7 +70,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
 
 ## Acceleration Methods
 
-For comprehensive illustration, please view cache-dit [User_Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
+For comprehensive illustration, please view cache-dit [User_Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
 
 ### 1. DBCache (Dual Block Cache)
 
@@ -222,6 +223,6 @@ To test another model, you can modify `--model` with the target model identifier
 
 ## Additional Resources
 
-- [Cache-DiT User Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
-- [Cache-DiT Benchmark](https://github.com/vipshop/cache-dit/tree/main/bench)
-- [DBCache Technical Details](https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md)
+- [Cache-DiT User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
+- [Cache-DiT Benchmark](https://cache-dit.readthedocs.io/en/latest/benchmark/HYBRID_CACHE/)
+- [DBCache Technical Details](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -23,12 +23,24 @@ The following table shows which models are currently supported by parallelism me
 | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
 | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
 | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
+| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
+
+
+!!! note "TP Limitations for Diffusion Models"
+    We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
+
+    - Good news: The text_encoder typically has minimal impact on overall inference performance.
+    - Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
+
+    We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
+
 
 !!! note "Why Z-Image is TP=2 only"
     Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
@@ -38,7 +50,7 @@ The following table shows which models are currently supported by parallelism me
 
 | Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
 |-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ |
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ❌ |
 
 ### Tensor Parallelism
 
@@ -56,10 +68,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,
+    ),
 )
 ```
 
@@ -72,6 +86,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
 
@@ -80,7 +95,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.
@@ -122,6 +140,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
 
@@ -130,7 +149,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.
@@ -172,6 +194,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
@@ -180,7 +203,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
@@ -210,100 +236,139 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 ##### How to parallelize a new model
 
-If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509).
+NOTE: "Terminology: SP vs CP"
+    Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py).
+    We use "Sequence Parallelism" to align with vLLM-Omni's terminology.
 
+---
 
-This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically:
+###### Non-intrusive `_sp_plan` (Recommended)
 
-- Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**.
-- Keep **text embeddings** (`encoder_hidden_states`) **replicated** on all SP ranks for correctness (unless you implement full joint-SP semantics and explicitly split text too).
-- Chunk **image RoPE freqs** to match the chunked image tokens.
-- `all_gather` the final output back along the sequence dimension.
+The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
 
-First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`:
+**Requirements for `forward()` function:**
 
-- Chunk `hidden_states` by SP world size.
-- Set `get_forward_context().split_text_embed_in_sp = False` because Qwen-Image uses joint attention and we would use a full text embedding in each rank.
-- After `pos_embed`, chunk `img_freqs` on `dim=0` (token axis) to match the chunked image tokens; only chunk `txt_freqs` if you explicitly split text embeddings in SP.
+- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations)
+- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule**
 
-Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example:
+**When to create a submodule:**
 
-```diff
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_sequence_parallel_rank,
-    get_sequence_parallel_world_size,
-    get_sp_group,
+```python
+# ❌ BAD: Inline operations - hooks cannot intercept
+def forward(self, x, cap_feats):
+    unified = torch.cat([x, cap_feats], dim=1)  # Cannot be sharded via _sp_plan
+    ...
+
+# ✅ GOOD: Extract into a submodule
+class UnifiedPrepare(nn.Module):
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)  # Now can be sharded via _sp_plan
+
+class MyModel(nn.Module):
+    def __init__(self):
+        self.unified_prepare = UnifiedPrepare()  # Submodule
+
+    def forward(self, x, cap_feats):
+        unified = self.unified_prepare(x, cap_feats)  # Hook can intercept here
+```
+
+---
+
+###### Defining `_sp_plan`
+
+**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference):
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # Corresponds to diffusers' ContextParallelInput
+    SequenceParallelOutput,  # Corresponds to diffusers' ContextParallelOutput
 )
-from vllm_omni.diffusion.forward_context import get_forward_context
-
-class QwenImageTransformer2DModel(...):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
-        encoder_hidden_states_mask: torch.Tensor = None,
-        timestep: torch.LongTensor = None,
-        img_shapes: list[tuple[int, int, int]] | None = None,
-        txt_seq_lens: list[int] | None = None,
-        ...
-    ):
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           # Chunk image tokens along sequence dimension.
-+           hidden_states = torch.chunk(
-+               hidden_states,
-+               get_sequence_parallel_world_size(),
-+               dim=-2,
-+           )[get_sequence_parallel_rank()]
-+
-+           # Qwen-Image uses *dual-stream* (text + image) and runs *joint attention*.
-+           # Text embeddings should be replicated across SP ranks for correctness.
-+           get_forward_context().split_text_embed_in_sp = False
-
-        hidden_states = self.img_in(hidden_states)
-
-        ...
-        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
-
-+       def get_rotary_emb_chunk(freqs):
-+           return torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[
-+               get_sequence_parallel_rank()
-+           ]
-+
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           img_freqs, txt_freqs = image_rotary_emb
-+           img_freqs = get_rotary_emb_chunk(img_freqs)
-+           if get_forward_context().split_text_embed_in_sp:
-+               txt_freqs = get_rotary_emb_chunk(txt_freqs)
-+           image_rotary_emb = (img_freqs, txt_freqs)
 ```
 
-Next, at the end of the `forward()` function, call `get_sp_group().all_gather` to gather the chunked outputs across devices and concatenate them along the sequence dimension (matching the earlier `dim=-2` chunking):
+| Parameter | Description |
+|-----------|-------------|
+| `split_dim` | Dimension to split/gather (usually `1` for sequence) |
+| `expected_dims` | Expected tensor rank for validation (optional) |
+| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors |
+| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) |
 
-```diff
-class QwenImageTransformer2DModel(...):
-    def forward(...):
-        ...
-        hidden_states = self.norm_out(hidden_states, temb)
-        output = self.proj_out(hidden_states)
+**Key naming convention:**
 
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           output = get_sp_group().all_gather(output, dim=-2)
-        return Transformer2DModelOutput(sample=output)
+| Key | Meaning | Python equivalent |
+|-----|---------|-------------------|
+| `""` | Root model | `model` |
+| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` |
+| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` |
+| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` |
+
+**Dictionary key types:**
+
+| Key type | `split_output` | Description |
+|----------|----------------|-------------|
+| `"param_name"` (str) | `False` | Shard **input parameter** by name |
+| `0`, `1` (int) | `True` | Shard **output tuple** by index |
+
+**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)):
+
+```python
+class MyTransformer(nn.Module):
+    _sp_plan = {
+        # Shard rope module OUTPUTS (returns tuple)
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
+        },
+        # Shard transformer block INPUT parameter
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at final projection
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
 ```
 
-Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with:
-```diff
-from vllm_omni import Omni
-+from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
+---
 
-omni = Omni(
-    model="Qwen/Qwen-Image",
-+    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
-)
+###### Hook flow
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
 ```
+Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ...
+                                                                              ↓
+... → [SequenceParallelGatherHook: post_forward] → Output
+```
+
+1. **SplitHook** shards tensors before/after the target module
+2. **Attention layers** handle Ulysses/Ring communication internally
+3. **GatherHook** collects sharded outputs
+
+The framework automatically applies these hooks when `sequence_parallel_size > 1`.
+
+---
+
+###### Method 2: Intrusive modification (For complex cases)
+
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`:
+
+```python
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+
+def forward(self, hidden_states, ...):
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_shard(hidden_states, dim=1)
+        # ... computation ...
+        output = sp_gather(output, dim=1)
+    return output
+```
+
+---
+
+###### Choosing the right approach
+
+| Scenario | Approach |
+|----------|----------|
+| Standard transformer | `_sp_plan` |
+| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` |
+| Dynamic/conditional sharding | Intrusive modification |
 
 
 ### CFG-Parallel
@@ -324,11 +389,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,
+    ),
 )
 ```
 

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
@@ -112,6 +124,7 @@ th {
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
+| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
 
 ### VideoGen
 

docs/user_guide/diffusion_acceleration.md

@@ -47,9 +47,9 @@ The following table shows which models are currently supported by each accelerat
 | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |❌ | ❌ | ❌ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ |
 
 ### VideoGen
 
@@ -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

docs/user_guide/examples/offline_inference/lora_inference.md

@@ -0,0 +1,107 @@
+# LoRA-Inference
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/lora_inference>.
+
+This contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
+
+## Overview
+
+Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
+
+- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
+- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
+
+Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
+
+## Usage
+
+### Pre-loaded LoRA (via --lora-path)
+
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_preloaded.png
+```
+
+**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
+
+### Per-request LoRA (via --lora-request-path)
+
+Load a LoRA adapter on-demand for each request:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-request-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_per_request.png
+```
+
+### No LoRA
+
+If no LoRA request is provided, we will use the base model without any LoRA adapters:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_no_lora.png
+```
+
+## Parameters
+
+### LoRA Parameters
+
+- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
+- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
+- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### Standard Parameters
+
+- `--prompt`: Text prompt for image generation (required)
+- `--seed`: Random seed for reproducibility (default: 42)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--output`: Output file path (default: `lora_output.png`)
+
+## How LoRA Works
+
+All LoRA adapters are handled uniformly:
+
+1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
+2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
+3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
+
+The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
+
+## LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
+## Example materials
+
+??? abstract "lora_inference.py"
+    ``````py
+    --8<-- "examples/offline_inference/lora_inference/lora_inference.py"
+    ``````

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -0,0 +1,94 @@
+# Qwen3-TTS Offline Inference
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_tts>.
+
+
+This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.
+
+## Model Overview
+
+Qwen3 TTS provides multiple task variants for speech generation:
+
+- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
+- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
+- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
+
+## Setup
+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.
+
+## Quick Start
+
+Run a single sample for a task:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Generated audio files are saved to `output_audio/` by default.
+
+## Task Usage
+
+### CustomVoice
+
+Single sample:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Batch sample (multiple prompts in one run):
+
+```
+python end2end.py --query-type CustomVoice --use-batch-sample
+```
+
+### VoiceDesign
+
+Single sample:
+
+```
+python end2end.py --query-type VoiceDesign
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type VoiceDesign --use-batch-sample
+```
+
+### Base (Voice Clone)
+
+Single sample:
+
+```
+python end2end.py --query-type Base
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type Base --use-batch-sample
+```
+
+Mode selection for Base:
+
+- `--mode-tag icl` (default): standard mode
+- `--mode-tag xvec_only`: enable `x_vector_only_mode` in the request
+
+Examples:
+
+```
+python end2end.py --query-type Base --mode-tag icl
+```
+
+## Notes
+
+- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
+- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen3_tts/end2end.py"
+    ``````

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

docs/user_guide/examples/online_serving/lora_inference.md

@@ -0,0 +1,69 @@
+# LoRA-Inference
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/lora_inference>.
+
+This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
+
+> Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+> Note: This example uses `/v1/chat/completions`. LoRA payloads for other OpenAI endpoints are not implemented here.
+
+## Start Server
+
+```bash
+# Pick a diffusion model (examples)
+# export MODEL=stabilityai/stable-diffusion-3.5-medium
+# export MODEL=Qwen/Qwen-Image
+
+bash run_server.sh
+```
+
+## Call API (curl)
+
+```bash
+# Required: local LoRA folder on the server
+export LORA_PATH=/path/to/lora_adapter
+
+# Optional
+export SERVER=http://localhost:8091
+export PROMPT="A piece of cheesecake"
+export LORA_NAME=my_lora
+export LORA_SCALE=1.0
+# Optional: if omitted, the server derives a stable id from LORA_PATH.
+# export LORA_INT_ID=123
+
+bash run_curl_lora_inference.sh
+```
+
+## Call API (Python)
+
+```bash
+python openai_chat_client.py \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora_adapter \
+  --lora-name my_lora \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+## LoRA Format
+
+LoRA adapters should be in PEFT format, for example:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/lora_inference/openai_chat_client.py"
+    ``````
+??? abstract "run_curl_lora_inference.sh"
+    ``````py
+    --8<-- "examples/online_serving/lora_inference/run_curl_lora_inference.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````py
+    --8<-- "examples/online_serving/lora_inference/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -15,6 +15,12 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
+If you want to open async chunking for qwen3-omni, launch the server with command below
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file

docs/user_guide/examples/online_serving/text_to_image.md

@@ -119,6 +119,7 @@ Use `extra_body` to pass generation parameters:
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |
+| `--cfg-parallel-size`.   | int   | 1       | Number of GPUs for CFG parallelism |
 
 ## Response Format
 

examples/offline_inference/bagel/end2end.py

@@ -0,0 +1,189 @@
+import argparse
+import os
+from typing import cast
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        default="ByteDance-Seed/BAGEL-7B-MoT",
+        help="Path to merged model directory.",
+    )
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    parser.add_argument("--prompt_type", default="text", choices=["text"])
+
+    parser.add_argument(
+        "--modality",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Modality mode to control stage execution.",
+    )
+
+    parser.add_argument(
+        "--image-path",
+        type=str,
+        default=None,
+        help="Path to input image for img2img.",
+    )
+
+    # OmniLLM init args
+    parser.add_argument("--enable-stats", action="store_true", default=False)
+    parser.add_argument("--init-sleep-seconds", type=int, default=20)
+    parser.add_argument("--batch-timeout", type=int, default=5)
+    parser.add_argument("--init-timeout", type=int, default=300)
+    parser.add_argument("--shm-threshold-bytes", type=int, default=65536)
+    parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"])
+    parser.add_argument("--ray-address", type=str, default=None)
+    parser.add_argument("--stage-configs-path", type=str, default=None)
+    parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
+
+    args = parser.parse_args()
+    return 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()]
+            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 not prompts:
+        # Default prompt for text2img test if none provided
+        prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        print(f"[Info] No prompts provided, using default: {prompts}")
+    omni_outputs = []
+
+    from PIL import Image
+
+    if args.modality == "img2img":
+        from PIL import Image
+
+        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+
+        print("[Info] Running in img2img mode (Stage 1 only)")
+        client = OmniDiffusion(model=model_name)
+
+        if args.image_path:
+            if os.path.exists(args.image_path):
+                loaded_image = Image.open(args.image_path).convert("RGB")
+                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(
+            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):
+            omni_outputs = [result]
+        else:
+            omni_outputs = result
+
+    else:
+        from vllm_omni.entrypoints.omni import Omni
+
+        omni_kwargs = {}
+        if args.stage_configs_path:
+            omni_kwargs["stage_configs_path"] = args.stage_configs_path
+
+        omni_kwargs.update(
+            {
+                "log_stats": args.enable_stats,
+                "init_sleep_seconds": args.init_sleep_seconds,
+                "batch_timeout": args.batch_timeout,
+                "init_timeout": args.init_timeout,
+                "shm_threshold_bytes": args.shm_threshold_bytes,
+                "worker_backend": args.worker_backend,
+                "ray_address": args.ray_address,
+            }
+        )
+
+        omni = Omni(model=model_name, **omni_kwargs)
+
+        formatted_prompts = []
+        for p in args.prompts:
+            if args.modality == "img2text":
+                if args.image_path:
+                    loaded_image = Image.open(args.image_path).convert("RGB")
+                    final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                    prompt_dict = {
+                        "prompt": final_prompt_text,
+                        "multi_modal_data": {"image": loaded_image},
+                        "modalities": ["text"],
+                    }
+                    formatted_prompts.append(prompt_dict)
+            elif args.modality == "text2text":
+                final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
+                formatted_prompts.append(prompt_dict)
+            else:
+                # text2img
+                final_prompt_text = f"<|im_start|>{p}<|im_end|>"
+                prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
+                formatted_prompts.append(prompt_dict)
+
+        params_list = omni.default_sampling_params_list
+        if args.modality == "text2img":
+            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  # type: ignore # The second stage is an OmniDiffusionSamplingParam
+
+        omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+
+    for i, req_output in enumerate(omni_outputs):
+        images = getattr(req_output, "images", None)
+        if not images and hasattr(req_output, "output"):
+            if isinstance(req_output.output, list):
+                images = req_output.output
+            else:
+                images = [req_output.output]
+
+        if images:
+            for j, img in enumerate(images):
+                img.save(f"output_{i}_{j}.png")
+
+        if hasattr(req_output, "request_output") and req_output.request_output:
+            for stage_out in req_output.request_output:
+                if hasattr(stage_out, "images") and stage_out.images:
+                    for k, img in enumerate(stage_out.images):
+                        save_path = f"output_{i}_stage_{getattr(stage_out, 'stage_id', '?')}_{k}.png"
+                        img.save(save_path)
+                        print(f"[Info] Saved stage output image to {save_path}")
+
+    print(omni_outputs)
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/bagel/run_single_prompt.sh

@@ -0,0 +1,6 @@
+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}

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
 
@@ -110,7 +111,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--negative_prompt",
         type=str,
-        default=" ",
+        default=None,
         required=False,
     )
     parser.add_argument(
@@ -181,6 +182,12 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for ring sequence parallelism.",
     )
+    parser.add_argument(
+        "--tensor_parallel_size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     parser.add_argument("--layers", type=int, default=4, help="Number of layers to decompose the input image into.")
     parser.add_argument(
         "--resolution",
@@ -273,6 +280,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        help="Enable CPU offloading for diffusion models.",
+    )
     return parser.parse_args()
 
 
@@ -301,7 +313,10 @@ def main():
     vae_use_slicing = is_npu()
     vae_use_tiling = is_npu()
     parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
     )
 
     # Configure cache based on backend type
@@ -335,9 +350,13 @@ def main():
         cache_config=cache_config,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        enable_cpu_offload=args.enable_cpu_offload,
     )
     print("Pipeline loaded")
 
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
     # Time profiling for generation
     print(f"\n{'=' * 60}")
     print("Generation Configuration:")
@@ -351,32 +370,56 @@ def main():
     else:
         print(f"  Input image size: {input_image.size}")
     print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
     )
     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,
-    }
 
-    outputs = omni.generate(**generate_kwargs)
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
+    # 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
 
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
     if not outputs:
         raise ValueError("No output generated from omni.generate()")
     logger.info("Outputs: %s", outputs)

examples/offline_inference/image_to_image/image_to_image.md

@@ -47,3 +47,4 @@ 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.
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.

examples/offline_inference/image_to_video/README.md

@@ -54,3 +54,4 @@ 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.
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.

examples/offline_inference/image_to_video/image_to_video.py

@@ -19,6 +19,7 @@ Usage:
 """
 
 import argparse
+import os
 from pathlib import Path
 
 import numpy as np
@@ -26,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
 
@@ -57,6 +59,11 @@ 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(
+        "--enable-cpu-offload",
+        action="store_true",
+        help="Enable CPU offloading for diffusion models.",
+    )
     return parser.parse_args()
 
 
@@ -95,26 +102,38 @@ def main():
     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,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
+        enable_cpu_offload=args.enable_cpu_offload,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
     # 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
@@ -178,6 +197,23 @@ def main():
     export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
 
 if __name__ == "__main__":
     main()

examples/offline_inference/lora_inference/README.md

@@ -0,0 +1,98 @@
+# LoRA Inference Examples
+
+This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
+
+## Overview
+
+Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
+
+- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
+- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
+
+Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
+
+## Usage
+
+### Pre-loaded LoRA (via --lora-path)
+
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_preloaded.png
+```
+
+**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
+
+### Per-request LoRA (via --lora-request-path)
+
+Load a LoRA adapter on-demand for each request:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-request-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_per_request.png
+```
+
+### No LoRA
+
+If no LoRA request is provided, we will use the base model without any LoRA adapters:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_no_lora.png
+```
+
+## Parameters
+
+### LoRA Parameters
+
+- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
+- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
+- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### Standard Parameters
+
+- `--prompt`: Text prompt for image generation (required)
+- `--seed`: Random seed for reproducibility (default: 42)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--output`: Output file path (default: `lora_output.png`)
+
+## How LoRA Works
+
+All LoRA adapters are handled uniformly:
+
+1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
+2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
+3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
+
+The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
+
+## LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```

examples/offline_inference/lora_inference/lora_inference.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+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
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate images with LoRA adapters.")
+    parser.add_argument("--model", default="stabilityai/stable-diffusion-3.5-medium", help="Model name or path.")
+    parser.add_argument("--prompt", required=True, help="Text prompt for image generation.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.")
+    parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
+    parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=50,
+        help="Number of denoising steps for the diffusion sampler.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="lora_output.png",
+        help="Path to save the generated image (PNG).",
+    )
+    parser.add_argument(
+        "--lora-path",
+        type=str,
+        default=None,
+        help="Path to LoRA adapter folder to pre-load at initialization (PEFT format). "
+        "Note: pre-loading populates the cache; you still need to pass a lora_request to activate it.",
+    )
+    parser.add_argument(
+        "--lora-request-path",
+        type=str,
+        default=None,
+        help="Path to LoRA adapter folder for per-request activation (dynamic LoRA). "
+        "If --lora-request-id is not provided, a stable ID will be derived from this path.",
+    )
+    parser.add_argument(
+        "--lora-request-id",
+        type=int,
+        default=None,
+        help="Integer ID for the LoRA adapter (for dynamic LoRA). "
+        "If not provided and --lora-request-path is set, will derive a stable ID from the path.",
+    )
+    parser.add_argument(
+        "--lora-scale",
+        type=float,
+        default=1.0,
+        help="Scale factor for LoRA weights (default: 1.0).",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    model = args.model
+
+    omni_kwargs = {}
+
+    if args.lora_path:
+        omni_kwargs["lora_path"] = args.lora_path
+        print(f"Using static LoRA from: {args.lora_path}")
+
+    omni = Omni(model=model, **omni_kwargs)
+
+    lora_request = None
+    if args.lora_request_path:
+        if args.lora_request_id is None:
+            lora_request_id = stable_lora_int_id(args.lora_request_path)
+        else:
+            lora_request_id = args.lora_request_id
+
+        lora_name = Path(args.lora_request_path).stem
+        lora_request = LoRARequest(
+            lora_name=lora_name,
+            lora_int_id=lora_request_id,
+            lora_path=args.lora_request_path,
+        )
+        print(f"Using per-request LoRA: name={lora_name}, id={lora_request_id}, scale={args.lora_scale}")
+    elif args.lora_path:
+        # pre-loaded LoRA
+        lora_request_id = stable_lora_int_id(args.lora_path)
+        lora_request = LoRARequest(
+            lora_name="preloaded",
+            lora_int_id=lora_request_id,
+            lora_path=args.lora_path,
+        )
+        print(f"Activating pre-loaded LoRA: id={lora_request_id}, scale={args.lora_scale}")
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+    )
+
+    if lora_request:
+        sampling_params.lora_request = lora_request
+        sampling_params.lora_scale = args.lora_scale
+
+    outputs = omni.generate(args.prompt, sampling_params)
+
+    if not outputs or len(outputs) == 0:
+        raise ValueError("No output generated from omni.generate()")
+
+    if isinstance(outputs, list):
+        first_output = outputs[0]
+    else:
+        first_output = outputs
+
+    images = None
+    if hasattr(first_output, "images") and first_output.images:
+        images = first_output.images
+    elif hasattr(first_output, "request_output") and first_output.request_output:
+        req_out = first_output.request_output
+        if isinstance(req_out, list) and len(req_out) > 0:
+            req_out = req_out[0]
+        if hasattr(req_out, "images") and req_out.images:
+            images = req_out.images
+
+    if not images:
+        raise ValueError("No images found in request_output")
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".png"
+    stem = output_path.stem or "lora_output"
+    if len(images) <= 1:
+        images[0].save(output_path)
+        print(f"Saved generated image to {output_path}")
+    else:
+        for idx, img in enumerate(images):
+            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+            img.save(save_path)
+            print(f"Saved generated image to {save_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -278,9 +278,9 @@ def get_audio_query(question: str = None, audio_path: str | None = None, samplin
 
 
 query_map = {
-    "mixed_modalities": get_mixed_modalities_query,
+    "use_mixed_modalities": get_mixed_modalities_query,
     "use_audio_in_video": get_use_audio_in_video_query,
-    "multi_audios": get_multi_audios_query,
+    "use_multi_audios": get_multi_audios_query,
     "use_image": get_image_query,
     "use_video": get_video_query,
     "use_audio": get_audio_query,
@@ -434,7 +434,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="mixed_modalities",
+        default="use_mixed_modalities",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -1,2 +1,2 @@
 python end2end.py --output-wav output_audio \
-                  --query-type mixed_modalities
+                  --query-type use_mixed_modalities

examples/offline_inference/qwen3_omni/end2end.py

@@ -12,6 +12,7 @@ from typing import NamedTuple
 import librosa
 import numpy as np
 import soundfile as sf
+import vllm
 from PIL import Image
 from vllm import SamplingParams
 from vllm.assets.audio import AudioAsset
@@ -225,18 +226,76 @@ def get_multi_audios_query() -> QueryResult:
     )
 
 
+# def get_use_audio_in_video_query(video_path: str | None = None) -> QueryResult:
+# question = (
+#     "Describe the content of the video in details, then convert what the "
+#     "baby say into text."
+# )
+# prompt = (
+#     f"<|im_start|>system\n{default_system}<|im_end|>\n"
+#     "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+#     f"{question}<|im_end|>\n"
+#     f"<|im_start|>assistant\n"
+# )
+# if video_path:
+#     if not os.path.exists(video_path):
+#         raise FileNotFoundError(f"Video file not found: {video_path}")
+#     video_frames = video_to_ndarrays(video_path, num_frames=16)
+# else:
+#     video_frames = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays
+# audio = extract_video_audio(video_path, sampling_rate=16000)
+# return QueryResult(
+#     inputs={
+#         "prompt": prompt,
+#         "multi_modal_data": {
+#             "video": video_frames,
+#             "audio": audio,
+#         },
+#         "mm_processor_kwargs": {
+#             "use_audio_in_video": True,
+#         },
+#     },
+#     limit_mm_per_prompt={"audio": 1, "video": 1},
+# )
+def get_use_audio_in_video_query() -> QueryResult:
+    question = "Describe the content of the video in details, then convert what the baby say into text."
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+    asset = VideoAsset(name="baby_reading", num_frames=16)
+    audio = asset.get_audio(sampling_rate=16000)
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "video": asset.np_ndarrays,
+                "audio": audio,
+            },
+            "mm_processor_kwargs": {
+                "use_audio_in_video": True,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1, "video": 1},
+    )
+
+
 query_map = {
     "text": get_text_query,
     "use_audio": get_audio_query,
     "use_image": get_image_query,
     "use_video": get_video_query,
-    "multi_audios": get_multi_audios_query,
-    "mixed_modalities": get_mixed_modalities_query,
+    "use_multi_audios": get_multi_audios_query,
+    "use_mixed_modalities": get_mixed_modalities_query,
+    "use_audio_in_video": get_use_audio_in_video_query,
 }
 
 
 def main(args):
     model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
 
     # Get paths from args
     video_path = getattr(args, "video_path", None)
@@ -259,6 +318,10 @@ def main(args):
             num_frames=getattr(args, "num_frames", 16),
             sampling_rate=getattr(args, "sampling_rate", 16000),
         )
+    elif args.query_type == "multi_audios":
+        query_result = query_func()
+    elif args.query_type == "use_audio_in_video":
+        query_result = query_func()
     else:
         query_result = query_func()
 
@@ -270,7 +333,7 @@ def main(args):
     )
 
     thinker_sampling_params = SamplingParams(
-        temperature=0.4,
+        temperature=0.9,
         top_p=0.9,
         top_k=-1,
         max_tokens=1200,
@@ -331,6 +394,8 @@ def main(args):
     total_requests = len(prompts)
     processed_count = 0
 
+    print(f"query type: {args.query_type}")
+
     for stage_outputs in omni_generator:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
@@ -385,7 +450,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="mixed_modalities",
+        default="use_mixed_modalities",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/offline_inference/qwen3_tts/README.md

@@ -0,0 +1,93 @@
+# Qwen3-TTS Offline Inference
+
+This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.
+
+## Model Overview
+
+Qwen3 TTS provides multiple task variants for speech generation:
+
+- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
+- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
+- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
+
+## Setup
+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.
+
+### ROCm Dependencies
+
+You will need to install these two dependencies `onnxruntime-rocm` and `sox`.
+
+```
+pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+pip install onnxruntime-rocm sox
+```
+
+## Quick Start
+
+Run a single sample for a task:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Generated audio files are saved to `output_audio/` by default.
+
+## Task Usage
+
+### CustomVoice
+
+Single sample:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Batch sample (multiple prompts in one run):
+
+```
+python end2end.py --query-type CustomVoice --use-batch-sample
+```
+
+### VoiceDesign
+
+Single sample:
+
+```
+python end2end.py --query-type VoiceDesign
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type VoiceDesign --use-batch-sample
+```
+
+### Base (Voice Clone)
+
+Single sample:
+
+```
+python end2end.py --query-type Base
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type Base --use-batch-sample
+```
+
+Mode selection for Base:
+
+- `--mode-tag icl` (default): standard mode
+- `--mode-tag xvec_only`: enable `x_vector_only_mode` in the request
+
+Examples:
+
+```
+python end2end.py --query-type Base --mode-tag icl
+```
+
+## Notes
+
+- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
+- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.

examples/offline_inference/qwen3_tts/end2end.py

@@ -0,0 +1,381 @@
+"""Offline inference demo for Qwen3 TTS via vLLM Omni.
+
+Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
+tasks, then runs Omni generation and saves output wav files.
+"""
+
+import os
+from typing import NamedTuple
+
+import soundfile as sf
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+
+
+class QueryResult(NamedTuple):
+    """Container for a prepared Omni request."""
+
+    inputs: dict
+    model_name: str
+
+
+def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
+    """Build CustomVoice sample inputs.
+
+    Args:
+        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
+
+    Returns:
+        QueryResult with Omni inputs and the CustomVoice model path.
+    """
+    task_type = "CustomVoice"
+    if use_batch_sample:
+        texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
+        instructs = ["", "Very happy."]
+        languages = ["Chinese", "English"]
+        speakers = ["Vivian", "Ryan"]
+        inputs = []
+        for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
+            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            inputs.append(
+                {
+                    "prompt": prompt,
+                    "additional_information": {
+                        "task_type": [task_type],
+                        "text": [text],
+                        "instruct": [instruct],
+                        "language": [language],
+                        "speaker": [speaker],
+                        "max_new_tokens": [2048],
+                    },
+                }
+            )
+    else:
+        text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。"
+        language = "Chinese"
+        speaker = "Vivian"
+        instruct = "用特别愤怒的语气说"
+        prompts = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        inputs = {
+            "prompt": prompts,
+            "additional_information": {
+                "task_type": [task_type],
+                "text": [text],
+                "language": [language],
+                "speaker": [speaker],
+                "instruct": [instruct],
+                "max_new_tokens": [2048],
+            },
+        }
+    return QueryResult(
+        inputs=inputs,
+        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    )
+
+
+def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
+    """Build VoiceDesign sample inputs.
+
+    Args:
+        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
+
+    Returns:
+        QueryResult with Omni inputs and the VoiceDesign model path.
+    """
+    task_type = "VoiceDesign"
+    if use_batch_sample:
+        texts = [
+            "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
+            "It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!",
+        ]
+        instructs = [
+            "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。",
+            "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.",
+        ]
+        languages = ["Chinese", "English"]
+        inputs = []
+        for text, instruct, language in zip(texts, instructs, languages):
+            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            inputs.append(
+                {
+                    "prompt": prompt,
+                    "additional_information": {
+                        "task_type": [task_type],
+                        "text": [text],
+                        "language": [language],
+                        "instruct": [instruct],
+                        "max_new_tokens": [2048],
+                        "non_streaming_mode": [True],
+                    },
+                }
+            )
+    else:
+        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
+        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
+        language = "Chinese"
+        prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        inputs = {
+            "prompt": prompt,
+            "additional_information": {
+                "task_type": [task_type],
+                "text": [text],
+                "language": [language],
+                "instruct": [instruct],
+                "max_new_tokens": [2048],
+                "non_streaming_mode": [True],
+            },
+        }
+    return QueryResult(
+        inputs=inputs,
+        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
+    )
+
+
+def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
+    """Build Base (voice clone) sample inputs.
+
+    Args:
+        use_batch_sample: When True, return a batch of prompts (Case 2).
+        mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior.
+
+    Returns:
+        QueryResult with Omni inputs and the Base model path.
+    """
+    task_type = "Base"
+    ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+    ref_audio_single = ref_audio_path_1
+    ref_text_single = (
+        "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+    )
+    syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."
+    syn_lang_single = "Auto"
+    x_vector_only_mode = mode_tag == "xvec_only"
+    if use_batch_sample:
+        syn_text_batch = [
+            "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.",
+            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
+        ]
+        syn_lang_batch = ["Chinese", "English"]
+        inputs = []
+        for text, language in zip(syn_text_batch, syn_lang_batch):
+            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            inputs.append(
+                {
+                    "prompt": prompt,
+                    "additional_information": {
+                        "task_type": [task_type],
+                        "ref_audio": [ref_audio_single],
+                        "ref_text": [ref_text_single],
+                        "text": [text],
+                        "language": [language],
+                        "x_vector_only_mode": [x_vector_only_mode],
+                        "max_new_tokens": [2048],
+                    },
+                }
+            )
+    else:
+        prompt = f"<|im_start|>assistant\n{syn_text_single}<|im_end|>\n<|im_start|>assistant\n"
+        inputs = {
+            "prompt": prompt,
+            "additional_information": {
+                "task_type": [task_type],
+                "ref_audio": [ref_audio_single],
+                "ref_text": [ref_text_single],
+                "text": [syn_text_single],
+                "language": [syn_lang_single],
+                "x_vector_only_mode": [x_vector_only_mode],
+                "max_new_tokens": [2048],
+            },
+        }
+    return QueryResult(
+        inputs=inputs,
+        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+    )
+
+
+def main(args):
+    """Run offline inference with Omni using prepared sample inputs.
+
+    Args:
+        args: Parsed CLI args from parse_args().
+    """
+    query_func = query_map[args.query_type]
+    if args.query_type in {"CustomVoice", "VoiceDesign"}:
+        query_result = query_func(use_batch_sample=args.use_batch_sample)
+    elif args.query_type == "Base":
+        query_result = query_func(
+            use_batch_sample=args.use_batch_sample,
+            mode_tag=args.mode_tag,
+        )
+    else:
+        query_result = query_func()
+
+    model_name = query_result.model_name
+    omni = Omni(
+        model=model_name,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.enable_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    sampling_params = SamplingParams(
+        temperature=0.9,
+        top_p=1.0,
+        top_k=50,
+        max_tokens=2048,
+        seed=42,
+        detokenize=False,
+        repetition_penalty=1.05,
+    )
+
+    sampling_params_list = [
+        sampling_params,
+    ]
+
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+
+    omni_generator = omni.generate(query_result.inputs, sampling_params_list)
+    for stage_outputs in omni_generator:
+        for output in stage_outputs.request_output:
+            request_id = output.request_id
+            audio_tensor = output.multimodal_output["audio"]
+            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            audio_samplerate = output.multimodal_output["sr"].item()
+            # Convert to numpy array and ensure correct format
+            audio_numpy = audio_tensor.float().detach().cpu().numpy()
+
+            # Ensure audio is 1D (flatten if needed)
+            if audio_numpy.ndim > 1:
+                audio_numpy = audio_numpy.flatten()
+
+            # Save audio file with explicit WAV format
+            sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV")
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+
+def parse_args():
+    """Parse CLI arguments for offline TTS inference.
+
+    Returns:
+        argparse.Namespace with CLI options.
+    """
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="CustomVoice",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing a single stage in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--output-wav",
+        default="output_audio",
+        help="[Deprecated] Output wav directory (use --output-dir).",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate.",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to a stage configs file.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading (default: 16000).",
+    )
+    parser.add_argument(
+        "--log-dir",
+        type=str,
+        default="logs",
+        help="Log directory (default: logs).",
+    )
+    parser.add_argument(
+        "--py-generator",
+        action="store_true",
+        default=False,
+        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
+    )
+    parser.add_argument(
+        "--use-batch-sample",
+        action="store_true",
+        default=False,
+        help="Use batch input sample for CustomVoice/VoiceDesign/Base query.",
+    )
+    parser.add_argument(
+        "--mode-tag",
+        type=str,
+        default="icl",
+        choices=["icl", "xvec_only"],
+        help="Mode tag for Base query x_vector_only_mode (default: icl).",
+    )
+
+    return parser.parse_args()
+
+
+query_map = {
+    "CustomVoice": get_custom_voice_query,
+    "VoiceDesign": get_voice_design_query,
+    "Base": get_base_query,
+}
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/text_to_audio/README.md

@@ -0,0 +1,37 @@
+# Text-To-Audio
+
+The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
+
+## Prerequisites
+
+If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you have access:
+
+1. **Accept Model License**: Visit the model page on Hugging Face (e.g., [stabilityai/stable-audio-open-1.0]) and accept the user agreement.
+2. **Authenticate**: Log in to Hugging Face locally to access the gated model.
+   ```bash
+   huggingface-cli login
+   ```
+
+## Local CLI Usage
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative_prompt "Low quality" \
+  --seed 42 \
+  --guidance_scale 7.0 \
+  --audio_length 10.0 \
+  --num_inference_steps 100 \
+  --output stable_audio_output.wav
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--seed`: integer seed for deterministic generation.
+- `--guidance_scale`: classifier-free guidance scale.
+- `--audio_length`: audio duration in seconds.
+- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--output`: path to save the generated WAV file.

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
 
 
@@ -142,17 +143,21 @@ def main():
     generation_start = time.perf_counter()
 
     # Generate audio
-    audio = 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,
+    outputs = omni.generate(
+        {
+            "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()
@@ -166,6 +171,21 @@ def main():
     suffix = output_path.suffix or ".wav"
     stem = output_path.stem or "stable_audio_output"
 
+    # Extract audio from omni.generate() outputs
+    if not outputs:
+        raise ValueError("No output generated from omni.generate()")
+
+    output = outputs[0]
+    if not hasattr(output, "request_output") or not output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
+    request_output = output.request_output[0]
+    if not hasattr(request_output, "multimodal_output"):
+        raise ValueError("No multimodal_output found in request_output")
+
+    audio = request_output.multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError("No audio output found in request_output")
+
     # Handle different output formats
     if isinstance(audio, torch.Tensor):
         audio = audio.cpu().float().numpy()

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
@@ -61,6 +96,7 @@ 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.
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
 > ℹ️ 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.
 

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:

examples/offline_inference/text_to_image/text_to_image.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
+import os
 import time
 from pathlib import Path
 
@@ -9,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
 
@@ -23,7 +25,9 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
     parser.add_argument(
-        "--negative_prompt", default="", help="negative prompt for classifier-free conditional guidance."
+        "--negative_prompt",
+        default=None,
+        help="negative prompt for classifier-free conditional guidance.",
     )
     parser.add_argument("--seed", type=int, default=142, help="Random seed for deterministic results.")
     parser.add_argument(
@@ -93,6 +97,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        help="Enable CPU offloading for diffusion models.",
+    )
     parser.add_argument(
         "--tensor_parallel_size",
         type=int,
@@ -148,6 +157,9 @@ def main():
         tensor_parallel_size=args.tensor_parallel_size,
     )
 
+    # 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,
@@ -156,8 +168,13 @@ def main():
         cache_config=cache_config,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        enable_cpu_offload=args.enable_cpu_offload,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
     # Time profiling for generation
     print(f"\n{'=' * 60}")
     print("Generation Configuration:")
@@ -173,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
@@ -189,6 +210,23 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
     # Extract images from OmniRequestOutput
     # omni.generate() returns list[OmniRequestOutput], extract images from the first output
     if not outputs or len(outputs) == 0:

examples/offline_inference/text_to_video/text_to_video.md

@@ -9,10 +9,11 @@ python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
-  --width 640 \
-  --num_frames 32 \
+  --width 832 \
+  --num_frames 33 \
   --guidance_scale 4.0 \
   --guidance_scale_high 3.0 \
+  --flow_shift 12.0 \
   --num_inference_steps 40 \
   --fps 16 \
   --output t2v_out.mp4
@@ -21,10 +22,11 @@ python text_to_video.py \
 Key arguments:
 
 - `--prompt`: text description (string).
-- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
 - `--num_frames`: Number of frames (Wan default is 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
+- `--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.
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.

examples/offline_inference/text_to_video/text_to_video.py

@@ -2,12 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
+import os
+import time
 from pathlib import Path
 
 import numpy as np
 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
 
@@ -34,6 +38,28 @@ 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(
+        "--enable-cpu-offload",
+        action="store_true",
+        help="Enable CPU offloading for diffusion models.",
+    )
+    parser.add_argument(
+        "--ulysses_degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ulysses sequence parallelism.",
+    )
+    parser.add_argument(
+        "--ring_degree",
+        type=int,
+        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()
 
 
@@ -46,25 +72,62 @@ def main():
     vae_use_slicing = is_npu()
     vae_use_tiling = is_npu()
 
+    # Configure parallel settings (only SP is supported for Wan)
+    # Note: cfg_parallel and tensor_parallel are not implemented for Wan models
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+    )
+
+    # 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,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
+        enable_cpu_offload=args.enable_cpu_offload,
+        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
+    # Print generation configuration
+    print(f"\n{'=' * 60}")
+    print("Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  Frames: {args.num_frames}")
+    print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}")
+    print(f"  Video size: {args.width}x{args.height}")
+    print(f"{'=' * 60}\n")
+
+    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
+
+    # Print profiling results
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
     # Extract video frames from OmniRequestOutput
     if isinstance(frames, list) and len(frames) > 0:
@@ -126,6 +189,23 @@ def main():
     export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
 
 if __name__ == "__main__":
     main()

examples/online_serving/bagel/openai_chat_client.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""
+Bagel OpenAI-compatible chat client for image generation and multimodal tasks.
+
+Usage:
+    python openai_chat_client.py --prompt "A cute cat" --output output.png
+    python openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    prompt: str,
+    server_url: str = "http://localhost:8091",
+    image_url: str | None = None,
+    height: int | None = None,
+    width: int | None = None,
+    steps: int | None = None,
+    seed: int | None = None,
+    negative_prompt: str | None = None,
+    modality: str = "text2img",  # "text2img" (default), "img2img", "img2text", "text2text"
+) -> bytes | str | None:
+    """Generate an image or text using the chat completions API.
+
+    Args:
+        prompt: Text description or prompt
+        server_url: Server URL
+        image_url: URL or path to input image (for img2img/img2text)
+        height: Image height in pixels
+        width: Image width in pixels
+        steps: Number of inference steps
+        seed: Random seed
+        negative_prompt: Negative prompt
+        modality: Task modality hint
+
+    Returns:
+        Image bytes (for image outputs) or Text string (for text outputs) or None if failed
+    """
+
+    # Construct Message Content
+    content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
+
+    if image_url:
+        # Check if local file
+        if Path(image_url).exists():
+            with open(image_url, "rb") as f:
+                b64_data = base64.b64encode(f.read()).decode("utf-8")
+                final_image_url = f"data:image/jpeg;base64,{b64_data}"
+        else:
+            final_image_url = image_url
+
+        content.append({"type": "image_url", "image_url": {"url": final_image_url}})
+
+    messages = [{"role": "user", "content": content}]
+
+    # Build request payload with all parameters at top level
+    # Note: vLLM ignores "extra_body", so we put parameters directly in the payload
+    payload = {"messages": messages}
+
+    # Set output modalities at top level
+    if modality == "text2img" or modality == "img2img":
+        payload["modalities"] = ["image"]
+    elif modality == "img2text" or modality == "text2text":
+        payload["modalities"] = ["text"]
+
+    # Add generation parameters directly to payload
+    if height is not None:
+        payload["height"] = height
+    if width is not None:
+        payload["width"] = width
+    if steps is not None:
+        payload["num_inference_steps"] = steps
+    if seed is not None:
+        payload["seed"] = seed
+    if negative_prompt:
+        payload["negative_prompt"] = negative_prompt
+
+    # Send request
+    try:
+        print(f"Sending request to {server_url} with modality {modality}...")
+        response = requests.post(
+            f"{server_url}/v1/chat/completions",
+            headers={"Content-Type": "application/json"},
+            json=payload,
+            timeout=300,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract content - check ALL choices since server may return multiple
+        # (e.g., text in choices[0], image in choices[1])
+        choices = data.get("choices", [])
+
+        # First pass: look for image output in any choice
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+
+            # Handle Image Output
+            if isinstance(choice_content, list) and len(choice_content) > 0:
+                first_item = choice_content[0]
+                if isinstance(first_item, dict) and "image_url" in first_item:
+                    img_url_str = first_item["image_url"].get("url", "")
+                    if img_url_str.startswith("data:image"):
+                        _, b64_data = img_url_str.split(",", 1)
+                        return base64.b64decode(b64_data)
+
+        # Second pass: look for text output if no image found
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+            if isinstance(choice_content, str) and choice_content:
+                return choice_content
+
+        print(f"Unexpected response format: {choices}")
+        return None
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
+    parser.add_argument("--prompt", "-p", default="<|im_start|>A cute cat<|im_end|>", help="Text prompt")
+    parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
+    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
+
+    # Modality Control
+    parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
+    parser.add_argument(
+        "--modality",
+        "-m",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Task modality",
+    )
+
+    # Generation Params
+    parser.add_argument("--height", type=int, default=512, help="Image height")
+    parser.add_argument("--width", type=int, default=512, help="Image width")
+    parser.add_argument("--steps", type=int, default=25, help="Inference steps")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--negative", help="Negative prompt")
+
+    args = parser.parse_args()
+
+    print(f"Mode: {args.modality}")
+    if args.image_url:
+        print(f"Input Image: {args.image_url}")
+
+    result = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        image_url=args.image_url,
+        height=args.height,
+        width=args.width,
+        steps=args.steps,
+        seed=args.seed,
+        negative_prompt=args.negative,
+        modality=args.modality,
+    )
+
+    if result:
+        if isinstance(result, bytes):
+            # It's an image
+            output_path = Path(args.output)
+            output_path.write_bytes(result)
+            print(f"Image saved to: {output_path}")
+            print(f"Size: {len(result) / 1024:.1f} KB")
+        elif isinstance(result, str):
+            # It's text
+            print("Response:")
+            print(result)
+    else:
+        print("Failed to generate response")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/bagel/run_server.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+# Bagel online serving startup script
+
+MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}"
+PORT="${PORT:-8091}"
+
+echo "Starting Bagel server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT"

examples/online_serving/lora_inference/README.md

@@ -0,0 +1,54 @@
+# Online LoRA Inference (Diffusion)
+
+This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
+
+> Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+> Note: This example uses `/v1/chat/completions`. LoRA payloads for other OpenAI endpoints are not implemented here.
+
+## Start Server
+
+```bash
+# Pick a diffusion model (examples)
+# export MODEL=stabilityai/stable-diffusion-3.5-medium
+# export MODEL=Qwen/Qwen-Image
+
+bash run_server.sh
+```
+
+## Call API (curl)
+
+```bash
+# Required: local LoRA folder on the server
+export LORA_PATH=/path/to/lora_adapter
+
+# Optional
+export SERVER=http://localhost:8091
+export PROMPT="A piece of cheesecake"
+export LORA_NAME=my_lora
+export LORA_SCALE=1.0
+# Optional: if omitted, the server derives a stable id from LORA_PATH.
+# export LORA_INT_ID=123
+
+bash run_curl_lora_inference.sh
+```
+
+## Call API (Python)
+
+```bash
+python openai_chat_client.py \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora_adapter \
+  --lora-name my_lora \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+## LoRA Format
+
+LoRA adapters should be in PEFT format, for example:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```

examples/online_serving/lora_inference/openai_chat_client.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""
+OpenAI-compatible chat client for diffusion LoRA inference.
+
+Example:
+  python openai_chat_client.py \
+    --server http://localhost:8091 \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora_adapter \
+    --lora-name my_lora \
+    --lora-scale 1.0 \
+    --output output.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    *,
+    prompt: str,
+    server_url: str,
+    height: int | None,
+    width: int | None,
+    num_inference_steps: int | None,
+    seed: int | None,
+    lora_name: str | None,
+    lora_path: str | None,
+    lora_scale: float | None,
+    lora_int_id: int | None,
+) -> bytes | None:
+    messages = [{"role": "user", "content": prompt}]
+
+    extra_body: dict = {}
+    if height is not None:
+        extra_body["height"] = height
+    if width is not None:
+        extra_body["width"] = width
+    if num_inference_steps is not None:
+        extra_body["num_inference_steps"] = num_inference_steps
+    if seed is not None:
+        extra_body["seed"] = seed
+
+    if lora_path:
+        lora_body: dict = {
+            "local_path": lora_path,
+            "name": lora_name or Path(lora_path).stem,
+        }
+        if lora_scale is not None:
+            lora_body["scale"] = float(lora_scale)
+        if lora_int_id is not None:
+            lora_body["int_id"] = int(lora_int_id)
+        extra_body["lora"] = lora_body
+
+    payload = {"messages": messages}
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    response = requests.post(
+        f"{server_url}/v1/chat/completions",
+        headers={"Content-Type": "application/json"},
+        json=payload,
+        timeout=300,
+    )
+    response.raise_for_status()
+    data = response.json()
+
+    content = data["choices"][0]["message"]["content"]
+    if isinstance(content, list) and content:
+        image_url = content[0].get("image_url", {}).get("url", "")
+        if image_url.startswith("data:image"):
+            _, b64_data = image_url.split(",", 1)
+            return base64.b64decode(b64_data)
+
+    raise RuntimeError(f"Unexpected response format: {content!r}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Diffusion LoRA OpenAI chat client")
+    parser.add_argument("--server", default="http://localhost:8091", help="Server URL")
+    parser.add_argument("--prompt", default="A piece of cheesecake", help="Text prompt")
+    parser.add_argument("--output", default="lora_online_output.png", help="Output image path")
+
+    parser.add_argument("--height", type=int, default=1024, help="Image height")
+    parser.add_argument("--width", type=int, default=1024, help="Image width")
+    parser.add_argument("--steps", type=int, default=50, help="num_inference_steps")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+
+    parser.add_argument("--lora-path", default=None, help="Server-local LoRA adapter folder (PEFT format)")
+    parser.add_argument("--lora-name", default=None, help="LoRA name (optional)")
+    parser.add_argument("--lora-scale", type=float, default=1.0, help="LoRA scale")
+    parser.add_argument(
+        "--lora-int-id",
+        type=int,
+        default=None,
+        help="LoRA integer id (cache key). If omitted, the server derives a stable id from lora_path.",
+    )
+
+    args = parser.parse_args()
+
+    image_bytes = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.steps,
+        seed=args.seed,
+        lora_name=args.lora_name,
+        lora_path=args.lora_path,
+        lora_scale=args.lora_scale if args.lora_path else None,
+        lora_int_id=args.lora_int_id if args.lora_path else None,
+    )
+
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_bytes(image_bytes)
+    print(f"Saved: {out_path} ({len(image_bytes) / 1024:.1f} KiB)")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/lora_inference/run_curl_lora_inference.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# Online diffusion LoRA inference via OpenAI-compatible chat API.
+
+SERVER="${SERVER:-http://localhost:8091}"
+PROMPT="${PROMPT:-A piece of cheesecake}"
+
+LORA_PATH="${LORA_PATH:-}"
+LORA_NAME="${LORA_NAME:-lora}"
+LORA_SCALE="${LORA_SCALE:-1.0}"
+LORA_INT_ID="${LORA_INT_ID:-}"
+
+HEIGHT="${HEIGHT:-1024}"
+WIDTH="${WIDTH:-1024}"
+NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-50}"
+SEED="${SEED:-42}"
+
+CURRENT_TIME=$(date +%Y%m%d%H%M%S)
+OUTPUT="${OUTPUT:-lora_online_output_${CURRENT_TIME}.png}"
+
+if [ -z "$LORA_PATH" ]; then
+  echo "ERROR: LORA_PATH is required (must be a server-local path)."
+  exit 1
+fi
+
+echo "Generating image with LoRA..."
+echo "Server: $SERVER"
+echo "Prompt: $PROMPT"
+echo "LoRA: name=$LORA_NAME id=${LORA_INT_ID:-auto} scale=$LORA_SCALE path=$LORA_PATH"
+echo "Output: $OUTPUT"
+
+LORA_INT_ID_FIELD=""
+if [ -n "$LORA_INT_ID" ]; then
+  LORA_INT_ID_FIELD=", \"int_id\": $LORA_INT_ID"
+fi
+
+curl -s "$SERVER/v1/chat/completions" \
+  -H "Content-Type: application/json" \
+  -d "{
+    \"messages\": [
+      {\"role\": \"user\", \"content\": \"$PROMPT\"}
+    ],
+    \"extra_body\": {
+      \"height\": $HEIGHT,
+      \"width\": $WIDTH,
+      \"num_inference_steps\": $NUM_INFERENCE_STEPS,
+      \"seed\": $SEED,
+      \"lora\": {
+        \"name\": \"$LORA_NAME\",
+        \"local_path\": \"$LORA_PATH\",
+        \"scale\": $LORA_SCALE$LORA_INT_ID_FIELD
+      }
+    }
+  }" | jq -r '.choices[0].message.content[0].image_url.url' | sed 's/^data:image[^,]*,\s*//' | base64 -d > "$OUTPUT"
+
+if [ -f "$OUTPUT" ]; then
+  echo "Image saved to: $OUTPUT"
+  echo "Size: $(du -h "$OUTPUT" | cut -f1)"
+else
+  echo "Failed to generate image"
+  exit 1
+fi

examples/online_serving/lora_inference/run_server.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+# Online diffusion serving with vLLM-Omni (OpenAI-compatible API).
+
+MODEL="${MODEL:-stabilityai/stable-diffusion-3.5-medium}"
+PORT="${PORT:-8091}"
+
+echo "Starting vLLM-Omni diffusion server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+
+if [ -z "${VLLM_BIN:-}" ]; then
+  if command -v vllm-omni >/dev/null 2>&1; then
+    VLLM_BIN="vllm-omni"
+  else
+    VLLM_BIN="vllm"
+  fi
+fi
+
+"$VLLM_BIN" serve "$MODEL" --omni \
+  --port "$PORT"

examples/online_serving/qwen3_omni/README.md

@@ -12,6 +12,12 @@ Please refer to [README.md](../../../README.md)
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
+If you want to open async chunking for qwen3-omni, launch the server with command below
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,4 +1,5 @@
 import base64
+import concurrent.futures
 import os
 from typing import NamedTuple
 
@@ -305,6 +306,26 @@ def get_multi_audios_query(custom_prompt: str | None = None):
     }
 
 
+def get_use_audio_in_video_query(
+    video_path: str | None = None,
+    audio_path: str | None = None,
+    custom_prompt: str | None = None,
+):
+    question = custom_prompt or (
+        "Describe the content of the video in details, then convert what the baby say into text."
+    )
+    video_url = get_video_url_from_path(video_path)
+    audio_url = get_audio_url_from_path(audio_path)
+    return {
+        "role": "user",
+        "content": [
+            {"type": "video_url", "video_url": {"url": video_url}},
+            {"type": "audio_url", "audio_url": {"url": audio_url}},
+            {"type": "text", "text": question},
+        ],
+    }
+
+
 query_map = {
     "text": get_text_query,
     "use_audio": get_audio_query,
@@ -312,6 +333,7 @@ query_map = {
     "use_video": get_video_query,
     "use_mixed_modalities": get_mixed_modalities_query,
     "use_multi_audios": get_multi_audios_query,
+    "use_audio_in_video": get_use_audio_in_video_query,
 }
 
 
@@ -372,6 +394,12 @@ def run_multimodal_generation(args) -> None:
         prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
     elif args.query_type == "text":
         prompt = query_func(custom_prompt=custom_prompt)
+    elif args.query_type == "use_audio_in_video":
+        prompt = query_func(
+            video_path=video_path,
+            audio_path=audio_path,
+            custom_prompt=custom_prompt,
+        )
     else:
         prompt = query_func()
 
@@ -387,51 +415,67 @@ def run_multimodal_generation(args) -> None:
     else:
         output_modalities = None
 
-    chat_completion = client.chat.completions.create(
-        messages=[
-            get_system_prompt(),
-            prompt,
-        ],
-        model=model_name,
-        modalities=output_modalities,
-        extra_body=extra_body,
-        stream=args.stream,
-    )
+    # Test multiple concurrent completions
+    num_concurrent_requests = args.num_concurrent_requests
 
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+        # Submit multiple completion requests concurrently
+        futures = [
+            executor.submit(
+                client.chat.completions.create,
+                messages=[
+                    get_system_prompt(),
+                    prompt,
+                ],
+                model=model_name,
+                modalities=output_modalities,
+                extra_body=extra_body,
+                stream=args.stream,
+            )
+            for _ in range(num_concurrent_requests)
+        ]
+
+        # Wait for all requests to complete and collect results
+        chat_completions = [future.result() for future in concurrent.futures.as_completed(futures)]
+
+    assert len(chat_completions) == num_concurrent_requests
     count = 0
     if not args.stream:
-        for choice in chat_completion.choices:
-            if choice.message.audio:
-                audio_data = base64.b64decode(choice.message.audio.data)
-                audio_file_path = f"audio_{count}.wav"
-                with open(audio_file_path, "wb") as f:
-                    f.write(audio_data)
-                print(f"Audio saved to {audio_file_path}")
-                count += 1
-            elif choice.message.content:
-                print("Chat completion output from text:", choice.message.content)
-    else:
-        printed_content = False
-        for chunk in chat_completion:
-            for choice in chunk.choices:
-                if hasattr(choice, "delta"):
-                    content = getattr(choice.delta, "content", None)
-                else:
-                    content = None
-
-                if getattr(chunk, "modality", None) == "audio" and content:
-                    audio_data = base64.b64decode(content)
+        # Verify all completions succeeded
+        for chat_completion in chat_completions:
+            for choice in chat_completion.choices:
+                if choice.message.audio:
+                    audio_data = base64.b64decode(choice.message.audio.data)
                     audio_file_path = f"audio_{count}.wav"
                     with open(audio_file_path, "wb") as f:
                         f.write(audio_data)
-                    print(f"\nAudio saved to {audio_file_path}")
+                    print(f"Audio saved to {audio_file_path}")
                     count += 1
+                elif choice.message.content:
+                    print("Chat completion output from text:", choice.message.content)
+    else:
+        printed_content = False
+        for chat_completion in chat_completions:
+            for chunk in chat_completion:
+                for choice in chunk.choices:
+                    if hasattr(choice, "delta"):
+                        content = getattr(choice.delta, "content", None)
+                    else:
+                        content = None
 
-                elif getattr(chunk, "modality", None) == "text":
-                    if not printed_content:
-                        printed_content = True
-                        print("\ncontent:", end="", flush=True)
-                    print(content, end="", flush=True)
+                    if getattr(chunk, "modality", None) == "audio" and content:
+                        audio_data = base64.b64decode(content)
+                        audio_file_path = f"audio_{count}.wav"
+                        with open(audio_file_path, "wb") as f:
+                            f.write(audio_data)
+                        print(f"\nAudio saved to {audio_file_path}")
+                        count += 1
+
+                    elif getattr(chunk, "modality", None) == "text":
+                        if not printed_content:
+                            printed_content = True
+                            print("\ncontent:", end="", flush=True)
+                        print(content, end="", flush=True)
 
 
 def parse_args():
@@ -490,6 +534,12 @@ def parse_args():
         action="store_true",
         help="Stream the response.",
     )
+    parser.add_argument(
+        "--num-concurrent-requests",
+        type=int,
+        default=1,
+        help="Number of concurrent requests to send. Default is 1.",
+    )
 
     return parser.parse_args()
 

examples/online_serving/qwen3_tts/README.md

@@ -0,0 +1,170 @@
+# Qwen3-TTS Online Serving
+
+This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
+
+## Supported Models
+
+| Model | Task Type | Description |
+|-------|-----------|-------------|
+| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+# CustomVoice model (default)
+./run_server.sh
+
+# Or specify task type
+./run_server.sh CustomVoice
+./run_server.sh VoiceDesign
+./run_server.sh Base
+```
+
+### 2. Run the Client
+
+```bash
+# CustomVoice: Use predefined speaker
+python openai_speech_client.py \
+    --text "你好，我是通义千问" \
+    --voice Vivian \
+    --language Chinese
+
+# CustomVoice with style instruction
+python openai_speech_client.py \
+    --text "今天天气真好" \
+    --voice Ryan \
+    --instructions "用开心的语气说"
+
+# VoiceDesign: Describe the voice style
+python openai_speech_client.py \
+    --task-type VoiceDesign \
+    --text "哥哥，你回来啦" \
+    --instructions "体现撒娇稚嫩的萝莉女声，音调偏高"
+
+# Base: Voice cloning
+python openai_speech_client.py \
+    --task-type Base \
+    --text "Hello, this is a cloned voice" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio"
+```
+
+### 3. Using curl
+
+```bash
+# Simple TTS request
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "Vivian",
+        "language": "English"
+    }' --output output.wav
+
+# With style instruction
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "I am so excited!",
+        "voice": "Vivian",
+        "instructions": "Speak with great enthusiasm"
+    }' --output excited.wav
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/speech
+```
+
+This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
+
+### Request Body
+
+```json
+{
+    "input": "Text to synthesize",
+    "voice": "Vivian",
+    "response_format": "wav",
+    "task_type": "CustomVoice",
+    "language": "Auto",
+    "instructions": "Optional style instructions",
+    "ref_audio": "URL or base64 for voice cloning",
+    "ref_text": "Reference audio transcript",
+    "x_vector_only_mode": false,
+    "max_new_tokens": 2048
+}
+```
+
+> **Note:** The `model` field is optional when serving a single model, as the server already knows which model is loaded.
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Parameters
+
+### Standard OpenAI Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | required | Text to synthesize |
+| `voice` | string | "Vivian" | Speaker/voice name |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| `model` | string | optional | Model name (optional when serving single model) |
+
+### Qwen3-TTS Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
+| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean |
+| `instructions` | string | "" | Voice style/emotion instructions |
+| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+
+### Voice Clone Parameters (Base task)
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) |
+| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) |
+
+## Python Usage
+
+```python
+import httpx
+
+# Simple request
+response = httpx.post(
+    "http://localhost:8000/v1/audio/speech",
+    json={
+        "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        "input": "Hello world",
+        "voice": "Vivian",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Limitations
+
+- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
+- **Single request**: Batch processing is not yet optimized for online serving.
+
+## Troubleshooting
+
+1. **Connection refused**: Make sure the server is running on the correct port
+2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh
+3. **Unsupported speaker**: Check supported speakers via model documentation
+4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning

examples/online_serving/qwen3_tts/openai_speech_client.py

@@ -0,0 +1,235 @@
+"""OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech endpoint.
+
+This script demonstrates how to use the OpenAI-compatible speech API
+to generate audio from text using Qwen3-TTS models.
+
+Examples:
+    # CustomVoice task (predefined speaker)
+    python openai_speech_client.py --text "Hello, how are you?" --voice Vivian
+
+    # CustomVoice with emotion instruction
+    python openai_speech_client.py --text "I'm so happy!" --voice Vivian \
+        --instructions "Speak with excitement"
+
+    # VoiceDesign task (voice from description)
+    python openai_speech_client.py --text "Hello world" \
+        --task-type VoiceDesign \
+        --instructions "A warm, friendly female voice"
+
+    # Base task (voice cloning)
+    python openai_speech_client.py --text "Hello world" \
+        --task-type Base \
+        --ref-audio "https://example.com/reference.wav" \
+        --ref-text "This is the reference transcript"
+"""
+
+import argparse
+import base64
+import os
+
+import httpx
+
+# Default server configuration
+DEFAULT_API_BASE = "http://localhost:8000"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file to base64 data URL."""
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Detect MIME type from extension
+    audio_path_lower = audio_path.lower()
+    if audio_path_lower.endswith(".wav"):
+        mime_type = "audio/wav"
+    elif audio_path_lower.endswith((".mp3", ".mpeg")):
+        mime_type = "audio/mpeg"
+    elif audio_path_lower.endswith(".flac"):
+        mime_type = "audio/flac"
+    elif audio_path_lower.endswith(".ogg"):
+        mime_type = "audio/ogg"
+    else:
+        mime_type = "audio/wav"  # Default
+
+    with open(audio_path, "rb") as f:
+        audio_bytes = f.read()
+    audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
+    return f"data:{mime_type};base64,{audio_b64}"
+
+
+def run_tts_generation(args) -> None:
+    """Run TTS generation via OpenAI-compatible /v1/audio/speech API."""
+
+    # Build request payload
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "voice": args.voice,
+        "response_format": args.response_format,
+    }
+
+    # Add optional parameters
+    if args.instructions:
+        payload["instructions"] = args.instructions
+    if args.task_type:
+        payload["task_type"] = args.task_type
+    if args.language:
+        payload["language"] = args.language
+    if args.max_new_tokens:
+        payload["max_new_tokens"] = args.max_new_tokens
+
+    # Voice clone parameters (Base task)
+    if args.ref_audio:
+        if args.ref_audio.startswith(("http://", "https://")):
+            payload["ref_audio"] = args.ref_audio
+        else:
+            payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
+    if args.ref_text:
+        payload["ref_text"] = args.ref_text
+    if args.x_vector_only:
+        payload["x_vector_only_mode"] = True
+
+    print(f"Task type: {args.task_type or 'CustomVoice'}")
+    print(f"Text: {args.text}")
+    print(f"Voice: {args.voice}")
+    print("Generating audio...")
+
+    # Make the API call
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    with httpx.Client(timeout=300.0) as client:
+        response = client.post(api_url, json=payload, headers=headers)
+
+    if response.status_code != 200:
+        print(f"Error: {response.status_code}")
+        print(response.text)
+        return
+
+    # Save audio response
+    output_path = args.output or "tts_output.wav"
+    with open(output_path, "wb") as f:
+        f.write(response.content)
+    print(f"Audio saved to: {output_path}")
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+
+    # Server configuration
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default=DEFAULT_API_BASE,
+        help=f"API base URL (default: {DEFAULT_API_BASE})",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_API_KEY,
+        help="API key (default: EMPTY)",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        help="Model name/path",
+    )
+
+    # Task configuration
+    parser.add_argument(
+        "--task-type",
+        "-t",
+        type=str,
+        default=None,
+        choices=["CustomVoice", "VoiceDesign", "Base"],
+        help="TTS task type (default: CustomVoice)",
+    )
+
+    # Input text
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Text to synthesize",
+    )
+
+    # Voice/speaker
+    parser.add_argument(
+        "--voice",
+        type=str,
+        default="Vivian",
+        help="Speaker/voice name (default: Vivian). Options: Vivian, Ryan, etc.",
+    )
+    parser.add_argument(
+        "--language",
+        type=str,
+        default=None,
+        help="Language: Auto, Chinese, English, etc.",
+    )
+    parser.add_argument(
+        "--instructions",
+        type=str,
+        default=None,
+        help="Voice style/emotion instructions",
+    )
+
+    # Base (voice clone) parameters
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio file path or URL for voice cloning (Base task)",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Reference audio transcript for voice cloning (Base task)",
+    )
+    parser.add_argument(
+        "--x-vector-only",
+        action="store_true",
+        help="Use x-vector only mode for voice cloning (no ICL)",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=None,
+        help="Maximum new tokens to generate",
+    )
+
+    # Output
+    parser.add_argument(
+        "--response-format",
+        type=str,
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio output format (default: wav)",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=None,
+        help="Output audio file path (default: tts_output.wav)",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_tts_generation(args)

examples/online_serving/qwen3_tts/run_server.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+# Launch vLLM-Omni server for Qwen3-TTS models
+#
+# Usage:
+#   ./run_server.sh                           # Default: CustomVoice model
+#   ./run_server.sh CustomVoice               # CustomVoice model
+#   ./run_server.sh VoiceDesign               # VoiceDesign model
+#   ./run_server.sh Base                      # Base (voice clone) model
+
+set -e
+
+TASK_TYPE="${1:-CustomVoice}"
+
+case "$TASK_TYPE" in
+    CustomVoice)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+        ;;
+    VoiceDesign)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
+        ;;
+    Base)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+        ;;
+    *)
+        echo "Unknown task type: $TASK_TYPE"
+        echo "Supported: CustomVoice, VoiceDesign, Base"
+        exit 1
+        ;;
+esac
+
+echo "Starting Qwen3-TTS server with model: $MODEL"
+
+vllm-omni serve "$MODEL" \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni

examples/online_serving/text_to_image/README.md

@@ -116,6 +116,7 @@ Use `extra_body` to pass generation parameters:
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |
+| `--cfg-parallel-size`.   | int   | 1       | Number of GPUs for CFG parallelism |
 
 ## Response Format
 

examples/online_serving/text_to_image/openai_chat_client.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Qwen-Image OpenAI-compatible chat client for image generation.
+Qwen-Image OpenAI-compatible image generation client.
 
 Usage:
     python openai_chat_client.py --prompt "A beautiful landscape" --output output.png
@@ -25,14 +25,14 @@ def generate_image(
     negative_prompt: str | None = None,
     num_outputs_per_prompt: int = 1,
 ) -> bytes | None:
-    """Generate an image using the chat completions API.
+    """Generate an image using the images generation API.
 
     Args:
         prompt: Text description of the image
         server_url: Server URL
         height: Image height in pixels
         width: Image width in pixels
-        steps: Number of inference steps
+        steps: Number of diffusion steps
         true_cfg_scale: Qwen-Image CFG scale
         seed: Random seed
         negative_prompt: Negative prompt
@@ -41,33 +41,31 @@ def generate_image(
     Returns:
         Image bytes or None if failed
     """
-    messages = [{"role": "user", "content": prompt}]
+    payload: dict[str, object] = {
+        "prompt": prompt,
+        "response_format": "b64_json",
+        "n": num_outputs_per_prompt,
+    }
+
+    if width is not None and height is not None:
+        payload["size"] = f"{width}x{height}"
+    elif width is not None:
+        payload["size"] = f"{width}x{width}"
+    elif height is not None:
+        payload["size"] = f"{height}x{height}"
 
-    # Build extra_body with generation parameters
-    extra_body = {}
-    if height is not None:
-        extra_body["height"] = height
-    if width is not None:
-        extra_body["width"] = width
     if steps is not None:
-        extra_body["num_inference_steps"] = steps
+        payload["num_inference_steps"] = steps
     if true_cfg_scale is not None:
-        extra_body["true_cfg_scale"] = true_cfg_scale
-    if seed is not None:
-        extra_body["seed"] = seed
+        payload["true_cfg_scale"] = true_cfg_scale
     if negative_prompt:
-        extra_body["negative_prompt"] = negative_prompt
-    extra_body["num_outputs_per_prompt"] = num_outputs_per_prompt
+        payload["negative_prompt"] = negative_prompt
+    if seed is not None:
+        payload["seed"] = seed
 
-    # Build request payload
-    payload = {"messages": messages}
-    if extra_body:
-        payload["extra_body"] = extra_body
-
-    # Send request
     try:
         response = requests.post(
-            f"{server_url}/v1/chat/completions",
+            f"{server_url}/v1/images/generations",
             headers={"Content-Type": "application/json"},
             json=payload,
             timeout=300,
@@ -75,15 +73,13 @@ def generate_image(
         response.raise_for_status()
         data = response.json()
 
-        # Extract image from response
-        content = data["choices"][0]["message"]["content"]
-        if isinstance(content, list) and len(content) > 0:
-            image_url = content[0].get("image_url", {}).get("url", "")
-            if image_url.startswith("data:image"):
-                _, b64_data = image_url.split(",", 1)
-                return base64.b64decode(b64_data)
+        items = data.get("data")
+        if isinstance(items, list) and items:
+            first = items[0].get("b64_json") if isinstance(items[0], dict) else None
+            if isinstance(first, str):
+                return base64.b64decode(first)
 
-        print(f"Unexpected response format: {content}")
+        print(f"Unexpected response format: {data}")
         return None
 
     except Exception as e:

examples/online_serving/text_to_image/run_curl_text_to_image.sh

@@ -1,35 +1,10 @@
 #!/bin/bash
 # Qwen-Image text-to-image curl example
 
-SERVER="${SERVER:-http://localhost:8091}"
-PROMPT="${PROMPT:-a good boy in the ocean}"
-CURRENT_TIME=$(date +%Y%m%d%H%M%S)
-OUTPUT="${OUTPUT:-qwen_image_output_${CURRENT_TIME}.png}"
-
-echo "Generating image..."
-echo "Prompt: $PROMPT"
-echo "Output: $OUTPUT"
-
-curl -s "$SERVER/v1/chat/completions" \
+curl -X POST http://localhost:8091/v1/images/generations \
   -H "Content-Type: application/json" \
-  -d "{
-    \"messages\": [
-      {\"role\": \"user\", \"content\": \"$PROMPT\"}
-    ],
-    \"extra_body\": {
-      \"height\": 1024,
-      \"width\": 1024,
-      \"num_inference_steps\": 100,
-      \"true_cfg_scale\": 4.0,
-      \"seed\": 42,
-      \"num_outputs_per_prompt\": 1
-    }
-  }" | jq -r '.choices[0].message.content[0].image_url.url' | sed 's/^data:image[^,]*,\s*//' | base64 -d > "$OUTPUT"
-
-if [ -f "$OUTPUT" ]; then
-    echo "Image saved to: $OUTPUT"
-    echo "Size: $(du -h "$OUTPUT" | cut -f1)"
-else
-    echo "Failed to generate image"
-    exit 1
-fi
+  -d '{
+    "prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
+    "size": "1024x1024",
+    "seed": 42
+  }' | jq -r '.data[0].b64_json' | base64 -d > dragon.png

mkdocs.yml

@@ -90,6 +90,7 @@ plugins:
       api_root_uri: "api"
       nav_item_prefix: ""  # No prefix in navigation tree (clean names)
       show_full_namespace: false  # Show only module name, not full path
+      on_implicit_namespace_package: skip  # Skip directories without __init__.py (e.g., assets)
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
         - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.12.0rc1"
+version = "0.14.0rc1"
 description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
@@ -36,10 +36,15 @@ dependencies = [
     "accelerate==1.12.0",
     "gradio==5.50",
     "soundfile>=0.13.1",
-    "cache-dit==1.1.8",
+    "cache-dit==1.2.0",
     "tqdm>=4.66.0",
     "torchsde>=0.2.6",  # Required for Stable Audio scheduler
-    # "vllm==0.12.0",  # TODO: fix the entrypoints overwrite problem
+    "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
 ]
 
 [project.optional-dependencies]
@@ -51,7 +56,10 @@ dev = [
     "mypy==1.11.1",
     "pre-commit==4.0.1",
     "openai-whisper>=20250625",
-    "psutil>=7.2.0"
+    "psutil>=7.2.0",
+    "soundfile>=0.13.1",
+    "imageio[ffmpeg]>=0.6.0",
+    "opencv-python>=4.12.0.88"
 ]
 
 docs = [
@@ -194,4 +202,5 @@ extend-ignore-identifiers-re = [
     ".*nothink.*",
     ".*NOTHINK.*",
     ".*nin.*",
+    "Ono_Anna",
 ]

tests/conftest.py

@@ -1,5 +1,15 @@
 import base64
+import datetime
+import io
+import math
 import os
+import random
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+# Set CPU device for CI environments without GPU
+if "VLLM_TARGET_DEVICE" not in os.environ:
+    os.environ["VLLM_TARGET_DEVICE"] = "cpu"
+
 import socket
 import subprocess
 import sys
@@ -7,44 +17,154 @@ import time
 from pathlib import Path
 from typing import Any
 
+import numpy as np
 import psutil
 import pytest
 import torch
-import whisper
 import yaml
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 from vllm.logger import init_logger
-from vllm.utils import get_open_port
+from vllm.utils.network_utils import get_open_port
 
 logger = init_logger(__name__)
 
 
 @pytest.fixture(autouse=True)
-def clean_gpu_memory_between_tests():
-    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
+def default_vllm_config():
+    """Set a default VllmConfig for all tests.
+
+    This fixture is auto-used for all tests to ensure that any test
+    that directly instantiates vLLM CustomOps (e.g., RMSNorm, LayerNorm)
+    or model components has the required VllmConfig context.
+
+    This fixture is required for vLLM 0.14.0+ where CustomOp initialization
+    requires a VllmConfig context set via set_current_vllm_config().
+    """
+    from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+    # Use CPU device if no GPU is available (e.g., in CI environments)
+    has_gpu = torch.cuda.is_available() and torch.cuda.device_count() > 0
+    device = "cuda" if has_gpu else "cpu"
+    device_config = DeviceConfig(device=device)
+
+    with set_current_vllm_config(VllmConfig(device_config=device_config)):
         yield
+
+
+@pytest.fixture(autouse=True)
+def clean_gpu_memory_between_tests():
+    print("\n=== PRE-TEST GPU CLEANUP ===")
+    _run_pre_test_cleanup()
+    yield
+    _run_post_test_cleanup()
+
+
+def _run_pre_test_cleanup():
+    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
+        print("GPU cleanup disabled")
         return
 
-    # Wait for GPU memory to be cleared before starting the test
-    import gc
-
-    from tests.utils import wait_for_gpu_memory_to_clear
+    print("Pre-test GPU status:")
+    _print_simple_gpu_status()
 
     num_gpus = torch.cuda.device_count()
     if num_gpus > 0:
         try:
+            from tests.utils import wait_for_gpu_memory_to_clear
+
             wait_for_gpu_memory_to_clear(
                 devices=list(range(num_gpus)),
                 threshold_ratio=0.1,
             )
-        except ValueError as e:
-            logger.info("Failed to clean GPU memory: %s", e)
+        except Exception as e:
+            print(f"Pre-test cleanup note: {e}")
 
-    yield
 
-    # Clean up GPU memory after the test
+def _run_post_test_cleanup():
+    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
+        return
+
+    import gc
+
     if torch.cuda.is_available():
-        torch.cuda.empty_cache()
         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"""
+
+    print("\n" + "=" * 80)
+    print("NVIDIA GPU Information (nvidia-smi)")
+    print("=" * 80)
+
+    try:
+        nvidia_result = subprocess.run(
+            ["nvidia-smi"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+
+        if nvidia_result.returncode == 0:
+            lines = nvidia_result.stdout.strip().split("\n")
+            for line in lines[:20]:
+                print(line)
+
+            if len(lines) > 20:
+                print(f"... (showing first 20 of {len(lines)} lines)")
+        else:
+            print("nvidia-smi command failed")
+
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        print("nvidia-smi not available or timed out")
+    except Exception as e:
+        print(f"Error running nvidia-smi: {e}")
+
+    print("\n" + "=" * 80)
+    print("Detailed GPU Processes (nvidia-smi pmon)")
+    print("=" * 80)
+
+    try:
+        pmon_result = subprocess.run(
+            ["nvidia-smi", "pmon", "-c", "1"],
+            capture_output=True,
+            text=True,
+            timeout=3,
+        )
+
+        if pmon_result.returncode == 0 and pmon_result.stdout.strip():
+            print(pmon_result.stdout)
+        else:
+            print("No active GPU processes found via nvidia-smi pmon")
+
+    except Exception:
+        print("nvidia-smi pmon not available")
+
+    print("\n" + "=" * 80)
+    print("System Processes with GPU keywords")
+    print("=" * 80)
 
 
 def dummy_messages_from_mix_data(
@@ -91,27 +211,397 @@ def dummy_messages_from_mix_data(
     return messages
 
 
-def cosine_similarity_text(s1, s2):
-    """
-        Calculate cosine similarity between two text strings.
-        Notes:
-    ------
-    - Higher score means more similar texts
-    - Score of 1.0 means identical word composition (bag-of-words)
-    - Score of 0.0 means completely different vocabulary
-    """
-    from sklearn.feature_extraction.text import CountVectorizer
-    from sklearn.metrics.pairwise import cosine_similarity
+def generate_synthetic_audio(
+    duration: int,  # seconds
+    num_channels: int,  # 1：Mono，2：Stereo 5：5.1 surround sound
+    sample_rate: int = 48000,  # Default use 48000Hz.
+    save_to_file: bool = False,
+) -> dict[str, Any]:
+    """ "Generate synthetic audio with rain."""
+    import soundfile as sf
 
-    vectorizer = CountVectorizer().fit_transform([s1, s2])
-    vectors = vectorizer.toarray()
-    return cosine_similarity([vectors[0]], [vectors[1]])[0][0]
+    # Initialize audio data array
+    num_samples = int(sample_rate * duration)
+    audio_data = np.zeros((num_samples, num_channels), dtype=np.float32)
+
+    # Configure parameters based on rain intensity
+    drop_density = 10  # Number of raindrops per second
+    drop_volume = 0.15  # Volume of individual raindrops
+    background_volume = 0.02  # Volume of background rain noise
+
+    # Pink noise sounds more natural than white noise for rain
+    white_noise = np.random.randn(num_samples)
+    pink_noise = np.convolve(white_noise, np.ones(8) / 8, mode="same")
+    pink_noise = pink_noise / np.max(np.abs(pink_noise)) if np.max(np.abs(pink_noise)) > 0 else pink_noise
+    bg_noise = pink_noise * background_volume
+
+    # Add background noise to all channels
+    for ch in range(num_channels):
+        audio_data[:, ch] += bg_noise
+
+    # Total number of raindrops = density × duration × channels for stereo effect
+    total_drops = int(drop_density * duration * num_channels)
+
+    for _ in range(total_drops):
+        # Random timing for raindrop
+        drop_time = random.uniform(0, duration)
+
+        # Random duration of raindrop sound (0.01-0.05 seconds)
+        drop_duration = random.uniform(0.01, 0.05)
+
+        # Random frequency gives variation in raindrop pitch
+        drop_freq = random.uniform(500, 5000)  # Hz
+
+        # Random channel selection for stereo positioning
+        channel = random.randint(0, num_channels - 1)
+
+        # Calculate sample positions for this raindrop
+        start_sample = int(drop_time * sample_rate)
+        drop_samples = int(drop_duration * sample_rate)
+        end_sample = min(start_sample + drop_samples, num_samples)
+
+        if start_sample < end_sample:
+            # Generate the raindrop sound
+            num_drop_samples = end_sample - start_sample
+            t = np.arange(num_drop_samples) / sample_rate
+
+            # Basic sine wave for raindrop sound
+            drop_sound = drop_volume * np.sin(2 * math.pi * drop_freq * t)
+
+            # Apply envelope for natural attack and decay
+            envelope = np.ones(num_drop_samples)
+            attack_samples = int(num_drop_samples * 0.1)  # 10% of samples for attack
+            decay_samples = num_drop_samples - attack_samples
+
+            if attack_samples > 0:
+                # Linear attack: volume increases from 0 to 1
+                envelope[:attack_samples] = np.linspace(0, 1, attack_samples)
+
+            if decay_samples > 0:
+                # Exponential decay for natural sound fade
+                decay = np.exp(-8 * t[attack_samples:] / drop_duration)
+                envelope[attack_samples:] = decay
+
+            # Apply envelope to raindrop sound
+            drop_sound *= envelope
+
+            # Add raindrop sound to selected channel
+            audio_data[start_sample:end_sample, channel] += drop_sound
+
+    # Step 3: Add simple reverb effect for realism
+    # Reverb simulates sound reflections in environment
+    if duration > 2:
+        # Single delay reverb (100ms delay)
+        delay_samples = int(0.1 * sample_rate)
+        if delay_samples < num_samples:
+            for ch in range(num_channels):
+                delayed = np.zeros(num_samples)
+                delayed[delay_samples:] = audio_data[:-delay_samples, ch] * 0.3
+                audio_data[:, ch] += delayed
+
+    # Step 4: Normalize audio to prevent clipping
+    # Find maximum amplitude and scale to 80% of maximum volume
+    max_amp = np.max(np.abs(audio_data))
+    if max_amp > 0:
+        audio_data = audio_data / max_amp * 0.8
+
+    # Handle file saving
+    audio_bytes = None
+
+    if save_to_file:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_path = f"audio_{num_channels}ch_{timestamp}.wav"
+
+        try:
+            sf.write(output_path, audio_data, sample_rate, format="WAV", subtype="PCM_16")
+            print(f"Audio saved: {output_path}")
+
+            with open(output_path, "rb") as f:
+                audio_bytes = f.read()
+        except Exception as e:
+            print(f"Save failed: {e}")
+            save_to_file = False
+
+    # If not saving or save failed, create in memory
+    if not save_to_file or audio_bytes is None:
+        buffer = io.BytesIO()
+        sf.write(buffer, audio_data, sample_rate, format="WAV", subtype="PCM_16")
+        buffer.seek(0)
+        audio_bytes = buffer.read()
+
+    # Return result
+    base64_audio = base64.b64encode(audio_bytes).decode("utf-8")
+    result = {
+        "base64": base64_audio,
+    }
+    if save_to_file and output_path:
+        result["file_path"] = output_path
+
+    return result
+
+
+def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> str:
+    """Generate synthetic video with bouncing balls and return base64 string."""
+
+    import cv2
+    import imageio
+
+    # Create random balls
+    num_balls = random.randint(3, 8)
+    balls = []
+
+    for _ in range(num_balls):
+        radius = min(width, height) // 8
+        if radius < 1:
+            raise ValueError(f"Video dimensions ({width}x{height}) are too small for synthetic video generation")
+        x = random.randint(radius, width - radius)
+        y = random.randint(radius, height - radius)
+
+        speed = random.uniform(3.0, 8.0)
+        angle = random.uniform(0, 2 * math.pi)
+        vx = speed * math.cos(angle)
+        vy = speed * math.sin(angle)
+
+        # OpenCV uses BGR format, but imageio expects RGB
+        # We'll create in BGR first, then convert to RGB later
+        color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))
+
+        balls.append({"x": x, "y": y, "vx": vx, "vy": vy, "radius": radius, "color_bgr": color_bgr})
+
+    # Generate video frames
+    video_frames = []
+
+    for frame_idx in range(num_frames):
+        # Create black background (BGR format)
+        frame_bgr = np.zeros((height, width, 3), dtype=np.uint8)
+
+        for ball in balls:
+            # Update position
+            ball["x"] += ball["vx"]
+            ball["y"] += ball["vy"]
+
+            # Boundary collision detection
+            if ball["x"] - ball["radius"] <= 0 or ball["x"] + ball["radius"] >= width:
+                ball["vx"] = -ball["vx"]
+                ball["x"] = max(ball["radius"], min(width - ball["radius"], ball["x"]))
+
+            if ball["y"] - ball["radius"] <= 0 or ball["y"] + ball["radius"] >= height:
+                ball["vy"] = -ball["vy"]
+                ball["y"] = max(ball["radius"], min(height - ball["radius"], ball["y"]))
+
+            # Use cv2 to draw circle
+            x, y = int(ball["x"]), int(ball["y"])
+            radius = ball["radius"]
+
+            # Draw solid circle (main circle)
+            cv2.circle(frame_bgr, (x, y), radius, ball["color_bgr"], -1)
+
+            # Add simple 3D effect: draw a brighter center
+            if radius > 3:  # Only add highlight when radius is large enough
+                highlight_radius = max(1, radius // 2)
+                highlight_x = max(highlight_radius, min(x - radius // 4, width - highlight_radius))
+                highlight_y = max(highlight_radius, min(y - radius // 4, height - highlight_radius))
+
+                # Create highlight color (brighter)
+                highlight_color = tuple(min(c + 40, 255) for c in ball["color_bgr"])
+                cv2.circle(frame_bgr, (highlight_x, highlight_y), highlight_radius, highlight_color, -1)
+
+        # Convert BGR to RGB for imageio
+        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+        video_frames.append(frame_rgb)
+
+    video_bytes = None
+    saved_file_path = None
+
+    buffer = io.BytesIO()
+    writer_kwargs = {
+        "format": "mp4",
+        "fps": 30,
+        "codec": "libx264",
+        "quality": 7,
+        "pixelformat": "yuv420p",
+        "macro_block_size": 16,
+        "ffmpeg_params": [
+            "-preset",
+            "medium",
+            "-crf",
+            "23",
+            "-movflags",
+            "+faststart",
+            "-pix_fmt",
+            "yuv420p",
+            "-vf",
+            f"scale={width}:{height}",
+        ],
+    }
+
+    if save_to_file:
+        import datetime
+
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_path = f"video_{width}x{height}_{timestamp}.mp4"
+        try:
+            with imageio.get_writer(output_path, **writer_kwargs) as writer:
+                for frame in video_frames:
+                    writer.append_data(frame)
+
+            saved_file_path = output_path
+            print(f"Video saved to: {saved_file_path}")
+            with open(output_path, "rb") as f:
+                video_bytes = f.read()
+
+        except Exception as e:
+            print(f"Warning: Failed to save video to file {output_path}: {e}")
+            save_to_file = False
+
+    if not save_to_file or video_bytes is None:
+        with imageio.get_writer(buffer, **writer_kwargs) as writer:
+            for frame in video_frames:
+                writer.append_data(frame)
+
+        buffer.seek(0)
+        video_bytes = buffer.read()
+
+    base64_video = base64.b64encode(video_bytes).decode("utf-8")
+
+    result = {
+        "base64": base64_video,
+    }
+    if save_to_file and saved_file_path:
+        result["file_path"] = saved_file_path
+
+    return result
+
+
+def generate_synthetic_image(width: int, height: int, save_to_file: bool = False) -> Any:
+    """Generate synthetic image with randomly colored squares and return base64 string."""
+    from PIL import Image, ImageDraw
+
+    # Create white background
+    image = Image.new("RGB", (width, height), (255, 255, 255))
+    draw = ImageDraw.Draw(image)
+
+    # Generate random number of squares
+    num_squares = random.randint(3, 8)
+
+    for _ in range(num_squares):
+        # Random square size
+        square_size = random.randint(min(width, height) // 8, min(width, height) // 4)
+
+        # Random position
+        x = random.randint(0, width - square_size - 1)
+        y = random.randint(0, height - square_size - 1)
+
+        # Random color
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+
+        # Random border width
+        border_width = random.randint(1, 5)
+
+        # Draw square
+        draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width)
+
+    # Handle file saving
+    image_bytes = None
+    saved_file_path = None
+
+    if save_to_file:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_path = f"image_{width}x{height}_{timestamp}.jpg"
+
+        try:
+            # Save image to file
+            image.save(output_path, format="JPEG", quality=85, optimize=True)
+            saved_file_path = output_path
+            print(f"Image saved to: {saved_file_path}")
+
+            # Read file for base64 encoding
+            with open(output_path, "rb") as f:
+                image_bytes = f.read()
+
+        except Exception as e:
+            print(f"Warning: Failed to save image to file {output_path}: {e}")
+            save_to_file = False
+
+    # If not saving or save failed, create in memory
+    if not save_to_file or image_bytes is None:
+        buffer = io.BytesIO()
+        image.save(buffer, format="JPEG", quality=85, optimize=True)
+        buffer.seek(0)
+        image_bytes = buffer.read()
+
+    # Generate base64
+    base64_image = base64.b64encode(image_bytes).decode("utf-8")
+
+    # Return result
+    result = {
+        "base64": base64_image,
+    }
+    if save_to_file and saved_file_path:
+        result["file_path"] = saved_file_path
+
+    return result
+
+
+def preprocess_text(text):
+    import re
+
+    word_to_num = {
+        "zero": "0",
+        "one": "1",
+        "two": "2",
+        "three": "3",
+        "four": "4",
+        "five": "5",
+        "six": "6",
+        "seven": "7",
+        "eight": "8",
+        "nine": "9",
+        "ten": "10",
+    }
+
+    for word, num in word_to_num.items():
+        pattern = r"\b" + re.escape(word) + r"\b"
+        text = re.sub(pattern, num, text, flags=re.IGNORECASE)
+
+    text = re.sub(r"[^\w\s]", "", text)
+    text = re.sub(r"\s+", " ", text)
+    return text.lower().strip()
+
+
+def cosine_similarity_text(text1, text2, n: int = 3):
+    from collections import Counter
+
+    if not text1 or not text2:
+        return 0.0
+
+    text1 = preprocess_text(text1)
+    text2 = preprocess_text(text2)
+
+    ngrams1 = [text1[i : i + n] for i in range(len(text1) - n + 1)]
+    ngrams2 = [text2[i : i + n] for i in range(len(text2) - n + 1)]
+
+    counter1 = Counter(ngrams1)
+    counter2 = Counter(ngrams2)
+
+    all_ngrams = set(counter1.keys()) | set(counter2.keys())
+    vec1 = [counter1.get(ng, 0) for ng in all_ngrams]
+    vec2 = [counter2.get(ng, 0) for ng in all_ngrams]
+
+    dot_product = sum(a * b for a, b in zip(vec1, vec2))
+    norm1 = sum(a * a for a in vec1) ** 0.5
+    norm2 = sum(b * b for b in vec2) ** 0.5
+
+    if norm1 == 0 or norm2 == 0:
+        return 0.0
+    return dot_product / (norm1 * norm2)
 
 
 def convert_audio_to_text(audio_data):
     """
     Convert base64 encoded audio data to text using speech recognition.
     """
+    import whisper
 
     audio_data = base64.b64decode(audio_data)
     output_path = f"./test_{int(time.time())}"
@@ -119,8 +609,14 @@ def convert_audio_to_text(audio_data):
         audio_file.write(audio_data)
 
     print(f"audio data is saved: {output_path}")
+
     model = whisper.load_model("base")
-    text = model.transcribe(output_path)["text"]
+    text = model.transcribe(
+        output_path,
+        temperature=0.0,
+        word_timestamps=True,
+        condition_on_previous_text=False,
+    )["text"]
     if text:
         return text
     else:
@@ -129,75 +625,238 @@ def convert_audio_to_text(audio_data):
 
 def modify_stage_config(
     yaml_path: str,
-    stage_updates: dict[int, dict[str, Any]],
+    updates: dict[str, Any],
+    deletes: dict[str, Any] = None,
 ) -> str:
     """
-    Batch modify configurations for multiple stages in a YAML file.
+    Modify configurations in a YAML file, supporting both top-level and stage-specific modifications,
+    including addition, modification, and deletion of configurations.
 
     Args:
         yaml_path: Path to the YAML configuration file.
-        stage_updates: Dictionary where keys are stage IDs and values are dictionaries of
-                      modifications for that stage. Each modification dictionary uses
-                      dot-separated paths as keys and new configuration values as values.
-                      Example: {
-                          0: {'engine_args.max_model_len': 5800},
-                          1: {'runtime.max_batch_size': 2}
-                      }
+        updates: Dictionary containing both top-level and stage-specific modifications to add or update.
+                Format: {
+                    'async_chunk': True,
+                    'stage_args': {
+                        0: {'engine_args.max_model_len': 5800},
+                        1: {'runtime.max_batch_size': 2}
+                    }
+                }
+        deletes: Dictionary containing configurations to delete.
+                Format: {
+                    'old_config': None,  # Delete entire key
+                    'stage_args': {
+                        0: ['engine_args.old_param'],
+                        1: ['runtime.unused_setting']
+                    }
+                }
 
     Returns:
         str: Path to the newly created modified YAML file with timestamp suffix.
-
-    Example:
-        >>> output_file = modify_stage_config(
-        ...     'config.yaml',
-        ...     {
-        ...         0: {'engine_args.max_model_len': 5800},
-        ...         1: {'runtime.max_batch_size': 2}
-        ...     }
-        ... )
-        >>> print(f"Modified configuration saved to: {output_file}")
-        Modified configuration saved to: config_1698765432.yaml
     """
     path = Path(yaml_path)
     if not path.exists():
         raise FileNotFoundError(f"yaml does not exist: {path}")
+
     try:
         with open(yaml_path, encoding="utf-8") as f:
             config = yaml.safe_load(f) or {}
     except Exception as e:
         raise ValueError(f"Cannot parse YAML file: {e}")
 
-    stage_args = config.get("stage_args", [])
-    if not stage_args:
-        raise ValueError("the stage_args does not exist")
+    # Helper function to apply update
+    def apply_update(config_dict: dict, key_path: str, value: Any) -> None:
+        """Apply update to dictionary using dot-separated path."""
+        # Handle direct list assignment (e.g., engine_input_source: [1, 2])
+        if "." not in key_path:
+            # Simple key, set directly
+            config_dict[key_path] = value
+            return
 
-    for stage_id, config_dict in stage_updates.items():
-        target_stage = None
-        for stage in stage_args:
-            if stage.get("stage_id") == stage_id:
-                target_stage = stage
-                break
+        current = config_dict
+        keys = key_path.split(".")
 
-        if target_stage is None:
-            available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
-            raise KeyError(f"Stage ID {stage_id} is not exist, available IDs: {available_ids}")
+        for i in range(len(keys) - 1):
+            key = keys[i]
 
-        for key_path, value in config_dict.items():
-            current = target_stage
-            keys = key_path.split(".")
-            for i in range(len(keys) - 1):
-                key = keys[i]
+            # Handle list indices
+            if key.isdigit() and isinstance(current, list):
+                index = int(key)
+                if index < 0:
+                    raise ValueError(f"Negative list index not allowed: {index}")
+                if index >= len(current):
+                    # Expand list if needed
+                    while len(current) <= index:
+                        # If we need to go deeper (more keys after this), create a dict
+                        # Otherwise, create None placeholder
+                        current.append({} if i < len(keys) - 2 else None)
+                current = current[index]
+            elif isinstance(current, dict):
+                # Handle dictionary keys
                 if key not in current:
-                    raise KeyError(f"the {'.'.join(keys[: i + 1])} does not exist")
-
-                elif not isinstance(current[key], dict) and i < len(keys) - 2:
-                    raise ValueError(f"{'.'.join(keys[: i + 1])}' cannot continue deeper because it's not a dict")
+                    # If there are more keys after this, create appropriate structure
+                    if i < len(keys) - 1:
+                        # Check if next key is a digit (list index) or string (dict key)
+                        if keys[i + 1].isdigit():
+                            current[key] = []
+                        else:
+                            current[key] = {}
+                    else:
+                        # This is the last key, create based on value type
+                        current[key] = [] if isinstance(value, list) else {}
+                elif not isinstance(current[key], (dict, list)) and i < len(keys) - 1:
+                    # If current value is not dict/list but we need to go deeper, replace it
+                    if keys[i + 1].isdigit():
+                        current[key] = []
+                    else:
+                        current[key] = {}
                 current = current[key]
-            current[keys[-1]] = value
+            else:
+                # Current is not a dict or list, cannot traverse further
+                raise TypeError(
+                    f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}"
+                )
+
+        # Set the final value
+        last_key = keys[-1]
+        if isinstance(current, list) and last_key.isdigit():
+            # Setting a value in a list by index
+            index = int(last_key)
+            if index < 0:
+                raise ValueError(f"Negative list index not allowed: {index}")
+            if index >= len(current):
+                # Expand list if needed
+                while len(current) <= index:
+                    current.append(None)
+            current[index] = value
+        elif isinstance(current, dict):
+            # Special case: if the value is a list and we're setting a top-level key
+            # Example: updating engine_input_source with [1, 2]
+            current[last_key] = value
+        else:
+            # Current is not a dict, cannot set key
+            raise TypeError(f"Cannot set value at {key_path}. Current type is {type(current).__name__}, expected dict.")
+
+    # Helper function to delete by path
+    def delete_by_path(config_dict: dict, path: str) -> None:
+        """Delete configuration by dot-separated path."""
+        if not path:
+            return
+
+        current = config_dict
+        keys = path.split(".")
+
+        # Traverse to the parent
+        for i in range(len(keys) - 1):
+            key = keys[i]
+
+            # Handle list indices
+            if key.isdigit() and isinstance(current, list):
+                index = int(key)
+                if index < 0 or index >= len(current):
+                    raise KeyError(f"List index {index} out of bounds")
+                current = current[index]
+            elif isinstance(current, dict):
+                if key not in current:
+                    raise KeyError(f"Path {'.'.join(keys[: i + 1])} does not exist")
+                current = current[key]
+            else:
+                raise TypeError(
+                    f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}"
+                )
+
+        # Delete the item
+        last_key = keys[-1]
+
+        if isinstance(current, list) and last_key.isdigit():
+            index = int(last_key)
+            if index < 0 or index >= len(current):
+                raise KeyError(f"List index {index} out of bounds")
+            del current[index]
+        elif isinstance(current, dict) and last_key in current:
+            del current[last_key]
+        else:
+            raise KeyError(f"Path {path} does not exist")
+
+    # Apply deletions first
+    if deletes:
+        for key, value in deletes.items():
+            if key == "stage_args":
+                if value and isinstance(value, dict):
+                    stage_args = config.get("stage_args", [])
+                    if not stage_args:
+                        raise ValueError("stage_args does not exist in config")
+
+                    for stage_id, delete_paths in value.items():
+                        if not delete_paths:
+                            continue
+
+                        # Find stage by ID
+                        target_stage = None
+                        for stage in stage_args:
+                            if stage.get("stage_id") == stage_id:
+                                target_stage = stage
+                                break
+
+                        if target_stage is None:
+                            available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
+                            raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
+
+                        # Delete specified paths in this stage
+                        for path in delete_paths:
+                            if path:  # Skip empty paths
+                                delete_by_path(target_stage, path)
+            elif "." in key:
+                # Delete using dot-separated path
+                delete_by_path(config, key)
+            elif value is None and key in config:
+                # Delete entire key
+                del config[key]
+
+    # Apply updates
+    for key, value in updates.items():
+        if key == "stage_args":
+            if value and isinstance(value, dict):
+                stage_args = config.get("stage_args", [])
+                if not stage_args:
+                    raise ValueError("stage_args does not exist in config")
+
+                for stage_id, stage_updates in value.items():
+                    # Find stage by ID
+                    target_stage = None
+                    for stage in stage_args:
+                        if stage.get("stage_id") == stage_id:
+                            target_stage = stage
+                            break
+
+                    if target_stage is None:
+                        available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
+                        raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
+
+                    # Apply updates to this stage
+                    for path, val in stage_updates.items():
+                        # Check if this is a simple key (not dot-separated)
+                        # Example: 'engine_input_source' vs 'engine_args.max_model_len'
+                        if "." not in path:
+                            # Direct key assignment (e.g., updating a list value)
+                            target_stage[path] = val
+                        else:
+                            # Dot-separated path (e.g., nested dict access)
+                            apply_update(target_stage, path, val)
+        elif "." in key:
+            # Apply using dot-separated path
+            apply_update(config, key, value)
+        else:
+            # Direct top-level key
+            config[key] = value
+
+    # Save to new file with timestamp
+    timestamp = int(time.time())
+    base_name = yaml_path.rsplit(".", 1)[0] if "." in yaml_path else yaml_path
+    output_path = f"{base_name}_{timestamp}.yaml"
 
-    output_path = f"{yaml_path.split('.')[0]}_{int(time.time())}.yaml"
     with open(output_path, "w", encoding="utf-8") as f:
-        yaml.dump(config, f, default_flow_style=False, sort_keys=False, allow_unicode=True, indent=2)
+        yaml.dump(config, f, default_flow_style=None, sort_keys=False, allow_unicode=True, indent=2)
 
     return output_path
 
@@ -264,24 +923,32 @@ class OmniServer:
         raise RuntimeError(f"Server failed to start within {max_wait} seconds")
 
     def _kill_process_tree(self, pid):
-        """kill process and its children"""
+        """kill process and its children with verification"""
         try:
             parent = psutil.Process(pid)
             children = parent.children(recursive=True)
+
+            # Get all PIDs first
+            all_pids = [pid] + [child.pid for child in children]
+
+            # Terminate children
             for child in children:
                 try:
                     child.terminate()
                 except psutil.NoSuchProcess:
                     pass
 
+            # Wait for children
             gone, still_alive = psutil.wait_procs(children, timeout=10)
 
+            # Kill remaining children
             for child in still_alive:
                 try:
                     child.kill()
                 except psutil.NoSuchProcess:
                     pass
 
+            # Terminate parent
             try:
                 parent.terminate()
                 parent.wait(timeout=10)
@@ -291,6 +958,24 @@ class OmniServer:
                 except psutil.NoSuchProcess:
                     pass
 
+            # VERIFICATION: Check if all processes are gone
+            time.sleep(1)  # Give system time
+            alive_processes = []
+            for check_pid in all_pids:
+                if psutil.pid_exists(check_pid):
+                    alive_processes.append(check_pid)
+
+            if alive_processes:
+                print(f"Warning: Processes still alive: {alive_processes}")
+                # Optional: Try system kill
+                import subprocess
+
+                for alive_pid in alive_processes:
+                    try:
+                        subprocess.run(["kill", "-9", str(alive_pid)], timeout=2)
+                    except Exception as e:
+                        print(f"Cleanup failed: {e}")
+
         except psutil.NoSuchProcess:
             pass
 
@@ -300,31 +985,6 @@ class OmniServer:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self.proc:
-            try:
-                parent = psutil.Process(self.proc.pid)
-                children = parent.children(recursive=True)
-                for child in children:
-                    try:
-                        child.terminate()
-                    except psutil.NoSuchProcess:
-                        pass
-
-                gone, still_alive = psutil.wait_procs(children, timeout=10)
-
-                for child in still_alive:
-                    try:
-                        child.kill()
-                    except psutil.NoSuchProcess:
-                        pass
-
-                try:
-                    parent.terminate()
-                    parent.wait(timeout=10)
-                except (psutil.NoSuchProcess, psutil.TimeoutExpired):
-                    try:
-                        parent.kill()
-                    except psutil.NoSuchProcess:
-                        pass
-
-            except psutil.NoSuchProcess:
-                pass
+            self._kill_process_tree(self.proc.pid)
+        _run_post_test_cleanup()
+        cleanup_dist_env_and_memory()

tests/diffusion/attention/test_attention_sp.py

@@ -170,6 +170,11 @@ def test_sequence_parallel(
     """Test Ulysses attention by comparing with and without SP enabled."""
     sequence_parallel_size = ulysses_degree * ring_degree
 
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < sequence_parallel_size:
+        pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")
+
     # Create temporary files to share results between processes
     with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
         baseline_output_file = f.name

tests/diffusion/distributed/test_comm.py

@@ -48,6 +48,11 @@ def test_4d_identity(
     use_sync: bool,
 ):
     """Test that two consecutive all-to-all operations return the original input."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     # Ensure num_heads is divisible by world_size
     if num_heads % world_size != 0:
         pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
@@ -177,6 +182,11 @@ def test_5d_identity(
     use_sync: bool,
 ):
     """Test that two consecutive all-to-all operations return the original input."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     # Ensure num_heads is divisible by world_size
     if num_heads % world_size != 0:
         pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
@@ -305,6 +315,11 @@ def test_ring_p2p(
     head_size: int,
 ):
     """Test Ring P2P communication (send_recv)."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     torch.multiprocessing.spawn(
         _test_ring_p2p_worker,
         args=(world_size, dtype, batch_size, num_heads, head_size),

tests/diffusion/lora/test_base_linear.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from vllm_omni.diffusion.lora.layers.base_linear import DiffusionBaseLinearLayerWithLoRA
+
+
+@dataclass
+class _DummyLoRAConfig:
+    fully_sharded_loras: bool = False
+
+
+class _DummyQuantMethod:
+    def __init__(self, weight: torch.Tensor):
+        self._weight = weight
+
+    def apply(self, _base_layer, x: torch.Tensor, bias: torch.Tensor | None):
+        y = x @ self._weight.t()
+        if bias is not None:
+            y = y + bias
+        return y
+
+
+def test_diffusion_base_linear_apply_multi_slice():
+    # Build a fake diffusion LoRA layer with 2 slices and rank=2.
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    in_dim = 3
+    out_slices = (2, 1)
+    rank = 2
+
+    # Base weight: identity-ish mapping to make base output easy to reason about.
+    base_weight = torch.tensor(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(base_weight)
+
+    # Allocate stacked weights: (max_loras=1, 1, rank, in_dim) and (1, 1, out_slice, rank)
+    a0 = torch.zeros((1, 1, rank, in_dim))
+    b0 = torch.zeros((1, 1, out_slices[0], rank))
+    a1 = torch.zeros((1, 1, rank, in_dim))
+    b1 = torch.zeros((1, 1, out_slices[1], rank))
+
+    # Slice 0: delta0 = (x @ A0.T) @ B0.T
+    A0 = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])  # (2, 3)
+    B0 = torch.tensor([[1.0, 0.0], [0.0, 1.0]])  # (2, 2)
+    a0[0, 0, :, :] = A0
+    b0[0, 0, :, :] = B0
+
+    # Slice 1: delta1 = (x @ A1.T) @ B1.T
+    A1 = torch.tensor([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])  # (2, 3)
+    B1 = torch.tensor([[2.0, 0.0]])  # (1, 2)
+    a1[0, 0, :, :] = A1
+    b1[0, 0, :, :] = B1
+
+    layer.lora_a_stacked = (a0, a1)
+    layer.lora_b_stacked = (b0, b1)
+    layer.output_slices = out_slices
+
+    x = torch.tensor([[1.0, 2.0, 3.0]])
+    out = layer.apply(x)
+
+    # Base output is identity: [1,2,3]
+    base_out = x @ base_weight.t()
+    # delta0:
+    # (x @ A0.T) = [1,2]
+    # [1,2] @ B0.T = [1,2]
+    delta0 = torch.tensor([[1.0, 2.0]])
+    # delta1:
+    # (x @ A1.T) = [3,1]
+    # [3,1] @ B1.T = [6]
+    delta1 = torch.tensor([[6.0]])
+    expected = torch.cat([base_out[:, :2] + delta0, base_out[:, 2:3] + delta1], dim=-1)
+    assert torch.allclose(out, expected)
+
+
+def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
+    # Verify that after reset_lora(), apply() skips LoRA matmuls even if the
+    # LoRA tensors are still allocated and non-empty.
+    from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    in_dim = 2
+    out_dim = 2
+    rank = 1
+
+    base_weight = torch.eye(in_dim)
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(base_weight)
+
+    a = torch.ones((1, 1, rank, in_dim))
+    b = torch.tensor([[[[1.0], [2.0]]]])  # (1,1,out_dim,rank)
+
+    layer.lora_a_stacked = (a,)
+    layer.lora_b_stacked = (b,)
+    layer.output_slices = (out_dim,)
+    layer._diffusion_lora_active_slices = (True,)
+
+    x = torch.tensor([[1.0, 2.0]])
+    out_active = layer.apply(x)
+    assert torch.allclose(out_active, torch.tensor([[4.0, 8.0]]))
+
+    monkeypatch.setattr(BaseLinearLayerWithLoRA, "reset_lora", lambda self, index: None)
+    layer.reset_lora(0)
+
+    assert layer._diffusion_lora_active_slices == (False,)
+    out_inactive = layer.apply(x)
+    assert torch.allclose(out_inactive, x)
+
+
+def test_diffusion_base_linear_apply_respects_inactive_slices():
+    # Build a fake diffusion LoRA layer with 2 slices and rank=2.
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    in_dim = 3
+    out_slices = (2, 1)
+    rank = 2
+
+    base_weight = torch.tensor(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(base_weight)
+
+    a0 = torch.zeros((1, 1, rank, in_dim))
+    b0 = torch.zeros((1, 1, out_slices[0], rank))
+    a1 = torch.zeros((1, 1, rank, in_dim))
+    b1 = torch.zeros((1, 1, out_slices[1], rank))
+
+    A0 = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])  # (2, 3)
+    B0 = torch.tensor([[1.0, 0.0], [0.0, 1.0]])  # (2, 2)
+    a0[0, 0, :, :] = A0
+    b0[0, 0, :, :] = B0
+
+    A1 = torch.tensor([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])  # (2, 3)
+    B1 = torch.tensor([[2.0, 0.0]])  # (1, 2)
+    a1[0, 0, :, :] = A1
+    b1[0, 0, :, :] = B1
+
+    layer.lora_a_stacked = (a0, a1)
+    layer.lora_b_stacked = (b0, b1)
+    layer.output_slices = out_slices
+    layer._diffusion_lora_active_slices = (True, False)
+
+    x = torch.tensor([[1.0, 2.0, 3.0]])
+    out = layer.apply(x)
+
+    # Only the first slice should be adapted.
+    expected = torch.tensor([[2.0, 4.0, 3.0]])
+    assert torch.allclose(out, expected)

tests/diffusion/lora/test_lora_manager.py

@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import torch
+from vllm.lora.lora_weights import LoRALayerWeights
+from vllm.lora.utils import get_supported_lora_modules
+from vllm.model_executor.layers.linear import LinearBase
+
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.lora.request import LoRARequest
+
+
+class _DummyLoRALayer:
+    def __init__(self, n_slices: int, output_slices: tuple[int, ...]):
+        self.n_slices = n_slices
+        self.output_slices = output_slices
+        self.set_calls: list[
+            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
+        ] = []
+        self.reset_calls: int = 0
+
+    def set_lora(self, index: int, lora_a, lora_b):
+        assert index == 0
+        self.set_calls.append((lora_a, lora_b))
+
+    def reset_lora(self, index: int):
+        assert index == 0
+        self.reset_calls += 1
+
+
+class _FakeLinearBase(LinearBase):
+    def __init__(self):
+        torch.nn.Module.__init__(self)
+
+
+def test_lora_manager_supported_modules_are_stable_with_wrapped_layers(monkeypatch):
+    # Simulate a pipeline that already contains LoRA wrappers where the original
+    # LinearBase is nested under ".base_layer".
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DummyBaseLayerWithLoRA(torch.nn.Module):
+        def __init__(self, base_layer: torch.nn.Module):
+            super().__init__()
+            self.base_layer = base_layer
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    pipeline = torch.nn.Module()
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.foo = _DummyBaseLayerWithLoRA(_FakeLinearBase())
+
+    # vLLM helper would see only the nested LinearBase and yield "base_layer".
+    assert get_supported_lora_modules(pipeline) == ["base_layer"]
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    assert "foo" in manager._supported_lora_modules
+    assert "base_layer" not in manager._supported_lora_modules
+
+
+def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DummyBaseLayerWithLoRA(torch.nn.Module):
+        def __init__(self, base_layer: torch.nn.Module):
+            super().__init__()
+            self.base_layer = base_layer
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        if isinstance(layer, _FakeLinearBase):
+            return _DummyBaseLayerWithLoRA(layer)
+        return layer
+
+    replace_calls: list[str] = []
+
+    def _fake_replace_submodule(root: torch.nn.Module, module_name: str, submodule: torch.nn.Module):
+        replace_calls.append(module_name)
+        setattr(root, module_name, submodule)
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
+
+    pipeline = torch.nn.Module()
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.foo = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    peft_helper = type("_PH", (), {"r": 1})()
+
+    manager._replace_layers_with_lora(peft_helper)
+    manager._replace_layers_with_lora(peft_helper)
+
+    # Only the top-level layer should have been replaced; nested ".base_layer"
+    # must be skipped to avoid nesting LoRA wrappers.
+    assert replace_calls == ["foo"]
+
+
+def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DummyBaseLayerWithLoRA(torch.nn.Module):
+        def __init__(self, base_layer: torch.nn.Module):
+            super().__init__()
+            self.base_layer = base_layer
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        return _DummyBaseLayerWithLoRA(layer)
+
+    replace_calls: list[str] = []
+
+    def _fake_replace_submodule(root: torch.nn.Module, module_name: str, submodule: torch.nn.Module):
+        replace_calls.append(module_name)
+        setattr(root, module_name, submodule)
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
+
+    pipeline = torch.nn.Module()
+    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.to_qkv = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    # Treat the dummy layer as a packed 3-slice projection so the manager uses
+    # `packed_modules_mapping` to decide replacement based on target_modules.
+    monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"])
+
+    peft_helper = type("_PH", (), {"r": 1, "target_modules": ["to_q"]})()
+    manager._replace_layers_with_lora(peft_helper)
+
+    assert replace_calls == ["to_qkv"]
+
+
+def test_lora_manager_activates_fused_lora_on_packed_layer():
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1))
+    manager._lora_modules = {"transformer.blocks.0.attn.to_qkv": packed_layer}
+
+    rank = 2
+    A = torch.ones((rank, 4))
+    B = torch.arange(0, sum(packed_layer.output_slices) * rank, dtype=torch.bfloat16).view(-1, rank)
+    lora = LoRALayerWeights(
+        module_name="transformer.blocks.0.attn.to_qkv",
+        rank=rank,
+        lora_alpha=rank,
+        lora_a=A,
+        lora_b=B,
+    )
+    manager._registered_adapters = {
+        7: type(
+            "LM",
+            (),
+            {
+                "id": 7,
+                "loras": {"transformer.blocks.0.attn.to_qkv": lora},
+                "get_lora": lambda self, k: self.loras.get(k),
+            },
+        )()
+    }
+    manager._adapter_scales = {7: 0.5}
+
+    manager._activate_adapter(7)
+
+    assert packed_layer.reset_calls == 0
+    assert len(packed_layer.set_calls) == 1
+    lora_a_list, lora_b_list = packed_layer.set_calls[0]
+    assert isinstance(lora_a_list, list)
+    assert isinstance(lora_b_list, list)
+    assert len(lora_a_list) == 3
+    assert len(lora_b_list) == 3
+    assert all(torch.allclose(a, A) for a in lora_a_list)
+    # B should be split into 3 slices and scaled.
+    b0, b1, b2 = lora_b_list
+    assert b0.shape[0] == 2 and b1.shape[0] == 1 and b2.shape[0] == 1
+    assert torch.allclose(torch.cat([b0, b1, b2], dim=0), B * 0.5)
+
+
+def test_lora_manager_activates_packed_lora_from_sublayers():
+    pipeline = torch.nn.Module()
+    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1))
+    manager._lora_modules = {"transformer.blocks.0.attn.to_qkv": packed_layer}
+
+    rank = 2
+    loras: dict[str, LoRALayerWeights] = {}
+    for name, out_dim in zip(["to_q", "to_k", "to_v"], [2, 1, 1]):
+        loras[f"transformer.blocks.0.attn.{name}"] = LoRALayerWeights(
+            module_name=f"transformer.blocks.0.attn.{name}",
+            rank=rank,
+            lora_alpha=rank,
+            lora_a=torch.ones((rank, 4)) * (1 if name == "to_q" else 2),
+            lora_b=torch.ones((out_dim, rank)) * (3 if name == "to_q" else 4),
+        )
+
+    manager._registered_adapters = {
+        1: type("LM", (), {"id": 1, "loras": loras, "get_lora": lambda self, k: self.loras.get(k)})()
+    }
+    manager._adapter_scales = {1: 2.0}
+
+    manager._activate_adapter(1)
+
+    assert packed_layer.reset_calls == 0
+    assert len(packed_layer.set_calls) == 1
+    lora_a_list, lora_b_list = packed_layer.set_calls[0]
+    assert isinstance(lora_a_list, list)
+    assert isinstance(lora_b_list, list)
+    assert len(lora_a_list) == 3
+    assert len(lora_b_list) == 3
+    # Scale should apply to B only.
+    assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0)
+    assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0)
+    assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0)
+
+
+def _dummy_lora_request(adapter_id: int) -> LoRARequest:
+    return LoRARequest(
+        lora_name=f"adapter_{adapter_id}",
+        lora_int_id=adapter_id,
+        lora_path=f"/tmp/adapter_{adapter_id}",
+    )
+
+
+def test_lora_manager_evicts_lru_adapter_when_cache_full(monkeypatch):
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+
+    req1 = _dummy_lora_request(1)
+    req2 = _dummy_lora_request(2)
+    req3 = _dummy_lora_request(3)
+
+    manager.set_active_adapter(req1, lora_scale=1.0)
+    manager.set_active_adapter(req2, lora_scale=1.0)
+
+    # Touch adapter 1 so adapter 2 becomes LRU.
+    manager.set_active_adapter(req1, lora_scale=1.0)
+
+    manager.set_active_adapter(req3, lora_scale=1.0)
+
+    assert set(manager.list_adapters()) == {1, 3}
+
+
+def test_lora_manager_does_not_evict_pinned_adapter(monkeypatch):
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+
+    manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
+    assert manager.pin_adapter(1)
+
+    manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
+    manager.set_active_adapter(_dummy_lora_request(3), lora_scale=1.0)
+
+    assert set(manager.list_adapters()) == {1, 3}
+
+
+def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+
+    manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
+    manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
+
+    assert manager.pin_adapter(1)
+    assert manager.pin_adapter(2)
+
+    manager.max_cached_adapters = 1
+    manager._evict_if_needed()
+
+    assert set(manager.list_adapters()) == {1, 2}

tests/distributed/omni_connectors/test_kv_flow.py

@@ -0,0 +1,194 @@
+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
+
+
+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."""
+
+    def __init__(self):
+        self.device = "cpu"
+        self.od_config = MagicMock()
+        self.bagel = MagicMock()
+        self.tokenizer = MagicMock()
+        self.language_model = MagicMock()
+        self.new_token_ids = {}
+
+
+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
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/e2e/offline_inference/conftest.py

@@ -7,10 +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
@@ -63,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.
 
@@ -81,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.
 
@@ -174,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:
@@ -188,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.
@@ -211,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,
@@ -248,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,
@@ -275,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,
@@ -302,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,
@@ -335,6 +337,7 @@ class OmniRunner:
         self.close()
         del self.omni
         cleanup_dist_env_and_memory()
+        _run_post_test_cleanup()
 
     def close(self):
         """Close and cleanup the Omni instance."""

tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml

@@ -78,6 +78,7 @@ stage_args:
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: audio
+      max_num_batched_tokens: 4069
     engine_input_source: [1]
     final_output: true
     final_output_type: audio

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -45,7 +45,7 @@ stage_args:
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
-       enforce_eager: false
+       enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
       #  tensor_parallel_size: 2
@@ -84,6 +84,7 @@ stage_args:
       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
     final_output: true

tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -84,6 +84,7 @@ stage_args:
       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
     final_output: true

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:
@@ -40,42 +42,51 @@ def test_cache_dit(model_name: str):
         "residual_diff_threshold": 0.24,
         "max_continuous_cached_steps": 3,
     }
+    m = None
+    try:
+        m = Omni(
+            model=model_name,
+            cache_backend="cache_dit",
+            cache_config=cache_config,
+        )
 
-    m = Omni(
-        model=model_name,
-        cache_backend="cache_dit",
-        cache_config=cache_config,
-    )
+        # Use minimal settings for fast testing
+        height = 256
+        width = 256
+        num_inference_steps = 4  # Minimal steps for fast test
 
-    # Use minimal settings for fast testing
-    height = 256
-    width = 256
-    num_inference_steps = 4  # Minimal steps for fast test
+        outputs = m.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            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]
+        assert first_output.final_output_type == "image"
+        if not hasattr(first_output, "request_output") or not first_output.request_output:
+            raise ValueError("No request_output found in OmniRequestOutput")
 
-    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
-    )
-    # Extract images from request_output[0]['images']
-    first_output = outputs[0]
-    assert first_output.final_output_type == "image"
-    if not hasattr(first_output, "request_output") or not first_output.request_output:
-        raise ValueError("No request_output found in OmniRequestOutput")
+        req_out = first_output.request_output[0]
+        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+            raise ValueError("Invalid request_output structure or missing 'images' key")
 
-    req_out = first_output.request_output[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+        images = req_out.images
 
-    images = req_out.images
-
-    # Verify generation succeeded
-    assert images is not None
-    assert len(images) == 1
-    # Check image size
-    assert images[0].width == width
-    assert images[0].height == height
+        # Verify generation succeeded
+        assert images is not None
+        assert len(images) == 1
+        # Check image size
+        assert images[0].width == width
+        assert images[0].height == height
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -1,11 +1,12 @@
 import sys
-import threading
-import time
 from pathlib import Path
 
 import pytest
 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
@@ -15,74 +16,44 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 
-
-class GPUMemoryMonitor:
-    """Poll global device memory usage via CUDA APIs."""
-
-    def __init__(self, device_index: int, interval: float = 0.05):
-        self.device_index = device_index
-        self.interval = interval
-        self.peak_used_mb = 0.0
-        self._stop_event = threading.Event()
-        self._thread: threading.Thread | None = None
-
-    def start(self) -> None:
-        def monitor_loop() -> None:
-            while not self._stop_event.is_set():
-                try:
-                    with torch.cuda.device(self.device_index):
-                        free_bytes, total_bytes = torch.cuda.mem_get_info()
-                    used_mb = (total_bytes - free_bytes) / (1024**2)
-                    self.peak_used_mb = max(self.peak_used_mb, used_mb)
-                except Exception:
-                    pass
-                time.sleep(self.interval)
-
-        self._thread = threading.Thread(target=monitor_loop, daemon=True)
-        self._thread.start()
-
-    def stop(self) -> None:
-        if self._thread is None:
-            return
-        self._stop_event.set()
-        self._thread.join(timeout=2.0)
-
-
 models = ["riverclouds/qwen_image_random"]
 
 
-@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
-@pytest.mark.parametrize("model_name", models)
-def test_cpu_offload_diffusion_model(model_name: str):
-    def inference(offload: bool = True):
-        torch.cuda.empty_cache()
-        device_index = torch.cuda.current_device()
-        monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
-        monitor.start()
-        m = Omni(model=model_name, enable_cpu_offload=offload)
-        torch.cuda.reset_peak_memory_stats(device=device_index)
-        height = 256
-        width = 256
+def inference(model_name: str, offload: bool = True):
+    torch.cuda.empty_cache()
+    device_index = torch.cuda.current_device()
+    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+    m = Omni(model=model_name, enable_cpu_offload=offload)
+    torch.cuda.reset_peak_memory_stats(device=device_index)
+    height = 256
+    width = 256
 
-        m.generate(
-            "a photo of a cat sitting on a laptop keyboard",
+    m.generate(
+        "a photo of a cat sitting on a laptop keyboard",
+        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()
 
-        monitor.stop()
-        torch.cuda.synchronize(device_index)
-        fallback_alloc = torch.cuda.max_memory_allocated(device=device_index) / (1024**2)
-        fallback_reserved = torch.cuda.max_memory_reserved(device=device_index) / (1024**2)
-        peak_memory_mb = max(monitor.peak_used_mb, fallback_alloc, fallback_reserved)
+    return peak
 
-        return peak_memory_mb
 
-    offload_peak_memory = inference(offload=True)
-    no_offload_peak_memory = inference(offload=False)
+@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
+@pytest.mark.parametrize("model_name", models)
+def test_cpu_offload_diffusion_model(model_name: str):
+    try:
+        no_offload_peak_memory = inference(model_name, offload=False)
+        cleanup_dist_env_and_memory()
+        offload_peak_memory = inference(model_name, offload=True)
+    except Exception:
+        pytest.fail("Inference failed")
     print(f"Offload peak memory: {offload_peak_memory} MB")
     print(f"No offload peak memory: {no_offload_peak_memory} MB")
     assert offload_peak_memory + 2500 < no_offload_peak_memory, (

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -0,0 +1,143 @@
+import json
+import os
+import sys
+from pathlib import Path
+
+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
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+
+# This test is specific to Z-Image LoRA behavior. Keep it focused on a single
+# 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):
+    def _extract_images(outputs: list[OmniRequestOutput]):
+        if not outputs:
+            raise ValueError("Empty outputs from Omni.generate()")
+        first_output = outputs[0]
+        assert first_output.final_output_type == "image"
+        if not hasattr(first_output, "request_output") or not first_output.request_output:
+            raise ValueError("No request_output found in OmniRequestOutput")
+
+        req_out = first_output.request_output[0]
+        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+            raise ValueError("Invalid request_output structure or missing 'images' key")
+        return req_out.images
+
+    def _write_zimage_lora(adapter_dir: Path) -> str:
+        adapter_dir.mkdir(parents=True, exist_ok=True)
+
+        # Z-Image transformer uses dim=3840 by default (see ZImageTransformer2DModel).
+        dim = 3840
+        module_name = "transformer.layers.0.attention.to_qkv"
+        rank = 1
+        lora_a = torch.zeros((rank, dim), dtype=torch.float32)
+        lora_a[0, 0] = 1.0
+
+        # QKVParallelLinear packs (Q, K, V). With tp=1 and n_kv_heads==n_heads in Z-Image,
+        # each slice is `dim`, so total out dim is `3 * dim`.
+        lora_b = torch.zeros((3 * dim, rank), dtype=torch.float32)
+        # Apply a visible delta to the Q slice only to keep the perturbation bounded.
+        lora_b[:dim, 0] = 0.1
+
+        save_file(
+            {
+                f"base_model.model.{module_name}.lora_A.weight": lora_a,
+                f"base_model.model.{module_name}.lora_B.weight": lora_b,
+            },
+            str(adapter_dir / "adapter_model.safetensors"),
+        )
+        (adapter_dir / "adapter_config.json").write_text(
+            json.dumps(
+                {
+                    "r": rank,
+                    "lora_alpha": rank,
+                    "target_modules": [module_name],
+                }
+            ),
+            encoding="utf-8",
+        )
+        return str(adapter_dir)
+
+    m = Omni(model=model_name)
+    try:
+        # high resolution may cause OOM on L4
+        height = 256
+        width = 256
+        prompt = "a photo of a cat sitting on a laptop keyboard"
+
+        outputs = m.generate(
+            prompt,
+            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)
+
+        assert len(images) == 1
+        # check image size
+        assert images[0].width == width
+        assert images[0].height == height
+
+        # Real LoRA E2E: generate again with a real on-disk PEFT adapter and
+        # verify that output changes.
+        if model_name == "Tongyi-MAI/Z-Image-Turbo":
+            from vllm_omni.lora.request import LoRARequest
+            from vllm_omni.lora.utils import stable_lora_int_id
+
+            lora_dir = _write_zimage_lora(tmp_path / "zimage_lora")
+            lora_request = LoRARequest(
+                lora_name="test",
+                lora_int_id=stable_lora_int_id(lora_dir),
+                lora_path=lora_dir,
+            )
+            outputs_lora = m.generate(
+                prompt,
+                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
+            assert images_lora[0].width == width
+            assert images_lora[0].height == height
+
+            import numpy as np
+
+            diff = np.abs(np.array(images[0], dtype=np.int16) - np.array(images_lora[0], dtype=np.int16)).mean()
+            assert diff > 0.0
+    finally:
+        m.close()

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?
@@ -233,7 +244,8 @@ def test_real_transformer_init_and_forward():
         }
     )
 
-    od_config = OmniDiffusionConfig(model="dummy-ovis", tf_model_config=tf_config, dtype=torch.float32, num_gpus=1)
+    od_config = OmniDiffusionConfig(model="dummy-ovis", tf_model_config=tf_config, dtype=torch.bfloat16, num_gpus=1)
+    torch.set_default_dtype(torch.bfloat16)
 
     # Mock distributed state for QKVParallelLinear initialization
     # We patch get_tp_group because get_tensor_model_parallel_rank calls it and asserts _TP is not None

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -35,7 +35,7 @@ test_params = [(model, stage_config) for model in models]
 
 @pytest.mark.core_model
 @pytest.mark.parametrize("test_config", test_params)
-@create_new_process_for_each_test()
+@create_new_process_for_each_test("spawn")
 def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
     """Test processing audio, image, and video together, generating audio output."""
     model, stage_config_path = test_config
@@ -95,7 +95,7 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t
 
 @pytest.mark.core_model
 @pytest.mark.parametrize("test_config", test_params)
-@create_new_process_for_each_test()
+@create_new_process_for_each_test("spawn")
 def test_mixed_modalities_to_text_only(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
     """Test processing audio, image, and video together, generating audio output."""
     model, stage_config_path = test_config

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -5,6 +5,10 @@ E2E offline tests for Omni model with video input and audio output.
 """
 
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
 from pathlib import Path
 
 import pytest
@@ -14,8 +18,6 @@ from vllm_omni.utils import is_rocm
 
 from .conftest import OmniRunner
 
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -9,6 +9,7 @@ correctly with diffusion models. It uses minimal settings to keep test time
 short for CI.
 """
 
+import gc
 import os
 import sys
 import time
@@ -20,8 +21,10 @@ 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[2]
+REPO_ROOT = Path(__file__).resolve().parents[3]
 if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
@@ -30,9 +33,6 @@ 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
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
-# Use random weights model for testing
 models = ["riverclouds/qwen_image_random"]
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
@@ -53,44 +53,149 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
     return abs_diff.mean().item(), abs_diff.max().item()
 
 
-def _get_images(output):
-    """Extract images from output, handling both dict and SimpleNamespace types.
+def _cleanup_distributed():
+    """Clean up distributed environment and GPU resources."""
+    if dist.is_initialized():
+        dist.destroy_process_group()
 
-    The output structure varies depending on serialization path:
-    - Direct memory: SimpleNamespace with .images attribute
-    - SHM serialization: dict with "images" key (dataclass converted via asdict)
-    - Wrapped output: SimpleNamespace(output=...) which needs unwrapping
-    """
-    # Check if output has direct images attribute (diffusion mode)
-    if hasattr(output, "images") and output.images:
-        return output.images
+    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+        os.environ.pop(key, None)
 
-    # Check request_output for pipeline mode
-    if output.request_output is None:
-        return None
+    gc.collect()
 
-    if isinstance(output.request_output, list) and len(output.request_output) == 0:
-        return None
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
 
-    item = output.request_output[0]
+    time.sleep(5)
 
-    # Handle wrapped SimpleNamespace (e.g. from omni_stage.py)
-    # Some items are wrapped as SimpleNamespace(request_id=..., output=...)
-    while hasattr(item, "output") and not hasattr(item, "images"):
-        item = item.output
-        if item is None:
-            return None
 
-    # Handle both dict (from SHM serialization) and object (direct) types
-    if isinstance(item, dict):
-        return item.get("images")
-    return getattr(item, "images", None)
+def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height: int, width: int, seed: int):
+    """Run baseline inference (no SP)."""
+    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
+    baseline = Omni(
+        model=model_name,
+        parallel_config=baseline_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+
+    try:
+        outputs = baseline.generate(
+            PROMPT,
+            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:
+        baseline.close()
+        _cleanup_distributed()
+
+
+def _run_sp(
+    model_name: str,
+    dtype: torch.dtype,
+    attn_backend: str,
+    height: int,
+    width: int,
+    seed: int,
+    ulysses_degree: int,
+    ring_degree: int,
+):
+    """Run SP inference."""
+    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    sp = Omni(
+        model=model_name,
+        parallel_config=sp_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+
+    try:
+        outputs = sp.generate(
+            PROMPT,
+            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:
+        sp.close()
+        _cleanup_distributed()
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test baseline inference only (no SP)."""
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_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")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=2, ring_degree=1)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP inference only (ring=2)."""
+    if device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=1, ring_degree=2)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
 
 
 @pytest.mark.parametrize("model_name", models)
 @pytest.mark.parametrize("ulysses_degree", [1, 2])
 @pytest.mark.parametrize("ring_degree", [1, 2])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])  # Only test bfloat16 to reduce CI time
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
 def test_sequence_parallel(
     model_name: str,
@@ -102,214 +207,60 @@ def test_sequence_parallel(
     """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs."""
     if ulysses_degree <= 1 and ring_degree <= 1:
         pytest.skip(
-            "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; provide ulysses_degree or ring_degree>1."
+            "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; "
+            "provide ulysses_degree or ring_degree>1."
         )
 
-    # Skip if not enough GPUs available for SP run
-    if device_count() < ulysses_degree * ring_degree:
-        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
+    sp_size = ulysses_degree * ring_degree
+    if device_count() < sp_size:
+        pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available")
 
-    # Use minimal settings for fast testing
     height = 256
     width = 256
-    num_inference_steps = 4  # Minimal steps for fast test
     seed = 42
 
-    # Step 1: Baseline (no Ulysses sequence parallel)
-    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
-    baseline = Omni(
-        model=model_name,
-        parallel_config=baseline_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = baseline.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        baseline_images = outputs[0].request_output[0].images
-    finally:
-        baseline.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)  # Wait for resources to release
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
 
-    assert baseline_images is not None
-    assert len(baseline_images) == 1
-    assert baseline_images[0].width == width
-    assert baseline_images[0].height == height
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
 
-    # Step 2: SP (Ulysses-SP + Ring-SP)
-    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-    sp = Omni(
-        model=model_name,
-        parallel_config=sp_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = sp.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        sp_images = outputs[0].request_output[0].images
-    finally:
-        sp.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)
-
-    assert sp_images is not None
-    assert len(sp_images) == 1
-    assert sp_images[0].width == width
-    assert sp_images[0].height == height
-
-    # Step 3: Compare outputs
     mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
 
-    # FP16/BF16 may differ slightly due to different computation order under parallelism.
-    if dtype in (torch.float16, torch.bfloat16):
-        mean_threshold = 2e-2
-        max_threshold = 2e-1
-    else:
-        mean_threshold = 1e-2
-        max_threshold = 1e-1
-
-    print(
-        "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
-    )
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
 
     assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
     )
 
 
 @pytest.mark.parametrize("model_name", models)
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_sequence_parallel_ulysses_sp_only(
-    model_name: str,
-    dtype: torch.dtype,
-    attn_backend: str,
-):
-    """Test sequence parallel with ulysses_degree=4, ring_degree=1, and the image size (332x332) where the sequence length is NOT divisible by sp_size."""
+def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP with ulysses_degree=4."""
     ulysses_degree = 4
     ring_degree = 1
 
-    # Skip if not enough GPUs available for SP run
     if device_count() < ulysses_degree * ring_degree:
         pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
 
-    # (272/8) * (272/8) = 17 * 17 = 289 Not divisible by sp_size=4
     height = 272
     width = 272
-    num_inference_steps = 4  # Minimal steps for fast test
     seed = 42
 
-    # Step 1: Baseline (no Ulysses sequence parallel)
-    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
-    baseline = Omni(
-        model=model_name,
-        parallel_config=baseline_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = baseline.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        baseline_images = outputs[0].request_output[0].images
-    finally:
-        baseline.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)  # Wait for resources to release
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
 
-    assert baseline_images is not None
-    assert len(baseline_images) == 1
-    assert baseline_images[0].width == width
-    assert baseline_images[0].height == height
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
 
-    # Step 2: SP (Ulysses-SP + Ring-SP)
-    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-    sp = Omni(
-        model=model_name,
-        parallel_config=sp_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = sp.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        sp_images = outputs[0].request_output[0].images
-    finally:
-        sp.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)
-
-    assert sp_images is not None
-    assert len(sp_images) == 1
-    assert sp_images[0].width == width
-    assert sp_images[0].height == height
-
-    # Step 3: Compare outputs
     mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
 
-    # FP16/BF16 may differ slightly due to different computation order under parallelism.
-    if dtype in (torch.float16, torch.bfloat16):
-        mean_threshold = 2e-2
-        max_threshold = 2e-1
-    else:
-        mean_threshold = 1e-2
-        max_threshold = 1e-1
-
-    print(
-        "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
-    )
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
 
     assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
     )

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -1,4 +1,3 @@
-import os
 import sys
 from pathlib import Path
 
@@ -6,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
@@ -15,8 +15,6 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
 # Use random weights model for CI testing (small, no authentication required)
 models = ["linyueqian/stable_audio_random"]
 
@@ -32,30 +30,33 @@ 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
     assert outputs is not None
     first_output = outputs[0]
-    assert first_output.final_output_type == "image"  # Generic output type
+    assert first_output.final_output_type == "image"
     assert hasattr(first_output, "request_output") and first_output.request_output
 
     req_out = first_output.request_output[0]
     assert isinstance(req_out, OmniRequestOutput)
-    assert hasattr(req_out, "images") and len(req_out.images) >= 1
-
-    # For stable audio, the "images" field contains audio numpy arrays
-    audio = req_out.images[0]
+    assert req_out.final_output_type == "audio"
+    assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output
+    audio = req_out.multimodal_output.get("audio")
     assert isinstance(audio, np.ndarray)
     # audio shape: (batch, channels, samples)
     # For stable-audio-open-1.0: sample_rate=44100, so 2 seconds = 88200 samples

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
 
@@ -34,33 +35,43 @@ elif is_rocm():
 
 @pytest.mark.parametrize("model_name", models)
 def test_diffusion_model(model_name: str):
-    m = Omni(model=model_name)
-    # high resolution may cause OOM on L4
-    height = 256
-    width = 256
-    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,
-    )
-    # Extract images from request_output[0]['images']
-    first_output = outputs[0]
-    assert first_output.final_output_type == "image"
-    if not hasattr(first_output, "request_output") or not first_output.request_output:
-        raise ValueError("No request_output found in OmniRequestOutput")
+    m = None
+    try:
+        m = Omni(model=model_name)
+        # high resolution may cause OOM on L4
+        height = 256
+        width = 256
+        outputs = m.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            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]
+        assert first_output.final_output_type == "image"
+        if not hasattr(first_output, "request_output") or not first_output.request_output:
+            raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+        req_out = first_output.request_output[0]
+        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+            raise ValueError("Invalid request_output structure or missing 'images' key")
 
-    images = req_out.images
+        images = req_out.images
 
-    assert len(images) == 2
-    # check image size
-    assert images[0].width == width
-    assert images[0].height == height
-    images[0].save("image_output.png")
+        assert len(images) == 2
+        # check image size
+        assert images[0].width == width
+        assert images[0].height == height
+        images[0].save("image_output.png")
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()

tests/e2e/offline_inference/test_t2v_model.py

@@ -5,14 +5,18 @@ 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:
     sys.path.insert(0, str(REPO_ROOT))
 
 from vllm_omni import Omni
+from vllm_omni.outputs import OmniRequestOutput
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 models = ["Wan-AI/Wan2.2-T2V-A14B-Diffusers"]
 
@@ -31,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"
@@ -45,10 +51,10 @@ def test_video_diffusion_model(model_name: str):
         raise ValueError("No request_output found in OmniRequestOutput")
 
     req_out = first_output.request_output[0]
-    if not isinstance(req_out, dict) or "images" not in req_out:
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 
-    frames = req_out["images"][0]
+    frames = req_out.images[0]
 
     assert frames is not None
     assert hasattr(frames, "shape")

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:
@@ -36,42 +38,51 @@ def test_teacache(model_name: str):
     cache_config = {
         "rel_l1_thresh": 0.2,  # Default threshold
     }
+    m = None
+    try:
+        m = Omni(
+            model=model_name,
+            cache_backend="tea_cache",
+            cache_config=cache_config,
+        )
 
-    m = Omni(
-        model=model_name,
-        cache_backend="tea_cache",
-        cache_config=cache_config,
-    )
+        # Use minimal settings for fast testing
+        height = 256
+        width = 256
+        num_inference_steps = 4  # Minimal steps for fast test
 
-    # Use minimal settings for fast testing
-    height = 256
-    width = 256
-    num_inference_steps = 4  # Minimal steps for fast test
+        outputs = m.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            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]
+        assert first_output.final_output_type == "image"
+        if not hasattr(first_output, "request_output") or not first_output.request_output:
+            raise ValueError("No request_output found in OmniRequestOutput")
 
-    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
-    )
-    # Extract images from request_output[0]['images']
-    first_output = outputs[0]
-    assert first_output.final_output_type == "image"
-    if not hasattr(first_output, "request_output") or not first_output.request_output:
-        raise ValueError("No request_output found in OmniRequestOutput")
+        req_out = first_output.request_output[0]
+        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+            raise ValueError("Invalid request_output structure or missing 'images' key")
 
-    req_out = first_output.request_output[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+        images = req_out.images
 
-    images = req_out.images
-
-    # Verify generation succeeded
-    assert images is not None
-    assert len(images) == 1
-    # Check image size
-    assert images[0].width == width
-    assert images[0].height == height
+        # Verify generation succeeded
+        assert images is not None
+        assert len(images) == 1
+        # Check image size
+        assert images[0].width == width
+        assert images[0].height == height
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()

tests/e2e/offline_inference/test_zimage_tensor_parallel.py

@@ -12,17 +12,21 @@ 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:
     sys.path.insert(0, str(REPO_ROOT))
 
+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
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
 
@@ -66,7 +70,12 @@ def _extract_single_image(outputs) -> Image.Image:
 
 def _run_zimage_generate(
     *, tp_size: int, height: int, width: int, num_inference_steps: int, seed: int
-) -> tuple[Image.Image, float]:
+) -> tuple[Image.Image, float, float]:
+    torch.cuda.empty_cache()
+    device_index = torch.cuda.current_device()
+    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+
     m = Omni(
         model=_get_zimage_model(),
         parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
@@ -83,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,
         )
 
@@ -107,9 +118,12 @@ def _run_zimage_generate(
             pass
 
         median_time_s = float(np.median(per_request_times_s))
-        return _extract_single_image([last_output]), median_time_s
+
+        peak_memory_mb = monitor.peak_used_mb
+
+        return _extract_single_image([last_output]), median_time_s, peak_memory_mb
     finally:
-        m.close()
+        monitor.stop()
         cleanup_dist_env_and_memory()
 
 
@@ -120,19 +134,19 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
         pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")
 
-    height = 256
-    width = 256
+    height = 512
+    width = 512
     num_inference_steps = 2
     seed = 42
 
-    tp1_img, tp1_time_s = _run_zimage_generate(
+    tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
         tp_size=1,
         height=height,
         width=width,
         num_inference_steps=num_inference_steps,
         seed=seed,
     )
-    tp2_img, tp2_time_s = _run_zimage_generate(
+    tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
         tp_size=2,
         height=height,
         width=width,
@@ -150,7 +164,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
 
     mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img)
     mean_threshold = 3e-2
-    max_threshold = 3.5e-1
+    max_threshold = 5e-1
     print(
         "Z-Image TP image diff stats (TP=1 vs TP=2): "
         f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
@@ -164,3 +178,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
 
     print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
     assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
+
+    print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+    assert tp2_peak_mem < tp1_peak_mem, (
+        f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
+    )

tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml

@@ -77,6 +77,7 @@ stage_args:
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
+      async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: audio  # Final output: audio waveform
       gpu_memory_utilization: 0.1

tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -84,6 +84,7 @@ stage_args:
       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
     final_output: true

tests/e2e/online_serving/test_async_omni.py

@@ -1,12 +1,17 @@
 import asyncio
+import os
+import sys
 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
 
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
 SEED = 42
 
 stage_config = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_thinker_ci.yaml")
@@ -21,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,
@@ -55,13 +60,19 @@ async def generate(
 @pytest.mark.asyncio
 async def test_abort():
     with ExitStack() as after:
-        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
+        # Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS.
+        engine = AsyncOmni(
+            model=model,
+            stage_configs_path=stage_config,
+            shm_threshold_bytes=sys.maxsize,
+        )
         after.callback(engine.shutdown)
 
-        NUM_REQUESTS = 5
-        NUM_EXPECTED_TOKENS = 100
-        NUM_EXPECTED_TOKENS_LONG = 1000
-        REQUEST_IDS_TO_ABORT = [1, 2, 3]
+        # Keep token counts modest to reduce flakiness on slow test hardware.
+        NUM_REQUESTS = 3
+        NUM_EXPECTED_TOKENS = 64
+        NUM_EXPECTED_TOKENS_LONG = 256
+        REQUEST_IDS_TO_ABORT = [1]
 
         prompt = "Hello my name is Robert and "
 
@@ -74,8 +85,10 @@ async def test_abort():
             tasks.append(asyncio.create_task(generate(engine, request_id, prompt, max_tokens)))
 
         # API server cancels requests when they disconnect.
+        # Explicitly abort in the engine to avoid orphaned requests hanging.
         for idx in REQUEST_IDS_TO_ABORT:
             tasks[idx].cancel()
+            await engine.abort(request_ids[idx])
             await asyncio.sleep(0.1)
 
         # Confirm the other requests are okay.
@@ -83,10 +96,10 @@ async def test_abort():
             # Confirm that it was actually canceled.
             if idx in REQUEST_IDS_TO_ABORT:
                 with pytest.raises((asyncio.CancelledError, GeneratorExit)):
-                    await task
+                    await asyncio.wait_for(task, timeout=60)
             else:
                 # Otherwise, make sure the request was not impacted.
-                num_generated_tokens, request_id = await task
+                num_generated_tokens, request_id = await asyncio.wait_for(task, timeout=180)
                 expected_tokens = NUM_EXPECTED_TOKENS
                 assert num_generated_tokens == expected_tokens, (
                     f"{request_id} generated {num_generated_tokens} but expected {expected_tokens}"
@@ -116,11 +129,16 @@ async def test_build_and_log_summary(monkeypatch):
     monkeypatch.setattr("vllm_omni.entrypoints.client_request_state.ClientRequestState", MockCRS)
 
     with ExitStack() as after:
-        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
+        # Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS.
+        engine = AsyncOmni(
+            model=model,
+            stage_configs_path=stage_config,
+            shm_threshold_bytes=sys.maxsize,
+        )
         after.callback(engine.shutdown)
         prompt = "Hello my name is Robert and "
-        NUM_EXPECTED_TOKENS = 100
-        NUM_REQUESTS = 5
+        NUM_EXPECTED_TOKENS = 64
+        NUM_REQUESTS = 3
         request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
 
         # Create concurrent requests.

tests/e2e/online_serving/test_i2i_multi_image_input.py

@@ -17,7 +17,7 @@ import openai
 import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
-from vllm.utils import get_open_port
+from vllm.utils.network_utils import get_open_port
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 

tests/e2e/online_serving/test_images_generations_lora.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online serving test for /v1/images/generations with per-request LoRA.
+
+This validates:
+- The API server accepts a per-request `lora` object in the Images API payload.
+- LoRA can be switched per request (adapter A -> adapter B -> no LoRA).
+- Output correctness is asserted using a small image slice with tolerance.
+"""
+
+import base64
+import json
+import os
+import signal
+import subprocess
+import sys
+import time
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import pytest
+import requests
+import torch
+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"
+
+
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+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():
+    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
+
+
+def _write_zimage_lora(adapter_dir: Path, *, q_scale: float = 0.0, k_scale: float = 0.0, v_scale: float = 0.0):
+    adapter_dir.mkdir(parents=True, exist_ok=True)
+
+    # Z-Image transformer uses dim=3840 by default.
+    dim = 3840
+    module_name = "transformer.layers.0.attention.to_qkv"
+    rank = 1
+
+    lora_a = torch.zeros((rank, dim), dtype=torch.float32)
+    lora_a[0, 0] = 1.0
+
+    # QKVParallelLinear packs (Q, K, V) => out dim is 3 * dim (tp=1).
+    lora_b = torch.zeros((3 * dim, rank), dtype=torch.float32)
+    if q_scale:
+        lora_b[:dim, 0] = q_scale
+    if k_scale:
+        lora_b[dim : 2 * dim, 0] = k_scale
+    if v_scale:
+        lora_b[2 * dim :, 0] = v_scale
+
+    save_file(
+        {
+            f"base_model.model.{module_name}.lora_A.weight": lora_a,
+            f"base_model.model.{module_name}.lora_B.weight": lora_b,
+        },
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps(
+            {
+                "r": rank,
+                "lora_alpha": rank,
+                "target_modules": [module_name],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+
+def _post_images(server: OmniServer, payload: dict) -> Image.Image:
+    url = f"http://{server.host}:{server.port}/v1/images/generations"
+    resp = requests.post(url, json=payload, headers={"Authorization": "Bearer EMPTY"}, timeout=900)
+    resp.raise_for_status()
+    data = resp.json()
+    b64 = data["data"][0]["b64_json"]
+    img_bytes = base64.b64decode(b64)
+    img = Image.open(BytesIO(img_bytes))
+    img.load()
+    return img.convert("RGB")
+
+
+def _image_blue_tail_slice(img: Image.Image) -> np.ndarray:
+    arr = np.asarray(img, dtype=np.uint8)
+    assert arr.ndim == 3 and arr.shape[-1] == 3
+    tail = arr[-3:, -3:, -1].astype(np.float32)
+    assert tail.shape == (3, 3)
+    return tail
+
+
+def _assert_slice_close(actual: np.ndarray, expected: np.ndarray, *, label: str) -> 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())
+    # 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()}"
+    )
+
+
+def _assert_slice_diff(actual: np.ndarray, baseline: np.ndarray, *, label: str) -> None:
+    assert actual.shape == (3, 3)
+    assert baseline.shape == (3, 3)
+    diff = np.abs(actual - baseline).mean()
+    assert diff > 0.1, f"{label} slice diff too small: {diff} ({actual.tolist()} vs {baseline.tolist()})"
+
+
+def _basic_payload() -> dict:
+    return {
+        "prompt": PROMPT,
+        "n": 1,
+        "size": SIZE,
+        "num_inference_steps": 2,
+        "guidance_scale": 0.0,
+        "seed": SEED,
+    }
+
+
+def test_images_generations_per_request_lora_switching(omni_server: OmniServer, tmp_path: Path) -> None:
+    # Base generation.
+    base_img = _post_images(omni_server, _basic_payload())
+    base_slice = _image_blue_tail_slice(base_img)
+
+    # Adapter A: apply delta to V slice only.
+    lora_a_dir = tmp_path / "zimage_lora_a"
+    _write_zimage_lora(lora_a_dir, v_scale=8.0)
+    payload_a = _basic_payload()
+    payload_a["lora"] = {"name": "a", "path": str(lora_a_dir), "scale": 64.0}
+    img_a = _post_images(omni_server, payload_a)
+    a_slice = _image_blue_tail_slice(img_a)
+    _assert_slice_diff(a_slice, base_slice, label="lora_a_vs_base")
+    a_vs_base = float(np.abs(a_slice - base_slice).mean())
+
+    # Adapter B: apply delta to K slice only (should differ from adapter A).
+    lora_b_dir = tmp_path / "zimage_lora_b"
+    _write_zimage_lora(lora_b_dir, k_scale=4.0)
+    payload_b = _basic_payload()
+    payload_b["lora"] = {"name": "b", "path": str(lora_b_dir), "scale": 64.0}
+    img_b = _post_images(omni_server, payload_b)
+    b_slice = _image_blue_tail_slice(img_b)
+    _assert_slice_diff(b_slice, base_slice, label="lora_b_vs_base")
+    _assert_slice_diff(b_slice, a_slice, label="lora_b_vs_lora_a")
+    b_vs_base = float(np.abs(b_slice - base_slice).mean())
+    b_vs_a = float(np.abs(b_slice - a_slice).mean())
+
+    # 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())
+
+    # Ensure LoRA effects are clearly above the baseline drift.
+    min_delta = base_reset + 0.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}"

tests/e2e/online_serving/test_qwen3_omni.py

@@ -4,141 +4,89 @@
 E2E Online tests for Qwen3-Omni model with video input and audio output.
 """
 
-import base64
-import concurrent.futures
-import ctypes
 import os
-import signal
-import socket
-import subprocess
-import sys
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import concurrent.futures
+import threading
 import time
 from pathlib import Path
 
 import openai
 import pytest
-from vllm.assets.video import VideoAsset
-from vllm.utils import get_open_port
 
+from tests.conftest import (
+    OmniServer,
+    convert_audio_to_text,
+    cosine_similarity_text,
+    dummy_messages_from_mix_data,
+    generate_synthetic_audio,
+    generate_synthetic_image,
+    generate_synthetic_video,
+    modify_stage_config,
+)
 from vllm_omni.utils import is_rocm
 
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
+
+def get_default_config():
+    return str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
+
+
+def get_chunk_config():
+    path = modify_stage_config(
+        get_default_config(),
+        updates={
+            "async_chunk": True,
+            "stage_args": {
+                0: {
+                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
+                },
+                1: {
+                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk"
+                },
+            },
+        },
+        deletes={"stage_args": {2: ["custom_process_input_func"]}},
+    )
+    return path
+
+
+CHUNK_CONFIG_PATH = get_chunk_config()
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
 if is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
-    stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
+    stage_configs = [get_default_config(), CHUNK_CONFIG_PATH]
 
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
 
 
-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:
-        """Start the vLLM-Omni server subprocess."""
-        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
-
-        # Helper to ensure child process dies when parent dies
-        libc = ctypes.CDLL("libc.so.6")
-
-        def preexec_fn():
-            # Ensure the child process receives SIGTERM when the parent (this test runner) dies.
-            # This prevents orphaned processes if the test is killed unexpectedly.
-            PR_SET_PDEATHSIG = 1
-            libc.prctl(PR_SET_PDEATHSIG, signal.SIGTERM)
-
-        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__))),  # Set working directory to vllm-omni root
-            start_new_session=True,
-            preexec_fn=preexec_fn,
-        )
-
-        # Wait for server to be ready
-        max_wait = 600  # 10 minutes
-        start_time = time.time()
-        while time.time() - start_time < max_wait:
-            try:
-                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-                    sock.settimeout(1)
-                    result = sock.connect_ex((self.host, self.port))
-                    if result == 0:
-                        print(f"Server ready on {self.host}:{self.port}")
-                        return
-            except Exception:
-                pass
-            time.sleep(2)
-
-        raise RuntimeError(f"Server failed to start within {max_wait} seconds")
-
-    def __enter__(self):
-        self._start_server()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.proc:
-            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()
+_omni_server_lock = threading.Lock()
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def omni_server(request):
     """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses module scope so the server starts only once for all tests.
+    Uses session scope so the server starts only once for the entire test session.
     Multi-stage initialization can take 10-20+ minutes.
     """
-    model, stage_config_path = request.param
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        yield server
+    with _omni_server_lock:
+        model, stage_config_path = request.param
+
+        print(f"Starting OmniServer with model: {model}")
+
+        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+            print("OmniServer started successfully")
+            yield server
+            print("OmniServer stopping...")
+
+        print("OmniServer stopped")
 
 
 @pytest.fixture
@@ -150,15 +98,6 @@ def client(omni_server):
     )
 
 
-@pytest.fixture(scope="session")
-def base64_encoded_video() -> str:
-    """Base64 encoded video for testing."""
-    video = VideoAsset(name="baby_reading", num_frames=4)
-    with open(video.video_path, "rb") as f:
-        content = f.read()
-        return base64.b64encode(content).decode("utf-8")
-
-
 def get_system_prompt():
     return {
         "role": "system",
@@ -192,71 +131,55 @@ def dummy_messages_from_video_data(
     ]
 
 
+def get_prompt(prompt_type="text_only"):
+    prompts = {
+        "text_only": "What is the capital of China?",
+        "mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
+    }
+    return prompts.get(prompt_type, prompts["text_only"])
+
+
+def get_max_batch_size(size_type="few"):
+    batch_sizes = {"few": 5, "medium": 100, "large": 256}
+    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.parametrize("omni_server", test_params, indirect=True)
-def test_video_to_audio_concurrent(
-    client: openai.OpenAI,
-    omni_server,
-    base64_encoded_video: str,
-) -> None:
-    """Test processing video with multiple concurrent completions, generating audio output via OpenAI API."""
-    # Create data URL for the base64 encoded video
-    video_data_url = f"data:video/mp4;base64,{base64_encoded_video}"
+def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
+    """
+    Test multi-modal input processing and text/audio output generation via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text + audio + video + image
+    Output Modal: text + audio
+    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")
 
-    messages = dummy_messages_from_video_data(video_data_url)
+    if param[1] == CHUNK_CONFIG_PATH:
+        pytest.skip("The current chunk scenario does not support multimodal.")
 
-    # Test multiple concurrent completions
-    num_concurrent_requests = 5
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
-        # Submit multiple completion requests concurrently
-        futures = [
-            executor.submit(
-                client.chat.completions.create,
-                model=omni_server.model,
-                messages=messages,
-            )
-            for _ in range(num_concurrent_requests)
-        ]
-
-        # Wait for all requests to complete and collect results
-        chat_completions = [future.result() for future in concurrent.futures.as_completed(futures)]
-
-    # Verify all completions succeeded
-    assert len(chat_completions) == num_concurrent_requests
-
-    for chat_completion in chat_completions:
-        assert len(chat_completion.choices) == 2  # 1 for text output, 1 for audio output
-
-        # Verify text output
-        text_choice = chat_completion.choices[0]
-        assert text_choice.finish_reason == "length"
-
-        # Verify we got a response
-        text_message = text_choice.message
-        assert text_message.content is not None and len(text_message.content) >= 10
-        assert text_message.role == "assistant"
-
-        # Verify audio output
-        audio_choice = chat_completion.choices[1]
-        assert audio_choice.finish_reason == "stop"
-        audio_message = audio_choice.message
-
-        # Check if audio was generated
-        if hasattr(audio_message, "audio") and audio_message.audio:
-            assert audio_message.audio.data is not None
-            assert len(audio_message.audio.data) > 0
-
-    # Test streaming completion
-    chat_completion = client.chat.completions.create(
-        model=omni_server.model,
-        messages=messages,
-        stream=True,
+    # Test single completion
+    e2e_list = list()
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        image_data_url=image_data_url,
+        audio_data_url=audio_data_url,
+        content_text=get_prompt("mix"),
     )
 
-    # Collect text and audio data from stream
+    # Test single completion
+    start_time = time.perf_counter()
+    chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True)
+
     text_content = ""
     audio_data = None
-
     for chunk in chat_completion:
         for choice in chunk.choices:
             if hasattr(choice, "delta"):
@@ -267,18 +190,94 @@ def test_video_to_audio_concurrent(
             modality = getattr(chunk, "modality", None)
 
             if modality == "audio" and content:
-                # Audio chunk - decode base64 content
+                # Audio chunk - content
                 if audio_data is None:
-                    audio_data = base64.b64decode(content)
+                    audio_data = content
                 else:
-                    audio_data += base64.b64decode(content)
+                    audio_data += content
             elif modality == "text" and content:
                 # Text chunk - accumulate text content
                 text_content += content if content else ""
 
-    # Verify text output
-    assert text_content is not None and len(text_content) >= 2
+    # Verify E2E
+    current_e2e = time.perf_counter() - start_time
+    print(f"the request e2e is: {current_e2e}")
+    # TODO: Verify the E2E latency after confirmation baseline.
+    e2e_list.append(current_e2e)
 
-    # Verify audio output
-    assert audio_data is not None
-    assert len(audio_data) > 0
+    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
+    # Verify all completions succeeded
+    assert audio_data is not None, "No audio output is generated"
+
+    # Verify text output success
+    assert text_content is not None and len(text_content) >= 2, "No text output is generated"
+    assert any(
+        keyword in text_content.lower() for keyword in ["square", "quadrate", "sphere", "globe", "circle", "round"]
+    ), "The output does not contain any of the keywords."
+
+    # Verify text output same as audio output
+    audio_content = convert_audio_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())
+    print(f"similarity is: {similarity}")
+    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.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
+    """
+    Test text input processing and text/audio output generation via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: text + audio
+    Datasets: few requests
+    """
+
+    num_concurrent_requests = get_max_batch_size()
+    messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
+
+    e2e_list = list()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+        # Submit multiple completion requests concurrently
+        futures = [
+            executor.submit(client.chat.completions.create, model=omni_server.model, messages=messages)
+            for _ in range(num_concurrent_requests)
+        ]
+        start_time = time.perf_counter()
+        # Wait for all requests to complete and collect results
+        chat_completions = list()
+        for future in concurrent.futures.as_completed(futures):
+            chat_completions.append(future.result())
+            # Verify E2E
+            current_e2e = time.perf_counter() - start_time
+            print(f"the request e2e is: {current_e2e}")
+            # TODO: Verify the E2E latency after confirmation baseline.
+            e2e_list.append(current_e2e)
+
+    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
+    # Verify all completions succeeded
+    assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
+    for chat_completion in chat_completions:
+        # Verify audio output success
+        audio_data = None
+        text_content = None
+        for choice in chat_completion.choices:
+            if choice.message.audio is not None:
+                audio_message = choice.message
+                audio_data = audio_message.audio.data
+                assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
+
+            if choice.message.content is not None:
+                # Verify text output success
+                text_content = choice.message.content
+                assert "beijing" in text_content.lower(), "The output do not contain keywords."
+
+        # Verify text output same as audio output
+        audio_content = convert_audio_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())
+        print(f"similarity is: {similarity}")
+        assert similarity > 0.9, "The audio content is not same as the text"