Address code review feedback

- Remove unused `field` import
- Move `json` import to top of file
- Add comment explaining local DiffusionParallelConfig import

Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>

.buildkite/pipeline.yml

@@ -29,6 +29,7 @@ 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:
@@ -92,28 +93,11 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion Sequence Parallelism Test"
-    timeout_in_minutes: 20
+  - label: "Diffusion Parallelism Test"
+    timeout_in_minutes: 25
     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

.buildkite/scripts/simple_test.sh

@@ -47,11 +47,9 @@ fi
 VENV_PYTHON="${VENV_DIR}/bin/python"
 [[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
 
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.14.0
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.12.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/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 }} --target final --progress plain ."
+          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"
         env:

README.md

@@ -15,6 +15,7 @@ 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:8091). Overrides host/port.",
+        help="Base URL of the server (e.g., http://localhost:30000). Overrides host/port.",
     )
     parser.add_argument("--host", type=str, default="localhost", help="Server host.")
-    parser.add_argument("--port", type=int, default=8091, help="Server port.")
+    parser.add_argument("--port", type=int, default=30000, 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.14.0
+ARG VLLM_BASE_TAG=v0.12.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.rocm

@@ -1,7 +1,9 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0
-FROM ${BASE_IMAGE} AS final
+ARG BASE_IMAGE=rocm/vllm-dev:nightly_main_20251205
+FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app
+ARG VLLM_VERSION=v0.12.0
+ARG PYTORCH_ROCM_ARCH="gfx942;gfx950"
 
 WORKDIR ${COMMON_WORKDIR}
 
@@ -11,18 +13,30 @@ 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 2: Copy vllm-omni code and install without uv
+# Step 3: Copy vllm-omni code and install without uv
 COPY . ${COMMON_WORKDIR}/vllm-omni
-RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+RUN cd ${COMMON_WORKDIR}/vllm-omni && python3 -m pip install --no-cache-dir ".[dev]"
 
+# 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"]

docs/.nav.yml

@@ -14,7 +14,6 @@ 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,6 +36,7 @@ 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][]
@@ -56,11 +57,38 @@ Core scheduling and caching components.
 - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
-- [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][]
+
+## 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][]
 
 ## Configuration
 
@@ -70,17 +98,6 @@ 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.14.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.12.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 

docs/contributing/profiling.md

@@ -1,30 +1,48 @@
 # Profiling vLLM-Omni
 
-> **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
+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.
 
-vLLM-Omni uses the PyTorch Profiler to analyze performance across both **Multi-Stage LLMs** and **Diffusion Models**.
+**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).
 
-### 1. Set the Output Directory
-Before running any script, set this environment variable. The system detects this and automatically saves traces here.
+**1. Enabling the Profiler**
 
-```bash
-export VLLM_TORCH_PROFILER_DIR=./profiles
+Before running your script, you must set the ```VLLM_TORCH_PROFILER_DIR``` environment
+variable.
+
+```Bash
+export VLLM_TORCH_PROFILER_DIR=/path/to/save/traces
 ```
 
-### 2. Start Profiling
+**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.
 
-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])
 ```
@@ -34,105 +52,13 @@ 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)
 
-
-**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)**
+**3. Online Inference(Async)**
 
 For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
 
@@ -151,16 +77,23 @@ async for output in async_omni.generate(prompt, sampling_params, request_id):
 await async_omni.stop_profile()
 ```
 
-### 3. Analyzing Omni Traces
+**4. Analyzing Omni Traces**
 
-Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
+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**
-**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
+```
+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 %)
+```
 
 **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. 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.  
+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)

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](https://github.com/vllm-project/vllm-omni/tree/main/examples).
+For more usages, please refer to [examples](../examples/README.md).

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

@@ -22,7 +22,7 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1 and
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.14.0 --torch-backend=auto
+uv pip install vllm==0.12.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.14.0 release version of the library
+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
 
 ```bash
-uv pip install vllm==0.14.0 --torch-backend=auto
+uv pip install vllm==0.12.0 --torch-backend=auto
 ```
 
-The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.12.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.14.0
+git checkout v0.12.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.14.0/vllm-0.14.0-cp38-abi3-manylinux_2_31_x86_64.whl
+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
 ```
 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.14.0 \
+    vllm/vllm-omni:v0.12.0rc1 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -34,28 +34,8 @@ DOCKER_BUILDKIT=1 docker build \
 
 #### Launch the docker image
 
-##### Launch with OpenAI API Server
-
 ```
-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 \
+docker run -it \
 --network=host \
 --group-add=video \
 --ipc=host \
@@ -64,9 +44,8 @@ docker run --rm -it \
 --device /dev/kfd \
 --device /dev/dri \
 -v <path/to/model>:/app/model \
--v ~/.cache/huggingface:/root/.cache/huggingface \
---entrypoint bash \
-vllm-omni-rocm
+vllm-omni-rocm \
+bash
 ```
 
 # --8<-- [end:build-docker]
@@ -78,25 +57,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 --rm \
-  --group-add=video \
-  --ipc=host \
-  --cap-add=SYS_PTRACE \
-  --security-opt seccomp=unconfined \
-  --device /dev/kfd \
-  --device /dev/dri \
-  -v <path/to/model>:/app/model \
-  -v ~/.cache/huggingface:/root/.cache/huggingface \
-  --env "HF_TOKEN=$HF_TOKEN" \
-  -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.14.0 \
-  --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 --rm -it \
+docker run -it \
   --network=host \
   --group-add=video \
   --ipc=host \
@@ -107,8 +68,26 @@ docker run --rm -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.14.0
+  vllm/vllm-omni-rocm:v0.12.0rc1 \
+  vllm serve --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 \
+  --network=host \
+  --group-add=video \
+  --ipc=host \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --device /dev/kfd \
+  --device /dev/dri \
+  -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
 ```
 
 # --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.14.0
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 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.14.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.12.0rc1 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,13 +17,7 @@ For installation on GPU from source:
 ```bash
 uv venv --python 3.12 --seed
 source .venv/bin/activate
-
-# 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
-
+uv pip install vllm==0.12.0 --torch-backend=auto
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .

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

@@ -34,9 +34,6 @@ th {
 |`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` |
 |`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

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

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -23,23 +23,13 @@ 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** | `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) |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) |
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
 
-
-!!! 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.
     For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**.
@@ -220,139 +210,100 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 ##### How to parallelize a new model
 
-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.
+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).
 
----
 
-###### Non-intrusive `_sp_plan` (Recommended)
+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:
 
-The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
+- 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.
 
-**Requirements for `forward()` function:**
+First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`:
 
-- 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**
+- 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.
 
-**When to create a submodule:**
+Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example:
 
-```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
+```diff
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+    get_sp_group,
 )
+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)
 ```
 
-| 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) |
+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):
 
-**Key naming convention:**
+```diff
+class QwenImageTransformer2DModel(...):
+    def forward(...):
+        ...
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
 
-| 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),
-    }
++       if self.parallel_config.sequence_parallel_size > 1:
++           output = get_sp_group().all_gather(output, dim=-2)
+        return Transformer2DModelOutput(sample=output)
 ```
 
----
+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
 
-###### Hook flow
+omni = Omni(
+    model="Qwen/Qwen-Image",
++    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
+)
 
+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

docs/user_guide/diffusion_acceleration.md

@@ -47,7 +47,7 @@ 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` | ❌ | ✅ | ❌ | ❌ | ❌ |
 

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -1,94 +0,0 @@
-# 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/online_serving/text_to_image.md

@@ -119,7 +119,6 @@ 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

@@ -1,179 +0,0 @@
-import argparse
-import os
-
-
-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
-    try:
-        # Preferred: load from txt file (one prompt per line)
-        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
-            with open(args.txt_prompts, encoding="utf-8") as f:
-                lines = [ln.strip() for ln in f.readlines()]
-            args.prompts = [ln for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
-    except Exception as e:
-        print(f"[Error] Failed to load prompts: {e}")
-        raise
-
-    if args.prompts is None:
-        # Default prompt for text2img test if none provided
-        args.prompts = ["<|im_start|>A cute cat<|im_end|>"]
-        print(f"[Info] No prompts provided, using default: {args.prompts}")
-    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)
-
-        generate_kwargs = {
-            "prompt": args.prompts,
-            "seed": 52,
-            "need_kv_receive": False,
-            "num_inference_steps": args.steps,
-        }
-
-        if args.image_path:
-            if os.path.exists(args.image_path):
-                loaded_image = Image.open(args.image_path).convert("RGB")
-                generate_kwargs["pil_image"] = loaded_image
-            else:
-                print(f"[Warning] Image path {args.image_path} does not exist.")
-
-        result = client.generate(**generate_kwargs)
-
-        # Ensure result is a list for iteration
-        if not isinstance(result, list):
-            omni_outputs = [result]
-        else:
-            omni_outputs = result
-
-    else:
-        import copy
-
-        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 = copy.deepcopy(omni.default_sampling_params_list)
-        if args.modality == "text2img":
-            params_list[0]["max_tokens"] = 1
-            if len(params_list) > 1:
-                params_list[1]["num_inference_steps"] = args.steps
-
-        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

@@ -1,6 +0,0 @@
-prompt="<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
-
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                                 --prompt_type text \
-                                 --init-sleep-seconds 0 \
-                                 --prompts ${prompt}

examples/offline_inference/image_to_image/image_edit.py

@@ -110,7 +110,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--negative_prompt",
         type=str,
-        default=None,
+        default=" ",
         required=False,
     )
     parser.add_argument(
@@ -181,12 +181,6 @@ 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",
@@ -307,10 +301,7 @@ 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,
-        tensor_parallel_size=args.tensor_parallel_size,
+        ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size
     )
 
     # Configure cache based on backend type
@@ -347,9 +338,6 @@ def main():
     )
     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:")
@@ -363,7 +351,7 @@ 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}, tensor_parallel_size={args.tensor_parallel_size}"
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
     )
     print(f"{'=' * 60}\n")
 
@@ -382,10 +370,6 @@ def main():
         "resolution": args.resolution,
     }
 
-    if profiler_enabled:
-        print("[Profiler] Starting profiling...")
-        omni.start_profile()
-
     outputs = omni.generate(**generate_kwargs)
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
@@ -393,23 +377,6 @@ 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.")
-
     if not outputs:
         raise ValueError("No output generated from omni.generate()")
     logger.info("Outputs: %s", outputs)

examples/offline_inference/image_to_video/image_to_video.py

@@ -19,7 +19,6 @@ Usage:
 """
 
 import argparse
-import os
 from pathlib import Path
 
 import numpy as np
@@ -96,9 +95,6 @@ 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,
@@ -107,10 +103,6 @@ def main():
         flow_shift=args.flow_shift,
     )
 
-    if profiler_enabled:
-        print("[Profiler] Starting profiling...")
-        omni.start_profile()
-
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
         args.prompt,
@@ -186,23 +178,6 @@ 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/qwen2_5_omni/end2end.py

@@ -278,9 +278,9 @@ def get_audio_query(question: str = None, audio_path: str | None = None, samplin
 
 
 query_map = {
-    "use_mixed_modalities": get_mixed_modalities_query,
+    "mixed_modalities": get_mixed_modalities_query,
     "use_audio_in_video": get_use_audio_in_video_query,
-    "use_multi_audios": get_multi_audios_query,
+    "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="use_mixed_modalities",
+        default="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 use_mixed_modalities
+                  --query-type mixed_modalities

examples/offline_inference/qwen3_omni/end2end.py

@@ -12,7 +12,6 @@ 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
@@ -226,76 +225,18 @@ 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,
-    "use_multi_audios": get_multi_audios_query,
-    "use_mixed_modalities": get_mixed_modalities_query,
-    "use_audio_in_video": get_use_audio_in_video_query,
+    "multi_audios": get_multi_audios_query,
+    "mixed_modalities": get_mixed_modalities_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)
@@ -318,10 +259,6 @@ 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()
 
@@ -333,7 +270,7 @@ def main(args):
     )
 
     thinker_sampling_params = SamplingParams(
-        temperature=0.9,
+        temperature=0.4,
         top_p=0.9,
         top_k=-1,
         max_tokens=1200,
@@ -394,8 +331,6 @@ 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:
@@ -450,7 +385,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="use_mixed_modalities",
+        default="mixed_modalities",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/offline_inference/qwen3_tts/README.md

@@ -1,84 +0,0 @@
-# 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.
-
-## 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

@@ -1,381 +0,0 @@
-"""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_image/text_to_image.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
-import os
 import time
 from pathlib import Path
 
@@ -24,9 +23,7 @@ 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=None,
-        help="negative prompt for classifier-free conditional guidance.",
+        "--negative_prompt", default="", 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(
@@ -151,9 +148,6 @@ 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,
@@ -164,10 +158,6 @@ def main():
         enforce_eager=args.enforce_eager,
     )
 
-    if profiler_enabled:
-        print("[Profiler] Starting profiling...")
-        omni.start_profile()
-
     # Time profiling for generation
     print(f"\n{'=' * 60}")
     print("Generation Configuration:")
@@ -199,23 +189,6 @@ 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

@@ -10,7 +10,7 @@ python text_to_video.py \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 640 \
-  --num_frames 33 \
+  --num_frames 32 \
   --guidance_scale 4.0 \
   --guidance_scale_high 3.0 \
   --num_inference_steps 40 \

examples/offline_inference/text_to_video/text_to_video.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
-import os
 from pathlib import Path
 
 import numpy as np
@@ -47,9 +46,6 @@ 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,
@@ -58,10 +54,6 @@ def main():
         flow_shift=args.flow_shift,
     )
 
-    if profiler_enabled:
-        print("[Profiler] Starting profiling...")
-        omni.start_profile()
-
     frames = omni.generate(
         args.prompt,
         negative_prompt=args.negative_prompt,
@@ -134,23 +126,6 @@ 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

@@ -1,184 +0,0 @@
-#!/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

@@ -1,12 +0,0 @@
-#!/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/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -305,26 +305,6 @@ 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,
@@ -332,7 +312,6 @@ 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,
 }
 
 
@@ -393,12 +372,6 @@ 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()
 

examples/online_serving/text_to_image/README.md

@@ -116,7 +116,6 @@ 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
 

mkdocs.yml

@@ -90,7 +90,6 @@ 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

@@ -36,13 +36,10 @@ dependencies = [
     "accelerate==1.12.0",
     "gradio==5.50",
     "soundfile>=0.13.1",
-    "cache-dit==1.2.0",
+    "cache-dit==1.1.8",
     "tqdm>=4.66.0",
     "torchsde>=0.2.6",  # Required for Stable Audio scheduler
-    "fa3-fwd", # flash attention 3, maintained by @ZJY0516
-    "openai-whisper>=20250625",
-    "imageio[ffmpeg]>=2.37.2",
-    # "vllm==0.14.0",  # TODO: fix the entrypoints overwrite problem
+    # "vllm==0.12.0",  # TODO: fix the entrypoints overwrite problem
 ]
 
 [project.optional-dependencies]
@@ -54,10 +51,7 @@ dev = [
     "mypy==1.11.1",
     "pre-commit==4.0.1",
     "openai-whisper>=20250625",
-    "psutil>=7.2.0",
-    "soundfile>=0.13.1",
-    "imageio[ffmpeg]>=0.6.0",
-    "opencv-python>=4.12.0.88"
+    "psutil>=7.2.0"
 ]
 
 docs = [

tests/conftest.py

@@ -1,15 +1,5 @@
 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
@@ -17,39 +7,17 @@ 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.logger import init_logger
-from vllm.utils.network_utils import get_open_port
+from vllm.utils import get_open_port
 
 logger = init_logger(__name__)
 
 
-@pytest.fixture(autouse=True)
-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():
     if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
@@ -123,397 +91,27 @@ def dummy_messages_from_mix_data(
     return messages
 
 
-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
-
-    # 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 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
+
+    vectorizer = CountVectorizer().fit_transform([s1, s2])
+    vectors = vectorizer.toarray()
+    return cosine_similarity([vectors[0]], [vectors[1]])[0][0]
 
 
 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())}"
@@ -521,15 +119,8 @@ 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,
-        temperature=0.0,
-        word_timestamps=True,
-        condition_on_previous_text=False,
-        initial_prompt="Please transcribe with proper word spacing.",
-    )["text"]
+    text = model.transcribe(output_path)["text"]
     if text:
         return text
     else:

tests/diffusion/attention/test_sequence_parallel.py

@@ -170,11 +170,6 @@ 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,11 +48,6 @@ 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})")
@@ -182,11 +177,6 @@ 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})")
@@ -315,11 +305,6 @@ 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/distributed/omni_connectors/test_kv_flow.py

@@ -1,189 +0,0 @@
-import unittest
-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.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
-        req = OmniDiffusionRequest(prompt="test")
-        from types import SimpleNamespace
-
-        req.past_key_values = SimpleNamespace(**layer_blocks)
-        req.kv_metadata = transfer_data.metadata
-
-        # 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.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)
-        req = OmniDiffusionRequest(prompt="integration_test")
-        req.past_key_values = data_dict["layer_blocks"]
-        req.kv_metadata = data_dict["metadata"]
-
-        # 4. Receiver (Injection Simulation)
-        # Use the logic verification again
-        pass
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml

@@ -78,7 +78,6 @@ 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: true
+       enforce_eager: false
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
       #  tensor_parallel_size: 2
@@ -84,7 +84,6 @@ 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,7 +84,6 @@ 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_diffusion_cpu_offload.py

@@ -1,10 +1,11 @@
 import sys
+import threading
+import time
 from pathlib import Path
 
 import pytest
 import torch
 
-from tests.utils import GPUMemoryMonitor
 from vllm_omni.utils.platform_utils import is_npu, is_rocm
 
 # ruff: noqa: E402
@@ -14,6 +15,39 @@ 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"]
 
 
@@ -38,10 +72,14 @@ def test_cpu_offload_diffusion_model(model_name: str):
             guidance_scale=0.0,
             generator=torch.Generator("cuda").manual_seed(42),
         )
-        peak = monitor.peak_used_mb
-        monitor.stop()
 
-        return peak
+        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_memory_mb
 
     offload_peak_memory = inference(offload=True)
     no_offload_peak_memory = inference(offload=False)

tests/e2e/offline_inference/test_ovis_image.py

@@ -233,8 +233,7 @@ def test_real_transformer_init_and_forward():
         }
     )
 
-    od_config = OmniDiffusionConfig(model="dummy-ovis", tf_model_config=tf_config, dtype=torch.bfloat16, num_gpus=1)
-    torch.set_default_dtype(torch.bfloat16)
+    od_config = OmniDiffusionConfig(model="dummy-ovis", tf_model_config=tf_config, dtype=torch.float32, num_gpus=1)
 
     # 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("spawn")
+@create_new_process_for_each_test()
 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("spawn")
+@create_new_process_for_each_test()
 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,8 +5,6 @@ E2E offline tests for Omni model with video input and audio output.
 """
 
 import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 from pathlib import Path
 
 import pytest
@@ -16,6 +14,8 @@ 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,7 +9,6 @@ correctly with diffusion models. It uses minimal settings to keep test time
 short for CI.
 """
 
-import gc
 import os
 import sys
 import time
@@ -22,7 +21,7 @@ import torch.distributed as dist
 from PIL import Image
 
 # ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[3]
+REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
@@ -31,6 +30,9 @@ 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"
@@ -51,145 +53,44 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
     return abs_diff.mean().item(), abs_diff.max().item()
 
 
-def _cleanup_distributed():
-    """Clean up distributed environment and GPU resources."""
-    if dist.is_initialized():
-        dist.destroy_process_group()
+def _get_images(output):
+    """Extract images from output, handling both dict and SimpleNamespace types.
 
-    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-        os.environ.pop(key, None)
+    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
 
-    gc.collect()
+    # Check request_output for pipeline mode
+    if output.request_output is None:
+        return None
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+    if isinstance(output.request_output, list) and len(output.request_output) == 0:
+        return None
 
-    time.sleep(5)
+    item = output.request_output[0]
 
+    # 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
 
-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,
-            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,
-            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
+    # Handle both dict (from SHM serialization) and object (direct) types
+    if isinstance(item, dict):
+        return item.get("images")
+    return getattr(item, "images", None)
 
 
 @pytest.mark.parametrize("model_name", models)
 @pytest.mark.parametrize("ulysses_degree", [1, 2])
 @pytest.mark.parametrize("ring_degree", [1, 2])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])  # Only test bfloat16 to reduce CI time
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
 def test_sequence_parallel(
     model_name: str,
@@ -201,60 +102,214 @@ 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."
         )
 
-    sp_size = ulysses_degree * ring_degree
-    if device_count() < sp_size:
-        pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available")
+    # 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")
 
+    # Use minimal settings for fast testing
     height = 256
     width = 256
+    num_inference_steps = 4  # Minimal steps for fast test
     seed = 42
 
-    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
-    assert baseline_images is not None and len(baseline_images) == 1
+    # 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
 
-    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
+    assert baseline_images is not None
+    assert len(baseline_images) == 1
+    assert baseline_images[0].width == width
+    assert baseline_images[0].height == height
 
+    # 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])
 
-    mean_threshold = 2e-2
-    max_threshold = 2e-1
+    # 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}"
+    )
 
     assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
+        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}"
     )
 
 
 @pytest.mark.parametrize("model_name", models)
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_backend: str):
-    """Test SP with ulysses_degree=4."""
+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."""
     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
 
-    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
-    assert baseline_images is not None and len(baseline_images) == 1
+    # 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
 
-    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
+    assert baseline_images is not None
+    assert len(baseline_images) == 1
+    assert baseline_images[0].width == width
+    assert baseline_images[0].height == height
 
+    # 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])
 
-    mean_threshold = 2e-2
-    max_threshold = 2e-1
+    # 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}"
+    )
 
     assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
+        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}"
     )

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -1,3 +1,4 @@
+import os
 import sys
 from pathlib import Path
 
@@ -14,6 +15,8 @@ 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"]
 

tests/e2e/offline_inference/test_t2v_model.py

@@ -11,10 +11,8 @@ 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_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
 models = ["Wan-AI/Wan2.2-T2V-A14B-Diffusers"]
 
@@ -47,10 +45,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, OmniRequestOutput) or not hasattr(req_out, "images"):
+    if not isinstance(req_out, dict) or "images" not in req_out:
         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_zimage_tensor_parallel.py

@@ -17,14 +17,12 @@ 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_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
 
@@ -68,12 +66,7 @@ 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, float]:
-    torch.cuda.empty_cache()
-    device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
-    monitor.start()
-
+) -> tuple[Image.Image, float]:
     m = Omni(
         model=_get_zimage_model(),
         parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
@@ -114,12 +107,9 @@ def _run_zimage_generate(
             pass
 
         median_time_s = float(np.median(per_request_times_s))
-
-        peak_memory_mb = monitor.peak_used_mb
-
-        return _extract_single_image([last_output]), median_time_s, peak_memory_mb
+        return _extract_single_image([last_output]), median_time_s
     finally:
-        monitor.stop()
+        m.close()
         cleanup_dist_env_and_memory()
 
 
@@ -130,19 +120,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 = 512
-    width = 512
+    height = 256
+    width = 256
     num_inference_steps = 2
     seed = 42
 
-    tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
+    tp1_img, tp1_time_s = _run_zimage_generate(
         tp_size=1,
         height=height,
         width=width,
         num_inference_steps=num_inference_steps,
         seed=seed,
     )
-    tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
+    tp2_img, tp2_time_s = _run_zimage_generate(
         tp_size=2,
         height=height,
         width=width,
@@ -160,7 +150,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 = 5e-1
+    max_threshold = 3.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}; "
@@ -174,8 +164,3 @@ 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,7 +77,6 @@ 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,7 +84,6 @@ 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,6 +1,4 @@
 import asyncio
-import os
-import sys
 from contextlib import ExitStack
 from pathlib import Path
 
@@ -9,8 +7,6 @@ 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")
@@ -59,19 +55,13 @@ async def generate(
 @pytest.mark.asyncio
 async def test_abort():
     with ExitStack() as after:
-        # 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,
-        )
+        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
         after.callback(engine.shutdown)
 
-        # 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]
+        NUM_REQUESTS = 5
+        NUM_EXPECTED_TOKENS = 100
+        NUM_EXPECTED_TOKENS_LONG = 1000
+        REQUEST_IDS_TO_ABORT = [1, 2, 3]
 
         prompt = "Hello my name is Robert and "
 
@@ -84,10 +74,8 @@ 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.
@@ -95,10 +83,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 asyncio.wait_for(task, timeout=60)
+                    await task
             else:
                 # Otherwise, make sure the request was not impacted.
-                num_generated_tokens, request_id = await asyncio.wait_for(task, timeout=180)
+                num_generated_tokens, request_id = await task
                 expected_tokens = NUM_EXPECTED_TOKENS
                 assert num_generated_tokens == expected_tokens, (
                     f"{request_id} generated {num_generated_tokens} but expected {expected_tokens}"
@@ -128,16 +116,11 @@ async def test_build_and_log_summary(monkeypatch):
     monkeypatch.setattr("vllm_omni.entrypoints.client_request_state.ClientRequestState", MockCRS)
 
     with ExitStack() as after:
-        # 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,
-        )
+        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
         after.callback(engine.shutdown)
         prompt = "Hello my name is Robert and "
-        NUM_EXPECTED_TOKENS = 64
-        NUM_REQUESTS = 3
+        NUM_EXPECTED_TOKENS = 100
+        NUM_REQUESTS = 5
         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.network_utils import get_open_port
+from vllm.utils import get_open_port
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 

tests/e2e/online_serving/test_qwen3_omni.py

@@ -4,13 +4,10 @@
 E2E Online tests for Qwen3-Omni model with video input and audio output.
 """
 
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 import base64
 import concurrent.futures
 import ctypes
+import os
 import signal
 import socket
 import subprocess
@@ -21,10 +18,12 @@ from pathlib import Path
 import openai
 import pytest
 from vllm.assets.video import VideoAsset
-from vllm.utils.network_utils import get_open_port
+from vllm.utils import get_open_port
 
 from vllm_omni.utils import is_rocm
 
+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/online_serving/test_qwen3_omni_expansion.py

@@ -6,8 +6,6 @@ E2E Online tests for Qwen3-Omni model.
 
 import concurrent.futures
 import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 import time
 from pathlib import Path
 
@@ -19,11 +17,11 @@ from tests.conftest import (
     convert_audio_to_text,
     cosine_similarity_text,
     dummy_messages_from_mix_data,
-    generate_synthetic_audio,
-    generate_synthetic_image,
     modify_stage_config,
 )
 
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2*H100-80G GPUs
@@ -70,20 +68,6 @@ def get_max_batch_size(size_type="few"):
     return batch_sizes.get(size_type, 5)
 
 
-def get_deploy_config(deploy_type="TP1"):
-    result = {
-        "TP1": {
-            0: {
-                "engine_args.gpu_memory_utilization": 0.95,
-                "engine_args.tensor_parallel_size": 1,
-                "runtime.devices": "0",
-            },
-            2: {"runtime.devices": "1"},
-        }
-    }
-    return result.get(deploy_type, result["TP1"])
-
-
 @pytest.mark.parametrize("test_config", test_params)
 def test_text_to_text_001(test_config: tuple[str, str]) -> None:
     """Test processing text, generating text output via OpenAI API."""
@@ -172,198 +156,3 @@ def test_text_to_text_audio_001(test_config: tuple[str, str]) -> None:
             assert cosine_similarity_text(audio_content.lower(), text_content.lower()) > 0.9, (
                 "The audio content is not same as the text"
             )
-
-
-@pytest.mark.parametrize("test_config", test_params)
-def test_audio_to_text_001(test_config: tuple[str, str]) -> None:
-    """Test processing text, generating text output via OpenAI API."""
-    model, stage_config_path = test_config
-    deploy_config = get_deploy_config()
-    deploy_config[0]["default_sampling_params.ignore_eos"] = True
-    stage_config_path = modify_stage_config(stage_config_path, deploy_config)
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 1)['base64']}"
-        messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
-        # Test single completion
-        api_client = client(server)
-        start_time = time.perf_counter()
-        chat_completion = api_client.chat.completions.create(
-            model=server.model, messages=messages, max_tokens=200, modalities=["text"]
-        )
-        # Verify only output text
-        assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
-
-        # Verify text output success
-        text_choice = chat_completion.choices[0]
-        assert text_choice.message.content is not None, "No text output is generated"
-        assert chat_completion.usage.completion_tokens == 200, (
-            "The output length differs from the requested max_tokens."
-        )
-
-        # Verify E2E
-        print(f"the request e2e is: {time.perf_counter() - start_time}")
-        # TODO: Verify the E2E latency after confirmation baseline.
-
-
-@pytest.mark.parametrize("test_config", test_params)
-def test_audio_to_text_audio_001(test_config: tuple[str, str]) -> None:
-    """Test processing text, generating audio output via OpenAI API."""
-
-    model, stage_config_path = test_config
-    num_concurrent_requests = get_max_batch_size()
-    stage_config_path = modify_stage_config(
-        stage_config_path,
-        {
-            0: {"runtime.max_batch_size": num_concurrent_requests},
-            1: {"runtime.max_batch_size": num_concurrent_requests},
-        },
-    )
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        audio_data_url = []
-        for _ in range(5):
-            audio_data_url.append(f"data:audio/wav;base64,{generate_synthetic_audio(1, 5)['base64']}")
-
-        messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
-
-        # Test single completion
-        api_client = client(server)
-        e2e_list = list()
-        with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
-            # Submit multiple completion requests concurrently
-            futures = [
-                executor.submit(api_client.chat.completions.create, model=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_message = chat_completion.choices[1].message
-            audio_data = audio_message.audio.data
-            assert audio_data is not None, "No audio output is generated"
-            assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
-
-            # Verify text output success
-            text_choice = chat_completion.choices[0]
-            text_content = text_choice.message.content
-            assert text_choice.message.content is not None, "No text output is generated"
-
-            # 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, text_content)
-            print(f"similarity between audio and text is: {similarity}")
-            assert similarity > 0.9, "The audio content is not same as the text"
-
-
-@pytest.mark.parametrize("test_config", test_params)
-def test_image_to_text_001(test_config: tuple[str, str]) -> None:
-    """Test processing text, generating text output via OpenAI API."""
-    model, stage_config_path = test_config
-    deploy_config = get_deploy_config()
-    stage_config_path = modify_stage_config(stage_config_path, deploy_config)
-
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
-        messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
-        # Test single completion
-        api_client = client(server)
-        start_time = time.perf_counter()
-        chat_completion = api_client.chat.completions.create(
-            model=server.model, messages=messages, max_tokens=100, modalities=["text"]
-        )
-        # Verify E2E
-        print(f"the request e2e is: {time.perf_counter() - start_time}")
-        # TODO: Verify the E2E latency after confirmation baseline.
-
-        # Verify only output text
-        assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
-
-        # Verify text output success
-        text_choice = chat_completion.choices[0]
-        text_content = text_choice.message.content
-        assert text_content is not None, "No text output is generated"
-        assert chat_completion.usage.completion_tokens <= 100, "The output length more than the requested max_tokens."
-        assert "square" in text_content.lower(), "The output do not contain keywords."
-
-
-@pytest.mark.parametrize("test_config", test_params)
-def test_image_to_text_audio_001(test_config: tuple[str, str]) -> None:
-    """Test processing text, generating audio output via OpenAI API."""
-
-    model, stage_config_path = test_config
-    num_concurrent_requests = 5
-    stage_config_path = modify_stage_config(
-        stage_config_path,
-        {
-            0: {"runtime.max_batch_size": num_concurrent_requests},
-            1: {"runtime.max_batch_size": num_concurrent_requests},
-        },
-    )
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
-        image_data_url = []
-        for _ in range(4):
-            image_data_url.append(f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}")
-
-        messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
-
-        # Test single completion
-        api_client = client(server)
-        e2e_list = list()
-        with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
-            # Submit multiple completion requests concurrently
-            futures = [
-                executor.submit(
-                    api_client.chat.completions.create,
-                    model=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_message = chat_completion.choices[1].message
-            audio_data = audio_message.audio.data
-            assert audio_data is not None, "No audio output is generated"
-            assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
-
-            # Verify text output success
-            text_choice = chat_completion.choices[0]
-            text_content = text_choice.message.content
-            assert text_content is not None, "No text output is generated"
-            assert "square" 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, text_content)
-            print(f"similarity between audio and text is: {similarity}")
-            assert similarity > 0.9, "The audio content is not same as the text"

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -15,7 +15,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.6
-    enforce_eager: false
+    enforce_eager: true
     trust_remote_code: true
     engine_output_type: latent  # Output hidden states for talker
     distributed_executor_backend: "mp"
@@ -46,7 +46,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
     enable_prefix_caching: false
@@ -58,7 +58,7 @@ stage_args:
   default_sampling_params:
     temperature: 0.9
     top_k: 50
-    max_tokens: 1000
+    max_tokens: 100
     seed: 42
     detokenize: False
     repetition_penalty: 1.05
@@ -81,7 +81,6 @@ stage_args:
     distributed_executor_backend: "mp"
     max_num_batched_tokens: 1000000
     hf_config_name: thinker_config
-    async_scheduling: false
   engine_input_source: [1]
   custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
   final_output: true
@@ -90,7 +89,7 @@ stage_args:
     temperature: 0.0
     top_p: 1.0
     top_k: -1
-    max_tokens: 2000
+    max_tokens: 200
     seed: 42
     detokenize: True
     repetition_penalty: 1.1

tests/entrypoints/test_omni_input_preprocessor.py

@@ -1,59 +0,0 @@
-from vllm_omni.inputs.preprocess import OmniInputPreprocessor
-
-
-def _make_preprocessor(monkeypatch):
-    preprocessor = object.__new__(OmniInputPreprocessor)
-    monkeypatch.setattr(preprocessor, "_truncate_inputs", lambda tokens, tokenization_kwargs=None: tokens)
-    monkeypatch.setattr(
-        preprocessor,
-        "_process_multimodal",
-        lambda *args, **kwargs: {"prompt_token_ids": [1, 2, 3]},
-    )
-    monkeypatch.setattr(preprocessor, "_tokenize_prompt", lambda prompt_text, tokenization_kwargs=None: [9, 8, 7])
-    return preprocessor
-
-
-def test_process_tokens_keeps_additional_information(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt_token_ids": [1, 2, 3],
-        "prompt_embeds": "embeds",
-        "additional_information": {"task": ["tts"], "lang": ["auto"]},
-    }
-
-    inputs = OmniInputPreprocessor._process_tokens(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [1, 2, 3]
-    assert inputs["prompt_embeds"] == "embeds"
-    assert inputs["additional_information"] == {"task": ["tts"], "lang": ["auto"]}
-
-
-def test_process_text_keeps_additional_information(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt": "hello",
-        "prompt_embeds": "embeds",
-        "additional_information": {"speaker": ["alice"]},
-    }
-
-    inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [9, 8, 7]
-    assert inputs["prompt_embeds"] == "embeds"
-    assert inputs["additional_information"] == {"speaker": ["alice"]}
-
-
-def test_process_text_multimodal_skips_empty_payloads(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt": "hello",
-        "multi_modal_data": {"image": "fake"},
-        "prompt_embeds": None,
-        "additional_information": None,
-    }
-
-    inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [1, 2, 3]
-    assert "prompt_embeds" not in inputs
-    assert "additional_information" not in inputs

tests/entrypoints/test_omni_new_request_data.py

@@ -1,49 +0,0 @@
-from types import SimpleNamespace
-
-import torch
-
-from vllm_omni.core.sched.output import OmniNewRequestData
-
-
-def test_omni_new_request_data_copies_payloads():
-    prompt_embeds = torch.randn(2, 3)
-    additional_information = {
-        "speaker": ["test"],
-        "codes": torch.tensor([1, 2], dtype=torch.int64),
-    }
-    request = SimpleNamespace(
-        request_id="req-1",
-        prompt_token_ids=[101, 102],
-        mm_features=None,
-        sampling_params=None,
-        pooling_params=None,
-        num_computed_tokens=0,
-        lora_request=None,
-        prompt_embeds=prompt_embeds,
-        additional_information=additional_information,
-    )
-
-    data = OmniNewRequestData.from_request(request, ([0, 1],), prefill_token_ids=[101, 102])
-
-    assert data.prompt_embeds is prompt_embeds
-    assert data.additional_information is additional_information
-    assert data.prefill_token_ids == [101, 102]
-
-
-def test_omni_new_request_data_allows_missing_payloads():
-    request = SimpleNamespace(
-        request_id="req-2",
-        prompt_token_ids=[201, 202],
-        mm_features=None,
-        sampling_params=None,
-        pooling_params=None,
-        num_computed_tokens=0,
-        lora_request=None,
-        prompt_embeds=None,
-        additional_information=None,
-    )
-
-    data = OmniNewRequestData.from_request(request, ([0],), prefill_token_ids=None)
-
-    assert data.prompt_embeds is None
-    assert data.additional_information is None

tests/entrypoints/test_stage_config.py

@@ -0,0 +1,396 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for structured stage configuration builder."""
+
+import json
+from unittest.mock import Mock
+
+import pytest
+from omegaconf import OmegaConf
+
+from vllm_omni.entrypoints.stage_config import (
+    StageConfig,
+    StageConfigFactory,
+    StageRuntimeConfig,
+)
+
+
+class TestStageRuntimeConfig:
+    """Tests for StageRuntimeConfig dataclass."""
+    
+    def test_runtime_config_creation(self):
+        """Test basic runtime config creation."""
+        config = StageRuntimeConfig(
+            process=True,
+            devices="0,1,2",
+            max_batch_size=4
+        )
+        
+        assert config.process is True
+        assert config.devices == "0,1,2"
+        assert config.max_batch_size == 4
+
+
+class TestStageConfig:
+    """Tests for StageConfig dataclass."""
+    
+    def test_stage_config_creation(self):
+        """Test basic stage config creation."""
+        runtime = StageRuntimeConfig(
+            process=True,
+            devices="0",
+            max_batch_size=1
+        )
+        
+        config = StageConfig(
+            stage_id=0,
+            stage_type="diffusion",
+            runtime=runtime,
+            engine_args={"model": "test-model"},
+            final_output=True,
+            final_output_type="image"
+        )
+        
+        assert config.stage_id == 0
+        assert config.stage_type == "diffusion"
+        assert config.runtime.devices == "0"
+        assert config.engine_args["model"] == "test-model"
+        assert config.final_output is True
+        assert config.final_output_type == "image"
+    
+    def test_to_omegaconf(self):
+        """Test conversion to OmegaConf DictConfig."""
+        runtime = StageRuntimeConfig(
+            process=True,
+            devices="0,1",
+            max_batch_size=2
+        )
+        
+        config = StageConfig(
+            stage_id=0,
+            stage_type="diffusion",
+            runtime=runtime,
+            engine_args={"cache_backend": "cache_dit"},
+            final_output=True,
+            final_output_type="image"
+        )
+        
+        omega_config = config.to_omegaconf()
+        
+        assert omega_config.stage_id == 0
+        assert omega_config.stage_type == "diffusion"
+        assert omega_config.runtime.process is True
+        assert omega_config.runtime.devices == "0,1"
+        assert omega_config.runtime.max_batch_size == 2
+        assert omega_config.engine_args.cache_backend == "cache_dit"
+        assert omega_config.final_output is True
+        assert omega_config.final_output_type == "image"
+
+
+class TestDeviceStringGeneration:
+    """Tests for device string generation helper."""
+    
+    def test_device_string_from_parallel_config(self):
+        """Test device string generation from parallel_config."""
+        # Mock parallel_config with world_size
+        parallel_config = Mock()
+        parallel_config.world_size = 4
+        
+        devices = StageConfigFactory._get_device_string(parallel_config)
+        
+        assert devices == "0,1,2,3"
+    
+    def test_device_string_from_num_devices(self):
+        """Test device string generation from num_devices parameter."""
+        devices = StageConfigFactory._get_device_string(None, num_devices=3)
+        
+        assert devices == "0,1,2"
+    
+    def test_device_string_default_single_device(self):
+        """Test default device string is single device."""
+        devices = StageConfigFactory._get_device_string(None, None)
+        
+        assert devices == "0"
+    
+    def test_device_string_single_device(self):
+        """Test device string with single device."""
+        parallel_config = Mock()
+        parallel_config.world_size = 1
+        
+        devices = StageConfigFactory._get_device_string(parallel_config)
+        
+        assert devices == "0"
+
+
+class TestCacheConfigNormalization:
+    """Tests for cache configuration normalization."""
+    
+    def test_cache_dit_defaults(self):
+        """Test default cache_dit configuration."""
+        config = StageConfigFactory._get_default_cache_config("cache_dit")
+        
+        assert config is not None
+        assert config["Fn_compute_blocks"] == 1
+        assert config["Bn_compute_blocks"] == 0
+        assert config["max_warmup_steps"] == 4
+        assert config["residual_diff_threshold"] == 0.24
+        assert config["max_continuous_cached_steps"] == 3
+        assert config["enable_taylorseer"] is False
+        assert config["taylorseer_order"] == 1
+        assert config["scm_steps_mask_policy"] is None
+        assert config["scm_steps_policy"] == "dynamic"
+    
+    def test_tea_cache_defaults(self):
+        """Test default tea_cache configuration."""
+        config = StageConfigFactory._get_default_cache_config("tea_cache")
+        
+        assert config is not None
+        assert config["rel_l1_thresh"] == 0.2
+    
+    def test_no_cache_defaults(self):
+        """Test no defaults for unknown cache backend."""
+        config = StageConfigFactory._get_default_cache_config("unknown")
+        
+        assert config is None
+    
+    def test_normalize_json_string_cache_config(self):
+        """Test normalization of JSON string cache_config."""
+        cache_config_str = '{"Fn_compute_blocks": 2, "max_warmup_steps": 6}'
+        
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            cache_config_str
+        )
+        
+        assert normalized is not None
+        assert normalized["Fn_compute_blocks"] == 2
+        assert normalized["max_warmup_steps"] == 6
+    
+    def test_normalize_invalid_json_uses_defaults(self):
+        """Test that invalid JSON falls back to defaults."""
+        cache_config_str = '{invalid json}'
+        
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            cache_config_str
+        )
+        
+        # Should fall back to defaults for cache_dit
+        assert normalized is not None
+        assert normalized["Fn_compute_blocks"] == 1
+    
+    def test_normalize_dict_cache_config(self):
+        """Test normalization of dict cache_config."""
+        cache_config_dict = {"Fn_compute_blocks": 3}
+        
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            cache_config_dict
+        )
+        
+        assert normalized == cache_config_dict
+    
+    def test_normalize_none_with_backend_uses_defaults(self):
+        """Test that None cache_config with backend uses defaults."""
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            None
+        )
+        
+        assert normalized is not None
+        assert normalized["Fn_compute_blocks"] == 1
+    
+    def test_normalize_none_without_backend_returns_none(self):
+        """Test that None cache_config without backend returns None."""
+        normalized = StageConfigFactory._normalize_cache_config(
+            "none", 
+            None
+        )
+        
+        assert normalized is None
+
+
+class TestDefaultDiffusionConfigBuilder:
+    """Tests for default diffusion stage configuration builder."""
+    
+    def test_default_diffusion_config_structure(self):
+        """Test that the builder produces correct schema structure."""
+        kwargs = {
+            "model": "test-model",
+            "dtype": "float16",
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        
+        assert len(configs) == 1
+        config = configs[0]
+        
+        # Verify structure
+        assert config.stage_id == 0
+        assert config.stage_type == "diffusion"
+        assert config.runtime.process is True
+        assert config.runtime.devices == "0"
+        assert config.runtime.max_batch_size == 1
+        assert config.final_output is True
+        assert config.final_output_type == "image"
+        
+        # Verify engine_args
+        assert config.engine_args.model == "test-model"
+        assert config.engine_args.dtype == "float16"  # Should be stringified
+        assert config.engine_args.model_stage == "diffusion"
+        assert config.engine_args.cache_backend == "none"
+    
+    def test_default_diffusion_with_parallel_config(self):
+        """Test builder with parallel configuration."""
+        parallel_config = Mock()
+        parallel_config.world_size = 4
+        
+        kwargs = {
+            "model": "test-model",
+            "parallel_config": parallel_config,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should generate device string for 4 devices
+        assert config.runtime.devices == "0,1,2,3"
+    
+    def test_default_diffusion_with_cache_backend(self):
+        """Test builder with cache backend."""
+        kwargs = {
+            "model": "test-model",
+            "cache_backend": "cache_dit",
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should have cache_dit defaults
+        assert config.engine_args.cache_backend == "cache_dit"
+        assert config.engine_args.cache_config is not None
+        assert config.engine_args.cache_config["Fn_compute_blocks"] == 1
+    
+    def test_default_diffusion_preserves_kwargs(self):
+        """Test that builder preserves additional kwargs in engine_args."""
+        kwargs = {
+            "model": "test-model",
+            "custom_param": "custom_value",
+            "another_param": 42,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should preserve custom parameters
+        assert config.engine_args.custom_param == "custom_value"
+        assert config.engine_args.another_param == 42
+    
+    def test_dtype_stringification(self):
+        """Test that dtype is converted to string."""
+        # Use a mock object for dtype
+        class DType:
+            def __str__(self):
+                return "torch.float16"
+        
+        kwargs = {
+            "model": "test-model",
+            "dtype": DType(),
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        assert config.engine_args.dtype == "torch.float16"
+
+
+class TestAsyncDiffusionConfigBuilder:
+    """Tests for async diffusion stage configuration builder."""
+    
+    def test_async_diffusion_with_parallel_config(self):
+        """Test async builder with parallel_config."""
+        parallel_config = Mock()
+        parallel_config.world_size = 2
+        
+        kwargs = {
+            "parallel_config": parallel_config,
+            "vae_use_slicing": True,
+            "vae_use_tiling": False,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion_async(kwargs)
+        config = configs[0]
+        
+        assert config.runtime.devices == "0,1"
+        assert config.engine_args.parallel_config == parallel_config
+        assert config.engine_args.vae_use_slicing is True
+        assert config.engine_args.vae_use_tiling is False
+    
+    def test_async_diffusion_builds_parallel_config(self):
+        """Test async builder builds parallel_config from parameters."""
+        kwargs = {
+            "tensor_parallel_size": 2,
+            "ulysses_degree": 2,
+            "ring_degree": 1,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion_async(kwargs)
+        config = configs[0]
+        
+        # Should build parallel_config and device string
+        # tensor_parallel_size=2 * sequence_parallel_size=2 = 4 devices
+        assert config.runtime.devices == "0,1,2,3"
+        assert config.engine_args.parallel_config is not None
+        assert config.engine_args.parallel_config.tensor_parallel_size == 2
+        assert config.engine_args.parallel_config.ulysses_degree == 2
+    
+    def test_async_diffusion_includes_async_specific_args(self):
+        """Test that async builder includes async-specific engine args."""
+        kwargs = {
+            "enable_cpu_offload": True,
+            "enforce_eager": True,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion_async(kwargs)
+        config = configs[0]
+        
+        assert config.engine_args.enable_cpu_offload is True
+        assert config.engine_args.enforce_eager is True
+        assert config.engine_args.model_stage == "diffusion"
+
+
+class TestCacheConfigIntegration:
+    """Integration tests for cache config normalization."""
+    
+    def test_cache_backend_integration(self):
+        """Test cache_backend arguments are correctly nested in config."""
+        kwargs = {
+            "model": "test-model",
+            "cache_backend": "cache_dit",
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Verify cache_dit defaults are in engine_args
+        assert config.engine_args.cache_backend == "cache_dit"
+        assert config.engine_args.cache_config is not None
+        assert "Fn_compute_blocks" in config.engine_args.cache_config
+        assert "max_warmup_steps" in config.engine_args.cache_config
+    
+    def test_custom_cache_config_override(self):
+        """Test custom cache_config overrides defaults."""
+        custom_config = {"Fn_compute_blocks": 5, "custom_key": "custom_value"}
+        
+        kwargs = {
+            "model": "test-model",
+            "cache_backend": "cache_dit",
+            "cache_config": custom_config,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should use custom config
+        assert config.engine_args.cache_config["Fn_compute_blocks"] == 5
+        assert config.engine_args.cache_config["custom_key"] == "custom_value"

tests/utils.py

@@ -7,7 +7,6 @@ import signal
 import subprocess
 import sys
 import tempfile
-import threading
 import time
 from collections.abc import Callable
 from contextlib import ExitStack, contextmanager, suppress
@@ -15,7 +14,6 @@ from typing import Any, Literal
 
 import cloudpickle
 import pytest
-import torch
 from typing_extensions import ParamSpec
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import cuda_device_count_stateless
@@ -476,44 +474,3 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
         return func
 
     return wrapper
-
-
-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=False)
-        self._thread.start()
-
-    def stop(self) -> None:
-        if self._thread is None:
-            return
-        self._stop_event.set()
-        self._thread.join(timeout=2.0)
-
-    @property
-    def peak_used_mb(self) -> float:
-        fallback_alloc = torch.cuda.max_memory_allocated(device=self.device_index) / (1024**2)
-        fallback_reserved = torch.cuda.max_memory_reserved(device=self.device_index) / (1024**2)
-        return max(self._peak_used_mb, fallback_alloc, fallback_reserved)
-
-    def __del__(self):
-        self.stop()

tests/worker/test_gpu_generation_model_runner.py

@@ -1,75 +0,0 @@
-import torch
-
-from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
-
-
-class _DummyInputBatch:
-    def __init__(self):
-        self.req_ids = ["req-1"]
-        self.req_id_to_index = {"req-1": 0}
-        self.num_reqs = 1
-        self.vocab_size = 10
-
-
-def _make_runner(multimodal_outputs):
-    runner = object.__new__(GPUGenerationModelRunner)
-    runner.execute_model_state = (
-        None,
-        None,
-        None,
-        None,
-        None,
-        None,
-        None,
-        None,
-        None,
-        multimodal_outputs,
-    )
-    runner.kv_connector_output = None
-    runner.input_batch = _DummyInputBatch()
-    runner.use_async_scheduling = False
-    runner.device = torch.device("cpu")
-    runner.supports_mm_inputs = False
-    return runner
-
-
-def test_sample_tokens_tensor_output():
-    multimodal_outputs = torch.randn(1, 2, 3)
-    runner = _make_runner(multimodal_outputs)
-
-    output = GPUGenerationModelRunner.sample_tokens(runner)
-
-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"].shape == (2, 3)
-
-
-def test_sample_tokens_list_output():
-    multimodal_outputs = [torch.randn(2, 1)]
-    runner = _make_runner(multimodal_outputs)
-
-    output = GPUGenerationModelRunner.sample_tokens(runner)
-
-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"].shape == (2, 1)
-
-
-def test_sample_tokens_list_allows_none_output():
-    multimodal_outputs = [None]
-    runner = _make_runner(multimodal_outputs)
-
-    output = GPUGenerationModelRunner.sample_tokens(runner)
-
-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"] is None
-
-
-def test_sample_tokens_dict_output():
-    multimodal_outputs = {"audio": torch.randn(1, 4), "unused": None}
-    runner = _make_runner(multimodal_outputs)
-
-    output = GPUGenerationModelRunner.sample_tokens(runner)
-
-    assert len(output.pooler_output) == 1
-    assert "audio" in output.pooler_output[0]
-    assert "unused" not in output.pooler_output[0]
-    assert output.pooler_output[0]["audio"].shape == (1, 4)

tests/worker/test_omni_gpu_model_runner.py

@@ -69,7 +69,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
 
     # Provide a minimal implementation that returns the expected 4-tuple.
     def _determine_batch_execution_and_padding(**kwargs):
-        return None, object(), None, None, None
+        return None, object(), None, None
 
     runner._determine_batch_execution_and_padding = _determine_batch_execution_and_padding
 

tools/pre_commit/check_pickle_imports.py

@@ -20,7 +20,7 @@ ALLOWED_FILES = {
     "tests/e2e/offline_inference/utils.py",
     "tests/utils.py",
     "vllm_omni/diffusion/distributed/group_coordinator.py",
-    "tests/diffusion/attention/test_attention_sp.py",
+    "tests/diffusion/attention/test_sequence_parallel.py",
 }
 
 PICKLE_RE = re.compile(

vllm_omni/__init__.py

@@ -20,7 +20,6 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
     # Allow importing vllm_omni without vllm (e.g., documentation builds)
     patch = None  # type: ignore
 
-
 from .config import OmniModelConfig
 from .entrypoints.async_omni import AsyncOmni
 

vllm_omni/assets/video.py

@@ -1,16 +0,0 @@
-import librosa
-import numpy as np
-from vllm.assets.video import VideoAsset
-
-
-def extract_video_audio(path: str = None, sampling_rate: int = 16000) -> np.ndarray:
-    """This function extracts the audio from a video file path and returns the audio as a numpy array.
-    Args:
-        path: The path to the video file.
-    Returns:
-        The audio as a numpy array.
-    """
-    if not path:
-        path = VideoAsset(name="baby_reading").video_path
-    audio_signal, sr = librosa.load(path, sr=sampling_rate)
-    return audio_signal

vllm_omni/config/model.py

@@ -1,12 +1,20 @@
 import warnings
+from importlib.util import find_spec
 from typing import Any
 
 import torch
+import vllm.envs as envs
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
+from vllm.attention.backends.registry import AttentionBackendEnum
 from vllm.config import ModelConfig, config
 from vllm.config.model import (
     _RUNNER_CONVERTS,
+    _RUNNER_TASKS,
+    ConvertOption,
+    ConvertType,
+    RunnerOption,
+    TaskOption,
     _get_and_verify_dtype,
     get_served_model_name,
 )
@@ -21,11 +29,12 @@ from vllm.transformers_utils.config import (
     get_pooling_config,
 )
 from vllm.transformers_utils.gguf_utils import (
-    is_gguf,
     maybe_patch_hf_config_from_gguf,
 )
-from vllm.transformers_utils.utils import maybe_model_redirect
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.transformers_utils.utils import (
+    is_gguf,
+    maybe_model_redirect,
+)
 
 import vllm_omni.model_executor.models as me_models
 
@@ -63,7 +72,6 @@ class OmniModelConfig(ModelConfig):
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: str | None = None
     hf_config_name: str | None = None
-    omni_kv_config: dict | None = None
 
     @property
     def registry(self):
@@ -138,6 +146,14 @@ class OmniModelConfig(ModelConfig):
 
         self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
 
+        if (backend := envs.VLLM_ATTENTION_BACKEND) and backend == "FLASHINFER" and find_spec("flashinfer") is None:
+            raise ValueError(
+                "VLLM_ATTENTION_BACKEND is set to FLASHINFER, but flashinfer "
+                "module was not found. See "
+                "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile "  # noqa: E501
+                "for instructions on how to install it."
+            )
+
         if self.override_attention_dtype is not None and not current_platform.is_rocm():
             warnings.warn(
                 "override-attention-dtype is set but not using ROCm platform",
@@ -167,16 +183,111 @@ class OmniModelConfig(ModelConfig):
         self.hf_text_config = self.draw_hf_text_config()
         self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
         self.encoder_config = self._get_encoder_config()
-        self.hf_image_processor_config = get_hf_image_processor_config(
-            self.model, hf_token=self.hf_token, revision=self.revision
-        )
-        self.model_arch_config = self.get_model_arch_config()
+        # Try to load image processor config, but allow it to fail for stages that don't need it
+        try:
+            self.hf_image_processor_config = get_hf_image_processor_config(
+                self.model, hf_token=self.hf_token, revision=self.revision
+            )
+        except (OSError, ValueError, IndexError) as e:
+            # Some stages (e.g., code2wav, talker) don't need image processor
+            # Log warning but allow initialization to continue
+            logger.warning(
+                f"Failed to load image processor config for model '{self.model}': {e}. "
+                "This is expected for stages that don't require image processing."
+            )
+            self.hf_image_processor_config = None
 
         architectures = self.architectures
         registry = self.registry
         is_generative_model = registry.is_text_generation_model(architectures, self)
         is_pooling_model = registry.is_pooling_model(architectures, self)
 
+        def _task_to_convert(task: TaskOption) -> ConvertType:
+            if task == "embedding" or task == "embed":
+                return "embed"
+            if task == "classify":
+                return "classify"
+            if task == "reward":
+                return "reward"
+            if task == "score":
+                new_task = self._get_default_pooling_task(architectures)
+                return "classify" if new_task == "classify" else "embed"
+
+            return "none"
+
+        if self.task is not None:
+            runner: RunnerOption = "auto"
+            convert: ConvertOption = "auto"
+            msg_prefix = (
+                "The 'task' option has been deprecated and will be removed in v0.13.0 or v1.0, whichever comes first."
+            )
+            msg_hint = "Please remove this option."
+
+            is_generative_task = self.task in _RUNNER_TASKS["generate"]
+            is_pooling_task = self.task in _RUNNER_TASKS["pooling"]
+
+            if is_generative_model and is_pooling_model:
+                if is_generative_task:
+                    runner = "generate"
+                    convert = "auto"
+                    msg_hint = (
+                        "Please replace this option with `--runner "
+                        "generate` to continue using this model "
+                        "as a generative model."
+                    )
+                elif is_pooling_task:
+                    runner = "pooling"
+                    convert = "auto"
+                    msg_hint = (
+                        "Please replace this option with `--runner "
+                        "pooling` to continue using this model "
+                        "as a pooling model."
+                    )
+                else:  # task == "auto"
+                    pass
+            elif is_generative_model or is_pooling_model:
+                if is_generative_task:
+                    runner = "generate"
+                    convert = "auto"
+                    msg_hint = "Please remove this option"
+                elif is_pooling_task:
+                    runner = "pooling"
+                    convert = _task_to_convert(self.task)
+                    msg_hint = (
+                        "Please replace this option with `--convert "
+                        f"{convert}` to continue using this model "
+                        "as a pooling model."
+                    )
+                else:  # task == "auto"
+                    pass
+            else:
+                # Neither generative nor pooling model - try to convert if possible
+                if is_pooling_task:
+                    runner = "pooling"
+                    convert = _task_to_convert(self.task)
+                    msg_hint = (
+                        "Please replace this option with `--runner pooling "
+                        f"--convert {convert}` to continue using this model "
+                        "as a pooling model."
+                    )
+                else:
+                    debug_info = {
+                        "architectures": architectures,
+                        "is_generative_model": is_generative_model,
+                        "is_pooling_model": is_pooling_model,
+                    }
+                    raise AssertionError(
+                        "The model should be a generative or "
+                        "pooling model when task is set to "
+                        f"{self.task!r}. Found: {debug_info}"
+                    )
+
+            self.runner = runner
+            self.convert = convert
+
+            msg = f"{msg_prefix} {msg_hint}"
+            warnings.warn(msg, DeprecationWarning, stacklevel=2)
+
         self.runner_type = self._get_runner_type(architectures, self.runner)
         self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert)
 
@@ -214,12 +325,9 @@ class OmniModelConfig(ModelConfig):
                     if getattr(self.pooler_config, k) is None:
                         setattr(self.pooler_config, k, v)
 
-            default_seq_pooling_type = self._model_info.default_seq_pooling_type
-            if self.pooler_config.seq_pooling_type is None:
-                self.pooler_config.seq_pooling_type = default_seq_pooling_type
-            default_tok_pooling_type = self._model_info.default_tok_pooling_type
-            if self.pooler_config.tok_pooling_type is None:
-                self.pooler_config.tok_pooling_type = default_tok_pooling_type
+            default_pooling_type = self._model_info.default_pooling_type
+            if self.pooler_config.pooling_type is None:
+                self.pooler_config.pooling_type = default_pooling_type
 
         self.dtype: torch.dtype = _get_and_verify_dtype(
             self.model,
@@ -231,11 +339,6 @@ class OmniModelConfig(ModelConfig):
 
         self.original_max_model_len = self.max_model_len
         self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
-
-        if self.is_encoder_decoder:
-            self.mm_processor_cache_gb = 0
-            logger.info("Encoder-decoder model detected, disabling mm processor cache.")
-
         # Init multimodal config if needed
         if self._model_info.supports_multimodal:
             if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data:
@@ -279,7 +382,7 @@ class OmniModelConfig(ModelConfig):
 
         # Avoid running try_verify_and_update_config multiple times
         self.config_updated = False
-        self._try_verify_and_update_model_config()
+
         self._verify_quantization()
         self._verify_cuda_graph()
         self._verify_bnb_config()

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -1,12 +1,8 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from dataclasses import asdict, dataclass
 from time import time
-from typing import Any
 
-import numpy as np
-from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
@@ -14,24 +10,10 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import check_stop, remove_all
 from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
-from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
-logger = init_logger(__name__)
-
-
-@dataclass
-class KVCacheTransferData:
-    request_id: str
-    layer_blocks: dict[str, Any]
-    block_ids: list[int]
-    metadata: dict[str, Any]
-
-    def to_dict(self) -> dict[str, Any]:
-        return asdict(self)
-
 
 class OmniARScheduler(VLLMScheduler):
     """
@@ -42,91 +24,8 @@ class OmniARScheduler(VLLMScheduler):
     specific to vLLM-Omni.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        # Track requests that need KV cache transfer when finished
-        # Value is {"seq_len": int, "block_ids": list[int]}
-        self.requests_needing_kv_transfer: dict[str, dict[str, Any]] = {}
-
-        # Track requests waiting for KV transfer (blocks not freed yet)
-        self.waiting_for_transfer_free: set[str] = set()
-
-        # Track ACTIVE transfers (submitted to runner but not yet acked via kv_extracted_req_ids)
-        self.active_kv_transfers: set[str] = set()
-
-        # [Omni] Pre-parse KV transfer criteria
-        self.kv_transfer_criteria = self._get_kv_transfer_criteria()
-
-        # Track requests that have already triggered prefill transfer to avoid duplicates
-        self.transfer_triggered_requests: set[str] = set()
-
-    def _get_kv_transfer_criteria(self) -> dict | None:
-        # Note: vllm_config is available in Scheduler after super().__init__
-        if not hasattr(self, "vllm_config"):
-            return None
-
-        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
-        if omni_kv_config:
-            if isinstance(omni_kv_config, dict):
-                return omni_kv_config.get("kv_transfer_criteria", None)
-            else:
-                return getattr(omni_kv_config, "kv_transfer_criteria", None)
-        return None
-
-    def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int]) -> bool:
-        """
-        Check triggers and process side effects (marking transfer).
-        Returns True if request should be STOPPED.
-        Returns False if request should continue (even if transfer was triggered).
-        """
-        if not self.kv_transfer_criteria:
-            return False
-
-        if request.request_id in self.waiting_for_transfer_free:
-            return False
-
-        criteria_type = self.kv_transfer_criteria.get("type")
-
-        # Universal duplicate check for once semantics
-        if request.request_id in self.transfer_triggered_requests:
-            return False
-
-        if criteria_type == "prefill_finished":
-            if request.num_computed_tokens >= request.num_prompt_tokens:
-                logger.debug(f"[Omni] Request {request.request_id} triggered prefill_finished transfer (Non-Stop)")
-                self.transfer_triggered_requests.add(request.request_id)
-                self._mark_request_for_kv_transfer(request.request_id, request.num_computed_tokens)
-
-                # Return False means "Do NOT stop the request" -> Continue Decoding
-                return False
-
-        elif criteria_type == "special_token":
-            target_token_id = self.kv_transfer_criteria.get("token_id")
-            if target_token_id is not None and target_token_id in new_token_ids:
-                logger.debug(f"[Omni] Request {request.request_id} triggered special_token criteria (Non-Stop)")
-
-                self.transfer_triggered_requests.add(request.request_id)
-
-                # Calculate precise snapshot length (trim to sentinel)
-                # Find the FIRST occurrence of the sentinel
-                try:
-                    idx = new_token_ids.index(target_token_id)
-                    # seq_len = tokens_before_this_step + idx + 1 (include sentinel)
-                    # request.num_computed_tokens already includes ALL new_token_ids
-                    # so we subtract (len(new_token_ids) - (idx + 1))
-                    tokens_to_exclude = len(new_token_ids) - (idx + 1)
-                    snapshot_len = request.num_computed_tokens - tokens_to_exclude
-                except ValueError:
-                    snapshot_len = request.num_computed_tokens
-
-                # Trigger Transfer
-                self._mark_request_for_kv_transfer(request.request_id, snapshot_len)
-
-                # Do NOT stop request
-                return False
-
-        return False
-
+    # Ensure scheduled_new_reqs carry omni-specific payloads
+    # (e.g., additional_information)
     def schedule(self) -> SchedulerOutput:  # type: ignore[override]
         scheduler_output = super().schedule()
         try:
@@ -156,13 +55,9 @@ class OmniARScheduler(VLLMScheduler):
                 new_list.append(omni_nr)
 
             scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-
-            # Add information about requests needing KV cache transfer
-            scheduler_output.finished_requests_needing_kv_transfer = self.get_finished_requests_needing_kv_transfer()
         except Exception:
             # If anything goes wrong, leave the original output unchanged
             init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
-            scheduler_output.finished_requests_needing_kv_transfer = {}
 
         return scheduler_output
 
@@ -178,11 +73,6 @@ class OmniARScheduler(VLLMScheduler):
         pooler_outputs = model_runner_output.pooler_output
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
-        cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
-
-        perf_stats: PerfStats | None = None
-        if self.perf_metrics and self.perf_metrics.is_enabled():
-            perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
@@ -241,14 +131,11 @@ class OmniARScheduler(VLLMScheduler):
                     spec_decoding_stats,
                     num_draft_tokens=num_draft_tokens,
                     num_accepted_tokens=num_accepted,
-                    num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
-                    request_id=req_id,
                 )
 
             stopped = False
             new_logprobs = None
             new_token_ids = generated_token_ids
-            pooler_output = pooler_outputs[req_index] if pooler_outputs else None
             kv_transfer_params = None
             status_before_stop = request.status
 
@@ -256,41 +143,14 @@ class OmniARScheduler(VLLMScheduler):
             if new_token_ids:
                 new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
 
-            # If criteria returns True, it means we must STOP the request (e.g. legacy behavior, not used now)
-            # If criteria returns False, it might have triggered a background transfer
-            # (e.g. prefill finished / special token) but continues decoding
-            if not stopped and self._process_kv_transfer_trigger(request, new_token_ids):
-                stopped = True
-
             # Stop checking for pooler models.
             pooler_output = None
             if pooler_outputs:
                 pooler_output = pooler_outputs[req_index]
                 if request.output_token_ids:
-                    stopped = check_stop(request, self.max_model_len)
-            routed_experts = None
+                    stopped = check_stop(request, self.max_model_len, pooler_output)
+
             if stopped:
-                # [Omni] Handle routed experts if enabled
-                if self.vllm_config.model_config.enable_return_routed_experts:
-                    kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
-                    block_ids = kv_blocks.get_block_ids()[0]
-                    num_tokens = request.num_tokens - 1
-
-                    # compute slot mapping
-                    block_ids_array = np.array(block_ids, dtype=np.int32)
-                    num_blocks = len(block_ids)
-                    block_size = self.block_size
-
-                    # generate block offsets
-                    block_offsets = np.arange(0, block_size)
-
-                    # compute slot mapping: slot = block_id * block_size + offset
-                    slot_mapping = (
-                        block_offsets.reshape((1, block_size)) + block_ids_array.reshape((num_blocks, 1)) * block_size
-                    ).flatten()[:num_tokens]
-
-                    routed_experts = self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
-
                 kv_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
@@ -305,13 +165,7 @@ class OmniARScheduler(VLLMScheduler):
                 struct_output_request = request.structured_output_request
                 assert struct_output_request is not None
                 assert struct_output_request.grammar is not None
-                ok = struct_output_request.grammar.accept_tokens(req_id, new_token_ids)
-                if not ok:
-                    logger.warning(
-                        "Unexpected: grammar rejected tokens %s for request %s.",
-                        new_token_ids,
-                        req_id,
-                    )
+                struct_output_request.grammar.accept_tokens(req_id, new_token_ids)
 
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
                 request.num_nans_in_logits = num_nans_in_logits[req_id]
@@ -333,7 +187,6 @@ class OmniARScheduler(VLLMScheduler):
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
                         num_cached_tokens=request.num_cached_tokens,
-                        routed_experts=routed_experts,
                         num_nans_in_logits=request.num_nans_in_logits,
                     )
                 )
@@ -348,37 +201,6 @@ class OmniARScheduler(VLLMScheduler):
             # This is a rare case and unlikely to impact performance.
             self.waiting.remove_requests(stopped_preempted_reqs)
 
-        # [Main] Handle failed KV load requests
-        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
-            requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
-            self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
-            for request in requests:
-                outputs[request.client_index].append(
-                    EngineCoreOutput(
-                        request_id=request.request_id,
-                        new_token_ids=[],
-                        finish_reason=request.get_finished_reason(),
-                        events=request.take_events(),
-                        trace_headers=request.trace_headers,
-                        num_cached_tokens=request.num_cached_tokens,
-                    )
-                )
-
-        # [Omni] Cleanup state for finished requests
-        for req in stopped_running_reqs:
-            if req.request_id not in self.waiting_for_transfer_free:
-                if req.request_id in self.transfer_triggered_requests:
-                    self.transfer_triggered_requests.remove(req.request_id)
-                if req.request_id in self.active_kv_transfers:
-                    self.active_kv_transfers.remove(req.request_id)
-
-        # Same for preempted
-        for req in stopped_preempted_reqs:
-            if req.request_id not in self.waiting_for_transfer_free:
-                if req.request_id in self.transfer_triggered_requests:
-                    self.transfer_triggered_requests.remove(req.request_id)
-                if req.request_id in self.active_kv_transfers:
-                    self.active_kv_transfers.remove(req.request_id)
         # KV Connector: update state for finished KV Transfers.
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
@@ -416,7 +238,7 @@ class OmniARScheduler(VLLMScheduler):
                     engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
             finished_req_ids.clear()
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None:
+        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
             # Return stats to only one of the front-ends.
             if (eco := next(iter(engine_core_outputs.values()), None)) is None:
                 # We must return the stats even if there are no request
@@ -424,200 +246,4 @@ class OmniARScheduler(VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
-        # This is where we free blocks that were held for transfer
-        try:
-            kv_extracted_ids = getattr(model_runner_output, "kv_extracted_req_ids", None)
-            if kv_extracted_ids:
-                for req_id in kv_extracted_ids:
-                    # Mark transfer as finished
-                    if req_id in self.active_kv_transfers:
-                        self.active_kv_transfers.remove(req_id)
-                        logger.debug(f"[Omni] KV Transfer finished for {req_id}")
-
-                    if req_id in self.waiting_for_transfer_free:
-                        # Now it's safe to free blocks
-                        req = self.requests.get(req_id)
-                        if req:
-                            self.kv_cache_manager.free(req)
-                            if req_id in self.requests:
-                                del self.requests[req_id]
-                            if req_id in self.transfer_triggered_requests:
-                                self.transfer_triggered_requests.remove(req_id)
-                            if req_id in self.active_kv_transfers:
-                                self.active_kv_transfers.remove(req_id)
-
-                            logger.debug(f"Freed blocks for {req_id} after transfer extraction")
-                        self.waiting_for_transfer_free.remove(req_id)
-        except Exception:
-            init_logger(__name__).exception("Failed to process finished transfer requests")
-
         return engine_core_outputs
-
-    def _free_request(self, request: Request) -> dict[str, Any] | None:
-        # TODO(wzliu)! for offline mode, we should not end process until all data is transferred
-        """Mark a request as finished and free its resources."""
-
-        # 1. Standard cleanup parts from base _free_request
-        delay_free_blocks = False
-        kv_xfer_params = None
-        if self.connector is not None:
-            delay_free_blocks, kv_xfer_params = self._connector_finished(request)
-
-        self.encoder_cache_manager.free(request)
-        request_id = request.request_id
-        self.finished_req_ids.add(request_id)
-        if self.finished_req_ids_dict is not None:
-            self.finished_req_ids_dict[request.client_index].add(request_id)
-
-        # 2. Omni Specific: Check if we need to transfer KV
-        if self._should_transfer_kv_for_request(request_id):
-            already_triggered = request_id in self.transfer_triggered_requests
-            is_active = request_id in self.active_kv_transfers
-
-            if already_triggered:
-                if is_active:
-                    # It triggered but hasn't finished yet. We MUST wait.
-                    logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
-                    self.waiting_for_transfer_free.add(request_id)
-                    # We do NOT mark for transfer again, just wait.
-                    kv_xfer_params = None  # No new transfer params
-                    return kv_xfer_params
-                else:
-                    logger.debug(
-                        f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "
-                        "Freeing immediately."
-                    )
-            else:
-                self.waiting_for_transfer_free.add(request_id)
-                self._mark_request_for_kv_transfer(request_id, request.num_computed_tokens)
-                # Return KV transfer metadata so it propagates to RequestOutput
-                if request_id in self.requests_needing_kv_transfer:
-                    transfer_data = self.requests_needing_kv_transfer[request_id]
-                    kv_xfer_params = {
-                        "past_key_values": transfer_data["block_ids"],
-                        "kv_metadata": {"seq_len": transfer_data["seq_len"], "block_ids": transfer_data["block_ids"]},
-                    }
-                    # Also update request.additional_information for good measure
-                    add_info = getattr(request, "additional_information", None)
-                    # If additional_information is an AdditionalInformationPayload-like object,
-                    # unpack list_data into a plain dict.
-                    if (
-                        add_info is not None
-                        and hasattr(add_info, "entries")
-                        and isinstance(getattr(add_info, "entries"), dict)
-                    ):
-                        request.additional_information = {
-                            k: getattr(v, "list_data")
-                            for k, v in getattr(add_info, "entries").items()
-                            if getattr(v, "list_data", None) is not None
-                        }
-                        add_info = request.additional_information
-                    if add_info is None:
-                        request.additional_information = {}
-                        add_info = request.additional_information
-                    if isinstance(add_info, dict):
-                        add_info.update(kv_xfer_params)
-
-                return kv_xfer_params
-
-        # 3. Standard Freeing
-        if not delay_free_blocks:
-            self._free_blocks(request)
-
-        return kv_xfer_params
-
-    def _free_blocks(self, request: Request):
-        # Helper to match base class structure if not directly available
-        # VLLMScheduler has _free_blocks
-        super()._free_blocks(request)
-
-    def _mark_request_for_kv_transfer(self, req_id: str, seq_len: int) -> None:
-        """Mark a request as needing KV cache transfer when it finishes."""
-        # Avoid duplicate marking (if already pending in queue)
-        if req_id in self.requests_needing_kv_transfer:
-            return
-
-        if self._should_transfer_kv_for_request(req_id):
-            # [Omni] Get block IDs from KVCacheManager
-            try:
-                block_ids_tuple = self.kv_cache_manager.get_block_ids(req_id)
-                if block_ids_tuple and len(block_ids_tuple) > 0:
-                    block_ids = block_ids_tuple[0]
-
-                    # [Omni] Fix: Truncate blocks to match seq_len snapshot
-                    # We need to know block_size. Usually in self.cache_config.block_size
-                    # Note: vllm_config might not be directly available, check scheduler_config or cache_config
-                    if hasattr(self, "cache_config") and hasattr(self.cache_config, "block_size"):
-                        block_size = self.cache_config.block_size
-                    elif hasattr(self, "scheduler_config") and hasattr(
-                        self.scheduler_config, "block_size"
-                    ):  # Some versions
-                        block_size = self.scheduler_config.block_size
-                    else:
-                        raise ValueError("Block size not found in cache_config or scheduler_config")
-
-                    # ceil(seq_len / block_size)
-                    num_blocks = (seq_len + block_size - 1) // block_size
-                    if len(block_ids) > num_blocks:
-                        logger.debug(
-                            f"[Omni] Truncating blocks for {req_id} from {len(block_ids)} "
-                            f"to {num_blocks} (seq_len={seq_len})"
-                        )
-                        block_ids = block_ids[:num_blocks]
-
-                else:
-                    block_ids = []
-            except Exception as e:
-                init_logger(__name__).warning(f"Failed to get block IDs for {req_id}: {e}")
-                block_ids = []
-
-            self.requests_needing_kv_transfer[req_id] = {"seq_len": seq_len, "block_ids": block_ids}
-            logger.debug(f"Marked request {req_id} for KV cache transfer (len={seq_len}, blocks={len(block_ids)})")
-
-    def _should_transfer_kv_for_request(self, req_id: str) -> bool:
-        """Determine if a request should trigger KV cache transfer."""
-        need_send = False
-        # Try to read from vLLM Config (where YAML config is typically loaded)
-        # Check for omni_kv_config attribute
-        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
-        if omni_kv_config:
-            # omni_kv_config could be an object or a dict
-            if isinstance(omni_kv_config, dict):
-                need_send = omni_kv_config.get("need_send_cache", False)
-            else:
-                need_send = getattr(omni_kv_config, "need_send_cache", False)
-        return need_send
-
-    def has_requests(self) -> bool:
-        """Check if there are any requests to process, including KV transfers."""
-        # [Omni] Also check for pending KV transfers
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_requests()
-
-    def has_finished_requests(self) -> bool:
-        """Check if there are any finished requests (including those needing KV transfer)."""
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_finished_requests()
-
-    def has_unfinished_requests(self) -> bool:
-        """Check if there are any unfinished requests (including those needing KV transfer)."""
-        # [Omni] Also check for pending KV transfers to ensure the engine loop continues
-        # MUST verify waiting_for_transfer_free and active_kv_transfers
-        # Otherwise engine loop might exit before transfer Ack is received.
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_unfinished_requests()
-
-    def get_finished_requests_needing_kv_transfer(self) -> dict[str, dict]:
-        """Get and clear the list of requests needing KV cache transfer.
-        Returns dict: {req_id: {"seq_len": int, "block_ids": list[int]}}
-        """
-        requests = self.requests_needing_kv_transfer.copy()
-
-        # Mark these requests as ACTIVE (sent to runner)
-        self.active_kv_transfers.update(requests.keys())
-
-        self.requests_needing_kv_transfer.clear()
-        return requests

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -2,7 +2,6 @@ import time
 from collections import defaultdict
 
 from vllm.distributed.kv_events import KVEventBatch
-from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.request_queue import create_request_queue
@@ -91,7 +90,7 @@ class OmniGenerationScheduler(VLLMScheduler):
             any_request = self.running[0]
             num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request.request_id)
 
-        # Assemble SchedulerOutput (align with v0.14.0)
+        # Assemble SchedulerOutput (align with v0.12.0)
         if self.use_v2_model_runner:
             # No resumed reqs in fast path; pass prefill_token_ids for new reqs.
             new_reqs_data = [
@@ -130,7 +129,7 @@ class OmniGenerationScheduler(VLLMScheduler):
             preempted_req_ids=set(),
         )
 
-        # Record the request ids scheduled in this step (v0.14.0 behavior).
+        # Record the request ids scheduled in this step (v0.12.0 behavior).
         self.prev_step_scheduled_req_ids.clear()
         self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
 
@@ -146,35 +145,6 @@ class OmniGenerationScheduler(VLLMScheduler):
         # Update internal state (advance num_computed_tokens, free encoder inputs,
         # etc.)
         self._update_after_schedule(scheduler_output)
-
-        try:
-            # Rewrap base NewRequestData entries with OmniNewRequestData,
-            # enriching with request-level payloads
-            new_list = []
-            for nr in scheduler_output.scheduled_new_reqs:
-                req_id = getattr(nr, "req_id", None)
-                request = self.requests.get(req_id) if req_id else None
-                # Build omni entry preserving all base fields
-                omni_nr = OmniNewRequestData(
-                    req_id=nr.req_id,
-                    prompt_token_ids=nr.prompt_token_ids,
-                    mm_features=nr.mm_features,
-                    sampling_params=nr.sampling_params,
-                    pooling_params=nr.pooling_params,
-                    block_ids=nr.block_ids,
-                    num_computed_tokens=nr.num_computed_tokens,
-                    lora_request=nr.lora_request,
-                    # Enrich with omni payloads from the live request object
-                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
-                    additional_information=(getattr(request, "additional_information", None) if request else None),
-                )
-                new_list.append(omni_nr)
-
-            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-        except Exception:
-            # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
-
         return scheduler_output
 
     """
@@ -206,7 +176,7 @@ class OmniGenerationScheduler(VLLMScheduler):
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
         kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
-        # Merge connector-side stats (align with v0.14.0)
+        # Merge connector-side stats (align with v0.12.0)
         if kv_connector_stats and self.connector:
             kv_stats = self.connector.get_kv_connector_stats()
             if kv_stats:
@@ -324,7 +294,7 @@ class OmniGenerationScheduler(VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Collect and publish KV cache events (align with v0.14.0)
+        # Collect and publish KV cache events (align with v0.12.0)
         events = self.kv_cache_manager.take_events()
         if self.connector is not None:
             connector_events = self.connector.take_events()

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import torch
-from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
@@ -14,6 +13,18 @@ from vllm_omni.diffusion.attention.backends.utils.fa import _pad_input, _unpad_i
 
 logger = init_logger(__name__)
 
+try:
+    # only tested with flash_attn v3
+    # from flash_attn_interface import flash_attn_func as flash_attn_3_func  # not available in flash-attn 2.8.1
+    from flash_attn import flash_attn_func, flash_attn_varlen_func  # can be FA2 or FA3
+except ImportError:
+    logger.warning(
+        "FlashAttentionBackend is not available. You may install flash-attn "
+        "by running `uv pip install flash-attn==2.8.1 --no-build-isolation`"
+        " or install pre-built flash-attn from https://github.com/Dao-AILab/flash-attention/releases"
+    )
+    raise ImportError
+
 
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True

vllm_omni/diffusion/attention/selector.py

@@ -5,14 +5,12 @@ import importlib
 import os
 from functools import cache
 
-import torch
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
 )
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
-from vllm_omni.utils.platform_utils import detect_device_type, is_rocm
 
 logger = init_logger(__name__)
 
@@ -65,22 +63,7 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
         The selected attention backend class
     """
     # Check environment variable
-
-    backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND", None)
-
-    if detect_device_type() == "cuda" and not is_rocm():
-        compute_capability = torch.cuda.get_device_capability()
-        major, minor = compute_capability
-        if 80 <= major * 10 + minor < 100:
-            if backend_name is None:
-                backend_name = "FLASH_ATTN"
-        else:
-            if backend_name == "FLASH_ATTN":
-                logger.warning(
-                    """Flash Attention requires GPU with compute capability >= 8.0 or < 10.0. "
-                               "Falling back to TORCH_SDPA backend."""
-                )
-                backend_name = "TORCH_SDPA"
+    backend_name: str | None = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
 
     if backend_name is not None:
         backend_name_upper = backend_name.upper()

vllm_omni/diffusion/cache/teacache/config.py

@@ -25,15 +25,6 @@ _MODEL_COEFFICIENTS = {
         3.20000000e00,
         -2.00000000e-02,
     ],
-    # Z-Image transformer coefficients
-    # Copied from Qwen-Image, need to be tuned specifically for Z-Image in future
-    "ZImageTransformer2DModel": [
-        -4.50000000e02,
-        2.80000000e02,
-        -4.50000000e01,
-        3.20000000e00,
-        -2.00000000e-02,
-    ],
 }
 
 

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -254,192 +254,6 @@ def extract_qwen_context(
     )
 
 
-def extract_zimage_context(
-    module: nn.Module,
-    x: list[torch.Tensor],
-    t: torch.Tensor,
-    cap_feats: list[torch.Tensor],
-    patch_size: int = 2,
-    f_patch_size: int = 1,
-    **kwargs: Any,
-) -> CacheContext:
-    """
-    Extract cache context for ZImageTransformer2DModel.
-
-    This is the ONLY Z-Image-specific code needed for TeaCache support.
-    It encapsulates preprocessing, modulated input extraction, transformer execution,
-    and postprocessing logic.
-
-    Args:
-        module: ZImageTransformer2DModel instance
-        x: List of image tensors per batch item
-        t: Timestep tensor
-        cap_feats: List of caption feature tensors per batch item
-        patch_size: Patch size for patchification (default: 2)
-        f_patch_size: Frame patch size (default: 1)
-        **kwargs: Additional keyword arguments ignored by this extractor
-
-    Returns:
-        CacheContext with all information needed for generic caching
-    """
-    from torch.nn.utils.rnn import pad_sequence
-
-    if not hasattr(module, "layers") or len(module.layers) == 0:
-        raise ValueError("Module must have main transformer layers")
-
-    bsz = len(x)
-    device = x[0].device
-
-    # ============================================================================
-    # PREPROCESSING (Z-Image specific)
-    # ============================================================================
-    # Scale timestep and create timestep embedding
-    t_scaled = t * module.t_scale
-    adaln_input = module.t_embedder(t_scaled)
-
-    # Patchify and embed inputs
-    (
-        x_patches,
-        cap_feats_processed,
-        x_size,
-        x_pos_ids,
-        cap_pos_ids,
-        x_inner_pad_mask,
-        cap_inner_pad_mask,
-    ) = module.patchify_and_embed(x, cap_feats, patch_size, f_patch_size)
-
-    # Process image patches through embedder and noise refiner
-    x_item_seqlens = [len(_) for _ in x_patches]
-    x_max_item_seqlen = max(x_item_seqlens)
-
-    x_embedded = torch.cat(x_patches, dim=0)
-    x_embedded = module.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_embedded)
-
-    # Match adaln_input dtype to x_embedded
-    adaln_input = adaln_input.type_as(x_embedded)
-
-    # Apply pad token
-    x_embedded[torch.cat(x_inner_pad_mask)] = module.x_pad_token
-    x_list = list(x_embedded.split(x_item_seqlens, dim=0))
-
-    # Compute rope embeddings for image patches
-    x_cos, x_sin = module.rope_embedder(torch.cat(x_pos_ids, dim=0))
-    x_cos = list(x_cos.split(x_item_seqlens, dim=0))
-    x_sin = list(x_sin.split(x_item_seqlens, dim=0))
-
-    # Pad sequences for batch processing
-    x_batched = pad_sequence(x_list, batch_first=True, padding_value=0.0)
-    x_cos_batched = pad_sequence(x_cos, batch_first=True, padding_value=0.0)
-    x_sin_batched = pad_sequence(x_sin, batch_first=True, padding_value=0.0)
-    x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device)
-    for i, seq_len in enumerate(x_item_seqlens):
-        x_attn_mask[i, :seq_len] = 1
-
-    # Run noise refiner blocks
-    for layer in module.noise_refiner:
-        x_batched = layer(x_batched, x_attn_mask, x_cos_batched, x_sin_batched, adaln_input)
-
-    # Process caption features through embedder and context refiner
-    cap_item_seqlens = [len(_) for _ in cap_feats_processed]
-    cap_max_item_seqlen = max(cap_item_seqlens)
-
-    cap_embedded = torch.cat(cap_feats_processed, dim=0)
-    cap_embedded = module.cap_embedder(cap_embedded)
-    cap_embedded[torch.cat(cap_inner_pad_mask)] = module.cap_pad_token
-    cap_list = list(cap_embedded.split(cap_item_seqlens, dim=0))
-
-    # Compute rope embeddings for caption
-    cap_cos, cap_sin = module.rope_embedder(torch.cat(cap_pos_ids, dim=0))
-    cap_cos = list(cap_cos.split(cap_item_seqlens, dim=0))
-    cap_sin = list(cap_sin.split(cap_item_seqlens, dim=0))
-
-    # Pad sequences for batch processing
-    cap_batched = pad_sequence(cap_list, batch_first=True, padding_value=0.0)
-    cap_cos_batched = pad_sequence(cap_cos, batch_first=True, padding_value=0.0)
-    cap_sin_batched = pad_sequence(cap_sin, batch_first=True, padding_value=0.0)
-    cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device)
-    for i, seq_len in enumerate(cap_item_seqlens):
-        cap_attn_mask[i, :seq_len] = 1
-
-    # Run context refiner blocks
-    for layer in module.context_refiner:
-        cap_batched = layer(cap_batched, cap_attn_mask, cap_cos_batched, cap_sin_batched)
-
-    # Create unified sequence (image + caption)
-    unified_list = []
-    unified_cos_list = []
-    unified_sin_list = []
-    for i in range(bsz):
-        x_len = x_item_seqlens[i]
-        cap_len = cap_item_seqlens[i]
-        unified_list.append(torch.cat([x_batched[i][:x_len], cap_batched[i][:cap_len]]))
-        unified_cos_list.append(torch.cat([x_cos_batched[i][:x_len], cap_cos_batched[i][:cap_len]]))
-        unified_sin_list.append(torch.cat([x_sin_batched[i][:x_len], cap_sin_batched[i][:cap_len]]))
-
-    unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)]
-    unified_max_item_seqlen = max(unified_item_seqlens)
-
-    unified = pad_sequence(unified_list, batch_first=True, padding_value=0.0)
-    unified_cos = pad_sequence(unified_cos_list, batch_first=True, padding_value=0.0)
-    unified_sin = pad_sequence(unified_sin_list, batch_first=True, padding_value=0.0)
-    unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device)
-    for i, seq_len in enumerate(unified_item_seqlens):
-        unified_attn_mask[i, :seq_len] = 1
-
-    # ============================================================================
-    # EXTRACT MODULATED INPUT (for cache decision)
-    # ============================================================================
-    # Use the first main transformer block's modulation
-    # The main layers have modulation=True and process the unified sequence
-    block = module.layers[0]
-    # Get modulation parameters: scale_msa, gate_msa, scale_mlp, gate_mlp
-    mod_params = block.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2)
-    scale_msa = 1.0 + mod_params[0]
-    # Extract modulated input: normalized hidden states scaled by modulation
-    modulated_input = block.attention_norm1(unified) * scale_msa
-
-    # ============================================================================
-    # DEFINE TRANSFORMER EXECUTION (Z-Image specific)
-    # ============================================================================
-    def run_transformer_blocks():
-        """Execute all Z-Image main transformer blocks."""
-        h = unified
-        for layer in module.layers:
-            h = layer(h, unified_attn_mask, unified_cos, unified_sin, adaln_input)
-        return (h,)
-
-    # ============================================================================
-    # DEFINE POSTPROCESSING (Z-Image specific)
-    # ============================================================================
-    def postprocess(h):
-        """Apply Z-Image specific output postprocessing."""
-        h = module.all_final_layer[f"{patch_size}-{f_patch_size}"](h, adaln_input)
-        h = list(h.unbind(dim=0))
-        output = module.unpatchify(h, x_size, patch_size, f_patch_size)
-        return output, {}
-
-    # ============================================================================
-    # RETURN CONTEXT
-    # ============================================================================
-    return CacheContext(
-        modulated_input=modulated_input,
-        hidden_states=unified,
-        encoder_hidden_states=None,  # Z-Image uses unified sequence, no separate encoder states
-        temb=adaln_input,
-        run_transformer_blocks=run_transformer_blocks,
-        postprocess=postprocess,
-        extra_states={
-            "unified_attn_mask": unified_attn_mask,
-            "unified_cos": unified_cos,
-            "unified_sin": unified_sin,
-            "x_size": x_size,
-            "x_item_seqlens": x_item_seqlens,
-            "patch_size": patch_size,
-            "f_patch_size": f_patch_size,
-        },
-    )
-
-
 # Registry for model-specific extractors
 # Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -448,7 +262,6 @@ def extract_zimage_context(
 # on the transformer module and multiple pipelines can share the same transformer.
 EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "QwenImageTransformer2DModel": extract_qwen_context,
-    "ZImageTransformer2DModel": extract_zimage_context,
     # Future models:
     # "FluxTransformer2DModel": extract_flux_context,
     # "CogVideoXTransformer3DModel": extract_cogvideox_context,

vllm_omni/diffusion/data.py

@@ -357,9 +357,6 @@ class OmniDiffusionConfig:
     # Logging
     log_level: str = "info"
 
-    # Omni configuration (injected from stage config)
-    omni_kv_config: dict[str, Any] = field(default_factory=dict)
-
     def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
         """
         Find an available port with retry logic.
@@ -456,12 +453,7 @@ class OmniDiffusionConfig:
         if "cache_backend" not in kwargs:
             cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
             kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
-
-        # Filter kwargs to only include valid fields
-        valid_fields = {f.name for f in fields(cls)}
-        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
-
-        return cls(**filtered_kwargs)
+        return cls(**kwargs)
 
 
 @dataclass

vllm_omni/diffusion/diffusion_engine.py

@@ -1,23 +1,26 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
+import multiprocessing as mp
 import time
-from collections.abc import Iterable
+import weakref
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
 from typing import Any
 
 import PIL.Image
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, OmniDiffusionConfig
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.scheduler import Scheduler, scheduler
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.platform_utils import get_diffusion_worker_class
 
 logger = init_logger(__name__)
 
@@ -29,6 +32,39 @@ def supports_image_input(model_class_name: str) -> bool:
     return bool(getattr(model_cls, "support_image_input", False))
 
 
+@dataclass
+class BackgroundResources:
+    """
+    Used as a finalizer for clean shutdown.
+    Create a BackgroundResources instance to encapsulate all background resources
+    (e.g., the scheduler and worker processes) that need explicit cleanup.
+    This object holds references to external system resources that are not managed
+    by Python's garbage collector (like OS processes, message queues, etc.),
+    so they must be cleaned up manually to avoid resource leaks or zombie processes.
+    """
+
+    scheduler: Scheduler | None = None
+    processes: list[mp.Process] | None = None
+
+    def __call__(self):
+        """Clean up background resources."""
+        if scheduler is not None:
+            try:
+                for _ in range(scheduler.num_workers):
+                    scheduler.mq.enqueue(SHUTDOWN_MESSAGE)
+                scheduler.close()
+            except Exception as exc:
+                logger.warning("Failed to send shutdown signal: %s", exc)
+        for proc in self.processes:
+            if not proc.is_alive():
+                continue
+            proc.join(30)
+            if proc.is_alive():
+                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
+                proc.terminate()
+                proc.join(30)
+
+
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
@@ -43,9 +79,9 @@ class DiffusionEngine:
         self.post_process_func = get_diffusion_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
 
-        executor_class = DiffusionExecutor.get_class(od_config)
-        self.executor = executor_class(od_config)
-
+        self._processes: list[mp.Process] = []
+        self._closed = False
+        self._make_client()
         try:
             self._dummy_run()
         except Exception as e:
@@ -163,127 +199,96 @@ class DiffusionEngine:
         """
         return DiffusionEngine(config)
 
-    def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
-        return self.executor.add_req(requests)
+    def _make_client(self):
+        # TODO rename it
+        scheduler.initialize(self.od_config)
 
-    def start_profile(self, trace_filename: str | None = None) -> None:
-        """
-        Start torch profiling on all diffusion workers.
+        # Get the broadcast handle from the initialized scheduler
+        broadcast_handle = scheduler.get_broadcast_handle()
 
-        Creates a directory (if needed) and sets up a base filename template
-        for per-rank profiler traces (typically saved as <template>_rank<N>.json).
+        processes, result_handle = self._launch_workers(
+            broadcast_handle=broadcast_handle,
+        )
 
-        Args:
-            trace_filename: Optional base filename (without extension or rank suffix).
-                            If None, generates one using current timestamp.
-        """
-        if trace_filename is None:
-            trace_filename = f"stage_0_diffusion_{int(time.time())}_rank"
-
-        trace_dir = os.environ.get("VLLM_TORCH_PROFILER_DIR", "./profiles")
-
-        # Expand ~ and ~user, then make absolute (robust against cwd changes)
-        trace_dir = os.path.expanduser(trace_dir)
-        trace_dir = os.path.abspath(trace_dir)
-
-        try:
-            os.makedirs(trace_dir, exist_ok=True)
-        except OSError as exc:
-            logger.error(f"Failed to create profiler directory {trace_dir}: {exc}")
-            raise
-
-        # Build final template path (without rank or extension — torch.profiler appends those)
-        full_template = os.path.join(trace_dir, trace_filename)
-
-        expected_pattern = f"{full_template}*.json"
-        logger.info(f"Starting diffusion profiling → {expected_pattern}")
-
-        # Also log the absolute directory once (useful in multi-node or containers)
-        logger.debug(f"Profiler output directory: {trace_dir}")
-
-        # Propagate to all workers
-        try:
-            self.collective_rpc(method="start_profile", args=(full_template,))
-        except Exception as e:
-            logger.error("Failed to start profiling on workers", exc_info=True)
-            raise RuntimeError(f"Could not start profiler: {e}") from e
-
-    def stop_profile(self) -> dict:
-        """
-        Stop profiling on all workers and collect the final trace/table paths.
-
-        The worker (torch_profiler.py) now handles trace export, compression to .gz,
-        and deletion of the original .json file. This method only collects and
-        reports the paths returned by the workers.
-
-        Returns:
-            dict with keys:
-            - "traces": list of final trace file paths (usually .json.gz)
-            - "tables": list of table strings (one per rank)
-        """
-        logger.info("Stopping diffusion profiling and collecting results...")
-
-        try:
-            # Give worker enough time — export + compression + table can be slow
-            results = self.collective_rpc(method="stop_profile", timeout=60000)
-        except Exception:
-            logger.error("Failed to stop profiling on workers", exc_info=True)
-            return {"traces": [], "tables": []}
-
-        output_files = {"traces": [], "tables": []}
-        successful_traces = 0
-
-        if not results:
-            logger.warning("No profiling results returned from any rank")
-            return output_files
-
-        for rank, res in enumerate(results):
-            if not isinstance(res, dict):
-                logger.warning(f"Rank {rank}: invalid result format (got {type(res)})")
-                continue
-
-            # 1. Trace file — should be .json.gz if compression succeeded
-            trace_path = res.get("trace")
-            if trace_path:
-                # We trust the worker — it created/compressed the file
-                logger.info(f"[Rank {rank}] Final trace: {trace_path}")
-                output_files["traces"].append(trace_path)
-                successful_traces += 1
-
-                # Optional: warn if path looks suspicious (e.g. still .json)
-                if not trace_path.endswith((".json.gz", ".json")):
-                    logger.warning(f"Rank {rank}: unusual trace path extension: {trace_path}")
-
-            # 2. Table file — plain text
-            table = res.get("table")
-            if table:
-                output_files["tables"].append(table)
-
-        # Final summary logging
-        num_ranks = len(results)
-        if successful_traces > 0:
-            final_paths_str = ", ".join(output_files["traces"][:3])
-            if len(output_files["traces"]) > 3:
-                final_paths_str += f" ... (+{len(output_files['traces']) - 3} more)"
-
-            logger.info(
-                f"Profiling stopped. Collected {successful_traces} trace file(s) "
-                f"from {num_ranks} rank(s). "
-                f"Final trace paths: {final_paths_str}"
-            )
-        elif output_files["traces"]:
-            logger.info(
-                f"Profiling stopped but no traces were successfully collected. "
-                f"Reported paths: {', '.join(output_files['traces'][:3])}"
-                f"{' ...' if len(output_files['traces']) > 3 else ''}"
-            )
+        if result_handle is not None:
+            scheduler.initialize_result_queue(result_handle)
         else:
-            logger.info("Profiling stopped — no trace files were collected from any rank.")
+            logger.error("Failed to get result queue handle from workers")
 
-        if output_files["tables"]:
-            logger.debug(f"Collected {len(output_files['tables'])} profiling table(s)")
+        self._processes = processes
 
-        return output_files
+        self.resources = BackgroundResources(scheduler=scheduler, processes=self._processes)
+        # Use weakref.finalize instead of __del__ or relying on self.close() at shutdown.
+        # During interpreter shutdown, global state (e.g., modules, built-ins) may already
+        # be cleared (set to None), so calling normal cleanup methods can fail with
+        # AttributeError: 'NoneType' object has no attribute '...'.
+        # weakref.finalize schedules cleanup *before* such destruction begins,
+        # ensuring resources are released while the runtime environment is still intact.
+        self._finalizer = weakref.finalize(self, self.resources)
+
+    def _launch_workers(self, broadcast_handle):
+        od_config = self.od_config
+        logger.info("Starting server...")
+
+        num_gpus = od_config.num_gpus
+        mp.set_start_method("spawn", force=True)
+        processes = []
+
+        # Get the appropriate worker class for current device
+        worker_proc = get_diffusion_worker_class()
+
+        # Launch all worker processes
+        scheduler_pipe_readers = []
+        scheduler_pipe_writers = []
+
+        for i in range(num_gpus):
+            reader, writer = mp.Pipe(duplex=False)
+            scheduler_pipe_writers.append(writer)
+            process = mp.Process(
+                target=worker_proc.worker_main,
+                args=(
+                    i,  # rank
+                    od_config,
+                    writer,
+                    broadcast_handle,
+                ),
+                name=f"DiffusionWorker-{i}",
+                daemon=True,
+            )
+            scheduler_pipe_readers.append(reader)
+            process.start()
+            processes.append(process)
+
+        # Wait for all workers to be ready
+        scheduler_infos = []
+        result_handle = None
+        for writer in scheduler_pipe_writers:
+            writer.close()
+
+        for i, reader in enumerate(scheduler_pipe_readers):
+            try:
+                data = reader.recv()
+            except EOFError:
+                logger.error(f"Rank {i} scheduler is dead. Please check if there are relevant logs.")
+                processes[i].join()
+                logger.error(f"Exit code: {processes[i].exitcode}")
+                raise
+
+            if data["status"] != "ready":
+                raise RuntimeError("Initialization failed. Please see the error messages above.")
+
+            if i == 0:
+                result_handle = data.get("result_handle")
+
+            scheduler_infos.append(data)
+            reader.close()
+
+        logger.debug("All workers are ready")
+
+        return processes, result_handle
+
+    def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
+        return scheduler.add_req(requests)
 
     def _dummy_run(self):
         """A dummy run to warm up the model."""
@@ -312,7 +317,7 @@ class DiffusionEngine:
 
     def collective_rpc(
         self,
-        method: str,
+        method: str | Callable,
         timeout: float | None = None,
         args: tuple = (),
         kwargs: dict | None = None,
@@ -321,7 +326,7 @@ class DiffusionEngine:
         """Call a method on worker processes and get results immediately.
 
         Args:
-            method: The method name (str) to execute on workers
+            method: The method name (str) or callable to execute on workers
             timeout: Optional timeout in seconds
             args: Positional arguments for the method
             kwargs: Keyword arguments for the method
@@ -330,18 +335,59 @@ class DiffusionEngine:
         Returns:
             Single result if unique_reply_rank is provided, otherwise list of results
         """
-        assert isinstance(method, str), "Only string method names are supported for now"
-        return self.executor.collective_rpc(
-            method=method,
-            timeout=timeout,
-            args=args,
-            kwargs=kwargs,
-            unique_reply_rank=unique_reply_rank,
-        )
+        if self._closed:
+            raise RuntimeError("DiffusionEngine is closed.")
+
+        deadline = None if timeout is None else time.monotonic() + timeout
+        kwargs = kwargs or {}
+
+        assert isinstance(method, str)
+        send_method = method
+
+        # Prepare RPC request message
+        rpc_request = {
+            "type": "rpc",
+            "method": send_method,
+            "args": args,
+            "kwargs": kwargs,
+            "output_rank": unique_reply_rank,
+        }
+
+        try:
+            # Broadcast RPC request to all workers via unified message queue
+            scheduler.mq.enqueue(rpc_request)
+
+            # Determine which workers we expect responses from
+            num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus
+
+            responses = []
+            for _ in range(num_responses):
+                dequeue_timeout = None if deadline is None else (deadline - time.monotonic())
+                try:
+                    if scheduler.result_mq is None:
+                        raise RuntimeError("Result queue not initialized")
+
+                    response = scheduler.result_mq.dequeue(timeout=dequeue_timeout)
+
+                    # Check if response indicates an error
+                    if isinstance(response, dict) and response.get("status") == "error":
+                        raise RuntimeError(
+                            f"Worker failed with error '{response.get('error')}', "
+                            "please check the stack trace above for the root cause"
+                        )
+
+                    responses.append(response)
+                except TimeoutError as e:
+                    raise TimeoutError(f"RPC call to {method} timed out.") from e
+
+            return responses[0] if unique_reply_rank is not None else responses
+
+        except Exception as e:
+            logger.error(f"RPC call failed: {e}")
+            raise
 
     def close(self) -> None:
-        if hasattr(self, "executor"):
-            self.executor.shutdown()
+        self._finalizer()
 
     def abort(self, request_id: str | Iterable[str]) -> None:
         # TODO implement it

vllm_omni/diffusion/distributed/__init__.py

@@ -1,38 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Distributed utilities for vLLM-Omni diffusion models."""
-
-from vllm_omni.diffusion.distributed.sp_plan import (
-    SequenceParallelConfig,
-    SequenceParallelInput,
-    SequenceParallelModelPlan,
-    SequenceParallelOutput,
-    SequenceParallelPartialInput,
-    get_sp_plan_from_model,
-    validate_sp_plan,
-)
-from vllm_omni.diffusion.distributed.sp_sharding import (
-    ShardingValidator,
-    get_sharding_validator,
-    sp_gather,
-    sp_shard,
-    sp_shard_with_padding,
-)
-
-__all__ = [
-    # Config
-    "SequenceParallelConfig",
-    # Plan types
-    "SequenceParallelInput",
-    "SequenceParallelOutput",
-    "SequenceParallelPartialInput",
-    "SequenceParallelModelPlan",
-    "validate_sp_plan",
-    "get_sp_plan_from_model",
-    # Sharding utilities
-    "sp_shard",
-    "sp_gather",
-    "sp_shard_with_padding",
-    "ShardingValidator",
-    "get_sharding_validator",
-]

vllm_omni/diffusion/distributed/sp_plan.py

@@ -1,457 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM and The HuggingFace Team
-# Type definitions in this module are adapted from HuggingFace diffusers library:
-#   diffusers/src/diffusers/models/_modeling_parallel.py
-"""Sequence Parallelism configuration and plan type definitions.
-
-This module provides:
-1. SequenceParallelConfig: Configuration for SP (ulysses_degree, ring_degree)
-2. SequenceParallelInput/Output: Type definitions for _sp_plan declarations
-3. Validation utilities for _sp_plan
-
-A _sp_plan is a dictionary that specifies how to shard/gather tensors at
-different points in a model's forward pass. This allows automatic handling
-of sequence parallelism without modifying the model's forward() method.
-
-NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP)
-in diffusers. We use "Sequence Parallelism" to align with vLLM-Omni terminology.
-
-Example:
-    class MyTransformer(nn.Module):
-        _sp_plan = {
-            # Split inputs before model forward
-            "": {
-                "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-                "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-            },
-            # Split RoPE embeddings after pos_embed layer
-            "pos_embed": {
-                0: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
-            },
-            # Gather output after proj_out layer
-            "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-        }
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch
-    import torch.nn as nn
-
-
-# =============================================================================
-# Sequence Parallel Configuration
-# =============================================================================
-
-
-@dataclass
-class SequenceParallelConfig:
-    """Configuration for Sequence Parallelism using vLLM-Omni's parallel state.
-
-    This class provides a unified interface for SP configuration that integrates
-    with vLLM-Omni's existing SequenceParallelGroupCoordinator. Unlike diffusers'
-    DeviceMesh-based approach (ContextParallelConfig), this uses the existing
-    parallel state management.
-
-    Note: This corresponds to `ContextParallelConfig` in diffusers library.
-
-    Args:
-        ulysses_degree: Number of devices for Ulysses (All-to-All) attention.
-            Sequence is split across devices, with Q/K/V redistributed via
-            All-to-All communication. Best for moderate sequences with good
-            interconnect bandwidth.
-        ring_degree: Number of devices for Ring attention. Sequence is split
-            across devices, with K/V passed in a ring topology. Best for long
-            sequences with limited memory/bandwidth.
-        convert_to_fp32: Whether to convert output and LSE to float32 for
-            numerical stability in ring attention.
-
-    Note:
-        ulysses_degree * ring_degree = sequence_parallel_size
-        vLLM-Omni supports hybrid Ulysses-Ring attention (both > 1).
-    """
-
-    ulysses_degree: int = 1
-    ring_degree: int = 1
-    convert_to_fp32: bool = True
-
-    # Internal state - populated by setup()
-    _rank: int | None = None
-    _world_size: int | None = None
-    _device: torch.device | None = None
-
-    def __post_init__(self) -> None:
-        if self.ulysses_degree < 1 or self.ring_degree < 1:
-            raise ValueError("`ulysses_degree` and `ring_degree` must be >= 1.")
-
-        if self.ulysses_degree == 1 and self.ring_degree == 1:
-            raise ValueError(
-                "At least one of `ulysses_degree` or `ring_degree` must be > 1 to use sequence parallelism."
-            )
-
-    @property
-    def sequence_parallel_size(self) -> int:
-        """Total sequence parallel world size."""
-        return self.ulysses_degree * self.ring_degree
-
-    def get_world_size(self) -> int:
-        """Get the sequence parallel world size from parallel state.
-
-        Returns:
-            The world size for sequence parallelism.
-
-        Raises:
-            RuntimeError: If parallel state is not initialized.
-        """
-        from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
-
-        return get_sequence_parallel_world_size()
-
-    def get_rank(self) -> int:
-        """Get the current rank in the sequence parallel group.
-
-        Returns:
-            The rank within the sequence parallel group.
-
-        Raises:
-            RuntimeError: If parallel state is not initialized.
-        """
-        from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_rank
-
-        return get_sequence_parallel_rank()
-
-    def get_ulysses_world_size(self) -> int:
-        """Get the Ulysses parallel world size.
-
-        Returns:
-            The world size for Ulysses (All-to-All) parallelism.
-        """
-        from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_world_size
-
-        return get_ulysses_parallel_world_size()
-
-    def get_ulysses_rank(self) -> int:
-        """Get the current rank in the Ulysses parallel group.
-
-        Returns:
-            The rank within the Ulysses parallel group.
-        """
-        from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_rank
-
-        return get_ulysses_parallel_rank()
-
-    def get_ring_world_size(self) -> int:
-        """Get the Ring parallel world size.
-
-        Returns:
-            The world size for Ring attention parallelism.
-        """
-        from vllm_omni.diffusion.distributed.parallel_state import get_ring_parallel_world_size
-
-        return get_ring_parallel_world_size()
-
-    def get_ring_rank(self) -> int:
-        """Get the current rank in the Ring parallel group.
-
-        Returns:
-            The rank within the Ring parallel group.
-        """
-        from vllm_omni.diffusion.distributed.parallel_state import get_ring_parallel_rank
-
-        return get_ring_parallel_rank()
-
-    def setup(self, rank: int, world_size: int, device: torch.device) -> None:
-        """Initialize the config with runtime parallel state.
-
-        This is called automatically when sequence parallelism is enabled.
-
-        Args:
-            rank: The global rank of this process.
-            world_size: Total world size.
-            device: The device for this rank.
-        """
-        self._rank = rank
-        self._world_size = world_size
-        self._device = device
-
-        expected_sp_size = self.ulysses_degree * self.ring_degree
-        actual_sp_size = self.get_world_size()
-
-        if expected_sp_size != actual_sp_size:
-            raise ValueError(
-                f"Configuration mismatch: ulysses_degree ({self.ulysses_degree}) * "
-                f"ring_degree ({self.ring_degree}) = {expected_sp_size}, but "
-                f"actual sequence parallel world size is {actual_sp_size}."
-            )
-
-    def is_initialized(self) -> bool:
-        """Check if the config has been initialized with runtime state.
-
-        Returns:
-            True if setup() has been called, False otherwise.
-        """
-        return self._rank is not None
-
-
-# =============================================================================
-# Sequence Parallel Plan Type Definitions
-# =============================================================================
-
-
-@dataclass(frozen=True)
-class SequenceParallelInput:
-    """Configuration for splitting an input tensor across sequence parallel ranks.
-
-    This specifies how to shard a tensor in the pre-forward or post-forward hook
-    of a layer. The tensor will be split along the specified dimension.
-
-    Note: This corresponds to `ContextParallelInput` in diffusers library.
-
-    Args:
-        split_dim: The dimension along which to split the tensor.
-        expected_dims: Expected number of dimensions. If provided, validates that
-            the tensor has this many dimensions before splitting. If the tensor
-            has a different number of dimensions, splitting is skipped with a warning.
-        split_output: If True, split the output of the layer instead of the input.
-            This is useful for layers whose outputs should be split after preprocessing
-            (e.g., RoPE embeddings).
-        auto_pad: If True, automatically pad the tensor if its size along split_dim
-            is not divisible by world_size. Creates an attention mask to indicate
-            valid vs padding positions. The mask is stored in ForwardContext.
-            Note: Ring attention does not support attention mask, so auto_pad
-            should only be used with Ulysses SP.
-
-    Example:
-        # Split hidden_states along sequence dimension (dim 1)
-        SequenceParallelInput(split_dim=1, expected_dims=3)
-
-        # Split RoPE output along sequence dimension (dim 0)
-        SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True)
-
-        # Split with auto-padding for variable-length sequences
-        SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
-    """
-
-    split_dim: int
-    expected_dims: int | None = None
-    split_output: bool = False
-    auto_pad: bool = False
-
-    def __repr__(self) -> str:
-        return (
-            f"SequenceParallelInput(split_dim={self.split_dim}, "
-            f"expected_dims={self.expected_dims}, split_output={self.split_output}, "
-            f"auto_pad={self.auto_pad})"
-        )
-
-
-@dataclass(frozen=True)
-class SequenceParallelOutput:
-    """Configuration for gathering an output tensor across sequence parallel ranks.
-
-    This specifies how to gather a tensor in the post-forward hook of a layer.
-    The tensor will be gathered along the specified dimension from all ranks.
-
-    Note: This corresponds to `ContextParallelOutput` in diffusers library.
-
-    Args:
-        gather_dim: The dimension along which to gather the tensor.
-        expected_dims: Expected number of dimensions. If provided, validates that
-            the tensor has this many dimensions before gathering.
-
-    Example:
-        # Gather output along sequence dimension (dim 1)
-        SequenceParallelOutput(gather_dim=1, expected_dims=3)
-    """
-
-    gather_dim: int
-    expected_dims: int | None = None
-
-    def __repr__(self) -> str:
-        return f"SequenceParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})"
-
-
-@dataclass(frozen=True)
-class SequenceParallelPartialInput:
-    """Configuration for partially splitting a tensor (e.g., split image part, keep text part).
-
-    This is designed for models like LongCat/Qwen where RoPE embeddings need special handling:
-    - Text portion: kept full across all ranks (for joint attention)
-    - Image portion: split across ranks
-
-    The tensor is assumed to be concatenated as [text_part, image_part] along split_dim.
-
-    Note: This is an extension beyond diffusers' standard ContextParallelInput,
-    designed for vLLM-Omni's dual-stream attention models.
-
-    Args:
-        split_dim: The dimension along which to split the image portion.
-        text_len_source: How to determine text length:
-            - str: Name of a forward parameter that contains text length
-            - int: Fixed text length value
-        expected_dims: Expected number of dimensions for validation.
-        split_output: If True, split the output instead of input.
-
-    Example:
-        # Split RoPE: text portion (from txt_ids.shape[0]) kept full, image portion split
-        SequenceParallelPartialInput(
-            split_dim=0,
-            text_len_source="txt_ids",  # Get text length from txt_ids.shape[0]
-            expected_dims=2,
-            split_output=True,
-        )
-
-        # Or with fixed text length
-        SequenceParallelPartialInput(
-            split_dim=0,
-            text_len_source=512,  # Fixed text length
-            expected_dims=2,
-            split_output=True,
-        )
-    """
-
-    split_dim: int
-    text_len_source: str | int
-    expected_dims: int | None = None
-    split_output: bool = False
-
-    def __repr__(self) -> str:
-        return (
-            f"SequenceParallelPartialInput(split_dim={self.split_dim}, "
-            f"text_len_source={self.text_len_source!r}, expected_dims={self.expected_dims}, "
-            f"split_output={self.split_output})"
-        )
-
-
-# =============================================================================
-# Type Aliases for _sp_plan Structure
-# =============================================================================
-
-# Any input config type
-AnySequenceParallelInput = SequenceParallelInput | SequenceParallelPartialInput
-
-# Input specification: maps parameter names (str) or output indices (int) to split config
-SequenceParallelInputType = dict[
-    str | int,
-    AnySequenceParallelInput | list[AnySequenceParallelInput] | tuple[AnySequenceParallelInput, ...],
-]
-
-# Output specification: single or multiple gather configs
-SequenceParallelOutputType = SequenceParallelOutput | list[SequenceParallelOutput] | tuple[SequenceParallelOutput, ...]
-
-# Full model plan: maps module names to input/output specifications
-# - Key "" refers to the model itself (root level)
-# - Key "module_name" refers to a submodule
-# - Key "module_name.*" refers to all children of a ModuleList
-#
-# Example of a complete _sp_plan:
-#
-#     _sp_plan = {
-#         # Root level: split model inputs before any submodule
-#         "": {
-#             "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-#         },
-#         # Submodule: split outputs of pos_embed (RoPE) layer
-#         "pos_embed": {
-#             0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
-#             1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
-#         },
-#         # Submodule: gather outputs of proj_out layer
-#         "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-#     }
-#
-SequenceParallelModelPlan = dict[str, SequenceParallelInputType | SequenceParallelOutputType]
-
-
-# =============================================================================
-# Validation Utilities
-# =============================================================================
-
-
-def _is_valid_input_config(value: object) -> bool:
-    """Check if a value is a valid input configuration type."""
-    return isinstance(value, (SequenceParallelInput, SequenceParallelPartialInput))
-
-
-def _is_valid_input_config_list(value: object) -> bool:
-    """Check if a value is a list/tuple of valid input configurations."""
-    if not isinstance(value, (list, tuple)):
-        return False
-    return all(_is_valid_input_config(x) for x in value)
-
-
-def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
-    """Validate a _sp_plan dictionary for correctness.
-
-    Args:
-        plan: The _sp_plan dictionary to validate.
-
-    Raises:
-        ValueError: If the plan is invalid.
-    """
-    if not isinstance(plan, dict):
-        raise ValueError(f"_sp_plan must be a dict, got {type(plan).__name__}")
-
-    for module_id, module_plan in plan.items():
-        if not isinstance(module_id, str):
-            raise ValueError(f"_sp_plan keys must be strings, got {type(module_id).__name__}")
-
-        # Check if it's an output specification (SequenceParallelOutput or list/tuple thereof)
-        if isinstance(module_plan, SequenceParallelOutput):
-            continue
-        if isinstance(module_plan, (list, tuple)):
-            if all(isinstance(x, SequenceParallelOutput) for x in module_plan):
-                continue
-            if _is_valid_input_config_list(module_plan):
-                # List of inputs for a specific parameter (when output is tuple)
-                continue
-
-        # Otherwise, should be an input specification dict
-        if isinstance(module_plan, dict):
-            for key, value in module_plan.items():
-                if not isinstance(key, (str, int)):
-                    raise ValueError(
-                        f"Input spec keys must be str or int, got {type(key).__name__} for module '{module_id}'"
-                    )
-                if isinstance(key, int) and not _is_valid_input_config(value):
-                    raise ValueError(
-                        f"Integer keys (output indices) must map to SequenceParallelInput/PartialInput, "
-                        f"got {type(value).__name__} for module '{module_id}'[{key}]"
-                    )
-                if _is_valid_input_config(value):
-                    if isinstance(key, int) and not value.split_output:
-                        raise ValueError(
-                            f"Integer keys (output indices) require split_output=True, "
-                            f"got split_output=False for module '{module_id}'[{key}]"
-                        )
-                elif _is_valid_input_config_list(value):
-                    pass  # Valid list of input configs
-                else:
-                    raise ValueError(
-                        f"Input spec values must be SequenceParallelInput/PartialInput or list thereof, "
-                        f"got {type(value).__name__} for module '{module_id}'['{key}']"
-                    )
-        else:
-            raise ValueError(
-                f"_sp_plan values must be dict (input spec) or SequenceParallelOutput, "
-                f"got {type(module_plan).__name__} for module '{module_id}'"
-            )
-
-
-def get_sp_plan_from_model(model: nn.Module) -> SequenceParallelModelPlan | None:
-    """Get the _sp_plan from a model if it exists.
-
-    Args:
-        model: The model to get the plan from.
-
-    Returns:
-        The _sp_plan dictionary, or None if not defined.
-    """
-    plan = getattr(model, "_sp_plan", None)
-    if plan is not None:
-        validate_sp_plan(plan)
-    return plan

vllm_omni/diffusion/distributed/sp_sharding.py

@@ -1,268 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and The HuggingFace Team
-"""Sequence Parallelism sharding utilities.
-
-This module provides low-level sharding and gathering functions for
-Sequence Parallelism. These can be used directly in model forward methods
-for semi-intrusive SP support, or internally by the SP hooks.
-"""
-
-from __future__ import annotations
-
-from contextlib import contextmanager
-from dataclasses import dataclass, field
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_sequence_parallel_rank,
-    get_sequence_parallel_world_size,
-    get_sp_group,
-)
-
-logger = init_logger(__name__)
-
-
-def sp_shard(
-    tensor: torch.Tensor,
-    dim: int,
-    validate: bool = True,
-) -> torch.Tensor:
-    """Shard a tensor along the specified dimension for sequence parallelism.
-
-    The tensor is split into world_size chunks along dim, and this rank
-    receives its corresponding chunk.
-
-    Args:
-        tensor: The tensor to shard.
-        dim: The dimension along which to split.
-        validate: If True, validate that the tensor size is divisible by world_size.
-
-    Returns:
-        The shard for this rank.
-
-    Raises:
-        ValueError: If validate=True and tensor size is not divisible by world_size.
-
-    Example:
-        # In model forward:
-        hidden_states = sp_shard(hidden_states, dim=1)
-    """
-    world_size = get_sequence_parallel_world_size()
-
-    if world_size == 1:
-        return tensor
-
-    rank = get_sequence_parallel_rank()
-    size = tensor.size(dim)
-
-    if validate and size % world_size != 0:
-        raise ValueError(
-            f"Tensor size along dim {dim} ({size}) must be divisible by "
-            f"world_size ({world_size}) for sequence parallel sharding."
-        )
-
-    return tensor.chunk(world_size, dim=dim)[rank]
-
-
-def sp_gather(
-    tensor: torch.Tensor,
-    dim: int,
-    validate: bool = True,
-) -> torch.Tensor:
-    """Gather a tensor along the specified dimension from all sequence parallel ranks.
-
-    The sharded tensors from all ranks are concatenated along dim.
-
-    Args:
-        tensor: The local shard to gather.
-        dim: The dimension along which to gather.
-        validate: If True, validate tensor consistency (currently unused).
-
-    Returns:
-        The full tensor gathered from all ranks.
-
-    Example:
-        # At end of model forward:
-        output = sp_gather(output, dim=1)
-    """
-    world_size = get_sequence_parallel_world_size()
-
-    if world_size == 1:
-        return tensor
-
-    sp_group = get_sp_group()
-    return sp_group.all_gather(tensor, dim=dim)
-
-
-def sp_shard_with_padding(
-    tensor: torch.Tensor,
-    dim: int,
-    pad_value: float = 0.0,
-) -> tuple[torch.Tensor, int]:
-    """Shard a tensor with automatic padding if not divisible by world_size.
-
-    This is useful for variable-length sequences where padding may be needed.
-
-    Args:
-        tensor: The tensor to shard.
-        dim: The dimension along which to split.
-        pad_value: Value to use for padding.
-
-    Returns:
-        Tuple of (sharded_tensor, padding_size). The padding_size indicates
-        how much padding was added to the original tensor before sharding.
-
-    Example:
-        sharded, pad_size = sp_shard_with_padding(hidden_states, dim=1)
-        # ... process ...
-        output = sp_gather(output, dim=1)
-        if pad_size > 0:
-            output = output[..., :-pad_size]  # Remove padding
-    """
-    world_size = get_sequence_parallel_world_size()
-
-    if world_size == 1:
-        return tensor, 0
-
-    size = tensor.size(dim)
-    remainder = size % world_size
-
-    if remainder == 0:
-        return sp_shard(tensor, dim, validate=False), 0
-
-    # Pad to make divisible
-    pad_size = world_size - remainder
-    pad_shape = list(tensor.shape)
-    pad_shape[dim] = pad_size
-    padding = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device)
-    tensor = torch.cat([tensor, padding], dim=dim)
-
-    return sp_shard(tensor, dim, validate=False), pad_size
-
-
-# NOTE: This class is a vLLM-Omni extension for
-# debugging intrusive SP implementations.
-# Purpose:
-# - Help developers detect bugs when implementing intrusive SP
-# - Verify that every sharded tensor is properly gathered
-# - Warn about common mistakes (double shard, gather without shard)
-#
-# When to use:
-# - During development/debugging of intrusive SP code
-# - In tests to verify shard/gather correctness
-@dataclass
-class ShardingValidator:
-    """Validator for tracking and verifying sharding operations.
-
-    This class helps ensure that sharding and gathering operations are
-    correctly paired in model forward passes. It tracks which tensors
-    have been sharded and verifies that they are properly gathered.
-
-    Usage:
-        validator = ShardingValidator()
-        with validator.track():
-            hidden_states = validator.shard(hidden_states, "hidden_states", dim=1)
-            # ... model computation ...
-            output = validator.gather(output, "hidden_states", dim=1)
-        validator.validate()  # Raises if any shard was not gathered
-
-    Attributes:
-        _sharded: Set of tensor names that have been sharded.
-        _gathered: Set of tensor names that have been gathered.
-        _enabled: Whether tracking is currently enabled.
-    """
-
-    _sharded: set[str] = field(default_factory=set)
-    _gathered: set[str] = field(default_factory=set)
-    _enabled: bool = False
-
-    def reset(self) -> None:
-        """Reset the validator state for a new forward pass."""
-        self._sharded.clear()
-        self._gathered.clear()
-
-    @contextmanager
-    def track(self):
-        """Context manager to enable tracking for a forward pass."""
-        self._enabled = True
-        self.reset()
-        try:
-            yield
-        finally:
-            self._enabled = False
-
-    def shard(
-        self,
-        tensor: torch.Tensor,
-        name: str,
-        dim: int,
-        validate_divisible: bool = True,
-    ) -> torch.Tensor:
-        """Shard a tensor and track the operation.
-
-        Args:
-            tensor: The tensor to shard.
-            name: A name to identify this tensor for validation.
-            dim: The dimension along which to split.
-            validate_divisible: If True, validate divisibility.
-
-        Returns:
-            The sharded tensor.
-        """
-        if self._enabled:
-            if name in self._sharded:
-                logger.warning(f"Tensor '{name}' sharded multiple times")
-            self._sharded.add(name)
-
-        return sp_shard(tensor, dim, validate=validate_divisible)
-
-    def gather(
-        self,
-        tensor: torch.Tensor,
-        name: str,
-        dim: int,
-    ) -> torch.Tensor:
-        """Gather a tensor and track the operation.
-
-        Args:
-            tensor: The local shard to gather.
-            name: The name used when sharding (for validation).
-            dim: The dimension along which to gather.
-
-        Returns:
-            The gathered tensor.
-        """
-        if self._enabled:
-            if name not in self._sharded:
-                logger.warning(f"Tensor '{name}' gathered without being sharded")
-            self._gathered.add(name)
-
-        return sp_gather(tensor, dim)
-
-    def validate(self) -> None:
-        """Validate that all sharded tensors were gathered.
-
-        Raises:
-            ValueError: If any sharded tensor was not gathered.
-        """
-        unmatched = self._sharded - self._gathered
-        if unmatched:
-            raise ValueError(
-                f"The following tensors were sharded but not gathered: {unmatched}. "
-                f"This may indicate a bug in the model's SP implementation."
-            )
-
-
-# Global validator instance for convenience
-_global_validator = ShardingValidator()
-
-
-def get_sharding_validator() -> ShardingValidator:
-    """Get the global sharding validator instance.
-
-    Returns:
-        The global ShardingValidator.
-    """
-    return _global_validator

vllm_omni/diffusion/executor/abstract.py

@@ -1,86 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Any
-
-from vllm.utils.import_utils import resolve_obj_by_qualname
-
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-
-
-class DiffusionExecutor(ABC):
-    """Abstract base class for Diffusion executors."""
-
-    uses_multiproc: bool = False
-
-    @staticmethod
-    def get_class(od_config: OmniDiffusionConfig) -> type["DiffusionExecutor"]:
-        executor_class: type[DiffusionExecutor]
-        distributed_executor_backend = od_config.distributed_executor_backend
-
-        if isinstance(distributed_executor_backend, type):
-            if not issubclass(distributed_executor_backend, DiffusionExecutor):
-                raise TypeError(
-                    "distributed_executor_backend must be a subclass of "
-                    f"DiffusionExecutor. Got {distributed_executor_backend}."
-                )
-            executor_class = distributed_executor_backend
-        elif distributed_executor_backend == "ray":
-            raise NotImplementedError("ray backend is not yet supported.")
-        elif distributed_executor_backend == "mp":
-            from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
-
-            executor_class = MultiprocDiffusionExecutor
-        elif distributed_executor_backend == "external_launcher":
-            raise NotImplementedError("external_launcher backend is not yet supported.")
-        elif isinstance(distributed_executor_backend, str):
-            try:
-                executor_class = resolve_obj_by_qualname(distributed_executor_backend)
-            except (ImportError, ValueError) as e:
-                raise ValueError(
-                    f"Failed to load executor backend '{distributed_executor_backend}'. "
-                    f"Ensure it is a valid python path. Error: {e}"
-                ) from e
-
-            if not issubclass(executor_class, DiffusionExecutor):
-                raise TypeError(
-                    f"distributed_executor_backend must be a subclass of DiffusionExecutor. Got {executor_class}."
-                )
-        else:
-            raise ValueError(f"Unknown distributed executor backend: {distributed_executor_backend}")
-        return executor_class
-
-    def __init__(self, od_config: OmniDiffusionConfig):
-        self.od_config = od_config
-        self._init_executor()
-
-    @abstractmethod
-    def _init_executor(self) -> None:
-        """Initialize the executor (e.g., launch workers, setup IPC)."""
-        pass
-
-    @abstractmethod
-    def add_req(self, requests: list[OmniDiffusionRequest]):
-        """Add requests to the execution queue."""
-        pass
-
-    @abstractmethod
-    def collective_rpc(
-        self,
-        method: str,
-        timeout: float | None = None,
-        args: tuple = (),
-        kwargs: dict | None = None,
-        unique_reply_rank: int | None = None,
-    ) -> Any:
-        """Execute a method on workers."""
-        pass
-
-    @abstractmethod
-    def check_health(self) -> None:
-        """Check if the executor and workers are healthy."""
-        pass
-
-    @abstractmethod
-    def shutdown(self) -> None:
-        """Shutdown the executor and release resources."""
-        pass

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -1,200 +0,0 @@
-import multiprocessing as mp
-import time
-import weakref
-from dataclasses import dataclass
-from typing import Any
-
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE
-from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.scheduler import Scheduler
-from vllm_omni.utils.platform_utils import get_diffusion_worker_class
-
-logger = init_logger(__name__)
-
-
-@dataclass
-class BackgroundResources:
-    """
-    Used as a finalizer for clean shutdown.
-    """
-
-    scheduler: Scheduler | None = None
-    processes: list[mp.Process] | None = None
-
-    def __call__(self):
-        """Clean up background resources."""
-        if self.scheduler is not None:
-            try:
-                for _ in range(self.scheduler.num_workers):
-                    self.scheduler.mq.enqueue(SHUTDOWN_MESSAGE)
-                self.scheduler.close()
-            except Exception as exc:
-                logger.warning("Failed to send shutdown signal: %s", exc)
-        if self.processes:
-            for proc in self.processes:
-                if not proc.is_alive():
-                    continue
-                proc.join(30)
-                if proc.is_alive():
-                    logger.warning("Terminating diffusion worker %s after timeout", proc.name)
-                    proc.terminate()
-                    proc.join(30)
-
-
-class MultiprocDiffusionExecutor(DiffusionExecutor):
-    uses_multiproc: bool = True
-
-    def _init_executor(self) -> None:
-        self._processes: list[mp.Process] = []
-        self._closed = False
-
-        # Initialize scheduler
-        self.scheduler = Scheduler()
-        self.scheduler.initialize(self.od_config)
-        broadcast_handle = self.scheduler.get_broadcast_handle()
-
-        # Launch workers
-        processes, result_handle = self._launch_workers(broadcast_handle)
-
-        if result_handle is not None:
-            self.scheduler.initialize_result_queue(result_handle)
-        else:
-            logger.error("Failed to get result queue handle from workers")
-
-        self._processes = processes
-
-        self.resources = BackgroundResources(scheduler=self.scheduler, processes=self._processes)
-        self._finalizer = weakref.finalize(self, self.resources)
-
-    def _launch_workers(self, broadcast_handle):
-        od_config = self.od_config
-        logger.info("Starting server...")
-
-        num_gpus = od_config.num_gpus
-        mp.set_start_method("spawn", force=True)
-        processes = []
-
-        # Get the appropriate worker class for current device
-        worker_proc = get_diffusion_worker_class()
-
-        # Launch all worker processes
-        scheduler_pipe_readers = []
-        scheduler_pipe_writers = []
-
-        for i in range(num_gpus):
-            reader, writer = mp.Pipe(duplex=False)
-            scheduler_pipe_writers.append(writer)
-            process = mp.Process(
-                target=worker_proc.worker_main,
-                args=(
-                    i,  # rank
-                    od_config,
-                    writer,
-                    broadcast_handle,
-                ),
-                name=f"DiffusionWorker-{i}",
-                daemon=True,
-            )
-            scheduler_pipe_readers.append(reader)
-            process.start()
-            processes.append(process)
-
-        # Wait for all workers to be ready
-        scheduler_infos = []
-        result_handle = None
-        for writer in scheduler_pipe_writers:
-            writer.close()
-
-        for i, reader in enumerate(scheduler_pipe_readers):
-            try:
-                data = reader.recv()
-            except EOFError:
-                logger.error(f"Rank {i} scheduler is dead. Please check if there are relevant logs.")
-                processes[i].join()
-                logger.error(f"Exit code: {processes[i].exitcode}")
-                raise
-
-            if data["status"] != "ready":
-                raise RuntimeError("Initialization failed. Please see the error messages above.")
-
-            if i == 0:
-                result_handle = data.get("result_handle")
-
-            scheduler_infos.append(data)
-            reader.close()
-
-        logger.debug("All workers are ready")
-
-        return processes, result_handle
-
-    def add_req(self, requests: list[OmniDiffusionRequest]):
-        return self.scheduler.add_req(requests)
-
-    def collective_rpc(
-        self,
-        method: str,
-        timeout: float | None = None,
-        args: tuple = (),
-        kwargs: dict | None = None,
-        unique_reply_rank: int | None = None,
-    ) -> Any:
-        if self._closed:
-            raise RuntimeError("DiffusionExecutor is closed.")
-
-        deadline = None if timeout is None else time.monotonic() + timeout
-        kwargs = kwargs or {}
-
-        # Prepare RPC request message
-        rpc_request = {
-            "type": "rpc",
-            "method": method,
-            "args": args,
-            "kwargs": kwargs,
-            "output_rank": unique_reply_rank,
-        }
-
-        try:
-            # Broadcast RPC request to all workers via unified message queue
-            self.scheduler.mq.enqueue(rpc_request)
-
-            # Determine which workers we expect responses from
-            num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus
-
-            responses = []
-            for _ in range(num_responses):
-                dequeue_timeout = None if deadline is None else (deadline - time.monotonic())
-                try:
-                    if self.scheduler.result_mq is None:
-                        raise RuntimeError("Result queue not initialized")
-
-                    response = self.scheduler.result_mq.dequeue(timeout=dequeue_timeout)
-
-                    # Check if response indicates an error
-                    if isinstance(response, dict) and response.get("status") == "error":
-                        raise RuntimeError(
-                            f"Worker failed with error '{response.get('error')}', "
-                            "please check the stack trace above for the root cause"
-                        )
-
-                    responses.append(response)
-                except TimeoutError as e:
-                    raise TimeoutError(f"RPC call to {method} timed out.") from e
-
-            return responses[0] if unique_reply_rank is not None else responses
-
-        except Exception as e:
-            logger.error(f"RPC call failed: {e}")
-            raise
-
-    def check_health(self) -> None:
-        # Simple check if processes are alive
-        for p in self._processes:
-            if not p.is_alive():
-                raise RuntimeError(f"Worker process {p.name} is dead")
-
-    def shutdown(self) -> None:
-        self._closed = True
-        self._finalizer()

vllm_omni/diffusion/forward_context.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from contextlib import contextmanager
 from dataclasses import dataclass
 
@@ -23,14 +21,6 @@ class ForwardContext:
     split_text_embed_in_sp: bool = False
     # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel
 
-    # Sequence Parallel padding support
-    # When sequence length is not divisible by SP world size, padding is added
-    # These values are used by SequenceParallelGatherHook to remove padding,
-    # and by attention layers to create attention masks dynamically
-    sp_padding_size: int = 0
-    # Original sequence length before padding (for removing padding in gather)
-    sp_original_seq_len: int | None = None
-
     def __post_init__(self):
         pass
 
@@ -96,14 +86,5 @@ def set_forward_context(
         attn_metadata=attn_metadata,
         split_text_embed_in_sp=split_text_embed_in_sp,
     )
-    # vLLM CustomOp dispatch (e.g. QKVParallelLinear) requires a global
-    # vLLM config set via set_current_vllm_config().
     with override_forward_context(forward_context):
-        if vllm_config is None:
-            yield
-        else:
-            # Local import to avoid importing vllm.config.vllm at module import time.
-            from vllm.config.vllm import set_current_vllm_config
-
-            with set_current_vllm_config(vllm_config):
-                yield
+        yield

vllm_omni/diffusion/hooks.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch.nn as nn
+
+
+class BaseState:
+    """Base class for hook state containers."""
+
+    def reset(self) -> None:  # pragma: no cover - default is no-op
+        pass
+
+
+class StateManager:
+    """Manage per-context hook state instances."""
+
+    def __init__(self, state_cls: Callable[[], BaseState]):
+        self._state_cls = state_cls
+        self._states: dict[str, BaseState] = {}
+        self._context: str = "default"
+
+    def set_context(self, name: str) -> None:
+        self._context = name or "default"
+
+    def get_state(self) -> BaseState:
+        if self._context not in self._states:
+            self._states[self._context] = self._state_cls()
+        return self._states[self._context]
+
+    def reset(self) -> None:
+        self._states.clear()
+
+
+class ModelHook:
+    """Base class for model hooks that can override a module's forward."""
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        raise NotImplementedError
+
+    def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any):
+        raise NotImplementedError
+
+    def reset_state(self, module: nn.Module) -> nn.Module:
+        raise NotImplementedError
+
+
+@dataclass
+class _WrappedForward:
+    module: nn.Module
+
+    def __call__(self, *args: Any, **kwargs: Any):
+        registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
+        if registry is None or not registry._hooks:
+            return self.module._original_forward(*args, **kwargs)
+        return registry.dispatch(*args, **kwargs)
+
+
+class HookRegistry:
+    """Registry of hooks attached to a module."""
+
+    def __init__(self, module: nn.Module):
+        self.module = module
+        self._hooks: dict[str, ModelHook] = {}
+
+    @classmethod
+    def get_or_create(cls, module: nn.Module) -> HookRegistry:
+        registry: HookRegistry | None = getattr(module, "_hook_registry", None)
+        if registry is None:
+            registry = cls(module)
+            setattr(module, "_hook_registry", registry)
+
+            # Wrap module.forward once so hooks can intercept calls.
+            if not hasattr(module, "_original_forward"):
+                module._original_forward = module.forward  # type: ignore[attr-defined]
+                module.forward = _WrappedForward(module)  # type: ignore[assignment]
+
+        return registry
+
+    def register_hook(self, name: str, hook: ModelHook) -> None:
+        hook.initialize_hook(self.module)
+        self._hooks[name] = hook
+
+    def get_hook(self, name: str) -> ModelHook | None:
+        return self._hooks.get(name)
+
+    def dispatch(self, *args: Any, **kwargs: Any):
+        # For now we support a single active hook and call it directly.
+        # This can be extended to a chain if needed.
+        if not self._hooks:
+            return self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+        # Deterministic order: sort by name.
+        name = sorted(self._hooks.keys())[0]
+        hook = self._hooks[name]
+        return hook.new_forward(self.module, *args, **kwargs)
+
+    def reset_hook(self, name: str) -> None:
+        hook = self._hooks.get(name)
+        if hook is not None:
+            hook.reset_state(self.module)

vllm_omni/diffusion/hooks/__init__.py

@@ -1,33 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Hook mechanism for model forward interception."""
-
-from vllm_omni.diffusion.hooks.base import (
-    BaseState,
-    HookRegistry,
-    ModelHook,
-    StateManager,
-)
-from vllm_omni.diffusion.hooks.sequence_parallel import (
-    SequenceParallelGatherHook,
-    SequenceParallelSplitHook,
-    apply_sequence_parallel,
-    disable_sequence_parallel_for_model,
-    enable_sequence_parallel_for_model,
-    remove_sequence_parallel,
-)
-
-__all__ = [
-    # Base hooks
-    "BaseState",
-    "StateManager",
-    "ModelHook",
-    "HookRegistry",
-    # Sequence parallel hooks (corresponds to diffusers' context_parallel)
-    "SequenceParallelSplitHook",
-    "SequenceParallelGatherHook",
-    "apply_sequence_parallel",
-    "remove_sequence_parallel",
-    "enable_sequence_parallel_for_model",
-    "disable_sequence_parallel_for_model",
-]

vllm_omni/diffusion/hooks/base.py

@@ -1,247 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Base hook classes for model forward interception.
-
-This module provides the foundational hook mechanism that allows intercepting
-and modifying model forward passes without invasive changes to model code.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Any
-
-import torch.nn as nn
-
-
-class BaseState:
-    """Base class for hook state containers."""
-
-    def reset(self) -> None:  # pragma: no cover - default is no-op
-        pass
-
-
-class StateManager:
-    """Manage per-context hook state instances."""
-
-    def __init__(self, state_cls: Callable[[], BaseState]):
-        self._state_cls = state_cls
-        self._states: dict[str, BaseState] = {}
-        self._context: str = "default"
-
-    def set_context(self, name: str) -> None:
-        self._context = name or "default"
-
-    def get_state(self) -> BaseState:
-        if self._context not in self._states:
-            self._states[self._context] = self._state_cls()
-        return self._states[self._context]
-
-    def reset(self) -> None:
-        self._states.clear()
-
-
-class ModelHook:
-    """Base class for model hooks that can override a module's forward.
-
-    Hooks can intercept the forward pass at two points:
-    - pre_forward: Called before the original forward, can modify args/kwargs
-    - post_forward: Called after the original forward, can modify output
-
-    Subclasses can override either or both methods. The default implementations
-    pass through args/kwargs/output unchanged.
-
-    For more complex behavior, override new_forward to completely replace
-    the forward logic.
-    """
-
-    def initialize_hook(self, module: nn.Module) -> nn.Module:
-        """Initialize the hook when it's registered to a module.
-
-        Args:
-            module: The module this hook is being attached to.
-
-        Returns:
-            The module (possibly modified).
-        """
-        return module
-
-    def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
-        """Called before the module's forward pass.
-
-        Args:
-            module: The module being called.
-            *args: Positional arguments to forward.
-            **kwargs: Keyword arguments to forward.
-
-        Returns:
-            Tuple of (args, kwargs) to pass to the forward method.
-        """
-        return args, kwargs
-
-    def post_forward(self, module: nn.Module, output: Any) -> Any:
-        """Called after the module's forward pass.
-
-        Args:
-            module: The module that was called.
-            output: The output from the forward method.
-
-        Returns:
-            The (possibly modified) output.
-        """
-        return output
-
-    def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> Any:
-        """Override the module's forward pass completely.
-
-        The default implementation calls pre_forward, then the original forward,
-        then post_forward. Override this method for more complex behavior.
-
-        Args:
-            module: The module being called.
-            *args: Positional arguments to forward.
-            **kwargs: Keyword arguments to forward.
-
-        Returns:
-            The output of the forward pass.
-        """
-        args, kwargs = self.pre_forward(module, *args, **kwargs)
-        output = module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
-        return self.post_forward(module, output)
-
-    def reset_state(self, module: nn.Module) -> nn.Module:
-        """Reset any state associated with this hook.
-
-        Args:
-            module: The module this hook is attached to.
-
-        Returns:
-            The module.
-        """
-        return module
-
-
-@dataclass
-class _WrappedForward:
-    """Wrapper that intercepts forward calls and dispatches to hooks."""
-
-    module: nn.Module
-
-    def __call__(self, *args: Any, **kwargs: Any):
-        registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
-        if registry is None or not registry._hooks:
-            return self.module._original_forward(*args, **kwargs)
-        return registry.dispatch(*args, **kwargs)
-
-
-class HookRegistry:
-    """Registry of hooks attached to a module.
-
-    Manages multiple hooks that can intercept a module's forward pass.
-    Hooks are called in sorted order by name for determinism.
-    """
-
-    def __init__(self, module: nn.Module):
-        self.module = module
-        self._hooks: dict[str, ModelHook] = {}
-
-    @classmethod
-    def get_or_create(cls, module: nn.Module) -> HookRegistry:
-        """Get existing registry or create a new one for the module.
-
-        Args:
-            module: The module to get/create a registry for.
-
-        Returns:
-            The HookRegistry for this module.
-        """
-        registry: HookRegistry | None = getattr(module, "_hook_registry", None)
-        if registry is None:
-            registry = cls(module)
-            setattr(module, "_hook_registry", registry)
-
-            # Wrap module.forward once so hooks can intercept calls.
-            if not hasattr(module, "_original_forward"):
-                module._original_forward = module.forward  # type: ignore[attr-defined]
-                module.forward = _WrappedForward(module)  # type: ignore[assignment]
-
-        return registry
-
-    def register_hook(self, name: str, hook: ModelHook) -> None:
-        """Register a hook with the given name.
-
-        Args:
-            name: Unique name for this hook.
-            hook: The hook instance to register.
-        """
-        hook.initialize_hook(self.module)
-        self._hooks[name] = hook
-
-    def remove_hook(self, name: str) -> None:
-        """Remove a hook by name.
-
-        Args:
-            name: The name of the hook to remove.
-        """
-        if name in self._hooks:
-            del self._hooks[name]
-
-    def get_hook(self, name: str) -> ModelHook | None:
-        """Get a hook by name.
-
-        Args:
-            name: The name of the hook.
-
-        Returns:
-            The hook if found, None otherwise.
-        """
-        return self._hooks.get(name)
-
-    def dispatch(self, *args: Any, **kwargs: Any) -> Any:
-        """Dispatch a forward call through registered hooks.
-
-        Currently supports a single active hook. Multiple hooks are called
-        in sorted order by name, with each hook's output passed to the next.
-
-        Args:
-            *args: Positional arguments to forward.
-            **kwargs: Keyword arguments to forward.
-
-        Returns:
-            The output of the forward pass.
-        """
-        if not self._hooks:
-            return self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
-
-        # For single hook case, call directly
-        if len(self._hooks) == 1:
-            hook = next(iter(self._hooks.values()))
-            return hook.new_forward(self.module, *args, **kwargs)
-
-        # For multiple hooks, chain them in sorted order
-        # Each hook can modify args/kwargs via pre_forward
-        sorted_hooks = sorted(self._hooks.items(), key=lambda x: x[0])
-
-        # Apply all pre_forward hooks
-        for _, hook in sorted_hooks:
-            args, kwargs = hook.pre_forward(self.module, *args, **kwargs)
-
-        # Call original forward
-        output = self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
-
-        # Apply all post_forward hooks in reverse order
-        for _, hook in reversed(sorted_hooks):
-            output = hook.post_forward(self.module, output)
-
-        return output
-
-    def reset_hook(self, name: str) -> None:
-        """Reset a hook's state by name.
-
-        Args:
-            name: The name of the hook to reset.
-        """
-        hook = self._hooks.get(name)
-        if hook is not None:
-            hook.reset_state(self.module)

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -1,703 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and the HuggingFace Team.
-# All rights reserved.
-#
-# This module is adapted from HuggingFace diffusers library:
-#   diffusers/src/diffusers/hooks/context_parallel.py
-#
-# NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
-# We use the term "Sequence Parallelism" to align with vLLM-Omni's existing terminology.
-#
-# Key adaptations from diffusers:
-#   - ModuleForwardMetadata: parameter lookup logic (adapted)
-#   - SequenceParallelSplitHook/GatherHook: hook structure (adapted from ContextParallel*)
-#   - apply_sequence_parallel: registration logic (adapted from apply_context_parallel)
-#
-# Key differences from diffusers:
-#   - Uses vLLM-Omni's SequenceParallelGroupCoordinator instead of DeviceMesh
-#   - Uses sp_shard/sp_gather from sp_sharding.py instead of funcol operations
-#   - Supports Ulysses + Ring hybrid parallelism
-#
-"""Sequence Parallelism hooks for non-intrusive SP support.
-
-This module implements the hook-based mechanism for applying sequence parallelism
-to models without modifying their forward() methods.
-
-Usage:
-    1. Define _sp_plan on your model class (corresponds to diffusers' _cp_plan)
-    2. Call apply_sequence_parallel(model, config, plan) to enable SP
-    3. Call remove_sequence_parallel(model, plan) to disable SP
-
-The hooks automatically shard inputs before forward and gather outputs after,
-based on the plan specification.
-"""
-
-from __future__ import annotations
-
-import inspect
-from dataclasses import dataclass
-from typing import Any
-
-import torch
-import torch.nn as nn
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.distributed.sp_plan import (
-    AnySequenceParallelInput,
-    SequenceParallelConfig,
-    SequenceParallelInput,
-    SequenceParallelModelPlan,
-    SequenceParallelOutput,
-    SequenceParallelPartialInput,
-)
-from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard
-from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook
-
-logger = init_logger(__name__)
-
-# Hook name templates for identifying SP hooks
-_SP_INPUT_HOOK_TEMPLATE = "sp_input---{}"
-_SP_OUTPUT_HOOK_TEMPLATE = "sp_output---{}"
-
-
-@dataclass
-class ModuleForwardMetadata:
-    """Metadata for mapping forward() parameter names to args/kwargs positions.
-
-    This caches the inspection of a module's forward signature to efficiently
-    locate parameters by name in subsequent calls.
-    """
-
-    cached_parameter_indices: dict[str, int] | None = None
-    _cls: type | None = None
-
-    def _get_parameter_from_args_kwargs(
-        self,
-        identifier: str,
-        args: tuple = (),
-        kwargs: dict | None = None,
-    ) -> tuple[Any, bool, int | None]:
-        """Get a parameter value from args or kwargs by name.
-
-        Args:
-            identifier: The parameter name to look up.
-            args: Positional arguments passed to forward.
-            kwargs: Keyword arguments passed to forward.
-
-        Returns:
-            Tuple of (value, is_kwarg, index).
-            - value: The parameter value (or None if not found)
-            - is_kwarg: True if found in kwargs
-            - index: Position in args if found there
-
-        Raises:
-            ValueError: If parameter not found in signature.
-        """
-        kwargs = kwargs or {}
-
-        # First check kwargs
-        if identifier in kwargs:
-            return kwargs[identifier], True, None
-
-        # Check cached indices
-        if self.cached_parameter_indices is not None:
-            index = self.cached_parameter_indices.get(identifier, None)
-            if index is None:
-                raise ValueError(f"Parameter '{identifier}' not found in cached indices.")
-            if index < len(args):
-                return args[index], False, index
-            return None, False, index
-
-        # Build cache from forward signature
-        if self._cls is None:
-            raise ValueError("Model class is not set for metadata.")
-
-        parameters = list(inspect.signature(self._cls.forward).parameters.keys())
-        parameters = parameters[1:]  # Skip `self`
-        self.cached_parameter_indices = {param: i for i, param in enumerate(parameters)}
-
-        if identifier not in self.cached_parameter_indices:
-            raise ValueError(f"Parameter '{identifier}' not found in function signature.")
-
-        index = self.cached_parameter_indices[identifier]
-
-        if index >= len(args):
-            return None, False, index
-
-        return args[index], False, index
-
-
-def _unwrap_module(module: nn.Module) -> nn.Module:
-    """Unwrap a module from any wrappers to get the original class.
-
-    Args:
-        module: Potentially wrapped module.
-
-    Returns:
-        The unwrapped module.
-    """
-    # Handle common wrappers
-    while hasattr(module, "_modules") and len(module._modules) == 1:
-        inner = next(iter(module._modules.values()))
-        if inner is not None:
-            module = inner
-        else:
-            break
-    return module
-
-
-class SequenceParallelSplitHook(ModelHook):
-    """Hook for splitting inputs before a module's forward pass.
-
-    This hook is registered to modules that need their inputs sharded
-    across sequence parallel ranks. It intercepts the forward call,
-    shards specified inputs according to the plan, and passes the
-    sharded inputs to the original forward.
-
-    For split_output=True inputs, it shards the output instead.
-
-    Supports both SequenceParallelInput (full split) and SequenceParallelPartialInput
-    (partial split for text/image separation).
-
-    Note: This corresponds to `ContextParallelSplitHook` in diffusers.
-    """
-
-    def __init__(
-        self,
-        metadata: dict[str | int, AnySequenceParallelInput | list[AnySequenceParallelInput]],
-        config: SequenceParallelConfig,
-    ) -> None:
-        super().__init__()
-        self.metadata = metadata
-        self.config = config
-        self.module_forward_metadata: ModuleForwardMetadata | None = None
-        # Cache for text lengths resolved from kwargs
-        self._text_len_cache: dict[str, int] = {}
-
-    def initialize_hook(self, module: nn.Module) -> nn.Module:
-        cls = _unwrap_module(module).__class__
-        self.module_forward_metadata = ModuleForwardMetadata(_cls=cls)
-        return module
-
-    def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
-        """Shard inputs before forward."""
-        args_list = list(args)
-        # Clear text length cache for this forward pass
-        self._text_len_cache.clear()
-
-        for name, spm in self.metadata.items():
-            # Skip if this is a split_output entry (handled in post_forward)
-            if isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) and spm.split_output:
-                continue
-
-            # Get the parameter value
-            input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs(
-                name, args_list, kwargs
-            )
-
-            if input_val is None:
-                continue
-
-            # Shard the input
-            if isinstance(input_val, torch.Tensor):
-                input_val = self._prepare_sp_input(input_val, spm, args_list, kwargs)
-            elif isinstance(input_val, (list, tuple)):
-                # Handle list/tuple of tensors with per-element config
-                if not isinstance(spm, (list, tuple)):
-                    raise ValueError(
-                        f"Expected list/tuple of SequenceParallelInput for parameter '{name}' "
-                        f"which is a list/tuple, but got {type(spm).__name__}"
-                    )
-                if len(input_val) != len(spm):
-                    raise ValueError(f"Expected {len(spm)} elements for parameter '{name}', got {len(input_val)}")
-                sharded_input_val = []
-                for i, x in enumerate(input_val):
-                    if torch.is_tensor(x) and not spm[i].split_output:
-                        x = self._prepare_sp_input(x, spm[i], args_list, kwargs)
-                    sharded_input_val.append(x)
-                input_val = type(input_val)(sharded_input_val)
-            else:
-                raise ValueError(f"Unsupported input type for sharding: {type(input_val).__name__}")
-
-            # Update args or kwargs
-            if is_kwarg:
-                kwargs[name] = input_val
-            elif index is not None and index < len(args_list):
-                args_list[index] = input_val
-            else:
-                raise ValueError(f"Failed to update parameter '{name}' after sharding.")
-
-        # Store kwargs for post_forward to resolve text lengths
-        self._last_kwargs = kwargs
-        self._last_args = tuple(args_list)
-
-        return tuple(args_list), kwargs
-
-    def post_forward(self, module: nn.Module, output: Any) -> Any:
-        """Shard outputs for split_output=True entries."""
-        is_tensor = isinstance(output, torch.Tensor)
-        is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)
-
-        if not is_tensor and not is_tensor_list:
-            # No tensor outputs to shard
-            return output
-
-        output_list = [output] if is_tensor else list(output)
-
-        for index, spm in self.metadata.items():
-            if not isinstance(index, int):
-                continue
-            if not isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) or not spm.split_output:
-                continue
-            if index >= len(output_list):
-                raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.")
-
-            output_list[index] = self._prepare_sp_input(output_list[index], spm, self._last_args, self._last_kwargs)
-
-        return output_list[0] if is_tensor else type(output)(output_list)
-
-    def _resolve_text_len(
-        self,
-        sp_input: SequenceParallelPartialInput,
-        args: tuple,
-        kwargs: dict,
-    ) -> int:
-        """Resolve text length from the source specification."""
-        source = sp_input.text_len_source
-
-        if isinstance(source, int):
-            return source
-
-        # String source - look up from kwargs or args
-        if source in self._text_len_cache:
-            return self._text_len_cache[source]
-
-        # Try to get from kwargs/args
-        try:
-            val, _, _ = self.module_forward_metadata._get_parameter_from_args_kwargs(source, args, kwargs)
-            if val is None:
-                raise ValueError(f"Parameter '{source}' is None, cannot determine text length.")
-            if isinstance(val, torch.Tensor):
-                # TODO: Currently assumes batch_size=1, where shape[0] is sequence length.
-                # For batch inference support, this should be updated to handle
-                # shape (batch_size, seq_len, ...) where text_len varies per sample.
-                text_len = val.shape[0]
-            elif isinstance(val, int):
-                text_len = val
-            else:
-                raise ValueError(f"Cannot determine text length from '{source}' of type {type(val).__name__}")
-            self._text_len_cache[source] = text_len
-            return text_len
-        except ValueError as e:
-            raise ValueError(f"Failed to resolve text_len_source '{source}': {e}") from e
-
-    def _prepare_sp_input(
-        self,
-        x: torch.Tensor,
-        sp_input: AnySequenceParallelInput,
-        args: tuple = (),
-        kwargs: dict | None = None,
-    ) -> torch.Tensor:
-        """Shard a tensor according to the input specification."""
-        kwargs = kwargs or {}
-
-        if sp_input.expected_dims is not None and x.dim() != sp_input.expected_dims:
-            logger.warning_once(f"Expected tensor with {sp_input.expected_dims} dims, got {x.dim()}. Skipping split.")
-            return x
-
-        if isinstance(sp_input, SequenceParallelInput):
-            # Full split with optional auto-padding
-            if sp_input.auto_pad:
-                return self._shard_with_auto_pad(x, sp_input.split_dim)
-            return sp_shard(x, sp_input.split_dim, validate=False)
-        elif isinstance(sp_input, SequenceParallelPartialInput):
-            # Partial split: keep text portion, split image portion
-            text_len = self._resolve_text_len(sp_input, args, kwargs)
-            dim = sp_input.split_dim
-
-            # Split tensor into text and image portions
-            text_part = x.narrow(dim, 0, text_len)
-            image_part = x.narrow(dim, text_len, x.size(dim) - text_len)
-
-            # Only shard the image portion
-            image_part_sharded = sp_shard(image_part, dim, validate=False)
-
-            # Concatenate back: [text_full, image_sharded]
-            return torch.cat([text_part, image_part_sharded], dim=dim)
-        else:
-            raise ValueError(f"Unsupported input config type: {type(sp_input).__name__}")
-
-    def _shard_with_auto_pad(self, x: torch.Tensor, dim: int) -> torch.Tensor:
-        """Shard tensor with automatic padding and attention mask creation.
-
-        When sequence length is not divisible by SP world size, this method:
-        1. Pads the tensor to make it divisible
-        2. Creates an attention mask indicating valid vs padding positions
-        3. Stores the mask and padding info in ForwardContext
-        """
-        from vllm_omni.diffusion.attention.selector import (
-            _BACKENDS_SUPPORT_ATTENTION_MASK,
-            get_attn_backend,
-        )
-        from vllm_omni.diffusion.distributed.parallel_state import (
-            get_ring_parallel_world_size,
-            get_sequence_parallel_rank,
-            get_sequence_parallel_world_size,
-        )
-        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
-
-        world_size = get_sequence_parallel_world_size()
-        if world_size == 1:
-            return x
-
-        seq_len = x.size(dim)
-        remainder = seq_len % world_size
-
-        if remainder == 0:
-            # No padding needed
-            return sp_shard(x, dim, validate=False)
-
-        # Check backend compatibility
-        attn_backend = get_attn_backend(-1)
-        if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK:
-            raise ValueError(
-                f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
-                f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "
-                f"Please switch to SDPA or Ascend attention backend."
-            )
-
-        # Ring attention does not support attention_mask
-        if get_ring_parallel_world_size() > 1:
-            raise ValueError(
-                f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
-                f"Cannot use Ring attention which does not support attention_mask. "
-                f"Please switch to Ulysses SP only."
-            )
-
-        # Calculate padding
-        pad_size = world_size - remainder
-        padded_seq_len = seq_len + pad_size
-
-        # Pad the tensor
-        pad_shape = list(x.shape)
-        pad_shape[dim] = pad_size
-        padding = torch.zeros(pad_shape, dtype=x.dtype, device=x.device)
-        x_padded = torch.cat([x, padding], dim=dim)
-
-        # Store padding info in forward context (only once, for primary tensor)
-        # Attention layers will create masks dynamically using this info
-        if is_forward_context_available():
-            ctx = get_forward_context()
-            # Only set if not already set (first auto_pad tensor wins)
-            if ctx.sp_original_seq_len is None:
-                ctx.sp_padding_size = pad_size
-                ctx.sp_original_seq_len = seq_len
-                logger.debug(
-                    f"Auto-padded sequence from {seq_len} to {padded_seq_len} "
-                    f"(pad_size={pad_size}, world_size={world_size}, dim={dim})"
-                )
-
-        # Shard the padded tensor
-        rank = get_sequence_parallel_rank()
-        return x_padded.chunk(world_size, dim=dim)[rank]
-
-
-class SequenceParallelGatherHook(ModelHook):
-    """Hook for gathering outputs after a module's forward pass.
-
-    This hook is registered to modules that need their outputs gathered
-    from all sequence parallel ranks. It intercepts the output and gathers
-    it according to the plan specification.
-
-    Note: This corresponds to `ContextParallelGatherHook` in diffusers.
-    """
-
-    def __init__(
-        self,
-        metadata: SequenceParallelOutput | list[SequenceParallelOutput],
-        config: SequenceParallelConfig,
-    ) -> None:
-        super().__init__()
-        if isinstance(metadata, SequenceParallelOutput):
-            metadata = [metadata]
-        self.metadata = metadata
-        self.config = config
-
-    def initialize_hook(self, module: nn.Module) -> nn.Module:
-        return module
-
-    def post_forward(self, module: nn.Module, output: Any) -> Any:
-        """Gather outputs after forward and remove padding if applied."""
-        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
-
-        is_tensor = isinstance(output, torch.Tensor)
-
-        if is_tensor:
-            output = [output]
-        elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)):
-            raise ValueError(f"Expected tensor or list/tuple of tensors, got {type(output).__name__}")
-
-        output = list(output)
-
-        if len(output) != len(self.metadata):
-            raise ValueError(f"Expected {len(self.metadata)} outputs, got {len(output)}.")
-
-        # Check if padding was applied during split
-        original_seq_len = None
-        if is_forward_context_available():
-            ctx = get_forward_context()
-            original_seq_len = ctx.sp_original_seq_len
-
-        for i, spm in enumerate(self.metadata):
-            if spm is None:
-                continue
-
-            x = output[i]
-            if spm.expected_dims is not None and x.dim() != spm.expected_dims:
-                logger.warning_once(
-                    f"Expected output tensor with {spm.expected_dims} dims, got {x.dim()}. Skipping gather."
-                )
-                continue
-
-            # Gather from all ranks
-            gathered = sp_gather(x, spm.gather_dim, validate=False)
-
-            # Remove padding if it was applied
-            if original_seq_len is not None and gathered.size(spm.gather_dim) > original_seq_len:
-                gathered = gathered.narrow(spm.gather_dim, 0, original_seq_len)
-                logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})")
-
-            output[i] = gathered
-
-        return output[0] if is_tensor else type(output)(output)
-
-
-def _get_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
-    """Get a submodule by dotted name, supporting wildcards.
-
-    Args:
-        model: The root module.
-        name: Dotted path to submodule. Use "*" to match all children
-            of a ModuleList.
-
-    Returns:
-        The submodule or list of submodules if wildcard used.
-
-    Raises:
-        ValueError: If the path is invalid or module not found.
-    """
-    if name.count("*") > 1:
-        raise ValueError("Wildcard '*' can only be used once in the name")
-    return _find_submodule_by_name(model, name)
-
-
-def _find_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
-    """Recursive helper for _get_submodule_by_name."""
-    if name == "":
-        return model
-
-    first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "")
-
-    if first_atom == "*":
-        if not isinstance(model, nn.ModuleList):
-            raise ValueError("Wildcard '*' can only be used with ModuleList")
-        submodules = []
-        for submodule in model:
-            subsubmodules = _find_submodule_by_name(submodule, remaining_name)
-            if not isinstance(subsubmodules, list):
-                subsubmodules = [subsubmodules]
-            submodules.extend(subsubmodules)
-        return submodules
-    else:
-        if hasattr(model, first_atom):
-            submodule = getattr(model, first_atom)
-            return _find_submodule_by_name(submodule, remaining_name)
-        else:
-            raise ValueError(f"'{first_atom}' is not a submodule of '{model.__class__.__name__}'")
-
-
-def apply_sequence_parallel(
-    module: nn.Module,
-    config: SequenceParallelConfig,
-    plan: SequenceParallelModelPlan,
-) -> None:
-    """Apply sequence parallel hooks to a model according to the plan.
-
-    This function registers hooks on the specified submodules to automatically
-    shard inputs and gather outputs for sequence parallelism.
-
-    Note: This corresponds to `apply_context_parallel` in diffusers.
-
-    The complete SP flow is:
-    1. Input sharding (SequenceParallelSplitHook): Split sequence across SP ranks
-    2. Attention parallelism (handled by vLLM-Omni's Attention layer):
-       - Ulysses: All-to-All over Q/K/V heads
-       - Ring: K/V circulation in ring topology
-       - Hybrid: Both (Ulysses handles head redistribution, Ring handles K/V)
-    3. Output gathering (SequenceParallelGatherHook): Gather sequence from SP ranks
-
-    Args:
-        module: The model to apply SP to.
-        config: The sequence parallel configuration.
-        plan: Dictionary mapping module names to input/output specifications.
-
-    Example:
-        config = SequenceParallelConfig(ulysses_degree=2)
-        plan = {
-            "": {"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3)},
-            "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-        }
-        apply_sequence_parallel(model, config, plan)
-
-    Note:
-        vLLM-Omni's Attention layer automatically handles the internal
-        parallelism (Ulysses All-to-All or Ring attention) based on the
-        forward_context configuration. This function only handles input/output
-        sharding for the model as a whole.
-    """
-    logger.debug(
-        f"Applying sequence parallel with config: ulysses={config.ulysses_degree}, "
-        f"ring={config.ring_degree}, plan keys: {list(plan.keys())}"
-    )
-
-    for module_id, sp_model_plan in plan.items():
-        submodule = _get_submodule_by_name(module, module_id)
-        if not isinstance(submodule, list):
-            submodule = [submodule]
-
-        logger.debug(f"Applying SP hooks to '{module_id}' ({len(submodule)} module(s))")
-
-        for m in submodule:
-            if isinstance(sp_model_plan, dict):
-                # Input specification
-                hook = SequenceParallelSplitHook(sp_model_plan, config)
-                hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
-            elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
-                # Output specification
-                if isinstance(sp_model_plan, SequenceParallelOutput):
-                    sp_model_plan = [sp_model_plan]
-                if not all(isinstance(x, SequenceParallelOutput) or x is None for x in sp_model_plan):
-                    raise ValueError(f"Expected SequenceParallelOutput elements, got {sp_model_plan}")
-                hook = SequenceParallelGatherHook(sp_model_plan, config)
-                hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
-            else:
-                raise ValueError(f"Unsupported plan type: {type(sp_model_plan).__name__}")
-
-            registry = HookRegistry.get_or_create(m)
-            registry.register_hook(hook_name, hook)
-
-
-def remove_sequence_parallel(
-    module: nn.Module,
-    plan: SequenceParallelModelPlan,
-) -> None:
-    """Remove sequence parallel hooks from a model.
-
-    Note: This corresponds to `remove_context_parallel` in diffusers.
-
-    Args:
-        module: The model to remove SP from.
-        plan: The same plan used when applying SP.
-    """
-    for module_id, sp_model_plan in plan.items():
-        submodule = _get_submodule_by_name(module, module_id)
-        if not isinstance(submodule, list):
-            submodule = [submodule]
-
-        for m in submodule:
-            registry = getattr(m, "_hook_registry", None)
-            if registry is None:
-                continue
-
-            if isinstance(sp_model_plan, dict):
-                hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
-            elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
-                hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
-            else:
-                continue
-
-            registry.remove_hook(hook_name)
-
-
-def enable_sequence_parallel_for_model(
-    model: nn.Module,
-    config: SequenceParallelConfig | None = None,
-) -> None:
-    """Enable sequence parallelism for a model using its _sp_plan.
-
-    This is a convenience function that reads the model's _sp_plan attribute
-    and applies sequence parallelism automatically.
-
-    Note: This corresponds to `enable_context_parallel_for_model` in diffusers,
-    but uses vLLM-Omni's _sp_plan instead of diffusers' _cp_plan.
-
-    The function performs two main tasks:
-    1. Applies _sp_plan hooks to shard inputs and gather outputs
-    2. Ensures Attention layers are configured for the correct parallel mode
-       (handled automatically by vLLM-Omni's forward_context mechanism)
-
-    Args:
-        model: The model to enable SP for. Must have a _sp_plan attribute.
-        config: Optional config. If None, uses default based on current
-            parallel state.
-
-    Raises:
-        ValueError: If model has no _sp_plan defined.
-
-    Note:
-        vLLM-Omni supports Ulysses + Ring hybrid parallelism:
-        - ulysses_degree > 1: Uses All-to-All communication over Q/K/V heads
-        - ring_degree > 1: Uses Ring attention with K/V passing
-        - Both > 1: Hybrid mode (Ulysses handles head redistribution,
-          Ring handles K/V circulation)
-    """
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        get_ring_parallel_world_size,
-        get_ulysses_parallel_world_size,
-    )
-    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
-
-    plan = get_sp_plan_from_model(model)
-    if plan is None:
-        raise ValueError(
-            f"Model {model.__class__.__name__} has no _sp_plan defined. "
-            f"Define _sp_plan as a class attribute or pass a plan explicitly."
-        )
-
-    if config is None:
-        # Create config from current parallel state
-        ulysses_degree = get_ulysses_parallel_world_size()
-        ring_degree = get_ring_parallel_world_size()
-        config = SequenceParallelConfig(
-            ulysses_degree=ulysses_degree,
-            ring_degree=ring_degree,
-        )
-        if ulysses_degree > 1 and ring_degree > 1:
-            mode = "hybrid"
-        elif ulysses_degree > 1:
-            mode = "ulysses"
-        else:
-            mode = "ring"
-        logger.info(
-            f"Created SP config from parallel state: "
-            f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
-            f"mode={mode}"
-        )
-
-    apply_sequence_parallel(model, config, plan)
-    logger.info(f"Enabled sequence parallelism for {model.__class__.__name__}")
-
-
-def disable_sequence_parallel_for_model(model: nn.Module) -> None:
-    """Disable sequence parallelism for a model.
-
-    Note: This corresponds to `disable_context_parallel_for_model` in diffusers.
-
-    Args:
-        model: The model to disable SP for.
-    """
-    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
-
-    plan = get_sp_plan_from_model(model)
-    if plan is not None:
-        remove_sequence_parallel(model, plan)

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -15,7 +15,8 @@ import numpy as np
 import torch
 from torch import nn
 from torch.nn.attention.flex_attention import flex_attention
-from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config as _Qwen2Config
 from transformers.models.qwen2.modeling_qwen2 import (
     Qwen2Attention,
     Qwen2MLP,
@@ -24,7 +25,6 @@ from transformers.models.qwen2.modeling_qwen2 import (
     Qwen2RotaryEmbedding,
 )
 from transformers.utils import ModelOutput
-from vllm.transformers_utils.configs.bagel import BagelConfig
 from vllm.vllm_flash_attn import flash_attn_varlen_func
 
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
@@ -71,7 +71,7 @@ torch._dynamo.config.accumulated_cache_size_limit = 4096
 flex_attention = torch.compile(flex_attention)
 
 
-class Qwen2MoTConfig(Qwen2Config):
+class Qwen2MoTConfig(_Qwen2Config):
     """Configuration for Qwen2MoT (Mixture of Tokens) model.
 
     This is fundamentally different from Qwen2, hence the distinct name.
@@ -639,6 +639,36 @@ def get_flattened_position_ids_extrapolate(img_h, img_w, patch_size, max_num_pat
     return pos_ids
 
 
+class BagelConfig(PretrainedConfig):
+    def __init__(
+        self,
+        visual_gen=True,
+        visual_und=True,
+        llm_config=None,
+        vit_config=None,
+        vae_config=None,
+        latent_patch_size=2,
+        max_latent_size=32,
+        vit_max_num_patch_per_side=70,
+        connector_act="gelu_pytorch_tanh",
+        interpolate_pos=False,
+        timestep_shift=1.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.visual_gen = visual_gen
+        self.visual_und = visual_und
+        self.llm_config = llm_config
+        self.vit_config = vit_config
+        self.vae_config = vae_config
+        self.latent_patch_size = latent_patch_size
+        self.max_latent_size = max_latent_size
+        self.vit_max_num_patch_per_side = vit_max_num_patch_per_side
+        self.connector_act = connector_act
+        self.interpolate_pos = interpolate_pos
+        self.timestep_shift = timestep_shift
+
+
 class Bagel(torch.nn.Module):
     config_class = BagelConfig
     base_model_prefix = "bagel"

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -20,7 +20,6 @@ from torch import nn
 from transformers import AutoTokenizer, SiglipImageProcessor, SiglipVisionConfig, SiglipVisionModel
 from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
-from vllm.transformers_utils.configs.bagel import BagelConfig
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -29,7 +28,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
 from .autoencoder import AutoEncoder, AutoEncoderParams
-from .bagel_transformer import Bagel, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM
+from .bagel_transformer import Bagel, BagelConfig, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM
 
 logger = init_logger(__name__)
 
@@ -305,84 +304,120 @@ class BagelPipeline(nn.Module):
         # Add text prompt (prefill) on gen context.
         # [Omni] Check for injected KV Cache from remote transfer
         injected_kv = getattr(req, "past_key_values", None)
-        if injected_kv is not None:
-            logger.info("Using injected KV Cache (direct)")
-            gen_context["past_key_values"] = injected_kv
+        injected_metadata = getattr(req, "kv_metadata", None)
 
-            # User requested: kv_lens and ropes set to [gen_context["past_key_values"].key_cache[0].shape[0]]
-            # Assuming injected_kv is compatible and has key_cache[0]
-            seq_len = injected_kv.key_cache[0].shape[0]
-            gen_context["kv_lens"] = [seq_len]
-            gen_context["ropes"] = [seq_len]
+        # Image input handling
+        image_input = getattr(req, "pil_image", None)
+        if image_input and not isinstance(image_input, list):
+            image_input = [image_input]
+
+        if image_input:
+            # If we have an image, we prefill with it
+            if self.image_processor and self.vae:
+
+                def vit_transforms(img):
+                    # SigLIP processor returns dict with pixel_values; we want the tensor
+                    return self.image_processor(images=img, return_tensors="pt").pixel_values[0]
+
+                def vae_transforms(img):
+                    if img.mode != "RGB":
+                        img = img.convert("RGB")
+                    # Convert to [-1, 1] tensor (H, W, C) -> (C, H, W)
+                    arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0
+                    return arr.permute(2, 0, 1)
+
+                # 1. Update VAE
+                gen_input_vae, newlens_vae, new_rope_vae = self.bagel.prepare_vae_images(
+                    curr_kvlens=gen_context["kv_lens"],
+                    curr_rope=gen_context["ropes"],
+                    images=image_input,
+                    transforms=vae_transforms,
+                    new_token_ids=self.new_token_ids,
+                )
+
+                for k, v in gen_input_vae.items():
+                    if torch.is_tensor(v):
+                        gen_input_vae[k] = v.to(self.device)
+
+                # VAE needs bfloat16 to match model strings usually, specifically encode
+                with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
+                    gen_context["past_key_values"] = self.bagel.forward_cache_update_vae(
+                        self.vae, gen_context["past_key_values"], **gen_input_vae
+                    )
+                gen_context["kv_lens"] = newlens_vae
+                gen_context["ropes"] = new_rope_vae
+
+                # 2. Update ViT
+                gen_input_img, newlens_img, new_rope_img = self.bagel.prepare_vit_images(
+                    curr_kvlens=gen_context["kv_lens"],
+                    curr_rope=gen_context["ropes"],
+                    images=image_input,
+                    transforms=vit_transforms,
+                    new_token_ids=self.new_token_ids,
+                )
+
+                for k, v in gen_input_img.items():
+                    if torch.is_tensor(v):
+                        gen_input_img[k] = v.to(self.device)
+
+                with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
+                    gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
+                        gen_context["past_key_values"], **gen_input_img
+                    )
+                gen_context["kv_lens"] = newlens_img
+                gen_context["ropes"] = new_rope_img
+            else:
+                logger.warning("Image provided but no image processor available.")
+
+        if injected_kv is not None and injected_metadata is not None:
+            logger.info("Using injected KV Cache from remote transfer")
+
+            # [Fix] Reconstruct NaiveCache if injected_kv is a dict of tensors
+            current_cache = gen_context["past_key_values"]
+            if isinstance(current_cache, NaiveCache) and isinstance(injected_kv, dict):
+                # injected_kv keys are like "0_k", "0_v", "1_k", ...
+                for key_name, tensor in injected_kv.items():
+                    try:
+                        # Parse layer index and type
+                        parts = key_name.split("_")
+                        if len(parts) < 2:
+                            continue
+
+                        layer_idx = int(parts[0])
+                        cache_type = parts[1]  # 'k' or 'v'
+
+                        # Ensure tensor is on correct device
+                        if tensor.device != self.device:
+                            tensor = tensor.to(self.device)
+
+                        if layer_idx in current_cache.key_cache:
+                            if cache_type == "k":
+                                current_cache.key_cache[layer_idx] = tensor
+                            elif cache_type == "v":
+                                current_cache.value_cache[layer_idx] = tensor
+                            elif cache_type == "kv":
+                                # Fallback if sender sent mixed/packed (less ideal)
+                                current_cache.key_cache[layer_idx] = tensor
+                                current_cache.value_cache[layer_idx] = tensor
+                    except Exception as e:
+                        logger.warning(f"Failed to load injected KV part {key_name}: {e}")
+
+            if "kv_lens" in injected_metadata:
+                val = injected_metadata["kv_lens"]
+                if isinstance(val, (int, float)):
+                    gen_context["kv_lens"] = [int(val)]
+                else:
+                    gen_context["kv_lens"] = list(val)
+
+            if "ropes" in injected_metadata:
+                val = injected_metadata["ropes"]
+                if isinstance(val, (int, float)):
+                    gen_context["ropes"] = [int(val)]
+                else:
+                    gen_context["ropes"] = list(val)
 
         else:
-            image_input = getattr(req, "pil_image", None)
-            if image_input and not isinstance(image_input, list):
-                image_input = [image_input]
-
-            if image_input:
-                # If we have an image, we prefill with it
-                if self.image_processor and self.vae:
-
-                    def vit_transforms(img):
-                        # SigLIP processor returns dict with pixel_values; we want the tensor
-                        return self.image_processor(images=img, return_tensors="pt").pixel_values[0]
-
-                    def vae_transforms(img):
-                        if img.mode != "RGB":
-                            img = img.convert("RGB")
-                        # Convert to [-1, 1] tensor (H, W, C) -> (C, H, W)
-                        arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0
-                        return arr.permute(2, 0, 1)
-
-                    # 1. Update VAE
-                    gen_input_vae, newlens_vae, new_rope_vae = self.bagel.prepare_vae_images(
-                        curr_kvlens=gen_context["kv_lens"],
-                        curr_rope=gen_context["ropes"],
-                        images=image_input,
-                        transforms=vae_transforms,
-                        new_token_ids=self.new_token_ids,
-                    )
-
-                    for k, v in gen_input_vae.items():
-                        if torch.is_tensor(v):
-                            gen_input_vae[k] = v.to(self.device)
-
-                    # VAE needs bfloat16 to match model strings usually, specifically encode
-                    with torch.autocast(
-                        device_type=self.device.type,
-                        enabled=self.device.type != "cpu",
-                        dtype=self.od_config.dtype,
-                    ):
-                        gen_context["past_key_values"] = self.bagel.forward_cache_update_vae(
-                            self.vae, gen_context["past_key_values"], **gen_input_vae
-                        )
-                    gen_context["kv_lens"] = newlens_vae
-                    gen_context["ropes"] = new_rope_vae
-
-                    # 2. Update ViT
-                    gen_input_img, newlens_img, new_rope_img = self.bagel.prepare_vit_images(
-                        curr_kvlens=gen_context["kv_lens"],
-                        curr_rope=gen_context["ropes"],
-                        images=image_input,
-                        transforms=vit_transforms,
-                        new_token_ids=self.new_token_ids,
-                    )
-
-                    for k, v in gen_input_img.items():
-                        if torch.is_tensor(v):
-                            gen_input_img[k] = v.to(self.device)
-
-                    with torch.autocast(
-                        device_type=self.device.type,
-                        enabled=self.device.type != "cpu",
-                        dtype=self.od_config.dtype,
-                    ):
-                        gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
-                            gen_context["past_key_values"], **gen_input_img
-                        )
-                    gen_context["kv_lens"] = newlens_img
-                    gen_context["ropes"] = new_rope_img
+            # Standard local prefill path
             generation_input, newlens, new_rope = self.bagel.prepare_prompts(
                 curr_kvlens=gen_context["kv_lens"],
                 curr_rope=gen_context["ropes"],
@@ -403,11 +438,7 @@ class BagelPipeline(nn.Module):
             for k, v in generation_input.items():
                 if torch.is_tensor(v):
                     generation_input[k] = v.to(self.device)
-            with torch.autocast(
-                device_type=self.device.type,
-                enabled=self.device.type != "cpu",
-                dtype=self.od_config.dtype,
-            ):
+            with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
                 gen_context["past_key_values"] = self.bagel.forward_cache_update_text(
                     gen_context["past_key_values"], **generation_input
                 )
@@ -452,11 +483,7 @@ class BagelPipeline(nn.Module):
             if torch.is_tensor(v):
                 generation_input[k] = v.to(self.device)
 
-        with torch.autocast(
-            device_type=self.device.type,
-            enabled=self.device.type != "cpu",
-            dtype=self.od_config.dtype,
-        ):
+        with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
             latents = self.bagel.generate_image(
                 past_key_values=gen_context["past_key_values"],
                 num_timesteps=gen_params.num_timesteps,

vllm_omni/diffusion/models/glm_image/__init__.py

@@ -1,21 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""GLM Image diffusion model components."""
-
-from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
-    GlmImageKVCache,
-    GlmImageTransformer2DModel,
-)
-from vllm_omni.diffusion.models.glm_image.pipeline_glm_image import (
-    GlmImagePipeline,
-    get_glm_image_post_process_func,
-    get_glm_image_pre_process_func,
-)
-
-__all__ = [
-    "GlmImageKVCache",
-    "GlmImagePipeline",
-    "GlmImageTransformer2DModel",
-    "get_glm_image_post_process_func",
-    "get_glm_image_pre_process_func",
-]

vllm_omni/diffusion/models/glm_image/glm_image_transformer.py

@@ -1,792 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections.abc import Iterable
-from enum import Enum
-from typing import Any
-
-import torch
-import torch.nn as nn
-from diffusers.models.attention import FeedForward
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings
-from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import QKVParallelLinear
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-
-from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.base import CachedTransformer
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding
-
-logger = init_logger(__name__)
-
-
-class GlmImageImageProjector(nn.Module):
-    """Projects latent image patches to transformer hidden dimension."""
-
-    def __init__(
-        self,
-        in_channels: int = 16,
-        hidden_size: int = 2560,
-        patch_size: int = 2,
-    ):
-        super().__init__()
-        self.patch_size = patch_size
-        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        batch_size, channel, height, width = hidden_states.shape
-        post_patch_height = height // self.patch_size
-        post_patch_width = width // self.patch_size
-
-        # Reshape: [B, C, H, W] -> [B, H', W', C*p*p] -> [B, H'*W', C*p*p]
-        hidden_states = hidden_states.reshape(
-            batch_size, channel, post_patch_height, self.patch_size, post_patch_width, self.patch_size
-        )
-        hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
-        hidden_states = self.proj(hidden_states)
-        return hidden_states
-
-
-class GlmImageRotaryPosEmbed(nn.Module):
-    """Rotary positional embedding for 2D image patches."""
-
-    def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
-        super().__init__()
-        self.dim = dim
-        self.patch_size = patch_size
-        self.theta = theta
-
-    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        batch_size, num_channels, height, width = hidden_states.shape
-        height, width = height // self.patch_size, width // self.patch_size
-
-        dim_h, dim_w = self.dim // 2, self.dim // 2
-        h_inv_freq = 1.0 / (
-            self.theta ** (torch.arange(0, dim_h, 2, dtype=torch.float32)[: (dim_h // 2)].float() / dim_h)
-        )
-        w_inv_freq = 1.0 / (
-            self.theta ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w)
-        )
-        h_seq = torch.arange(height, device=hidden_states.device)
-        w_seq = torch.arange(width, device=hidden_states.device)
-        h_inv_freq = h_inv_freq.to(hidden_states.device)
-        w_inv_freq = w_inv_freq.to(hidden_states.device)
-
-        freqs_h = torch.outer(h_seq, h_inv_freq)
-        freqs_w = torch.outer(w_seq, w_inv_freq)
-
-        # Create position matrices: [height, 1, dim//4] and [1, width, dim//4]
-        freqs_h = freqs_h.unsqueeze(1).expand(height, width, -1)
-        freqs_w = freqs_w.unsqueeze(0).expand(height, width, -1)
-
-        # Concatenate: [height, width, dim//2] -> [height, width, dim]
-        freqs = torch.cat([freqs_h, freqs_w], dim=-1)
-        freqs = torch.cat([freqs, freqs], dim=-1)
-        freqs = freqs.reshape(height * width, -1)
-        return (freqs.cos(), freqs.sin())
-
-
-class GlmImageAdaLayerNormZero(nn.Module):
-    """Adaptive LayerNorm with zero initialization for both image and text streams."""
-
-    def __init__(self, embedding_dim: int, dim: int) -> None:
-        super().__init__()
-        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.linear = nn.Linear(embedding_dim, 12 * dim, bias=True)
-
-    def forward(
-        self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
-    ) -> tuple[torch.Tensor, ...]:
-        dtype = hidden_states.dtype
-        norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
-        norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype)
-
-        emb = self.linear(temb)
-        (
-            shift_msa,
-            c_shift_msa,
-            scale_msa,
-            c_scale_msa,
-            gate_msa,
-            c_gate_msa,
-            shift_mlp,
-            c_shift_mlp,
-            scale_mlp,
-            c_scale_mlp,
-            gate_mlp,
-            c_gate_mlp,
-        ) = emb.chunk(12, dim=1)
-
-        hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
-        encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa.unsqueeze(1)) + c_shift_msa.unsqueeze(1)
-
-        return (
-            hidden_states,
-            gate_msa,
-            shift_mlp,
-            scale_mlp,
-            gate_mlp,
-            encoder_hidden_states,
-            c_gate_msa,
-            c_shift_mlp,
-            c_scale_mlp,
-            c_gate_mlp,
-        )
-
-
-class GlmImageAdaLayerNormContinuous(nn.Module):
-    """Final AdaLN for output projection (no activation before Linear)."""
-
-    def __init__(
-        self,
-        embedding_dim: int,
-        conditioning_embedding_dim: int,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        bias: bool = True,
-    ):
-        super().__init__()
-        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
-        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
-
-    def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
-        # NO SiLU here
-        emb = self.linear(conditioning_embedding.to(x.dtype))
-        scale, shift = torch.chunk(emb, 2, dim=1)
-        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
-        return x
-
-
-class KVCacheMode(Enum):
-    """Mode for KV cache operations.
-
-    - WRITE: Store the K/V tensors from condition images
-    - READ: Concatenate cached K/V with current K/V
-    - SKIP: Do not use cache (pass-through)
-    """
-
-    WRITE = "write"
-    READ = "read"
-    SKIP = "skip"
-
-
-class GlmImageLayerKVCache:
-    """KV cache for a single attention layer.
-
-    Stores key and value tensors for image editing. The cache accumulates
-    KV pairs during write mode and provides them during read mode.
-
-    Shape convention (vllm-omni):
-        key/value: [batch_size, seq_length, num_heads, head_dim]
-    """
-
-    def __init__(self):
-        self.k_cache: torch.Tensor | None = None
-        self.v_cache: torch.Tensor | None = None
-
-    def store(self, key: torch.Tensor, value: torch.Tensor) -> None:
-        """Store or accumulate KV tensors.
-
-        If cache is empty, stores the tensors directly.
-        If cache is not empty, concatenates new tensors along seq_length dim.
-
-        Args:
-            key: Key tensor of shape [B, S, H, D]
-            value: Value tensor of shape [B, S, H, D]
-        """
-        if self.k_cache is None:
-            self.k_cache = key
-            self.v_cache = value
-        else:
-            # Concatenate along sequence dimension (dim=1 for [B, S, H, D])
-            self.k_cache = torch.cat([self.k_cache, key], dim=1)
-            self.v_cache = torch.cat([self.v_cache, value], dim=1)
-
-    def get(self) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-        """Get cached KV tensors.
-
-        Returns:
-            Tuple of (k_cache, v_cache), both may be None if cache is empty.
-        """
-        return self.k_cache, self.v_cache
-
-    def clear(self) -> None:
-        """Clear the cache."""
-        self.k_cache = None
-        self.v_cache = None
-
-    @property
-    def is_empty(self) -> bool:
-        """Check if cache is empty."""
-        return self.k_cache is None
-
-    def __repr__(self) -> str:
-        if self.is_empty:
-            return "GlmImageLayerKVCache(empty)"
-        return f"GlmImageLayerKVCache(k_shape={self.k_cache.shape}, v_shape={self.v_cache.shape})"
-
-
-class GlmImageKVCache:
-    """Container for all layers' KV caches.
-
-    Manages KV cache for all transformer layers in GLM-Image model.
-    Provides a unified interface for setting mode and clearing cache.
-
-    Args:
-        num_layers: Number of transformer layers in the model.
-
-    Example:
-        kv_cache = GlmImageKVCache(num_layers=28)
-        kv_cache.set_mode(KVCacheMode.WRITE)
-        # ... process condition image ...
-        kv_cache.set_mode(KVCacheMode.READ)
-        # ... process target image ...
-        kv_cache.clear()
-    """
-
-    def __init__(self, num_layers: int):
-        self.num_layers = num_layers
-        self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
-        self._mode: KVCacheMode | None = None
-
-    def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
-        """Get cache for a specific layer.
-
-        Args:
-            layer_idx: Index of the layer (0-indexed).
-
-        Returns:
-            GlmImageLayerKVCache for the specified layer.
-
-        Raises:
-            IndexError: If layer_idx is out of range.
-        """
-        if layer_idx < 0 or layer_idx >= self.num_layers:
-            raise IndexError(f"Layer index {layer_idx} out of range [0, {self.num_layers})")
-        return self.caches[layer_idx]
-
-    def __len__(self) -> int:
-        """Return number of layers."""
-        return self.num_layers
-
-    @property
-    def mode(self) -> KVCacheMode | None:
-        """Get current cache mode."""
-        return self._mode
-
-    def set_mode(self, mode: KVCacheMode | str | None) -> None:
-        """Set cache mode for all layers.
-
-        Args:
-            mode: Cache mode (WRITE, READ, SKIP) or string ("write", "read", "skip").
-                  Use None to disable cache operations.
-
-        Raises:
-            ValueError: If mode is an invalid string.
-        """
-        if mode is None:
-            self._mode = None
-        elif isinstance(mode, str):
-            try:
-                self._mode = KVCacheMode(mode.lower())
-            except ValueError:
-                raise ValueError(f"Invalid mode: '{mode}', must be one of 'write', 'read', 'skip'")
-        else:
-            self._mode = mode
-
-    def clear(self) -> None:
-        """Clear cache for all layers and reset mode."""
-        for cache in self.caches:
-            cache.clear()
-        self._mode = None
-
-    @property
-    def is_empty(self) -> bool:
-        """Check if all layer caches are empty."""
-        return all(cache.is_empty for cache in self.caches)
-
-    def __repr__(self) -> str:
-        mode_str = self._mode.value if self._mode else "None"
-        return f"GlmImageKVCache(num_layers={self.num_layers}, mode={mode_str}, is_empty={self.is_empty})"
-
-
-class GlmImageAttention(nn.Module):
-    """
-    Joint attention for GLM-Image model using vllm-omni's optimized attention.
-
-    This combines text and image streams for joint attention computation.
-    Supports KV caching for image editing workflows via external cache.
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        head_dim: int,
-        out_bias: bool = True,
-        eps: float = 1e-5,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.head_dim = head_dim
-        self.inner_dim = num_heads * head_dim
-
-        # QKV projection (fused for efficiency)
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=head_dim,
-            total_num_heads=num_heads,
-            disable_tp=True,
-            bias=True,
-        )
-
-        # QK normalization (LayerNorm, not RMSNorm for GLM-Image)
-        self.norm_q = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps)
-        self.norm_k = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps)
-
-        # Output projection
-        self.to_out = nn.Sequential(
-            nn.Linear(self.inner_dim, dim, bias=out_bias),
-            nn.Dropout(0.0),
-        )
-
-        # RoPE and attention
-        self.rope = RotaryEmbedding(is_neox_style=False)
-        self.attn = Attention(
-            num_heads=num_heads,
-            head_size=head_dim,
-            softmax_scale=1.0 / (head_dim**0.5),
-            causal=False,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-        attention_mask: torch.Tensor | None = None,
-        kv_cache: GlmImageLayerKVCache | None = None,
-        kv_cache_mode: KVCacheMode | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for joint attention.
-
-        Args:
-            hidden_states: Image hidden states [B, img_seq_len, D]
-            encoder_hidden_states: Text hidden states [B, text_seq_len, D]
-            image_rotary_emb: Tuple of (cos, sin) for RoPE
-            attention_mask: Optional attention mask for text tokens
-            kv_cache: Optional layer KV cache for image editing
-            kv_cache_mode: Cache mode (WRITE, READ, SKIP)
-
-        Returns:
-            Tuple of (image_hidden_states, text_hidden_states)
-        """
-        dtype = encoder_hidden_states.dtype
-        batch_size, text_seq_length, _ = encoder_hidden_states.shape
-
-        # Concatenate text and image: [text, image]
-        hidden_states_combined = torch.cat([encoder_hidden_states, hidden_states], dim=1)
-
-        # QKV projection
-        qkv, _ = self.to_qkv(hidden_states_combined)
-        query, key, value = qkv.chunk(3, dim=-1)
-
-        # Reshape: [B, S, H*D] -> [B, S, H, D]
-        query = query.unflatten(-1, (self.num_heads, -1))
-        key = key.unflatten(-1, (self.num_heads, -1))
-        value = value.unflatten(-1, (self.num_heads, -1))
-
-        # QK normalization
-        query = self.norm_q(query).to(dtype=dtype)
-        key = self.norm_k(key).to(dtype=dtype)
-
-        # Apply RoPE only to image tokens (not text tokens)
-        if image_rotary_emb is not None:
-            cos, sin = image_rotary_emb
-            cos = cos.to(query.dtype)
-            sin = sin.to(query.dtype)
-            # Only apply RoPE to image part (after text_seq_length)
-            query_img = query[:, text_seq_length:, :, :]
-            key_img = key[:, text_seq_length:, :, :]
-            from diffusers.models.embeddings import apply_rotary_emb
-
-            query_img = apply_rotary_emb(query_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
-            # key_img = self.rope(key_img, cos, sin)
-            key_img = apply_rotary_emb(key_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
-            query = torch.cat([query[:, :text_seq_length, :, :], query_img], dim=1)
-            key = torch.cat([key[:, :text_seq_length, :, :], key_img], dim=1)
-
-        # Handle KV cache for image editing
-        if kv_cache is not None and kv_cache_mode is not None:
-            if kv_cache_mode == KVCacheMode.WRITE:
-                kv_cache.store(key, value)
-            elif kv_cache_mode == KVCacheMode.READ:
-                k_cached, v_cached = kv_cache.get()
-                if k_cached is not None:
-                    key = torch.cat([k_cached, key], dim=1)
-                    value = torch.cat([v_cached, value], dim=1)
-            # KVCacheMode.SKIP: do nothing
-
-        # Attention computation
-        hidden_states_out = self.attn(query, key, value)
-        hidden_states_out = hidden_states_out.flatten(2, 3)
-        hidden_states_out = hidden_states_out.to(dtype)
-
-        # Output projection
-        hidden_states_out = self.to_out(hidden_states_out)
-
-        # Split back to text and image
-        encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :]
-        hidden_states_out = hidden_states_out[:, text_seq_length:, :]
-
-        return hidden_states_out, encoder_hidden_states_out
-
-
-class GlmImageTransformerBlock(nn.Module):
-    """Single transformer block for GLM-Image."""
-
-    def __init__(
-        self,
-        dim: int = 2560,
-        num_attention_heads: int = 64,
-        attention_head_dim: int = 40,
-        time_embed_dim: int = 512,
-    ) -> None:
-        super().__init__()
-
-        # 1. Attention with AdaLN
-        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
-        self.attn1 = GlmImageAttention(
-            dim=dim,
-            num_heads=num_attention_heads,
-            head_dim=attention_head_dim,
-        )
-
-        # 2. Feedforward
-        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-        attention_mask: torch.Tensor | None = None,
-        attention_kwargs: dict[str, Any] | None = None,
-        kv_cache: GlmImageLayerKVCache | None = None,
-        kv_cache_mode: KVCacheMode | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for transformer block.
-
-        Args:
-            hidden_states: Image hidden states
-            encoder_hidden_states: Text hidden states
-            temb: Timestep embedding
-            image_rotary_emb: RoPE embeddings
-            attention_mask: Text attention mask
-            attention_kwargs: Additional attention arguments
-            kv_cache: Layer-specific KV cache for image editing
-            kv_cache_mode: Cache mode (WRITE, READ, SKIP)
-
-        Returns:
-            Tuple of (image_hidden_states, text_hidden_states)
-        """
-        # 1. Timestep conditioning via AdaLN
-        (
-            norm_hidden_states,
-            gate_msa,
-            shift_mlp,
-            scale_mlp,
-            gate_mlp,
-            norm_encoder_hidden_states,
-            c_gate_msa,
-            c_shift_mlp,
-            c_scale_mlp,
-            c_gate_mlp,
-        ) = self.norm1(hidden_states, encoder_hidden_states, temb)
-
-        # 2. Attention
-        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            image_rotary_emb=image_rotary_emb,
-            attention_mask=attention_mask,
-            kv_cache=kv_cache,
-            kv_cache_mode=kv_cache_mode,
-        )
-        hidden_states = hidden_states + attn_hidden_states * gate_msa.unsqueeze(1)
-        encoder_hidden_states = encoder_hidden_states + attn_encoder_hidden_states * c_gate_msa.unsqueeze(1)
-
-        # 3. Feedforward
-        norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
-        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) * (
-            1 + c_scale_mlp.unsqueeze(1)
-        ) + c_shift_mlp.unsqueeze(1)
-
-        ff_output = self.ff(norm_hidden_states)
-        ff_output_context = self.ff(norm_encoder_hidden_states)
-        hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
-        encoder_hidden_states = encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)
-
-        return hidden_states, encoder_hidden_states
-
-
-class GlmImageTransformer2DModel(CachedTransformer):
-    """
-    GLM-Image Transformer model for 2D image generation.
-
-    This is the vllm-omni optimized version of the GLM-Image DiT model.
-
-    Args:
-        od_config: OmniDiffusionConfig containing model configuration. The
-            transformer hyper-parameters (e.g. patch size / channels / heads)
-            are read from `od_config.tf_model_config`.
-    """
-
-    def __init__(
-        self,
-        od_config: OmniDiffusionConfig,
-    ):
-        super().__init__()
-
-        patch_size = od_config.tf_model_config.patch_size
-        in_channels = od_config.tf_model_config.in_channels
-        out_channels = od_config.tf_model_config.out_channels
-        num_attention_heads = od_config.tf_model_config.num_attention_heads
-        attention_head_dim = od_config.tf_model_config.attention_head_dim
-        time_embed_dim = od_config.tf_model_config.time_embed_dim
-        condition_dim = od_config.tf_model_config.condition_dim
-        prior_vq_quantizer_codebook_size = od_config.tf_model_config.prior_vq_quantizer_codebook_size
-        text_embed_dim = od_config.tf_model_config.text_embed_dim
-
-        # Get num_layers from config if available
-        model_config = od_config.tf_model_config
-        if model_config is not None and hasattr(model_config, "num_layers"):
-            num_layers = model_config.num_layers
-
-        self.od_config = od_config
-        self.patch_size = patch_size
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-
-        # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
-        pooled_projection_dim = 2 * 2 * condition_dim
-        inner_dim = num_attention_heads * attention_head_dim
-
-        # 1. RoPE
-        self.rope = GlmImageRotaryPosEmbed(attention_head_dim, patch_size, theta=10000.0)
-
-        # 2. Patch & Text-timestep embedding
-        self.image_projector = GlmImageImageProjector(in_channels, inner_dim, patch_size)
-        self.glyph_projector = FeedForward(text_embed_dim, inner_dim, inner_dim=inner_dim, activation_fn="gelu")
-        self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim)
-        self.prior_projector = FeedForward(inner_dim, inner_dim, inner_dim=inner_dim, activation_fn="linear-silu")
-
-        self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
-            embedding_dim=time_embed_dim,
-            condition_dim=condition_dim,
-            pooled_projection_dim=pooled_projection_dim,
-            timesteps_dim=time_embed_dim,
-        )
-
-        # 3. Transformer blocks
-        self.transformer_blocks = nn.ModuleList(
-            [
-                GlmImageTransformerBlock(inner_dim, num_attention_heads, attention_head_dim, time_embed_dim)
-                for _ in range(num_layers)
-            ]
-        )
-
-        # 4. Output projection
-        self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, time_embed_dim, elementwise_affine=False)
-        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        prior_token_id: torch.Tensor,
-        prior_token_drop: torch.Tensor,
-        timestep: torch.LongTensor,
-        target_size: torch.Tensor,
-        crop_coords: torch.Tensor,
-        attention_kwargs: dict[str, Any] | None = None,
-        return_dict: bool = True,
-        attention_mask: torch.Tensor | None = None,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-        kv_cache: GlmImageKVCache | None = None,
-    ) -> torch.Tensor | Transformer2DModelOutput:
-        """
-        Forward pass of the GLM-Image Transformer.
-
-        Args:
-            hidden_states: Input latent tensor of shape [B, C, H, W].
-            encoder_hidden_states: Text embeddings of shape [B, S, D].
-            prior_token_id: Prior VQ token IDs.
-            prior_token_drop: Mask for dropping prior tokens (CFG).
-            timestep: Diffusion timestep.
-            target_size: Target image size for conditioning.
-            crop_coords: Crop coordinates for conditioning.
-            attention_kwargs: Additional attention arguments.
-            return_dict: Whether to return a dataclass.
-            attention_mask: Optional attention mask for text tokens.
-            image_rotary_emb: Pre-computed rotary embeddings.
-            kv_cache: Optional KV cache for image editing. When provided,
-                      the cache's mode determines behavior:
-                      - WRITE: Store KV from condition images
-                      - READ: Use cached KV during generation
-                      - SKIP: No caching (same as None)
-
-        Returns:
-            Output tensor or Transformer2DModelOutput.
-        """
-        batch_size, num_channels, height, width = hidden_states.shape
-
-        # Get KV cache mode
-        kv_cache_mode = kv_cache.mode if kv_cache is not None else None
-
-        # 1. RoPE
-        if image_rotary_emb is None:
-            image_rotary_emb = self.rope(hidden_states)
-            # Move to correct device
-            image_rotary_emb = (
-                image_rotary_emb[0].to(hidden_states.device),
-                image_rotary_emb[1].to(hidden_states.device),
-            )
-
-        # 2. Patch & Timestep embeddings
-        p = self.patch_size
-        post_patch_height = height // p
-        post_patch_width = width // p
-
-        hidden_states = self.image_projector(hidden_states)
-        encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
-
-        # Prior embedding with dropout
-        prior_embedding = self.prior_token_embedding(prior_token_id)
-        prior_embedding[prior_token_drop] *= 0.0
-        prior_hidden_states = self.prior_projector(prior_embedding)
-        hidden_states = hidden_states + prior_hidden_states
-
-        # Timestep conditioning
-        temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype)
-
-        # 3. Transformer blocks
-        for layer_idx, block in enumerate(self.transformer_blocks):
-            # Get layer-specific KV cache if available
-            layer_kv_cache = kv_cache[layer_idx] if kv_cache is not None else None
-
-            hidden_states, encoder_hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                temb,
-                image_rotary_emb,
-                attention_mask,
-                attention_kwargs,
-                kv_cache=layer_kv_cache,
-                kv_cache_mode=kv_cache_mode,
-            )
-
-        # 4. Output norm & projection
-        hidden_states = self.norm_out(hidden_states, temb)
-        hidden_states = self.proj_out(hidden_states)
-
-        # 5. Unpatchify: [B, H'*W', C*p*p] -> [B, C, H, W]
-        hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p)
-        output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
-
-        if not return_dict:
-            return (output,)
-        return Transformer2DModelOutput(sample=output)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """
-        Load weights from pretrained checkpoint.
-
-        This method handles the mapping from diffusers weight names to vllm-omni weight names,
-        especially for fused QKV projections.
-        """
-        stacked_params_mapping = [
-            # Fused QKV projection: to_q, to_k, to_v -> to_qkv
-            (".to_qkv", ".to_q", "q"),
-            (".to_qkv", ".to_k", "k"),
-            (".to_qkv", ".to_v", "v"),
-        ]
-
-        params_dict = dict(self.named_parameters())
-
-        # Also include buffers (for any beta/eps parameters)
-        for name, buffer in self.named_buffers():
-            params_dict[name] = buffer
-
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            # Handle fused QKV projections
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-
-                # Map diffusers name to vllm-omni name
-                name = name.replace(weight_name, param_name)
-
-                if name not in params_dict:
-                    logger.warning(f"Skipping weight {name} - not found in model")
-                    break
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Standard weight loading (not fused)
-                if name not in params_dict:
-                    logger.warning(f"Skipping weight {name} - not found in model")
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-
-            loaded_params.add(name)
-
-        return loaded_params
-
-    def create_kv_cache(self) -> GlmImageKVCache:
-        """
-        Create a KV cache for image editing.
-
-        Returns a new GlmImageKVCache instance sized for this model's
-        number of transformer layers. Use this for image editing workflows.
-
-        Example:
-            kv_cache = transformer.create_kv_cache()
-            kv_cache.set_mode("write")
-            transformer(condition_image, kv_cache=kv_cache)
-            kv_cache.set_mode("read")
-            for t in timesteps:
-                transformer(noisy_target, kv_cache=kv_cache)
-            kv_cache.clear()
-
-        Returns:
-            GlmImageKVCache instance with correct number of layers.
-        """
-        return GlmImageKVCache(num_layers=len(self.transformer_blocks))
-
-    @property
-    def num_layers(self) -> int:
-        """Return number of transformer layers."""
-        return len(self.transformer_blocks)
-
-    @property
-    def dtype(self) -> torch.dtype:
-        """Return dtype of model parameters."""
-        return next(self.parameters()).dtype

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -1,981 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GlmImagePipeline implementation for vLLM-Omni.
-
-This pipeline implements GLM-Image text-to-image generation with:
-- AR stage: GlmImageForConditionalGeneration generates prior tokens
-- DiT stage: GlmImageTransformer2DModel performs diffusion denoising
-- VAE: AutoencoderKL decodes latents to images
-"""
-
-from __future__ import annotations
-
-import inspect
-import json
-import logging
-import os
-import re
-from collections.abc import Iterable
-
-import numpy as np
-import PIL.Image
-import torch
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
-from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler,
-)
-from diffusers.utils.torch_utils import randn_tensor
-from torch import nn
-from transformers import (
-    ByT5Tokenizer,
-    GlmImageForConditionalGeneration,
-    GlmImageProcessor,
-    T5EncoderModel,
-)
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_cfg_group,
-    get_classifier_free_guidance_rank,
-    get_classifier_free_guidance_world_size,
-)
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
-    GlmImageKVCache,
-    GlmImageTransformer2DModel,
-)
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.model_executor.model_loader.weight_utils import (
-    download_weights_from_hf_specific,
-)
-
-logger = logging.getLogger(__name__)
-
-
-def get_glm_image_pre_process_func(od_config: OmniDiffusionConfig):
-    """Get pre-processing function for GLM-Image pipeline.
-
-    Pre-processes condition images before they are sent to the pipeline.
-    This is called by DiffusionEngine before batching requests.
-    """
-    model_name = od_config.model
-    if os.path.exists(model_name):
-        model_path = model_name
-    else:
-        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
-
-    vae_config_path = os.path.join(model_path, "vae/config.json")
-    with open(vae_config_path) as f:
-        vae_config = json.load(f)
-        block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512])
-        vae_scale_factor = 2 ** (len(block_out_channels) - 1)
-
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
-    # GLM-Image uses patch_size=2 for transformer
-    patch_size = 2
-
-    def pre_process_func(requests: list[OmniDiffusionRequest]):
-        """Pre-process condition images for Image Edit mode."""
-        for req in requests:
-            images = req.pil_image
-            if images is None:
-                # Text-to-image mode, no preprocessing needed
-                continue
-
-            if not isinstance(images, list):
-                images = [images]
-
-            preprocessed = []
-            height, width = None, None
-
-            for img in images:
-                if isinstance(img, PIL.Image.Image):
-                    img_h, img_w = img.size[::-1]  # PIL is (width, height)
-                else:
-                    img_h, img_w = img.shape[:2]
-
-                # Align to multiple of vae_scale_factor * patch_size
-                multiple_of = vae_scale_factor * patch_size
-                img_h = (img_h // multiple_of) * multiple_of
-                img_w = (img_w // multiple_of) * multiple_of
-
-                processed = image_processor.preprocess(img, height=img_h, width=img_w)
-                preprocessed.append(processed)
-
-                # Use first image dimensions as default
-                if height is None:
-                    height, width = img_h, img_w
-
-            # Store in request
-            req.preprocessed_image = preprocessed
-            req.prompt_image = images  # Keep original PIL images
-            if req.height is None:
-                req.height = height
-            if req.width is None:
-                req.width = width
-
-        return requests
-
-    return pre_process_func
-
-
-def get_glm_image_post_process_func(od_config: OmniDiffusionConfig):
-    """Get post-processing function for GLM-Image pipeline."""
-    model_name = od_config.model
-    if os.path.exists(model_name):
-        model_path = model_name
-    else:
-        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
-
-    vae_config_path = os.path.join(model_path, "vae/config.json")
-    with open(vae_config_path) as f:
-        vae_config = json.load(f)
-        block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512])
-        vae_scale_factor = 2 ** (len(block_out_channels) - 1)
-
-    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
-
-    def post_process_func(images: torch.Tensor) -> list[PIL.Image.Image]:
-        return image_processor.postprocess(images, output_type="pil")
-
-    return post_process_func
-
-
-def calculate_shift(
-    image_seq_len: int,
-    base_seq_len: int = 256,
-    base_shift: float = 0.25,
-    max_shift: float = 0.75,
-) -> float:
-    """Calculate timestep shift based on image sequence length."""
-    m = (image_seq_len / base_seq_len) ** 0.5
-    mu = m * max_shift + base_shift
-    return mu
-
-
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: int | None = None,
-    device: str | torch.device | None = None,
-    timesteps: list[int] | None = None,
-    sigmas: list[float] | None = None,
-    **kwargs,
-) -> tuple[torch.Tensor, int]:
-    """
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps.
-    Handles custom timesteps and sigmas schedules.
-    """
-    accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-    accepts_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-
-    if timesteps is not None and sigmas is not None:
-        # Both provided - check if scheduler supports both
-        if not accepts_timesteps and not accepts_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep or sigma schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif timesteps is not None:
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        if not accepts_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigma schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-
-    return timesteps, num_inference_steps
-
-
-def retrieve_latents(
-    encoder_output: torch.Tensor,
-    generator: torch.Generator | None = None,
-    sample_mode: str = "sample",
-) -> torch.Tensor:
-    """Extract latents from VAE encoder output."""
-    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
-        return encoder_output.latent_dist.sample(generator)
-    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
-        return encoder_output.latent_dist.mode()
-    elif hasattr(encoder_output, "latents"):
-        return encoder_output.latents
-    else:
-        raise AttributeError("Could not access latents of provided encoder_output")
-
-
-class GlmImagePipeline(nn.Module):
-    """
-    GLM-Image Pipeline for text-to-image and image-to-image generation.
-
-    This pipeline integrates:
-    - AR model (GlmImageForConditionalGeneration): Generates prior image tokens
-    - Text encoder (T5EncoderModel): Encodes glyph/text embeddings
-    - DiT model (GlmImageTransformer2DModel): Diffusion transformer
-    - VAE (AutoencoderKL): Encodes/decodes images to/from latent space
-
-    The pipeline flow:
-    1. AR generates prior_token_ids from text prompt
-    2. T5 encodes glyph text for text rendering
-    3. DiT performs iterative denoising conditioned on prior tokens
-    4. VAE decodes final latents to image
-    """
-
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.od_config = od_config
-        self.parallel_config = od_config.parallel_config
-        self.device = get_local_device()
-
-        model = od_config.model
-        local_files_only = os.path.exists(model)
-
-        if local_files_only:
-            model_path = model
-        else:
-            model_path = download_weights_from_hf_specific(model, od_config.revision, ["*"])
-
-        # Load scheduler
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            model_path, subfolder="scheduler", local_files_only=True
-        )
-
-        # Load AR model (vision_language_encoder)
-        logger.info("Loading GlmImageForConditionalGeneration (AR model)...")
-        self.vision_language_encoder = GlmImageForConditionalGeneration.from_pretrained(
-            model_path,
-            subfolder="vision_language_encoder",
-            local_files_only=True,
-            torch_dtype=torch.bfloat16,
-        ).to(self.device)
-        self.vision_language_encoder.eval()
-
-        # Load processor for AR model
-        self.processor = GlmImageProcessor.from_pretrained(model_path, subfolder="processor", local_files_only=True)
-
-        # Load text encoder (T5 for glyph embeddings)
-        logger.info("Loading T5EncoderModel (glyph encoder)...")
-        self.text_encoder = T5EncoderModel.from_pretrained(
-            model_path,
-            subfolder="text_encoder",
-            local_files_only=True,
-            torch_dtype=torch.bfloat16,
-        ).to(self.device)
-        self.text_encoder.eval()
-
-        # Load tokenizer for glyph encoding
-        self.tokenizer = ByT5Tokenizer.from_pretrained(model_path, subfolder="tokenizer", local_files_only=True)
-
-        # Load VAE
-        logger.info("Loading AutoencoderKL (VAE)...")
-        self.vae = AutoencoderKL.from_pretrained(
-            model_path, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16
-        ).to(self.device)
-        self.vae.eval()
-
-        # Load transformer (DiT)
-        logger.info("Loading GlmImageTransformer2DModel (DiT)...")
-        self.transformer = GlmImageTransformer2DModel(od_config=od_config)
-
-        # Weight sources for DiT loading
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder="transformer",
-                revision=od_config.revision,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            )
-        ]
-
-        # Configure scale factors
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
-        self.default_sample_size = 128
-
-        # Get transformer config for patch size
-        self._patch_size = getattr(self.transformer, "patch_size", 2)
-
-    # ==================== Input Validation ====================
-
-    def check_inputs(
-        self,
-        prompt: str | list[str] | None,
-        height: int | None,
-        width: int | None,
-        prompt_embeds: torch.Tensor | None = None,
-    ) -> None:
-        """Validate input arguments before generation."""
-        # Check dimension alignment
-        multiple_of = self.vae_scale_factor * self._patch_size
-        if height is not None and height % multiple_of != 0:
-            logger.warning(
-                f"`height` should be divisible by {multiple_of} but is {height}. "
-                "Dimensions will be adjusted accordingly."
-            )
-        if width is not None and width % multiple_of != 0:
-            logger.warning(
-                f"`width` should be divisible by {multiple_of} but is {width}. Dimensions will be adjusted accordingly."
-            )
-
-        # Check prompt/prompt_embeds mutual exclusivity
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. "
-                "Please provide only one of the two."
-            )
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Provide either `prompt` or `prompt_embeds`. Cannot leave both undefined.")
-
-        # Check prompt type
-        if prompt is not None and not isinstance(prompt, (str, list)):
-            raise ValueError(f"`prompt` must be of type `str` or `list` but is {type(prompt)}")
-
-    # ==================== AR Stage Methods ====================
-
-    @staticmethod
-    def _compute_generation_params(
-        image_grid_thw: torch.Tensor,
-        is_text_to_image: bool,
-    ) -> tuple[int, int, int, int]:
-        """
-        Compute AR generation parameters from image grid.
-
-        Args:
-            image_grid_thw: Image grid tensor of shape [N, 3] where each row is [t, h, w]
-            is_text_to_image: Whether this is text-to-image (vs image-to-image)
-
-        Returns:
-            Tuple of (max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w)
-        """
-        grid_sizes = []
-        grid_hw = []
-
-        for i in range(image_grid_thw.shape[0]):
-            t, h, w = image_grid_thw[i].tolist()
-            grid_sizes.append(int(h * w))
-            grid_hw.append((int(h), int(w)))
-
-        if not is_text_to_image:
-            # Image-to-image: only generate target image tokens
-            max_new_tokens = grid_sizes[-1] + 1
-            large_image_start_offset = 0
-            target_grid_h, target_grid_w = grid_hw[-1]
-        else:
-            # Text-to-image: generate both small preview and large target
-            total_tokens = sum(grid_sizes)
-            max_new_tokens = total_tokens + 1
-            large_image_start_offset = sum(grid_sizes[1:])
-            target_grid_h, target_grid_w = grid_hw[0]
-
-        return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w
-
-    @staticmethod
-    def _extract_large_image_tokens(
-        outputs: torch.Tensor, input_length: int, large_image_start_offset: int, large_image_tokens: int
-    ) -> torch.Tensor:
-        """Extract large image tokens from AR output."""
-        generated_tokens = outputs[0][input_length:]
-        large_image_start = large_image_start_offset
-        large_image_end = large_image_start + large_image_tokens
-        return generated_tokens[large_image_start:large_image_end]
-
-    @staticmethod
-    def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> torch.Tensor:
-        """Upsample token IDs by 2x using nearest neighbor interpolation."""
-        token_ids = token_ids.view(1, 1, token_h, token_w)
-        token_ids = torch.nn.functional.interpolate(token_ids.float(), scale_factor=2, mode="nearest").to(
-            dtype=torch.long
-        )
-        token_ids = token_ids.view(1, -1)
-        return token_ids
-
-    @torch.inference_mode()
-    def generate_prior_tokens(
-        self,
-        prompt: str,
-        height: int,
-        width: int,
-        image: list[PIL.Image.Image] | None = None,
-        factor: int = 32,
-    ) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
-        """
-        Generate prior tokens using the AR model.
-
-        Args:
-            prompt: Text prompt for generation
-            height: Target image height
-            width: Target image width
-            image: Optional condition images for image-to-image
-            factor: Token factor (default 32)
-
-        Returns:
-            Tuple of (prior_token_ids, prior_token_image_ids)
-            prior_token_image_ids is a list of tensors, one per condition image
-        """
-        device = self.vision_language_encoder.device
-        height = (height // factor) * factor
-        width = (width // factor) * factor
-        is_text_to_image = image is None or len(image) == 0
-
-        # Build message content
-        content = []
-        if image is not None:
-            for img in image:
-                content.append({"type": "image", "image": img})
-        content.append({"type": "text", "text": prompt})
-        messages = [{"role": "user", "content": content}]
-
-        # Apply chat template - processor will handle target dimensions and build grid
-        inputs = self.processor.apply_chat_template(
-            messages,
-            tokenize=True,
-            target_h=height,
-            target_w=width,
-            return_dict=True,
-            return_tensors="pt",
-        ).to(device)
-
-        image_grid_thw = inputs.get("image_grid_thw")
-
-        # Compute generation parameters from the full grid
-        max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params(
-            image_grid_thw=image_grid_thw, is_text_to_image=is_text_to_image
-        )
-
-        # Process condition images if provided
-        # Use image_grid_thw[:-1] to exclude the target image grid (last entry)
-        prior_token_image_ids = None
-        if image is not None and image_grid_thw is not None and len(image_grid_thw) > 1:
-            # Get features only for condition images (exclude target image grid)
-            condition_grid = image_grid_thw[:-1]
-            prior_token_image_embed = self.vision_language_encoder.get_image_features(
-                inputs["pixel_values"], condition_grid
-            )
-            prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0)
-            flat_prior_token_image_ids = self.vision_language_encoder.get_image_tokens(
-                prior_token_image_embed, condition_grid
-            )
-            # Split by image grid sizes and convert to list
-            split_sizes = (condition_grid.prod(dim=-1)).tolist()
-            prior_token_image_ids_list = torch.split(flat_prior_token_image_ids, split_sizes, dim=0)
-            # Convert to list with upsampling
-            prior_token_image_ids = []
-            for i, token_ids in enumerate(prior_token_image_ids_list):
-                grid_t, grid_h, grid_w = condition_grid[i].tolist()
-                token_ids = token_ids.view(1, -1)
-                # Upsample 2x (from d32 to d64)
-                token_ids_upsampled = self._upsample_token_ids(token_ids, grid_h, grid_w)
-                prior_token_image_ids.append(token_ids_upsampled)
-
-        # Generate with AR model
-        outputs = self.vision_language_encoder.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            do_sample=True,
-        )
-
-        # Extract and upsample tokens
-        large_image_tokens = token_h * token_w
-        prior_token_ids_d32 = self._extract_large_image_tokens(
-            outputs, inputs["input_ids"].shape[-1], large_image_offset, large_image_tokens
-        )
-        prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w)
-
-        return prior_token_ids, prior_token_image_ids
-
-    # ==================== Text Encoding Methods ====================
-
-    def get_glyph_texts(self, prompt: str | list[str]) -> list[str]:
-        """Extract text within quotes for glyph rendering."""
-        prompt = prompt[0] if isinstance(prompt, list) else prompt
-        ocr_texts = (
-            re.findall(r"'([^']*)'", prompt)
-            + re.findall(r"“([^“”]*)”", prompt)
-            + re.findall(r'"([^"]*)"', prompt)
-            + re.findall(r"「([^「」]*)」", prompt)
-        )
-        return ocr_texts
-
-    def _get_glyph_embeds(
-        self,
-        prompt: str | list[str],
-        max_sequence_length: int = 2048,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> torch.Tensor:
-        """Get glyph embeddings from T5 encoder for text rendering."""
-        device = device or self.device
-        dtype = dtype or self.text_encoder.dtype
-
-        glyph_texts = self.get_glyph_texts(prompt)
-        input_ids = self.tokenizer(
-            glyph_texts if len(glyph_texts) > 0 else [""],
-            max_length=max_sequence_length,
-            truncation=True,
-        ).input_ids
-
-        # Pad to even length
-        input_ids = [[self.tokenizer.pad_token_id] * ((len(ids) + 1) % 2) + ids for ids in input_ids]
-        max_length = max(len(ids) for ids in input_ids)
-
-        attention_mask = torch.tensor(
-            [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids],
-            device=device,
-        )
-        input_ids = torch.tensor(
-            [ids + [self.tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids],
-            device=device,
-        )
-
-        outputs = self.text_encoder(input_ids, attention_mask=attention_mask)
-        glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0)
-
-        return glyph_embeds.to(device=device, dtype=dtype)
-
-    def encode_prompt(
-        self,
-        prompt: str | list[str],
-        do_classifier_free_guidance: bool = True,
-        num_images_per_prompt: int = 1,
-        prompt_embeds: torch.Tensor | None = None,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-        max_sequence_length: int = 2048,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Encode prompt into glyph embeddings for text rendering."""
-        device = device or self.device
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            prompt_embeds = self._get_glyph_embeds(prompt, max_sequence_length, device, dtype)
-
-        seq_len = prompt_embeds.size(1)
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        negative_prompt_embeds = None
-        if do_classifier_free_guidance:
-            negative_prompt = [""] * batch_size
-            negative_prompt_embeds = self._get_glyph_embeds(negative_prompt, max_sequence_length, device, dtype)
-            seq_len = negative_prompt_embeds.size(1)
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        return prompt_embeds, negative_prompt_embeds
-
-    # ==================== Latent Preparation ====================
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: int,
-        height: int,
-        width: int,
-        dtype: torch.dtype,
-        device: torch.device,
-        generator: torch.Generator | None,
-        latents: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """Prepare random noise latents."""
-        if latents is not None:
-            return latents.to(device)
-
-        shape = (
-            batch_size,
-            num_channels_latents,
-            int(height) // self.vae_scale_factor,
-            int(width) // self.vae_scale_factor,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(f"Passed {len(generator)} generators but batch size is {batch_size}.")
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        return latents
-
-    def diffuse(
-        self,
-        latents: torch.Tensor,
-        prior_token_id: torch.Tensor,
-        prompt_embeds: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor | None,
-        timesteps: torch.Tensor,
-        target_size: torch.Tensor,
-        crop_coords: torch.Tensor,
-        guidance_scale: float,
-        do_classifier_free_guidance: bool,
-        kv_caches: GlmImageKVCache | None = None,
-    ) -> torch.Tensor:
-        """
-        Denoising loop for diffusion process with CFG-Parallel support.
-
-        Args:
-            latents: Initial noise latents
-            prior_token_id: Prior tokens generated by AR model
-            prompt_embeds: Encoded positive prompt embeddings (glyph embeddings)
-            negative_prompt_embeds: Encoded negative prompt embeddings
-            timesteps: Denoising timesteps
-            target_size: Target image size tensor [[height, width]]
-            crop_coords: Crop coordinates tensor
-            guidance_scale: CFG scale
-            do_classifier_free_guidance: Whether to apply CFG
-            kv_caches: Optional KV cache for Image Edit mode
-
-        Returns:
-            Denoised latents ready for VAE decode
-        """
-        # Prepare conditional/unconditional drop flags
-        prior_token_drop_cond = torch.full_like(prior_token_id, False, dtype=torch.bool)
-        prior_token_drop_uncond = torch.full_like(prior_token_id, True, dtype=torch.bool)
-
-        transformer_dtype = self.transformer.dtype
-
-        # Enable CFG-parallel: rank0 computes positive, rank1 computes negative
-        cfg_parallel_ready = do_classifier_free_guidance and get_classifier_free_guidance_world_size() > 1
-
-        for i, t in enumerate(timesteps):
-            latent_model_input = latents.to(transformer_dtype)
-            timestep = t.expand(latents.shape[0]) - 1
-
-            if cfg_parallel_ready:
-                cfg_group = get_cfg_group()
-                cfg_rank = get_classifier_free_guidance_rank()
-
-                if cfg_rank == 0:
-                    # Rank 0: Compute positive (conditional) prediction
-                    local_pred = self.transformer(
-                        hidden_states=latent_model_input,
-                        encoder_hidden_states=prompt_embeds,
-                        prior_token_id=prior_token_id,
-                        prior_token_drop=prior_token_drop_cond,
-                        timestep=timestep,
-                        target_size=target_size,
-                        crop_coords=crop_coords,
-                        kv_cache=kv_caches,
-                        return_dict=False,
-                    )[0].float()
-                else:
-                    # Rank 1: Compute negative (unconditional) prediction
-                    local_pred = self.transformer(
-                        hidden_states=latent_model_input,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        prior_token_id=prior_token_id,
-                        prior_token_drop=prior_token_drop_uncond,
-                        timestep=timestep,
-                        target_size=target_size,
-                        crop_coords=crop_coords,
-                        kv_cache=kv_caches,
-                        return_dict=False,
-                    )[0].float()
-
-                # All-gather predictions from all ranks
-                gathered = cfg_group.all_gather(local_pred, separate_tensors=True)
-
-                if cfg_rank == 0:
-                    # Rank 0: Combine predictions and apply CFG
-                    noise_pred_cond = gathered[0]
-                    noise_pred_uncond = gathered[1]
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
-                    # Scheduler step
-                    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-
-                # Broadcast updated latents to all ranks
-                cfg_group.broadcast(latents, src=0)
-
-            else:
-                # Sequential CFG (single GPU or no CFG)
-                # Conditional forward pass
-                noise_pred_cond = self.transformer(
-                    hidden_states=latent_model_input,
-                    encoder_hidden_states=prompt_embeds,
-                    prior_token_id=prior_token_id,
-                    prior_token_drop=prior_token_drop_cond,
-                    timestep=timestep,
-                    target_size=target_size,
-                    crop_coords=crop_coords,
-                    kv_cache=kv_caches,
-                    return_dict=False,
-                )[0].float()
-
-                if do_classifier_free_guidance:
-                    # Unconditional forward pass
-                    noise_pred_uncond = self.transformer(
-                        hidden_states=latent_model_input,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        prior_token_id=prior_token_id,
-                        prior_token_drop=prior_token_drop_uncond,
-                        timestep=timestep,
-                        target_size=target_size,
-                        crop_coords=crop_coords,
-                        kv_cache=kv_caches,
-                        return_dict=False,
-                    )[0].float()
-
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
-                else:
-                    noise_pred = noise_pred_cond
-
-                # Scheduler step
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-
-        return latents
-
-    # ==================== Main Forward Pass ====================
-
-    def _prepare_condition_image_kv_cache(
-        self,
-        condition_images: list[torch.Tensor],
-        prior_token_image_ids: list[torch.Tensor],
-        prompt_embeds: torch.Tensor,
-        generator: torch.Generator | None = None,
-    ) -> GlmImageKVCache:
-        """
-        Prepare KV cache by running condition images through transformer at timestep 0.
-
-        This is used for Image Edit mode where we need to cache the condition image's
-        KV states for cross-attention during denoising.
-
-        Args:
-            condition_images: List of preprocessed condition images
-            prior_token_image_ids: Prior token IDs for each condition image from AR model
-            prompt_embeds: Prompt embeddings (used to get dtype)
-            generator: Optional random generator
-
-        Returns:
-            GlmImageKVCache with cached KV states from condition images
-        """
-        kv_caches = self.transformer.create_kv_cache()
-        kv_caches.set_mode("write")
-
-        # Prepare VAE normalization parameters
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.latent_channels, 1, 1)
-            .to(device=self.device, dtype=prompt_embeds.dtype)
-        )
-        latents_std = (
-            torch.tensor(self.vae.config.latents_std)
-            .view(1, self.vae.config.latent_channels, 1, 1)
-            .to(device=self.device, dtype=prompt_embeds.dtype)
-        )
-
-        # Process each condition image through transformer to populate KV cache
-        for condition_image, condition_prior_token_id in zip(condition_images, prior_token_image_ids):
-            condition_image = condition_image.to(device=self.device, dtype=prompt_embeds.dtype)
-
-            # Encode condition image to latent space
-            # Use argmax (mode) for deterministic encoding of condition images
-            condition_latent = retrieve_latents(
-                self.vae.encode(condition_image), generator=generator, sample_mode="argmax"
-            )
-            condition_latent = (condition_latent - latents_mean) / latents_std
-
-            # Run forward pass at timestep 0 to cache KV states
-            # Empty encoder_hidden_states since we only want to cache image features
-            _ = self.transformer(
-                hidden_states=condition_latent,
-                encoder_hidden_states=torch.zeros_like(prompt_embeds)[:1, :0, ...],
-                prior_token_id=condition_prior_token_id,
-                prior_token_drop=torch.full_like(condition_prior_token_id, False, dtype=torch.bool),
-                timestep=torch.zeros((1,), device=self.device),
-                target_size=torch.tensor([condition_image.shape[-2:]], device=self.device, dtype=prompt_embeds.dtype),
-                crop_coords=torch.zeros((1, 2), device=self.device, dtype=prompt_embeds.dtype),
-                kv_cache=kv_caches,
-                return_dict=False,
-            )
-
-        return kv_caches
-
-    @torch.inference_mode()
-    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
-        """
-        Main generation forward pass.
-
-        Args:
-            req: OmniDiffusionRequest with generation parameters
-
-        Returns:
-            DiffusionOutput containing generated image
-        """
-        prompt = req.prompt or ""
-        if isinstance(prompt, list):
-            prompt = prompt[0] if prompt else ""
-
-        # Get pre-computed prompt embeddings if provided
-        prompt_embeds = req.prompt_embeds if isinstance(req.prompt_embeds, torch.Tensor) else None
-
-        # Get condition images for Image Edit mode
-        # Use pre-processed images from pre_process_func
-        preprocessed_images = req.preprocessed_image
-        condition_images = getattr(req, "prompt_image", None)
-        img_height = req.height
-        img_width = req.width
-
-        is_image_edit = preprocessed_images is not None
-
-        # Use image dimensions as default if available
-        height = req.height or img_height or self.default_sample_size * self.vae_scale_factor
-        width = req.width or img_width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.num_inference_steps or 50
-        guidance_scale = req.guidance_scale or 1.5
-
-        # 0. Validate inputs
-        self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds)
-
-        batch_size = 1
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # Set seed if provided
-        generator = None
-        if req.seed is not None:
-            generator = torch.Generator(device=self.device).manual_seed(req.seed)
-
-        # 1. Get prior tokens - either from external source (multistage) or generate internally
-        # Check if prior_token_ids are provided externally (from AR stage in multistage mode)
-        external_prior_tokens = req.extra.get("prior_token_ids") if req.extra else None
-        external_prior_image_ids = req.extra.get("prior_token_image_ids") if req.extra else None
-
-        if external_prior_tokens is not None:
-            # Multistage mode: use externally provided prior tokens from vLLM AR stage
-            logger.info("Using externally provided prior tokens from AR stage...")
-            prior_token_id = external_prior_tokens
-            if isinstance(prior_token_id, list):
-                prior_token_id = torch.tensor(prior_token_id, dtype=torch.long, device=self.device)
-            elif isinstance(prior_token_id, torch.Tensor):
-                prior_token_id = prior_token_id.to(device=self.device, dtype=torch.long)
-            # Ensure shape is [1, num_tokens] for batch processing
-            if prior_token_id.dim() == 1:
-                prior_token_id = prior_token_id.unsqueeze(0)
-            prior_token_image_ids = external_prior_image_ids
-        else:
-            # Single-stage mode: generate prior tokens with internal AR model
-            logger.info("Generating prior tokens with AR model...")
-            prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
-                prompt=prompt,
-                image=condition_images,
-                height=height,
-                width=width,
-            )
-
-        # 2. Encode prompt for glyph embeddings
-        logger.info("Encoding prompt...")
-        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-            prompt,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            num_images_per_prompt=1,
-            prompt_embeds=prompt_embeds,
-            device=self.device,
-            dtype=self.transformer.dtype,
-        )
-
-        # 3. Prepare KV cache for Image Edit mode
-        kv_caches = None
-        if is_image_edit and prior_token_image_ids is not None:
-            logger.info("Preparing KV cache for Image Edit mode...")
-            kv_caches = self._prepare_condition_image_kv_cache(
-                condition_images=preprocessed_images,
-                prior_token_image_ids=prior_token_image_ids,
-                prompt_embeds=prompt_embeds,
-                generator=generator,
-            )
-            # Switch to read mode for denoising
-            kv_caches.set_mode("read")
-
-        # 4. Prepare latents
-        latent_channels = self.transformer.in_channels
-        latents = self.prepare_latents(
-            batch_size=batch_size,
-            num_channels_latents=latent_channels,
-            height=height,
-            width=width,
-            dtype=prompt_embeds.dtype,
-            device=self.device,
-            generator=generator,
-        )
-
-        # 5. Prepare timesteps
-        image_seq_len = ((height // self.vae_scale_factor) * (width // self.vae_scale_factor)) // (self._patch_size**2)
-        timesteps_array = np.linspace(self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1)[:-1]
-        timesteps_array = timesteps_array.astype(np.int64).astype(np.float32)
-        sigmas = timesteps_array / self.scheduler.config.num_train_timesteps
-
-        mu = calculate_shift(
-            image_seq_len,
-            self.scheduler.config.get("base_image_seq_len", 256),
-            self.scheduler.config.get("base_shift", 0.25),
-            self.scheduler.config.get("max_shift", 0.75),
-        )
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler, num_inference_steps, self.device, timesteps_array.tolist(), sigmas.tolist(), mu=mu
-        )
-
-        # 6. Prepare conditioning tensors
-        target_size = torch.tensor([[height, width]], dtype=prompt_embeds.dtype, device=self.device)
-        crop_coords = torch.zeros((1, 2), dtype=prompt_embeds.dtype, device=self.device)
-
-        # 7. Denoising loop with CFG-parallel support
-        logger.info(f"Starting denoising loop with {num_inference_steps} steps...")
-        latents = self.diffuse(
-            latents=latents,
-            prior_token_id=prior_token_id,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            timesteps=timesteps,
-            target_size=target_size,
-            crop_coords=crop_coords,
-            guidance_scale=guidance_scale,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            kv_caches=kv_caches,
-        )
-
-        # 8. VAE decode
-        logger.info("Decoding latents with VAE...")
-        latents = latents.to(self.vae.dtype)
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.latent_channels, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents_std = (
-            torch.tensor(self.vae.config.latents_std)
-            .view(1, self.vae.config.latent_channels, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents = latents * latents_std + latents_mean
-        image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
-
-        # 9. Leave post-process to vllm-omni pipeline
-
-        return DiffusionOutput(output=image)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load transformer weights."""
-        # Filter weights for transformer only
-        transformer_weights = (
-            (name.replace("transformer.", ""), weight) for name, weight in weights if name.startswith("transformer.")
-        )
-        return self.transformer.load_weights(transformer_weights)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -642,8 +642,8 @@ class QwenImagePipeline(
     def forward(
         self,
         req: OmniDiffusionRequest,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
+        prompt: str | list[str] = "",
+        negative_prompt: str | list[str] = "",
         true_cfg_scale: float = 4.0,
         height: int | None = None,
         width: int | None = None,
@@ -665,8 +665,8 @@ class QwenImagePipeline(
         # # TODO: only support single prompt now
         # if req.prompt is not None:
         #     prompt = req.prompt[0] if isinstance(req.prompt, list) else req.prompt
-        prompt = req.prompt
-        negative_prompt = req.negative_prompt
+        prompt = req.prompt if req.prompt is not None else prompt
+        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
         height = req.height or self.default_sample_size * self.vae_scale_factor
         width = req.width or self.default_sample_size * self.vae_scale_factor
         num_inference_steps = req.num_inference_steps or num_inference_steps

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -695,8 +695,8 @@ class QwenImageEditPipeline(
     def forward(
         self,
         req: OmniDiffusionRequest,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
+        prompt: str | list[str] = "",
+        negative_prompt: str | list[str] = "",
         image: PIL.Image.Image | torch.Tensor | None = None,
         true_cfg_scale: float = 4.0,
         height: int | None = None,
@@ -717,8 +717,8 @@ class QwenImageEditPipeline(
         max_sequence_length: int = 512,
     ) -> DiffusionOutput:
         """Forward pass for image editing."""
-        prompt = req.prompt
-        negative_prompt = req.negative_prompt
+        prompt = req.prompt if req.prompt is not None else prompt
+        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
 
         # Get preprocessed image from request (pre-processing is done in DiffusionEngine)
         if hasattr(req, "preprocessed_image"):