docs: restore BitsAndBytes navigation

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
docs: remove GGUF from active navigation
2026-08-12 00:10:38 +08:00 · 2026-08-11 17:22:19 +08:00 · 2026-08-11 17:07:49 +08:00 · 2026-08-11 16:55:39 +08:00 · 2026-08-11 16:16:46 +08:00
48 changed files with 375 additions and 1657 deletions
@@ -345,7 +345,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
@@ -115,7 +115,7 @@ steps:
          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -40,30 +40,12 @@ nav:
  - Models:
    - models/supported_models.md
  - Features:
-    - Sleep Mode: features/sleep_mode.md
+    - Overview: features/README.md
-    - Diffusion Features:
+    - Runtime and Stage Execution:
-      - Overview: user_guide/diffusion_features.md
+      - Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
-      - Feature Compatibility: user_guide/feature_compatibility.md
+      - Sleep Mode: features/sleep_mode.md
      - Regional Compilation: user_guide/diffusion/regional_compilation.md
      - Cache Acceleration:
        - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
        - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
      - Attention Backends: user_guide/diffusion/attention_backends.md
      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
      - Parallelism:
        - Overview: user_guide/diffusion/parallelism/overview.md
        - CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
        - Expert Parallel: user_guide/diffusion/parallelism/expert_parallel.md
        - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
        - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
        - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
        - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
      - CPU Offloading: user_guide/diffusion/cpu_offload.md
      - LoRA: user_guide/diffusion/lora.md
      - Custom Pipeline: features/custom_pipeline.md
      - Execution Modes: user_guide/diffusion/execution_modes.md
    - Quantization:
-      - Quantization: user_guide/quantization/overview.md
+      - Overview: user_guide/quantization/overview.md
      - Online Quantization: user_guide/quantization/online.md
      - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
      - FP8 W8A8: user_guide/quantization/fp8.md
@@ -72,9 +54,33 @@ nav:
      - MXFP8 W8A8: user_guide/quantization/mxfp8.md
      - MXFP4 W4A4: user_guide/quantization/mxfp4.md
      - ModelOpt: user_guide/quantization/modelopt.md
      - GGUF: user_guide/quantization/gguf.md
      - AutoRound: user_guide/quantization/autoround.md
      - msModelSlim: user_guide/quantization/msmodelslim.md
    - Diffusion Acceleration:
      - Overview and Compatibility:
        - Overview: user_guide/diffusion_features.md
        - Feature Compatibility: user_guide/feature_compatibility.md
      - CPU Offloading: user_guide/diffusion/cpu_offload.md
      - Cache Acceleration:
        - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
        - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
      - Parallelism:
        - Overview: user_guide/diffusion/parallelism/overview.md
        - CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
        - Expert Parallel: user_guide/diffusion/parallelism/expert_parallel.md
        - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
        - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
        - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
        - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
      - Attention Backends: user_guide/diffusion/attention_backends.md
      - Regional Compilation: user_guide/diffusion/regional_compilation.md
      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
      - Startup and Loading: user_guide/diffusion/startup_and_loading.md
      - LoRA: user_guide/diffusion/lora.md
    - Experimental:
      - Session State Manager: features/session_state_manager.md
  - Integrations:
    - Custom Diffusion Pipeline: features/custom_pipeline.md
    - ComfyUI: features/comfyui.md
  - Fault Tolerance:
    - Failure Modes: user_guide/fault_injection_reliability_matrix.md
@@ -108,6 +114,7 @@ nav:
          - Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md
          - Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md
          - Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
      - Quantization: design/feature/quantization.md
      - Diffusion Acceleration:
        - Parallelism:
          - CFG-Parallel: design/feature/cfg_parallel.md
@@ -117,9 +124,8 @@ nav:
          - Sequence Parallel: design/feature/sequence_parallel.md
          - Tensor Parallel: design/feature/tensor_parallel.md
          - VAE Patch Parallelism: design/feature/vae_parallel.md
-        - Attention Backends:
+        - Attention Optimization:
          - Skip-Softmax: design/feature/skip_softmax.md
        - Quantization: design/feature/quantization.md
        - Cache-DiT: design/feature/cache_dit.md
        - TeaCache: design/feature/teacache.md
        - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
@@ -11,6 +11,10 @@ in the active navigation.
 ## Feature Design Documents
 For user-facing configuration and current compatibility, see the
 [Features overview](../features/README.md). A design document defines an
 implementation contract; it is not, by itself, a general support claim.
 ### Runtime and stage execution
 - [Disaggregated Inference](feature/disaggregated_inference.md)
@@ -30,6 +34,10 @@ in the active navigation.
 - [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md)
 - [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md)
 ### Quantization
 - [Quantization](feature/quantization.md)
 ### Diffusion acceleration
 #### Parallelism
@@ -42,14 +50,14 @@ in the active navigation.
 - [Tensor Parallel](feature/tensor_parallel.md)
 - [VAE Patch Parallelism](feature/vae_parallel.md)
-#### Attention Backends
+#### Attention optimization
 The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md)
 guide lists all selectable backends and their platform defaults. Skip-Softmax
 is the backend-related optimization with a standalone feature-design contract:
 - [Skip-Softmax](feature/skip_softmax.md)
 #### Quantization
 - [Quantization](feature/quantization.md)
 - [Cache-DiT](feature/cache_dit.md)
 - [TeaCache](feature/teacache.md)
 - [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
@@ -0,0 +1,67 @@
 # Features
 Use this section to choose and configure vLLM-Omni runtime and optimization
 features. The navigation follows the active
 [Feature Design](../design/index.md#feature-design-documents) taxonomy whenever
 an implementation contract has a user-facing workflow.
 !!! note "User guides, design documents, and recipes"
    User guides explain how to enable a feature and where it is supported.
    Design documents define internal contracts and do not, by themselves,
    imply general or production support. For model-specific launch flags and
    hardware requirements, use the linked examples and validated recipes from
    [Supported Models](../models/supported_models.md).
 ## Runtime and Stage Execution
 | Goal | User guide | Related design contract |
 | --- | --- | --- |
 | Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
 | Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
 Some runtime designs are deliberately not promoted as standalone User Guide
 features yet:
 - [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
  [OmniConnector implementations](../design/index.md#communication) describe
  topology and transport contracts. Practical deployment configuration remains
  under [Pipeline and deploy configurations](../configuration/stage_configs.md).
 - [Async Chunk](../design/feature/async_chunk.md) and
  [Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
  are model- and pipeline-dependent. Use the selected model's deploy
  configuration and recipe for the supported settings.
 - [Automatic Prefix Caching](../design/feature/prefix_caching.md) remains a
  design-level contract until its user-facing configuration and compatibility
  surface is consolidated.
 ## Quantization
 Quantization is a cross-model feature rather than a diffusion-only
 optimization. The unified [`quantization_config` guide](../user_guide/quantization/overview.md)
 covers diffusion-only models, multi-stage omni/TTS models, and multi-stage
 diffusion models. Its [design contract](../design/feature/quantization.md)
 defines the shared configuration and backend extension points.
 ## Diffusion Acceleration
 | Goal | User guide | Related design contract |
 | --- | --- | --- |
 | Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) |
 | Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) |
 | Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) |
 | Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
 | Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) |
 | Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
 | Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
 | Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
 | Apply diffusion adapters | [LoRA](../user_guide/diffusion/lora.md) | User-facing extension guide |
 The [Pipeline Parallelism guide](../user_guide/diffusion/parallelism/pipeline_parallel.md)
 remains available by direct link, but it is not promoted in the primary User
 Guide navigation until its user-facing support and placement are settled.
 ## Experimental
 [Session State Manager](session_state_manager.md) is opt-in and experimental.
 Its APIs and compatibility may change without notice.
@@ -1,18 +1,14 @@
 # LoRA (Low-Rank Adaptation) Guide
-LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni supports two LoRA backends: **PEFT** for PEFT-style adapters, and **Distill** usually for few-steps inference. **PEFT** backend allows you to customize model behavior without modifying the base model weights. **Distill** backend fuses LoRA weights into the base model at initialization.
+LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni currently supports PEFT-style LoRA adapters, allowing you to customize model behavior without modifying the base model weights.
 ## Overview
-LoRA adapters are lightweight, model-specific fine-tuning weights that can be applied to diffusion models in two ways:
+LoRA adapters are lightweight, model-specific fine-tuning weights that can be dynamically loaded and applied to diffusion models. vLLM-Omni uses a unified LoRA handling mechanism similar to vLLM with LRU cache management.
 - **PEFT backend** (`--lora-backend peft`, default): Loads a PEFT-format adapter folder via `DiffusionLoRAManager`. Adapters are cached (LRU) and activated per request via `LoRARequest`. It uses a unified LoRA handling mechanisms similar to vLLM with LRU cache management.
 - **Distill backend** (`--lora-backend distill`): Calls `pipeline.load_lora_weights` once at initialization to fuse one or more concrete checkpoint files directly into the base weights. Typically used for distilled few-step LoRAs (e.g. Lightning, LightX2V).
 ## LoRA Adapter Format
-### PEFT (Parameter-Efficient Fine-Tuning) format (default)
+LoRA adapters must be in **PEFT (Parameter-Efficient Fine-Tuning)** format. A typical LoRA adapter directory structure:
 A typical PEFT-format LoRA adapter directory structure:
 ```
 lora_adapter/
@@ -29,20 +25,19 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi
 ### Offline Inference
-#### PEFT backend: pre-loaded LoRA
+#### Pre-loaded LoRA
-Load a PEFT-format LoRA adapter at initialization. The adapter is pre-loaded into the cache and can be activated per request:
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
 ```python
 from vllm_omni import Omni
 from vllm_omni.lora.request import LoRARequest
-lora_path = "/path/to/lora_adapter"
+lora_path="/path/to/lora_adapter"
 omni = Omni(
    model="stabilityai/stable-diffusion-3.5-medium",
-    lora_path=lora_path,
+    lora_path=lora_path
    lora_backend="peft",  # default, can be omitted
 )
 lora_request = LoRARequest(
@@ -58,98 +53,6 @@ outputs = omni.generate(
 )
 ```
 #### Distill backend: fuse distilled LoRA at init
 For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or more concrete `.safetensors` files. The weights are fused into the base model once at init; subsequent `generate()` calls do not need a `LoRARequest`.
 ##### Supported pipelines
 For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization.
 | Pipeline | Supported distilled LoRA repo | Notes |
 |----------|--------------------------------|-------|
 | `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. |
 | `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
 | `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
 Other diffusion pipelines are not currently listed as supporting distilled LoRA. Use the PEFT backend for request-time adapters, or bake converted weights into a local Diffusers directory before serving.
 Single-file example (Qwen-Image-Lightning):
 ```python
 from vllm_omni import Omni
 omni = Omni(
    model="Qwen/Qwen-Image-2512",
    lora_path="/path/to/Qwen-Image-2512-Lightning.safetensors",
    lora_backend="distill",
 )
 outputs = omni.generate(prompt="A piece of cheesecake")
 ```
 Multi-file example (Wan2.2 MoE, high + low noise):
 ```python
 from vllm_omni import Omni
 omni = Omni(
    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
    lora_path=[
        "/path/to/wan2.2_high_noise_lora.safetensors",   # -> transformer
        "/path/to/wan2.2_low_noise_lora.safetensors",    # -> transformer_2
    ],
    lora_backend="distill",
 )
 ```
 The CLI examples under `examples/offline_inference/` accept the same flags, e.g.:
 ```bash
 python examples/offline_inference/text_to_video/text_to_video.py \
  --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
  --lora-backend distill \
  --lora-path /path/to/high.safetensors /path/to/low.safetensors \
  --prompt "A cat playing with yarn"
 ```
 ##### Online serving
 Distilled LoRAs can also be fused when an online diffusion server starts. The
 adapter remains active for the lifetime of that server, so requests do not pass
 a per-request `lora` field.
 ```bash
 vllm serve Qwen/Qwen-Image-2512 \
  --omni \
  --port 8091 \
  --lora-backend distill \
  --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors
 ```
 Then send a normal generation request with the sampling settings expected by
 the distilled checkpoint:
 ```bash
 curl -s http://localhost:8091/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [
      {"role": "user", "content": "A Chinese female college student, around 20 years old, with a very short haircut that conveys a gentle, artistic vibe. Her hair naturally falls to partially cover her cheeks, projecting a tomboyish yet charming demeanor. She has cool-toned fair skin and delicate features, with a slightly shy yet subtly confident expression—her mouth crooked in a playful, youthful smirk. She wears an off-shoulder top, revealing one shoulder, with a well-proportioned figure. The image is framed as a close-up selfie: she dominates the foreground, while the background clearly shows her dormitory—a neatly made bed with white linens on the top bunk, a tidy study desk with organized stationery, and wooden cabinets and drawers. The photo is captured on a smartphone under soft, even ambient lighting, with natural tones, high clarity, and a bright, lively atmosphere full of youthful, everyday energy."}
    ],
    "extra_body": {
      "height": 1024,
      "width": 1024,
      "num_inference_steps": 4,
      "true_cfg_scale": 1.0,
      "seed": 42
    }
  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > qwen_image_distill_lora_output.png
 ```
 For Wan2.2 MoE serving, pass the high-noise checkpoint first and the low-noise
 checkpoint second to `--lora-path`.
 !!! note "Server-side Path Requirement"
    The LoRA adapter path (`local_path`) must be readable on the **server** machine. If your client and server are on different machines, ensure the LoRA adapter is accessible via a shared mount or copied to the server.
@@ -238,7 +141,6 @@ Notes:
 - This route avoids runtime LoRA loading changes in vLLM-Omni when you choose to bake converted weights into a local Diffusers directory.
 - Output quality and speed depend on the replacement checkpoints and sampling params you choose.
 - If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path <high>.safetensors <low>.safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above.
 ## See Also
@@ -0,0 +1,60 @@
 # Diffusion Startup and Loading
 Large diffusion models can take several minutes to load at startup. vLLM-Omni
 loads safetensors shards in parallel to reduce this initialization time.
 Multi-thread weight loading is enabled by default with four threads. No
 configuration is needed for the default behavior.
 ## Configuration
 | Parameter | CLI flag | Default | Description |
 | --- | --- | --- | --- |
 | `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` | Pass the flag to disable multi-thread loading |
 | `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of parallel weight-loading threads |
 !!! tip
    The default balances startup speed and disk I/O contention. Fast NVMe
    storage may benefit from more threads, while network storage or hard disks
    may not.
 ## Online Serving
 ```bash
 # Default: multi-thread loading with four threads
 vllm serve Qwen/Qwen-Image --omni --port 8091
 # Increase the thread count
 vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni \
  --num-weight-load-threads 8
 # Disable multi-thread loading
 vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
 ```
 ## Offline Inference
 ```python
 from vllm_omni import Omni
 # Default: multi-thread loading with four threads
 omni = Omni(model="Qwen/Qwen-Image")
 # Increase the thread count
 omni = Omni(
    model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
    num_weight_load_threads=8,
 )
 ```
 ## Reference Benchmarks
 The following measurements were collected on NVIDIA H800 hardware. Treat them
 as reference results rather than a guarantee for other storage or hardware
 configurations.
 | Model | Sequential loading | Multi-thread loading | Speedup |
 | --- | ---: | ---: | ---: |
 | **Qwen/Qwen-Image** (53.7 GiB) | 168 s | 27 s | **6.2x** |
 | **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283 s | 56 s | **5.1x** |
@@ -49,7 +49,7 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
 | Method | Description | Best For |
 |--------|-------------|----------|
-| **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
+| **[Multi-Thread Weight Loading](diffusion/startup_and_loading.md)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
 **Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
@@ -100,7 +100,9 @@ The following tables show which models support each feature:
 - **🔀SP (Ulysses & Ring)**: Includes both Ulysses-SP and Ring-Attention methods
 - ✅ = Fully supported
 - ✅* = Supported with the constraint listed below the table
 - ❌ = Not supported
 - ❓ = Not verified; not recommended
 > Notes:
@@ -119,7 +121,7 @@ The following tables show which models support each feature:
 | **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ✅           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ✅*        |
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -145,6 +147,7 @@ The following tables show which models support each feature:
 > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
 > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
 > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
 ### VideoGen
@@ -155,13 +158,16 @@ The following tables show which models support each feature:
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **LongCat-Video-Avatar-1.5** |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |
 > **Step execution note:** Helios supports single-request step execution only;
 > use `max_num_seqs=1`.
 **Frame Interpolation Support**
 - **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
@@ -190,7 +196,7 @@ The following tables show which models support each feature:
 | **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | | |
 | **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
 | **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
-| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | | |
+| **🔀HSDP** | ❓ | ❓ | ✅ | ❓ | ✅ | ❌ | | | | | | | | |
 | **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | | |
 | **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | | |
 | **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | | |
@@ -201,97 +207,43 @@ The following tables show which models support each feature:
 !!! info
-    1. Tensor Parallel and HSDP are not compatible.
+    1. HSDP can be combined with Ulysses-SP or CFG-Parallel. Tensor Parallel
       and HSDP are not compatible; other HSDP combinations in the table
       remain unverified.
    2. TeaCache and Cache-DiT are not compatible.
    3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
-    4. CPU Offloading (Layerwise) supports single-card for now.
+    4. The CPU Offloading (Layerwise) row describes local layerwise offload.
-    5. Using FP8-Quant as an example of qunatization methods.
+       Multi-device Distributed Layerwise Offload has a separate topology and
       compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading).
    5. The compatibility matrix uses FP8 as the representative quantization method.
    6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
 ## Multi-Thread Weight Loading
-Large diffusion models can take several minutes to load weights at startup (e.g., ~3 min for Qwen-Image, ~5 min for Wan2.2 I2V 14B). Multi-thread weight loading speeds up this process by loading safetensors shards in parallel using a thread pool instead of sequentially.
+The loading guide now lives at [Diffusion Startup and
-
+Loading](diffusion/startup_and_loading.md). This heading remains so existing
-This optimization is **enabled by default** with 4 threads. No configuration is needed for the default behavior.
+links to this section continue to work.
 ### Configuration
 | Parameter | CLI Flag | Default | Description |
 |-----------|----------|---------|-------------|
 | `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` (enabled) | Pass the flag to disable multi-thread loading |
 | `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of threads for parallel weight loading |
 !!! tip
    The default of 4 threads balances speed and disk I/O contention. On fast NVMe storage you may benefit from more threads (e.g., 8). On HDD or network storage, the default of 4 avoids saturating I/O bandwidth.
 ### Online Serving
 ```bash
 # Default (multi-thread enabled, 4 threads)
 vllm serve Qwen/Qwen-Image --omni --port 8091
 # Custom thread count
 vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --num-weight-load-threads 8
 # Disable multi-thread loading
 vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
 ```
 ### Offline Inference
 ```python
 from vllm_omni import Omni
 # Default (multi-thread enabled, 4 threads)
 omni = Omni(model="Qwen/Qwen-Image")
 # Custom thread count
 omni = Omni(
    model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
    num_weight_load_threads=8,
 )
 ```
 ### Benchmarks
 Measured on NVIDIA H800:
 | Model | Before | After | Speedup |
 |-------|--------|-------|---------|
 | **Qwen/Qwen-Image** (53.7 GiB) | 168s | 27s | **6.2x** |
 | **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283s | 56s | **5.1x** |
 ## Learn More
-**Cache Acceleration:**
+The Diffusion Acceleration navigation groups the remaining guides as follows:
- **[TeaCache Configuration Guide](diffusion/cache_acceleration/teacache.md)** - Parameter tuning, performance tips, troubleshooting
+| Area | Guide |
- **[Cache-DiT Advanced Guide](diffusion/cache_acceleration/cache_dit.md)** - DBCache, TaylorSeer, SCM techniques and optimization
+| --- | --- |
 | Compatibility | [Feature Compatibility](feature_compatibility.md) |
 | CPU offloading | [CPU Offloading](diffusion/cpu_offload.md) |
 | Cache acceleration | [TeaCache](diffusion/cache_acceleration/teacache.md), [Cache-DiT](diffusion/cache_acceleration/cache_dit.md) |
 | Parallelism | [Parallelism Overview](diffusion/parallelism/overview.md) |
 | Attention | [Attention Backends](diffusion/attention_backends.md) |
 | Compilation | [Regional Compilation](diffusion/regional_compilation.md) |
 | Video extension | [Frame Interpolation](diffusion/frame_interpolation.md) |
 | Startup | [Startup and Loading](diffusion/startup_and_loading.md) |
 | Adapters | [LoRA](diffusion/lora.md) |
-**Parallelism Methods:**
+Related cross-model and runtime features are documented separately:
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
+- [Quantization](quantization/overview.md) covers diffusion-only models,
-
+  multi-stage omni/TTS models, and multi-stage diffusion models.
-**Memory Optimization:**
+- [Execution Modes and Streaming](diffusion/execution_modes.md) covers the
-
+  diffusion runtime, including batching, step execution, and streaming output.
 - **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage
 - **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
 - **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
 **Extensions:**
 - **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
 - **[Frame Interpolation Guide](diffusion/frame_interpolation.md)** - Worker-side post-generation video frame interpolation for smoother motion
 **Execution Modes:**
 - **[Diffusion Execution Modes](diffusion/execution_modes.md)** - Configure request batching, step execution, continuous batching, and streaming output
 **Startup Optimization:**
 - **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** - Speed up model startup by loading safetensors shards in parallel
 **Advanced Topics:**
 - **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance
@@ -88,9 +88,8 @@ python text_to_image.py \
 | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
 | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file |
+| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
 | `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
 **NextStep-1.1 specific arguments:**
@@ -125,9 +125,6 @@ python text_to_video.py \
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
 - `--flow-shift`: scheduler flow_shift parameter.
 - `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
 - `--lora-scale`: scale factor for LoRA weights.
 - `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
 ### Wan2.2-specific
@@ -60,8 +60,6 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import (
    build_image_to_video_prompt,
    get_extra_body_params,
@@ -296,30 +294,6 @@ def parse_args() -> argparse.Namespace:
        default=None,
        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
    )
    parser.add_argument(
        "--lora-path",
        type=str,
        nargs="+",
        default=None,
        help="Path to a LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. "
        "For Wan2.2 MoE models, pass the high-noise checkpoint first and the low-noise checkpoint second.",
    )
    parser.add_argument(
        "--lora-scale",
        type=float,
        default=1.0,
        help="Scale factor for PEFT LoRA weights (default: 1.0).",
    )
    parser.add_argument(
        "--lora-backend",
        type=str,
        default="peft",
        choices=["peft", "distill"],
        help=(
            "LoRA backend. 'peft' loads a PEFT-format adapter for request-time activation; "
            "'distill' fuses one or more concrete LoRA checkpoint files at initialization."
        ),
    )
    parser.add_argument(
        "--extra-body",
        type=functools.partial(parse_json_object, flag_name="--extra-body"),
@@ -513,12 +487,6 @@ def main():
        omni_kwargs["flow_shift"] = flow_shift
    if args.quantization is not None:
        omni_kwargs["quantization"] = args.quantization
    if args.lora_path is not None:
        lora_path = args.lora_path
        if len(lora_path) == 1:
            lora_path = lora_path[0]
        omni_kwargs["lora_path"] = lora_path
        omni_kwargs["lora_backend"] = args.lora_backend
    # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
    # gate is an engine-level config (offline analog of the server's --no-guardrails).
    if args.extra_body and "guardrails" in args.extra_body:
@@ -549,18 +517,6 @@ def main():
    print(f"  Video size: {width}x{height}")
    print(f"{'=' * 60}\n")
    lora_request = None
    if args.lora_path and args.lora_backend == "peft":
        if len(args.lora_path) != 1:
            raise ValueError("Only one LoRA path is expected for PEFT backend.")
        lora_path = args.lora_path[0]
        lora_request = LoRARequest(
            lora_name=Path(lora_path).stem,
            lora_int_id=stable_lora_int_id(lora_path),
            lora_path=lora_path,
        )
    negative_prompt = args.negative_prompt
    if negative_prompt is None and not is_ltx2:
        # Preserve the historical empty-prompt behavior for non-LTX examples.
@@ -574,11 +530,6 @@ def main():
        width=width,
        num_frames=num_frames,
    )
    sampling_extra_args = {"sample_solver": args.sample_solver}
    if lora_request is not None:
        sampling_extra_args["lora_request"] = lora_request
        sampling_extra_args["lora_scale"] = args.lora_scale
    sampling_params = OmniDiffusionSamplingParams(
        height=height,
        width=width,
@@ -589,7 +540,7 @@ def main():
        num_inference_steps=num_inference_steps,
        num_frames=num_frames,
        frame_rate=frame_rate,
-        extra_args=sampling_extra_args,
+        extra_args={"sample_solver": args.sample_solver},
    )
    if flow_shift is not None:
        sampling_params.extra_args["flow_shift"] = flow_shift
@@ -109,9 +109,8 @@ python text_to_image.py \
 | `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
 | `--ulysses-mode` | str | `"strict"` | Ulysses SP mode: `"strict"` or `"advanced_uaa"` |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file |
+| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
 | `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
 | `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
 | `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
 | `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). |
@@ -260,7 +260,6 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument(
        "--lora-path",
        type=str,
        nargs="+",
        default=None,
        help="Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.",
    )
@@ -270,15 +269,6 @@ def parse_args() -> argparse.Namespace:
        default=1.0,
        help="Scale factor for LoRA weights (default: 1.0).",
    )
    parser.add_argument(
        "--lora-backend",
        type=str,
        default="peft",
        choices=["peft", "distill"],
        help="LoRA backend for loading LoRA adapters. Default: peft"
        "'peft' loads a PEFT-format adapter folder, used e.g. for RL"
        "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
    )
    parser.add_argument(
        "--vae-patch-parallel-size",
        type=int,
@@ -414,12 +404,8 @@ def main():
    # Prepare LoRA kwargs for Omni initialization
    lora_args: dict[str, Any] = {}
    if args.lora_path:
-        lora_path = args.lora_path
+        lora_args["lora_path"] = args.lora_path
-        if len(lora_path) == 1:
+        print(f"Using LoRA from: {args.lora_path}")
            lora_path = lora_path[0]
        lora_args["lora_path"] = lora_path
        lora_args["lora_backend"] = args.lora_backend
        print(f"Using LoRA from: {lora_path} with backend: {args.lora_backend}")
    # Build quantization kwargs: use quantization_config dict when
    # ignored_layers is specified so the list flows through OmniDiffusionConfig
@@ -508,16 +494,12 @@ def main():
    # Build LoRA request when --lora-path is set
    lora_request = None
-    if args.lora_path and args.lora_backend == "peft":
+    if args.lora_path:
-        if len(args.lora_path) != 1:
+        lora_request_id = stable_lora_int_id(args.lora_path)
            raise ValueError("Only one LoRA path is expected for PEFT backend.")
        lora_path = args.lora_path[0]
        lora_request_id = stable_lora_int_id(lora_path)
        lora_request = LoRARequest(
-            lora_name=Path(lora_path).stem,
+            lora_name=Path(args.lora_path).stem,
            lora_int_id=lora_request_id,
-            lora_path=lora_path,
+            lora_path=args.lora_path,
        )
    generation_start = time.perf_counter()
@@ -218,9 +218,6 @@ python text_to_video.py \
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (such as `fp8` for FP8).
 - `--flow-shift`: scheduler flow_shift parameter.
 - `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
 - `--lora-scale`: scale factor for LoRA weights.
 - `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
 ### Wan2.2-specific
@@ -14,8 +14,6 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -336,29 +334,6 @@ def parse_args() -> argparse.Namespace:
        action="store_true",
        help="Enable expert parallelism for MoE layers.",
    )
    parser.add_argument(
        "--lora-path",
        type=str,
        nargs="+",
        default=None,
        help="Path to LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. Loaded at initialization and used for generation."
        "Note: for Wan2.2 MoE models, two checkpoints positionally mapped to high-noise and low-noise modules.",
    )
    parser.add_argument(
        "--lora-scale",
        type=float,
        default=1.0,
        help="Scale factor for LoRA weights (default: 1.0).",
    )
    parser.add_argument(
        "--lora-backend",
        type=str,
        default="peft",
        choices=["peft", "distill"],
        help="LoRA backend for loading LoRA adapters. Default: peft"
        "'peft' loads a PEFT-format adapter folder, used e.g. for RL"
        "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
    )
    parser.add_argument(
        "--use-hsdp",
        action="store_true",
@@ -465,12 +440,6 @@ def main():
        omni_kwargs["cache_backend"] = args.cache_backend
        omni_kwargs["cache_config"] = cache_config
        omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
    if args.lora_path is not None:
        lora_path = args.lora_path
        if len(lora_path) == 1:
            lora_path = lora_path[0]
        omni_kwargs["lora_path"] = lora_path
        omni_kwargs["lora_backend"] = args.lora_backend
    # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
    # gate is an engine-level config (offline analog of the server's --no-guardrails).
@@ -500,19 +469,6 @@ def main():
    print(f"  Video size: {args.width}x{args.height}")
    print(f"{'=' * 60}\n")
    lora_request = None
    if args.lora_path and args.lora_backend == "peft":
        if len(args.lora_path) != 1:
            raise ValueError("Only one LoRA path is expected for PEFT backend.")
        lora_path = args.lora_path[0]
        lora_request_id = stable_lora_int_id(lora_path)
        lora_request = LoRARequest(
            lora_name=Path(lora_path).stem,
            lora_int_id=lora_request_id,
            lora_path=lora_path,
        )
    prompt_dict = {"prompt": args.prompt}
    if args.negative_prompt is not None:
        prompt_dict["negative_prompt"] = args.negative_prompt
@@ -520,11 +476,6 @@ def main():
        # Preserve the historical empty-prompt behavior for non-LTX examples.
        prompt_dict["negative_prompt"] = ""
    extra_args = {}
    if lora_request:
        extra_args["lora_request"] = lora_request
        extra_args["lora_scale"] = args.lora_scale
    sampling_kwargs = dict(
        height=args.height,
        width=args.width,
@@ -532,7 +483,6 @@ def main():
        guidance_scale=args.guidance_scale,
        num_inference_steps=args.num_inference_steps,
        num_frames=args.num_frames,
        extra_args=extra_args,
    )
    if args.guidance_scale_high is not None:
        sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
@@ -115,23 +115,7 @@ python gradio_demo.py
 ## LoRA
-This example supports both startup-time distilled LoRAs and request-time
+This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
 PEFT-compatible adapters. The LoRA path must be readable on the **server**
 machine (usually a local path or a mounted directory).
 ### Start Server with a Distilled LoRA
 ```bash
 vllm serve Qwen/Qwen-Image-2512 \
  --omni \
  --port 8091 \
  --lora-backend distill \
  --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors
 ```
 The distilled LoRA is fused once during server initialization and applies to
 every request. Use the checkpoint's few-step sampling settings, for example
 `num_inference_steps=4` and `true_cfg_scale=1.0`.
 ### Using Python Client with LoRA
@@ -82,7 +82,7 @@ example and the DiT sections below.
 - Optional environment variables:
 ```bash
-export CACHE_DIT_VERSION=1.5.0
+export CACHE_DIT_VERSION=1.3.0
 ```
 HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
@@ -16,7 +16,7 @@ diffusers==0.38.0
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
-cache-dit==1.5.0
+cache-dit==1.3.0
 tqdm>=4.66.0
 torchsde>=0.2.6
 openai-whisper>=20250625
@@ -2206,24 +2206,6 @@ class TestBaseConfigInheritance:
 class TestPlatformOverrides:
    """Test platform-specific deploy config overrides."""
    def test_qwen3_omni_musa_clears_parent_quantization_for_audio_stages(self):
        pipeline = resolve_pipeline_config(
            "qwen3_omni_moe",
            Q3_OMNI_ALL_STAGES_HF_CONFIG,
        )
        assert isinstance(pipeline, PipelineConfig)
        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
        base_stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))
        assert "hf_overrides" not in base_stages[1].yaml_engine_args
        assert "hf_overrides" not in base_stages[2].yaml_engine_args
        musa = _apply_platform_overrides(load_deploy_config(deploy_path), platform="musa")
        musa_stages = merge_pipeline_deploy(pipeline, musa)
        expected = {"quantization_config": None}
        assert musa_stages[1].yaml_engine_args["hf_overrides"] == expected
        assert musa_stages[2].yaml_engine_args["hf_overrides"] == expected
    def test_qwen3_tts_rocm_disables_code2wav_outer_cudagraph(self):
        deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml"))
@@ -1,438 +0,0 @@
 from collections import defaultdict
 import pytest
 import torch
 import torch.nn as nn
 from diffusers.loaders.lora_conversion_utils import (
    _convert_non_diffusers_qwen_lora_to_diffusers,
 )
 from pytest_mock import MockerFixture
 from torch.testing import assert_close
 from vllm_omni.diffusion.lora.loader import (
    LoraLoaderMixin,
    QwenImageLoraLoaderMixin,
    WanLoraLoaderMixin,
    _prepare_lora_delta,
    _remap_state_dict_keys,
 )
 HEAD_DIM = 24
 RANK = 2
 NUM_LAYERS = 10
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 class DummyTransformerBlock(nn.Module):
    def __init__(
        self,
        d_model: int,
        bias: bool = False,
        stacked_params: bool = False,
    ):
        super().__init__()
        self.attn = nn.Module()
        if stacked_params:
            self.attn.to_qkv = nn.Linear(d_model * 3, d_model, bias=bias)
        else:
            self.attn.to_q = nn.Linear(d_model, d_model, bias=bias)
            self.attn.to_k = nn.Linear(d_model, d_model, bias=bias)
            self.attn.to_v = nn.Linear(d_model, d_model, bias=bias)
        self.attn.to_out = nn.Linear(d_model, d_model, bias=bias)
 class DummyTransformer(nn.Module):
    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
        super().__init__()
        self.blocks = nn.ModuleList(
            [DummyTransformerBlock(d_model, bias=bias, stacked_params=stacked_params) for _ in range(num_layers)]
        )
        if stacked_params:
            self.stacked_params_mapping = [
                (".attn.to_qkv", ".attn.to_q", "q"),
                (".attn.to_qkv", ".attn.to_k", "k"),
                (".attn.to_qkv", ".attn.to_v", "v"),
            ]
 class DummyPipeline(nn.Module, LoraLoaderMixin):
    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
        super().__init__()
        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=stacked_params)
 # ======================================================
 # helper functions for test
 # ======================================================
 def _filter_parameters(named_params, mode: str = "all"):
    if mode == "all":
        return named_params
    filtered = []
    for name, param in named_params:
        if name.endswith(mode):
            filtered.append((name, param))
    return filtered
 def assert_params_equal(module, original_params, mode: str = "all"):
    for name, param in _filter_parameters(module.named_parameters(), mode):
        assert_close(param, original_params[name])
 def assert_params_not_equal(module, original_params, mode: str = "all"):
    for name, param in _filter_parameters(module.named_parameters(), mode):
        assert not torch.allclose(param, original_params[name])
 def make_dummy_lora_state_dict(
    base_key,
    weight_a: torch.Tensor | None = None,
    weight_b: torch.Tensor | None = None,
    weight_bias: torch.Tensor | None = None,
    bias: bool = False,
    lora_a_suffix: str = "lora_A.weight",
    lora_b_suffix: str = "lora_B.weight",
    lora_bias_suffix: str = "bias",
 ):
    """
    Create a dummy LoRA state dict.
    Returns:
        tuple: (state_dict, expected_delta)
    """
    key_a = f"{base_key}.{lora_a_suffix}"
    key_b = f"{base_key}.{lora_b_suffix}"
    key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
    if weight_a is None:
        weight_a = torch.randn(RANK, HEAD_DIM)
    if weight_b is None:
        weight_b = torch.randn(HEAD_DIM, RANK)
    state_dict = {}
    expected_delta = None
    if bias:
        state_dict[key_a] = weight_a
        state_dict[key_b] = weight_b
        weight_bias = torch.randn(HEAD_DIM)
        state_dict[key_bias] = weight_bias
        expected_delta = weight_bias
    else:
        state_dict[key_a] = weight_a
        state_dict[key_b] = weight_b
        expected_delta = weight_b @ weight_a
    return state_dict, expected_delta
 def make_lora_state_dict_for_module(
    module,
    lora_rank=RANK,
    bias=False,
    lora_a_suffix: str = "lora_A.weight",
    lora_b_suffix: str = "lora_B.weight",
    lora_bias_suffix: str = "bias",
    prefix: str = "transformer",
    remap_proj_out: bool = True,
 ):
    param_to_weight_names = defaultdict(list)
    if hasattr(module, "stacked_params_mapping"):
        for param_name, weight_name, _ in module.stacked_params_mapping:
            param_to_weight_names[param_name].append(weight_name)
    state_dict = {}
    for name, param in module.named_parameters(prefix=prefix):
        if name.endswith(".weight"):
            base_key = name[: -len(".weight")]
            if remap_proj_out and base_key.endswith(".to_out"):
                base_key = base_key.replace(".to_out", ".to_out.0")
            d_in, d_out = param.shape
            is_stacked_param = False
            for param_name, weight_names in param_to_weight_names.items():
                if param_name not in name:
                    continue
                is_stacked_param = True
                target_in = d_in // len(weight_names)
                for weight_name in weight_names:
                    lora_key = base_key.replace(param_name, weight_name)
                    weight_a = torch.randn(lora_rank, d_out)
                    weight_b = torch.randn(target_in, lora_rank)
                    state_dict[f"{lora_key}.{lora_a_suffix}"] = weight_a
                    state_dict[f"{lora_key}.{lora_b_suffix}"] = weight_b
            if is_stacked_param:
                continue
            weight_a = torch.randn(lora_rank, d_out)
            weight_b = torch.randn(d_in, lora_rank)
            key_a = f"{base_key}.{lora_a_suffix}"
            key_b = f"{base_key}.{lora_b_suffix}"
            state_dict[key_a] = weight_a
            state_dict[key_b] = weight_b
        elif name.endswith(".bias") and bias:
            base_key = name[: -len(".bias")]
            if remap_proj_out and base_key.endswith(".to_out"):
                base_key = base_key.replace(".to_out", ".to_out.0")
            d_bias = param.shape[0]
            is_stacked_param = False
            for param_name, weight_names in param_to_weight_names.items():
                if param_name not in name:
                    continue
                is_stacked_param = True
                target_bias = d_bias // len(weight_names)
                for weight_name in weight_names:
                    lora_key = base_key.replace(param_name, weight_name)
                    weight_bias = torch.randn(target_bias)
                    state_dict[f"{lora_key}.{lora_bias_suffix}"] = weight_bias
            if is_stacked_param:
                continue
            key_bias = f"{base_key}.{lora_bias_suffix}"
            state_dict[key_bias] = torch.randn_like(param)
    return state_dict
 # ======================================================
 # Test loader helper functions
 # ======================================================
 def test_prepare_lora_delta_with_stacked_params():
    base_key = "transformer.attn.to_qkv"
    sd0, delta0 = make_dummy_lora_state_dict("transformer.attn.to_q")
    sd1, delta1 = make_dummy_lora_state_dict("transformer.attn.to_k")
    sd2, delta2 = make_dummy_lora_state_dict("transformer.attn.to_v")
    lora_state_dict = {}
    lora_state_dict.update(sd0)
    lora_state_dict.update(sd1)
    lora_state_dict.update(sd2)
    expected_delta = torch.concat([delta0, delta1, delta2])
    stacked_mapping = {".attn.to_qkv": [".attn.to_q", ".attn.to_k", ".attn.to_v"]}
    actual_delta, used_keys = _prepare_lora_delta(
        lora_state_dict,
        base_key,
        stacked_mapping,
    )
    # validate the delta
    assert actual_delta is not None
    assert_close(actual_delta, expected_delta)
    # validate the used keys
    key_q_a = "transformer.attn.to_q.lora_A.weight"
    key_q_b = "transformer.attn.to_q.lora_B.weight"
    key_k_a = "transformer.attn.to_k.lora_A.weight"
    key_k_b = "transformer.attn.to_k.lora_B.weight"
    key_v_a = "transformer.attn.to_v.lora_A.weight"
    key_v_b = "transformer.attn.to_v.lora_B.weight"
    for key in [key_q_a, key_q_b, key_k_a, key_k_b, key_v_a, key_v_b]:
        assert key in used_keys
 def test_remap_state_dict_keys_matched_two():
    lora_state_dict, _ = make_dummy_lora_state_dict(
        "transformer.attn.to_out.0", bias=True, lora_bias_suffix="lora_B.bias"
    )
    lora_state_dict = _remap_state_dict_keys(
        lora_state_dict,
        [(".to_out.0.", ".to_out."), (".lora_B.bias", ".bias")],
    )
    assert len(lora_state_dict) == 3
    assert "transformer.attn.to_out.bias" in lora_state_dict
 # ======================================================
 # Test LoraLoaderMixin
 # ======================================================
 class TestLoraLoaderMixin:
    @pytest.mark.parametrize(
        "bias, stacked_params",
        [
            (False, False),
            (True, False),
            (False, True),
            (True, True),
        ],
        ids=[
            "basic",
            "bias",
            "stacked_params",
            "bias_stacked_params",
        ],
    )
    def test_load_lora_into_module(self, bias: bool, stacked_params: bool):
        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM, bias=bias, stacked_params=stacked_params)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer, bias=bias)
        assert len(lora_state_dict) > 0
        # remap the keys to vllm-omni format
        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
        # validate the parameters are updated
        assert_params_not_equal(pipeline.transformer, original_params)
        # validate all keys are used
        missing_keys = set(lora_state_dict.keys()) - used_keys
        assert len(missing_keys) == 0
    def test_unload_module_lora(self):
        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
        # validate the state dict is valid
        assert len(lora_state_dict) > 0
        # remap the keys to vllm-omni format
        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
        # validate the weights are updated
        assert_params_not_equal(pipeline.transformer, original_params)
        # validate all keys are used
        missing_keys = set(lora_state_dict.keys()) - used_keys
        assert len(missing_keys) == 0
        unload_used_keys = pipeline.unload_module_lora(lora_state_dict, pipeline.transformer)
        missing_keys = set(lora_state_dict.keys()) - unload_used_keys
        assert len(missing_keys) == 0
        assert len(unload_used_keys - used_keys) == 0
        # validate the weights are restored
        assert_params_equal(pipeline.transformer, original_params)
 # ======================================================
 # Test QwenImageLoraLoaderMixin
 # ======================================================
 class DummyQwenImagePipeline(nn.Module, QwenImageLoraLoaderMixin):
    def __init__(self, num_layers: int, d_model: int, bias: bool = False):
        super().__init__()
        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=True)
 class TestQwenImageLoraLoaderMixin:
    def test_load_lora_weights(self, mocker: MockerFixture):
        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
        assert len(lora_state_dict) > 0
        mocker.patch(
            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
        )
        pipeline.load_lora_weights(lora_state_dict, "adapter0")
        # validate the weights are updated after lora loaded
        assert_params_not_equal(pipeline.transformer, original_params)
        # validate lora_loaded map is updated after lora loaded
        assert "adapter0" in pipeline.lora_loaded
    def test_unload_lora_weights(self, mocker: MockerFixture):
        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
        assert len(lora_state_dict) > 0
        mocker.patch(
            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
        )
        pipeline.load_lora_weights(lora_state_dict, "adapter0")
        # validate the weights are updated after lora loaded
        assert_params_not_equal(pipeline.transformer, original_params)
        # validate lora_loaded map is updated after lora loaded
        assert "adapter0" in pipeline.lora_loaded
        pipeline.unload_lora_weights("adapter0")
        # validate the weights are restored after lora unloaded
        assert_params_equal(pipeline.transformer, original_params)
        # validate lora_loaded map is updated
        assert "adapter0" not in pipeline.lora_loaded
    def test_load_diffusers_lora_applies_non_unit_alpha(self, mocker: MockerFixture):
        pipeline = DummyQwenImagePipeline(1, HEAD_DIM)
        to_out = pipeline.transformer.blocks[0].attn.to_out
        original_weight = to_out.weight.detach().clone()
        base_key = "blocks.0.attn.to_out.0"
        lora_a = torch.randn(RANK, HEAD_DIM)
        lora_b = torch.randn(HEAD_DIM, RANK)
        alpha = torch.tensor(6.0)
        state_dict = {
            f"{base_key}.lora_A.weight": lora_a,
            f"{base_key}.lora_B.weight": lora_b,
            f"{base_key}.alpha": alpha,
        }
        mocker.patch(
            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
        )
        pipeline.load_lora_weights(state_dict, "adapter0")
        expected_weight = original_weight + (lora_b @ lora_a) * (alpha.item() / RANK)
        assert_close(to_out.weight, expected_weight)
        pipeline.unload_lora_weights("adapter0")
        assert_close(to_out.weight, original_weight)
 # ======================================================
 # Test WanLoraLoaderMixin
 # ======================================================
 class DummyWanPipeline(nn.Module, WanLoraLoaderMixin):
    def __init__(self, num_layers: int, d_model: int):
        super().__init__()
        self.has_transformer_2 = True
        self.transformer = DummyTransformer(num_layers, d_model, stacked_params=True)
        self.transformer_2 = DummyTransformer(num_layers, d_model, stacked_params=True)
 def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
    pipeline = DummyWanPipeline(NUM_LAYERS, HEAD_DIM)
    original_high = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
    original_low = {name: param.clone() for name, param in pipeline.transformer_2.named_parameters()}
    high_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
    low_state_dict = make_lora_state_dict_for_module(pipeline.transformer_2)
    mocker.patch(
        "vllm_omni.diffusion.lora.loader._load_lora_state_dict",
        side_effect=[high_state_dict, low_state_dict],
    )
    pipeline.load_lora_weights(["high.safetensors", "low.safetensors"], "adapter0")
    assert_params_not_equal(pipeline.transformer, original_high)
    assert_params_not_equal(pipeline.transformer_2, original_low)
    assert set(pipeline.lora_loaded["adapter0"]) == {"transformer", "transformer_2"}
    pipeline.unload_lora_weights("adapter0")
    assert_params_equal(pipeline.transformer, original_high)
    assert_params_equal(pipeline.transformer_2, original_low)
    assert "adapter0" not in pipeline.lora_loaded
@@ -19,11 +19,38 @@ DEPLOY_CONFIG = modify_stage_config(
    get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
    updates={
        "base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
-        # Talker context is 4096 (tts_config.max_position_embeddings); KV sizing
+        # Cap per-stage KV so Thinker/Talker/Code2Wav share one GPU with
-        # is left automatic so duplex matches the simplex deploy profiles.
+        # max_sessions=2. Talker must stay within its 4096 context so the
        # 0.5 GiB budget passes vLLM's min-KV check at init.
        "stages": {
-            1: {"max_model_len": 4096},
+            0: {"kv_cache_memory_bytes": 6 * 1024 * 1024 * 1024},
            1: {
                "max_model_len": 4096,
                "kv_cache_memory_bytes": 512 * 1024 * 1024,
            },
            2: {"kv_cache_memory_bytes": 256 * 1024 * 1024},
        },
        # Platform overrides apply after ordinary stage settings and would
        # otherwise reinstate the base CUDA 2 GiB Talker default.
        "platforms": {
            "cuda": {
                "stages": [
                    {
                        "stage_id": 1,
                        "kv_cache_memory_bytes": 512 * 1024 * 1024,
                    }
                ]
            }
        },
    },
 )
 CORE_DEPLOY_CONFIG = modify_stage_config(
    DEPLOY_CONFIG,
    updates={
        "stages": {
            0: {"enforce_eager": True},
            1: {"enforce_eager": True},
        }
    },
 )
 ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
@@ -44,6 +71,17 @@ SERVER_PARAMS = [
        id="three-stage-single-gpu",
    )
 ]
 CORE_SERVER_PARAMS = [
    pytest.param(
        OmniServerParams(
            model=MODEL,
            stage_config_path=CORE_DEPLOY_CONFIG,
            use_stage_cli=False,
            server_args=["--trust-remote-code"],
        ),
        id="three-stage-single-gpu",
    )
 ]
 def validated_wav(path: Path, expected_sha256: str) -> Path:
@@ -74,11 +112,9 @@ def validated_soft_interrupt_wav() -> Path:
 def resolve_ref_audio(model_prefix: str) -> Path:
-    # ``MODEL`` may be a local checkpoint directory instead of a Hub repo id
+    if model_prefix:
-    # (``Path("/prefix") / "/abs/path"`` collapses to the absolute path), so only
+        model_root = Path(model_prefix) / MODEL
-    # fall back to the Hub cache when no such directory exists on disk.
+    else:
    model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
    if not model_root.is_dir():
        model_root = Path(snapshot_download(MODEL, local_files_only=True))
    ref_audio = model_root / REF_AUDIO_RELATIVE_PATH
    if not ref_audio.is_file():
@@ -14,15 +14,47 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
 NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
 MODEL = "black-forest-labs/FLUX.1-Kontext-dev"
-BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
 def _get_diffusion_feature_cases(model: str):
    return [
        pytest.param(
-            OmniServerParams(model=model),
+            OmniServerParams(
-            id="base",
+                model=model,
-            marks=BASE_FEATURE_MARKS,
+                server_args=[
                    "--tensor-parallel-size",
                    "2",
                    "--enable-cpu-offload",
                ],
            ),
            id="parallel_tp_2",
            marks=PARALLEL_FEATURE_MARKS,
        ),
        pytest.param(
            OmniServerParams(
                model=model,
                server_args=[
                    "--enable-cpu-offload",
                    "--cfg-parallel-size",
                    "2",
                ],
            ),
            id="parallel_cfg_2",
            marks=PARALLEL_FEATURE_MARKS,
        ),
        pytest.param(
            OmniServerParams(
                model=model,
                server_args=[
                    "--tensor-parallel-size",
                    "2",
                    "--enable-cpu-offload",
                    "--cache-backend",
                    "cache_dit",
                ],
            ),
            id="parallel_002",
        ),
    ]
@@ -1,7 +1,10 @@
-"""E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
+"""
 E2E Online tests for MiniCPM-o 4.5 model with multimodal input and audio / text output.
-Exercises async chunk streaming (``--async-chunk``) across separate Thinker,
+MiniCPM-o 4.5 is a three-stage pipeline: Stage 0 Thinker (text), Stage 1 Talker
-Talker, and Code2Wav stages.
+(TTS codec), Stage 2 Code2Wav (vocoder). The talker -> Code2Wav bridge supports
 both async_chunk (streaming) and full-payload (``--no-async-chunk``) modes, so the
 tests below exercise the sync path as well.
 """
 import os
@@ -11,17 +14,12 @@ import pytest
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
 from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from tests.helpers.stage_config import get_deploy_config_path
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 _MODEL = "openbmb/MiniCPM-o-4_5"
-_CI_DEPLOY = modify_stage_config(
+_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
    get_deploy_config_path("minicpmo_4_5.yaml"),
    updates={
        "stages": {0: {"default_sampling_params.max_tokens": 64}, 1: {"default_sampling_params.max_tokens": 1024}}
    },
 )
 test_params = [
    pytest.param(
@@ -29,10 +27,12 @@ test_params = [
            model=_MODEL,
            stage_config_path=_CI_DEPLOY,
            use_stage_cli=False,
-            server_args=["--trust-remote-code", "--async-chunk"],
+            server_args=[
                "--trust-remote-code",
            ],
        ),
-        id="async_chunk",
+        id="default",
-    ),
+    )
 ]
@@ -98,7 +98,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
 def test_text_to_audio_001(omni_server, openai_client) -> None:
    """
    Test text-only input generating text + audio output via OpenAI API.
-    This exercises Talker TTS region detection and the Code2Wav stage.
+    This exercises the talker TTS region detection and in-process token2wav vocoder.
    Deploy Setting: default single GPU
    Input Modal: text
    Output Modal: text + audio
@@ -132,21 +132,13 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
    messages = dummy_messages_from_mix_data(
        system_prompt=get_system_prompt(),
        audio_data_url=audio_data_url,
-        content_text=get_prompt(),
+        content_text=get_prompt("mix"),
    )
    request_config = {
        "model": omni_server.model,
        "messages": messages,
        "stream": True,
        "key_words": {"text": ["Beijing"]},
        "modalities": ["text", "audio"],
        "extra_body": {
            "chat_template_kwargs": {
                "use_tts_template": True,
                "enable_thinking": False,
            }
        },
    }
    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@@ -11,6 +11,7 @@ import pytest
 import websockets
 from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
    CORE_SERVER_PARAMS,
    SERVER_PARAMS,
    demo_args,
    multi_session_args,
@@ -106,7 +107,7 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[
@pytest.mark.core_model
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+@pytest.mark.parametrize("omni_server", CORE_SERVER_PARAMS, indirect=True)
 def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None:
    ref_audio = resolve_ref_audio(model_prefix)
    events = asyncio.run(
@@ -133,12 +134,8 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str,
        output_dir=tmp_path / "single_session",
    )
    args.turns = 2
    # Every turn replays the same active-speech window as the first one. The default
    # shorter follow-up window is a different mid-utterance slice, which the native
    # duplex model may legitimately answer with "listen" instead of a response.
    args.turn_duration_ms = [args.first_turn_ms] * args.turns
    result = asyncio.run(run_demo(args))
-    assert result["ok"] is True
+    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
    assert result["audio_delta_count"] > 0
    assert result["done_count"] == 2
    assert result["error_count"] == 0
@@ -163,7 +160,7 @@ def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str,
            )
        )
    )
-    assert result["ok"] is True
+    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
    assert result["session_count"] == 2
    assert result["resume"]["ok"] is True
    assert result["takeover"]["ok"] is True
@@ -2,9 +2,6 @@
 These cover modality combinations, separate Code2Wav behavior, sequential
 request isolation, longer audio output, and Chinese speech.
 Like Qwen3-Omni expansion, ``default`` exercises ``--no-async-chunk`` and
 ``async_chunk`` exercises ``--async-chunk``.
 """
 import os
@@ -23,27 +20,18 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 _MODEL = "openbmb/MiniCPM-o-4_5"
 _CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
 # Deploy YAML defaults to ``async_chunk: true``; keep sync and async as
 # separate parametrizations so both handoff paths are exercised.
 test_params = [
    pytest.param(
        OmniServerParams(
            model=_MODEL,
            stage_config_path=_CI_DEPLOY,
            use_stage_cli=False,
-            server_args=["--trust-remote-code", "--no-async-chunk"],
+            server_args=[
                "--trust-remote-code",
            ],
        ),
        id="default",
-    ),
+    )
    pytest.param(
        OmniServerParams(
            model=_MODEL,
            stage_config_path=_CI_DEPLOY,
            use_stage_cli=False,
            server_args=["--trust-remote-code", "--async-chunk"],
        ),
        id="async_chunk",
    ),
 ]
@@ -284,37 +284,6 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
    assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
 def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
    """Ensure startup distilled LoRA options reach the online diffusion stage."""
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    OmniServeCommand().subparser_init(subparsers)
    args = parser.parse_args(
        [
            "serve",
            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
            "--omni",
            "--lora-backend",
            "distill",
            "--lora-path",
            "/models/high.safetensors",
            "/models/low.safetensors",
        ]
    )
    explicit_kwargs = args.get_explicit_kwargs_dict()
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
    engine_args = stage_cfg["engine_args"]
    assert explicit_kwargs["lora_backend"] == "distill"
    assert engine_args["lora_backend"] == "distill"
    assert engine_args["lora_path"] == [
        "/models/high.safetensors",
        "/models/low.safetensors",
    ]
 def test_serve_cli_forwards_distributed_offload_residency():
    """Ensure the two-GPU DLO placement controls reach the diffusion stage."""
    parser = TrackingArgumentParser()
@@ -37,19 +37,6 @@ def tiny_ltx2_builder() -> str:
    return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
 def _shrink_flux_clip_text_encoder(config: dict) -> dict:
    config["num_hidden_layers"] = 2
    return config
 def _shrink_flux_t5_text_encoder(config: dict) -> dict:
    config["num_layers"] = 2
    config["num_decoder_layers"] = 2
    config["d_ff"] = 64
    config["num_heads"] = 4
    return config
 def tiny_qwen_image_builder() -> str:
    def shrink_text_encoder(config: dict) -> dict:
        config["num_hidden_layers"] = 2
@@ -70,24 +57,23 @@ def tiny_qwen_image_builder() -> str:
 def tiny_flux_builder() -> str:
    def shrink_clip_text_encoder(config: dict) -> dict:
        config["num_hidden_layers"] = 2
        return config
    def shrink_t5_text_encoder(config: dict) -> dict:
        config["num_layers"] = 2
        config["num_decoder_layers"] = 2
        config["d_ff"] = 64
        config["num_heads"] = 4
        return config
    return build_tiny_from_configs(
        "FluxPipeline",
        "black-forest-labs/FLUX.1-schnell",
        transform={
-            "text_encoder": _shrink_flux_clip_text_encoder,
+            "text_encoder": shrink_clip_text_encoder,
-            "text_encoder_2": _shrink_flux_t5_text_encoder,
+            "text_encoder_2": shrink_t5_text_encoder,
            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
        },
    )
 def tiny_flux_kontext_builder() -> str:
    return build_tiny_from_configs(
        "FluxKontextPipeline",
        "black-forest-labs/FLUX.1-Kontext-dev",
        transform={
            "text_encoder": _shrink_flux_clip_text_encoder,
            "text_encoder_2": _shrink_flux_t5_text_encoder,
            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
        },
    )
@@ -59,16 +59,6 @@ DIFFUSION_TEST_SETTINGS = {
        builder=diff_model_builders.tiny_flux_builder,
        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
    ),
    "FluxKontextPipeline": DiffusionModelTestOpts(
        model="black-forest-labs/FLUX.1-Kontext-dev",
        builder=diff_model_builders.tiny_flux_kontext_builder,
        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
        extra_test_groups=[
            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD, DiffusionAccs.CACHE_DIT],
        ],
    ),
    "Flux2Pipeline": DiffusionModelTestOpts(
        model="black-forest-labs/FLUX.2-dev",
        builder=diff_model_builders.tiny_flux2_builder,
@@ -48,6 +48,7 @@ EXCLUDED_MODELS = [
    "InternVLAA1Pipeline",
    "LongCatImageEditPipeline",
    "StableDiffusion3Pipeline",
    "FluxKontextPipeline",
    "HunyuanImage3ForCausalMM",
    "ErnieImagePipeline",
    "NextStep11Pipeline",
@@ -3,88 +3,13 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
-from vllm_omni.worker.gpu_model_runner import (
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner, _filter_mrope_kwargs_for_model
    OmniGPUModelRunner,
    _filter_mrope_kwargs_for_model,
 )
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def _runner_for_talker_graph_init(
    *,
    talker_mtp_graph_safe: bool | None,
    has_separate_talker: bool = True,
    model_stage: str = "talker",
 ) -> OmniGPUModelRunner:
    runner = object.__new__(OmniGPUModelRunner)
    runner.model = SimpleNamespace(
        talker=object() if has_separate_talker else None,
        talker_mtp=object(),
        model_stage=model_stage,
    )
    if talker_mtp_graph_safe is not None:
        runner.model.talker_mtp_graph_safe = talker_mtp_graph_safe
    runner.model_config = SimpleNamespace(hf_text_config=SimpleNamespace(hidden_size=4))
    runner.compilation_config = SimpleNamespace(
        cudagraph_mode=CUDAGraphMode.FULL,
        max_cudagraph_capture_size=1,
    )
    runner.vllm_config = object()
    runner.max_num_reqs = 1
    runner.dtype = torch.float32
    runner._make_buffer = lambda *args, **kwargs: SimpleNamespace(args=args, kwargs=kwargs)
    return runner
 def test_talker_mtp_skips_graph_when_model_declares_unsafe(monkeypatch):
    runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=False)
    talker_mtp = runner.model.talker_mtp
    monkeypatch.setattr(
        "vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
        lambda: pytest.fail("graph wrapper must not be selected"),
    )
    OmniGPUModelRunner._init_talker_mtp(runner)
    assert runner.talker_mtp is talker_mtp
@pytest.mark.parametrize("model_stage", ["thinker", "code2wav"])
 def test_non_talker_stage_does_not_use_talker_mtp_graph(monkeypatch, model_stage: str):
    runner = _runner_for_talker_graph_init(
        talker_mtp_graph_safe=None,
        has_separate_talker=False,
        model_stage=model_stage,
    )
    talker_mtp = runner.model.talker_mtp
    monkeypatch.setattr(
        "vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
        lambda: pytest.fail("graph wrapper must not be selected"),
    )
    OmniGPUModelRunner._init_talker_mtp(runner)
    assert runner.talker_mtp is talker_mtp
@pytest.mark.parametrize("talker_mtp_graph_safe", [None, True])
 def test_talker_mtp_uses_graph_for_legacy_or_explicit_safe_model(monkeypatch, talker_mtp_graph_safe):
    runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=talker_mtp_graph_safe)
    wrapped = object()
    monkeypatch.setattr(
        "vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
        lambda: lambda *args, **kwargs: wrapped,
    )
    OmniGPUModelRunner._init_talker_mtp(runner)
    assert runner.talker_mtp is wrapped
 class DummyBuffer:
    """A minimal buffer wrapper that exposes the `.gpu` attribute."""
@@ -563,9 +563,8 @@ class _DiffusionConfigProjection:
    diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
    diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
    diffusers_pipeline_cls: Any = None
-    lora_path: str | list[str] | None = None
+    lora_path: str | None = None
    lora_scale: float = 1.0
    lora_backend: str = "peft"
    max_cpu_loras: int | None = None
    output_type: str = "pil"
    enable_cpu_offload: bool = False
@@ -385,9 +385,6 @@ class StageDeployConfig:
    # Diffusion model loading and adapter construction.
    model_class_name: str | None = None
    diffusion_load_format: str | None = None
    lora_path: str | list[str] | None = None
    lora_backend: str | None = None
    lora_scale: float | None = None
    diffusers_load_kwargs: dict[str, Any] | None = None
    diffusers_call_kwargs: dict[str, Any] | None = None
    diffusion_quantization_config: str | None = None
@@ -69,18 +69,6 @@ stages:
      repetition_penalty: 1.1
 platforms:
  # The supported Qwen3-Omni ModelOpt checkpoint quantizes the thinker only. Clear the
  # checkpoint-level metadata for the BF16 audio stages on the verified MUSA
  # profile before vLLM performs its normal quantization auto-detection.
  musa:
    stages:
      - stage_id: 1
        hf_overrides:
          quantization_config: null
      - stage_id: 2
        hf_overrides:
          quantization_config: null
  npu:
    stages:
      - stage_id: 0
@@ -27,7 +27,6 @@ from vllm_omni.diffusion.diffusion_kv.config import (
    DiffusionKVCacheMode,
    parse_diffusion_kv_cache_mode,
 )
 from vllm_omni.diffusion.lora.manager import LoRABackend
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.errors import client_error_metadata
@@ -683,9 +682,8 @@ class OmniDiffusionConfig:
    # pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)
    # LoRA parameters
-    lora_path: str | list[str] | None = None
+    lora_path: str | None = None
-    lora_scale: float | list[float] = 1.0
+    lora_scale: float = 1.0
    lora_backend: LoRABackend = LoRABackend.PEFT  # available choices: ["peft", "distill"]
    max_cpu_loras: int | None = None
    output_type: str = "pil"
@@ -1,461 +0,0 @@
 import os
 from collections import OrderedDict, defaultdict
 from collections.abc import Callable
 import torch
 from diffusers.loaders.lora_conversion_utils import (
    _convert_non_diffusers_qwen_lora_to_diffusers,
    _convert_non_diffusers_wan_lora_to_diffusers,
 )
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from vllm.logger import init_logger
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
 logger = init_logger(__name__)
 lora_convert_mapping: dict[str, Callable] = {
    "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
    "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
    "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
    "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers,
    "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers,
 }
 def get_converter_by_pipeline(pipeline):
    for pipeline_cls in pipeline.__class__.__mro__:
        converter = lora_convert_mapping.get(pipeline_cls.__name__)
        if converter is not None:
            return converter
    return None
 def _prepare_lora_delta(
    lora_state_dict,
    base_key: str,
    param_to_weight_names: dict[str, list[str]] | None = None,
    is_bias: bool = False,
    lora_a_suffix: str = "lora_A.weight",
    lora_b_suffix: str = "lora_B.weight",
    lora_bias_suffix: str = "bias",
 ):
    used_keys = set()
    # stacked_params_mapping                       param_to_weight_names
    # [(".to_qkv", ".to_q.", "q")
    # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
    # (".to_qkv", ".to_v.", "v")]
    # stacked_sd is to store packed parameter deltas (format: [str, list[torch.Tensor]])
    # example: {".to_qkv": [delta_q, delta_k, delta_v]}
    stacked_deltas = []
    is_stacked_param = False
    if param_to_weight_names is None:
        param_to_weight_names = defaultdict(list)
    for param_name, weight_names in param_to_weight_names.items():
        if param_name not in base_key:
            continue
        is_stacked_param = True
        # handle lora_a_key and lora_b_key together
        if is_bias:
            for weight_name in weight_names:
                lora_bias_key = f"{base_key.replace(param_name, weight_name)}.{lora_bias_suffix}"
                if lora_bias_key not in lora_state_dict:
                    return None, used_keys
                delta = lora_state_dict[lora_bias_key]
                stacked_deltas.append(delta)
                used_keys.add(lora_bias_key)
        else:
            for weight_name in weight_names:
                lora_a_key = f"{base_key.replace(param_name, weight_name)}.{lora_a_suffix}"
                lora_b_key = lora_a_key.replace(lora_a_suffix, lora_b_suffix)
                if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
                    return None, used_keys
                a = lora_state_dict[lora_a_key]
                b = lora_state_dict[lora_b_key]
                delta = torch.matmul(b, a)
                stacked_deltas.append(delta)
                used_keys.add(lora_a_key)
                used_keys.add(lora_b_key)
        continue
    if is_stacked_param:
        return torch.concat(stacked_deltas), used_keys
    if is_bias:
        lora_bias_key = f"{base_key}.{lora_bias_suffix}"
        if lora_bias_key not in lora_state_dict:
            return None, used_keys
        used_keys.add(lora_bias_key)
        return lora_state_dict[lora_bias_key], used_keys
    lora_a_key = f"{base_key}.{lora_a_suffix}"
    lora_b_key = f"{base_key}.{lora_b_suffix}"
    if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
        return None, used_keys
    a = lora_state_dict[lora_a_key]
    b = lora_state_dict[lora_b_key]
    used_keys.add(lora_a_key)
    used_keys.add(lora_b_key)
    return torch.matmul(b, a), used_keys
 def _load_lora_state_dict(
    pretrained_model_name_or_path: str,
    weights_name: str | None = None,
    use_safetensors: bool = True,
    subfolder: str | None = None,
    cache_dir: str | None = None,
    force_download: bool = False,
    local_files_only: bool | None = None,
 ):
    # first we try to load it from local disk
    if use_safetensors and pretrained_model_name_or_path.endswith(".safetensors"):
        return load_file(pretrained_model_name_or_path)
    if os.path.isdir(pretrained_model_name_or_path):
        model_file = pretrained_model_name_or_path
        if subfolder:
            model_file = os.path.join(model_file, subfolder)
        if weights_name is None:
            raise ValueError("weights_name is required when loading from a directory")
        model_file = os.path.join(model_file, weights_name)
        return load_file(model_file)
    # finally, we try to load it from the internet
    try:
        model_file = hf_hub_download(
            pretrained_model_name_or_path,
            filename=weights_name,
            subfolder=subfolder,
            cache_dir=cache_dir,
            force_download=force_download,
            local_files_only=local_files_only,
        )
        return load_file(model_file)
    except Exception as e:
        logger.error(f"Failed to download {weights_name} from {pretrained_model_name_or_path}: {e}")
        raise e
 def _remap_state_dict_keys(sd, rules):
    """
    Remap keys in a state_dict by sequentially applying a list of substring substitution rules.
    This utility function transforms parameter names by applying a sequence of
    (old_substring, new_substring) pairs. For each key in the input dictionary,
    **all rules are evaluated in the given order**, and every rule whose
    `old_substring` is found within the current key triggers a replacement.
    Substitutions are cumulative: the key is updated after each match, and
    subsequent rules operate on the result of previous replacements.
    The order of rules matters; users are responsible for arranging them to
    achieve the desired final key. Keys that do not match any rule are copied
    unchanged.
    Args:
        sd (dict[str, Any]): Original state_dict mapping parameter names to tensors.
        rules (List[Tuple[str, str]]): A list of (old_substring, new_substring) pairs.
            Rules are applied sequentially, and **all matching rules are executed**.
    Returns:
        dict[str, Any]: A new state_dict with remapped keys. Values are unchanged.
    """
    new_sd = OrderedDict()
    for k, v in sd.items():
        new_key = k
        matched = False
        for p1, p2 in rules:
            if p1 not in k:
                continue
            new_key = new_key.replace(p1, p2)
            matched = True
        if matched:
            new_sd[new_key] = v
        else:
            new_sd[k] = v
    return new_sd
 def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Fold Diffusers-format LoRA alpha values into their B matrices."""
    scaled_state_dict = dict(state_dict)
    for alpha_key in (key for key in state_dict if key.endswith(".alpha")):
        base_key = alpha_key[: -len(".alpha")]
        lora_a_key = f"{base_key}.lora_A.weight"
        lora_b_key = f"{base_key}.lora_B.weight"
        if lora_a_key not in state_dict or lora_b_key not in state_dict:
            raise ValueError(f"LoRA alpha key {alpha_key!r} does not have matching lora_A and lora_B weights.")
        rank = state_dict[lora_a_key].shape[0]
        alpha = state_dict[alpha_key]
        scaled_state_dict[lora_b_key] = state_dict[lora_b_key] * (alpha.item() / rank)
        scaled_state_dict.pop(alpha_key)
    return scaled_state_dict
 class LoraLoaderMixin:
    transformer_name = "transformer"
    # Lazy initialization to avoid MRO issues: __init__ may not be called
    # when mixin with nn.Module
    @property
    def lora_loaded(self):
        if not hasattr(self, "_lora_loaded"):
            self._lora_loaded = {}
        return self._lora_loaded
    @lora_loaded.setter
    def lora_loaded(self, value):
        self._lora_loaded = value
    @classmethod
    def load_lora_into_module(
        cls,
        state_dict,
        module,
        prefix: str = "transformer",
        lora_a_suffix: str = "lora_A.weight",
        lora_b_suffix: str = "lora_B.weight",
        lora_bias_suffix: str = "bias",
    ):
        lora_loaded_keys = set()
        # prepare stacked parameters mapping for later use
        # stacked_params_mapping                       param_to_weight_names
        # [(".to_qkv", ".to_q.", "q")
        # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
        # (".to_qkv", ".to_v.", "v")]
        param_to_weight_names = defaultdict(list)
        if hasattr(module, "stacked_params_mapping"):
            for param_name, weight_name, _ in module.stacked_params_mapping:
                param_to_weight_names[param_name].append(weight_name)
        for name, params in module.named_parameters(prefix):
            is_bias = False
            if name.endswith(".bias"):
                base_key = name[: -len(".bias")]
                is_bias = True
            elif name.endswith(".weight"):
                base_key = name[: -len(".weight")]
            else:
                continue
            delta, used_keys = _prepare_lora_delta(
                state_dict,
                base_key,
                param_to_weight_names,
                is_bias,
                lora_a_suffix,
                lora_b_suffix,
                lora_bias_suffix,
            )
            if delta is None:
                continue
            lora_loaded_keys.update(used_keys)
            delta = delta.to(device=params.device, dtype=params.dtype)
            with torch.no_grad():
                params.add_(delta)
            del delta
        missing_keys = set(state_dict.keys()) - lora_loaded_keys
        for k in missing_keys:
            logger.warning(f"Missing loading lora key: {k}")
        logger.info(f"{len(lora_loaded_keys)} lora keys loaded into {module.__class__.__name__}.")
        return lora_loaded_keys
    @classmethod
    def unload_module_lora(
        cls,
        state_dict,
        module,
        prefix: str = "transformer",
        lora_a_suffix: str = "lora_A.weight",
        lora_b_suffix: str = "lora_B.weight",
        lora_bias_suffix: str = "bias",
    ):
        lora_unloaded_keys = set()
        param_to_weight_names = defaultdict(list)
        if hasattr(module, "stacked_params_mapping"):
            for param_name, weight_name, _ in module.stacked_params_mapping:
                param_to_weight_names[param_name].append(weight_name)
        for name, param in module.named_parameters(prefix):
            is_bias = False
            if name.endswith(".bias"):
                base_key = name[: -len(".bias")]
                is_bias = True
            elif name.endswith(".weight"):
                base_key = name[: -len(".weight")]
            else:
                continue
            delta, used_keys = _prepare_lora_delta(
                state_dict,
                base_key,
                param_to_weight_names,
                is_bias,
                lora_a_suffix,
                lora_b_suffix,
                lora_bias_suffix,
            )
            if delta is None:
                continue
            lora_unloaded_keys.update(used_keys)
            delta = delta.to(device=param.device, dtype=param.dtype)
            with torch.no_grad():
                param.sub_(delta)
            del delta
        missing_keys = set(state_dict.keys()) - lora_unloaded_keys
        for k in missing_keys:
            logger.warning(f"Missing unloading lora key: {k}")
        logger.info(f"{len(lora_unloaded_keys)} lora keys unloaded from {module.__class__.__name__}.")
        return lora_unloaded_keys
 class QwenImageLoraLoaderMixin(LoraLoaderMixin):
    def load_lora_weights(
        self,
        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
        adapter_name: str | None = None,
    ):
        if adapter_name in self.lora_loaded:
            return
        if isinstance(pretrained_model_name_or_path_or_dict, dict):
            state_dict = pretrained_model_name_or_path_or_dict
        else:
            state_dict = _load_lora_state_dict(pretrained_model_name_or_path_or_dict)
        has_alpha = any(k.endswith(".alpha") for k in state_dict)
        is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
        is_diffusers_lora_format = any(k.endswith(".lora_A.weight") for k in state_dict)
        if has_alpha and is_diffusers_lora_format:
            state_dict = _apply_diffusers_lora_alpha_scaling(state_dict)
        if has_alpha or is_non_diffusers_format:
            converter = get_converter_by_pipeline(self)
            if converter is None:
                raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")
            state_dict = converter(state_dict)
        state_dict = _remap_state_dict_keys(state_dict, [(".to_out.0.", ".to_out.")])
        self.load_lora_into_module(
            state_dict,
            self.transformer,
            prefix=self.transformer_name,
        )
        self.lora_loaded[adapter_name] = state_dict
    def unload_lora_weights(self, adapter_name: str):
        if adapter_name not in self.lora_loaded:
            return
        state_dict = self.lora_loaded[adapter_name]
        transformer = get_transformer_from_pipeline(self)
        self.unload_module_lora(state_dict, transformer, prefix=self.transformer_name)
        del self.lora_loaded[adapter_name]
 class WanLoraLoaderMixin(LoraLoaderMixin):
    transformer_2_name = "transformer_2"
    def load_lora_weights(
        self,
        pretrained_model_name_or_path: str | list[str],
        adapter_name: str | None = None,
    ):
        """Load LoRA weights into the pipeline's transformer(s).
        Args:
            pretrained_model_name_or_path: Path(s) to concrete .safetensors LoRA files.
                - str: a single .safetensors file. Only valid for Wan2.1
                  (single-transformer pipelines).
                - list[str]: one file per transformer, matched **by position** to
                  [transformer, transformer_2]. Required for Wan2.2 MoE.
            adapter_name: Name to register the adapter under.
        """
        if adapter_name in self.lora_loaded:
            return
        # Normalize to list[(lora_path, target_module_name)].
        lora_paths = None
        if isinstance(pretrained_model_name_or_path, str):
            lora_paths = [pretrained_model_name_or_path]
        else:
            lora_paths = list(pretrained_model_name_or_path)
        target_modules = (
            [self.transformer_name, self.transformer_2_name] if self.has_transformer_2 else [self.transformer_name]
        )
        if len(lora_paths) != len(target_modules):
            raise ValueError(
                f"{self.__class__.__name__} expects {len(target_modules)} LoRA file(s) "
                f"(target_modules={target_modules}, paths={lora_paths}), got {len(lora_paths)}. Pass a list of concrete "  # noqa
                f".safetensors files, one per transformer."
            )
        module_to_lora_sd = dict()
        for lora_path, target_module_name in zip(lora_paths, target_modules):
            state_dict = _load_lora_state_dict(lora_path)
            is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
            if is_non_diffusers_format:
                converter = get_converter_by_pipeline(self)
                if converter is None:
                    raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")
                state_dict = converter(state_dict)
            state_dict = _remap_state_dict_keys(
                state_dict,
                [
                    (".ffn.net.0.", ".ffn.net_0."),
                    (".ffn.net.2.", ".ffn.net_2."),
                    (
                        ".to_out.0.",
                        ".to_out.",
                    ),
                    # '.bias_B.bias' is a horrible name in diffuser's.
                    # So we remap it to '.bias'
                    (".lora_B.bias", ".bias"),
                ],
            )
            module = get_transformer_from_pipeline(self, transformer_name=target_module_name)
            if module is None:
                logger.warning(
                    f"Skip LoRA {lora_path}: target module '{target_module_name}' "
                    f"is not loaded in pipeline {self.__class__.__name__} "
                    f"(e.g. disabled by boundary_ratio)."
                )
                continue
            self.load_lora_into_module(state_dict, module, prefix=self.transformer_name, lora_bias_suffix="bias")
            module_to_lora_sd[target_module_name] = state_dict
        self.lora_loaded[adapter_name] = module_to_lora_sd
    def unload_lora_weights(self, adapter_name: str):
        if adapter_name not in self.lora_loaded:
            return
        module_to_lora_sd = self.lora_loaded[adapter_name]
        for module_name, lora_sd in module_to_lora_sd.items():
            module = get_transformer_from_pipeline(self, transformer_name=module_name)
            self.unload_module_lora(lora_sd, module, prefix=self.transformer_name)
        del self.lora_loaded[adapter_name]
@@ -4,7 +4,6 @@
 import math
 import time
 from collections import OrderedDict
 from enum import Enum
 from typing import get_args
 import torch
@@ -34,11 +33,6 @@ from vllm_omni.lora.utils import stable_lora_int_id
 logger = init_logger(__name__)
 class LoRABackend(str, Enum):
    PEFT = "peft"
    DISTILL = "distill"
 class DiffusionLoRAManager:
    """Manager for LoRA adapters in diffusion models.
@@ -25,7 +25,6 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.lora.loader import QwenImageLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -265,11 +264,7 @@ def apply_rotary_emb_qwen(
 class QwenImagePipeline(
-    nn.Module,
+    nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
    QwenImageCFGParallelMixin,
    DiffusionPipelineProfilerMixin,
    SupportsComponentDiscovery,
    QwenImageLoraLoaderMixin,
 ):
    supports_request_batch = True
    _dit_modules: ClassVar[list[str]] = ["transformer"]
@@ -25,7 +25,6 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -286,7 +285,6 @@ class Wan22Pipeline(
    DenoiseProgressMixin,
    DiffusionPipelineProfilerMixin,
    SupportsComponentDiscovery,
    WanLoraLoaderMixin,
 ):
    _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
@@ -25,7 +25,6 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -148,7 +147,6 @@ class Wan22I2VPipeline(
    DiffusionPipelineProfilerMixin,
    DenoiseProgressMixin,
    SupportsComponentDiscovery,
    WanLoraLoaderMixin,
 ):
    """
    Wan2.2 Image-to-Video Pipeline.
@@ -70,10 +70,9 @@ def find_module_with_attr(model, attr_name="transformer"):
    return None
-def get_transformer_from_pipeline(pipeline: Any, transformer_name: str = "transformer"):
+def get_transformer_from_pipeline(pipeline: Any):
-    """Return a named transformer from a pipeline or pipeline wrapper."""
+    pipe = find_module_with_attr(pipeline, attr_name="transformer")
    pipe = find_module_with_attr(pipeline, attr_name=transformer_name)
    if pipe is not None:
-        return getattr(pipe, transformer_name)
+        return pipe.transformer
    return None
@@ -49,7 +49,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
-from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
@@ -388,31 +388,14 @@ class DiffusionWorker:
        """Initialize the LoRA manager for this worker."""
        if self.model_runner.pipeline is None:
            return
-
+        self.lora_manager = DiffusionLoRAManager(
-        lora_path = self.od_config.lora_path
+            pipeline=self.model_runner.pipeline,
-        if isinstance(lora_path, list) and len(lora_path) == 1:
+            device=self.device,
-            lora_path = lora_path[0]
+            dtype=self.od_config.dtype,
-
+            max_cached_adapters=self.od_config.max_cpu_loras,
-        lora_backend = self.od_config.lora_backend
+            lora_path=self.od_config.lora_path,
-        if lora_backend == LoRABackend.PEFT:
+            lora_scale=self.od_config.lora_scale,
-            self.lora_manager = DiffusionLoRAManager(
+        )
                pipeline=self.model_runner.pipeline,
                device=self.device,
                dtype=self.od_config.dtype,
                max_cached_adapters=self.od_config.max_cpu_loras,
                lora_path=lora_path,
                lora_scale=self.od_config.lora_scale,
            )
        elif lora_backend == LoRABackend.DISTILL:
            pipeline = self.model_runner.pipeline
            if hasattr(pipeline, "load_lora_weights"):
                if self.od_config.lora_scale > 1.0:
                    logger.warning("lora_scale > 1.0 may not take any effect when using distilled LoRA backend.")
                pipeline.load_lora_weights(lora_path)
            else:
                logger.warning("Pipeline does not support loading distilled LoRA weights for now.")
        else:
            raise ValueError(f"Unknown LoRA backend: {lora_backend}. Available choices: {LoRABackend.__members__}")
    def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
        """Start or stop profiling for this GPU worker.
@@ -471,9 +471,6 @@ class OrchestratorArgs:
    num_gpus: int | None = None
    model_class_name: str | None = None
    diffusion_load_format: str | None = None
    lora_path: list[str] | None = None
    lora_backend: str | None = None
    lora_scale: float | None = None
    diffusers_load_kwargs: str = "{}"
    diffusers_call_kwargs: str = "{}"
    ulysses_degree: int | None = None
@@ -1048,9 +1048,6 @@ class AsyncOmniEngine:
            "boundary_ratio": kwargs.get("boundary_ratio", None),
            "flow_shift": kwargs.get("flow_shift", None),
            "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
            "lora_path": kwargs.get("lora_path", None),
            "lora_scale": kwargs.get("lora_scale", 1.0),
            "lora_backend": kwargs.get("lora_backend", "peft"),
            "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
            "worker_extension_cls": kwargs.get("worker_extension_cls", None),
            "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]),
@@ -1254,9 +1251,6 @@ class AsyncOmniEngine:
                if lora_scale is not None:
                    if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
                        cfg.engine_args.lora_scale = lora_scale
                if kwargs.get("lora_backend") is not None:
                    if not hasattr(cfg.engine_args, "lora_backend") or cfg.engine_args.lora_backend is None:
                        cfg.engine_args.lora_backend = kwargs["lora_backend"]
                if (
                    kwargs.get("diffusion_attention_config") is not None
                    or kwargs.get("diffusion_attention_backend") is not None
@@ -164,13 +164,6 @@ class OmniServeCommand(CLISubcommand):
                    "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)."
                )
        lora_backend = getattr(args, "lora_backend", None)
        lora_path = getattr(args, "lora_path", None)
        if lora_backend == "distill" and not lora_path:
            raise ValueError("--lora-backend distill requires --lora-path")
        if (lora_backend or "peft") == "peft" and isinstance(lora_path, list) and len(lora_path) > 1:
            raise ValueError("--lora-backend peft accepts only one startup LoRA path")
        # --omni-lb-policy is validated against the LoadBalancingPolicy enum.
        omni_lb_policy = getattr(args, "omni_lb_policy", None)
        if omni_lb_policy is not None:
@@ -434,32 +427,6 @@ class OmniServeCommand(CLISubcommand):
                "custom_pipeline, dummy, or diffusers for the HF diffusers adapter."
            ),
        )
        omni_config_group.add_argument(
            "--lora-path",
            type=str,
            nargs="+",
            default=None,
            help=(
                "LoRA checkpoint path(s) loaded when the diffusion server starts. "
                "The distill backend accepts one file per pipeline transformer."
            ),
        )
        omni_config_group.add_argument(
            "--lora-backend",
            type=str,
            choices=["peft", "distill"],
            default=None,
            help=(
                "Diffusion LoRA loading backend. 'distill' fuses checkpoint files "
                "into the base model at server startup; 'peft' uses the adapter manager."
            ),
        )
        omni_config_group.add_argument(
            "--lora-scale",
            type=float,
            default=None,
            help="Scale for a startup PEFT LoRA. Distilled LoRAs are fused at their checkpoint scale.",
        )
        omni_config_group.add_argument(
            "--diffusion-compile-granularity",
            choices=["regional", "full"],
@@ -152,9 +152,6 @@ class Qwen3OmniMoeForConditionalGeneration(
                dtype=torch.long,
            )
        elif self.model_stage == "talker":
            # The outer wrapper exposes talker_mtp for every stage, but only
            # the talker stage owns the module that the method invokes.
            self.talker_mtp_graph_safe = current_omni_platform.supports_talker_mtp_graph_capture()
            multimodal_config.skip_mm_profiling = True
            self.has_preprocess = True
            self.has_postprocess = True
@@ -155,11 +155,6 @@ class OmniPlatform(Platform):
        """Check if the platform supports torch.compile with inductor backend."""
        raise NotImplementedError
    @classmethod
    def supports_talker_mtp_graph_capture(cls) -> bool:
        """Whether a model may capture its dedicated talker MTP graph."""
        return True
    @classmethod
    def has_flash_attn_package(cls) -> bool:
        """Check if a Flash Attention package is available and usable on this platform."""
@@ -3,8 +3,6 @@
 import torch
 from vllm.config import VllmConfig
 from vllm.config.kernel import IrOpPriorityConfig
 from vllm.logger import init_logger
 from vllm.platforms.interface import DeviceCapability
 from vllm_musa.platform import MUSAPlatformBase
@@ -32,6 +30,10 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
    def get_omni_generation_worker_cls(cls) -> str:
        return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
    @classmethod
    def get_default_stage_config_path(cls) -> str:
        return "vllm_omni/deploy"
    @classmethod
    def has_flash_attn_package(cls) -> bool:
        from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed
@@ -113,15 +115,6 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
        """MUSA supports torch.compile with inductor backend."""
        return True
    @classmethod
    def get_default_stage_config_path(cls) -> str:
        return "vllm_omni/deploy"
    @classmethod
    def supports_talker_mtp_graph_capture(cls) -> bool:
        """MUSA keeps Qwen3 talker MTP outside its dedicated FULL graph."""
        return False
    @classmethod
    def supports_float64(cls) -> bool:
        """MUSA does not support float64 yet."""
@@ -144,7 +137,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
    @classmethod
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
        """Get the compute capability of the MUSA device."""
-        major, minor = torch.musa.get_device_capability(device_id)
+        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)
    @classmethod
@@ -185,26 +178,3 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        return torch.musa.get_device_name(device_id)
    @classmethod
    def set_device_control_env_var(cls, devices: str | int | None) -> None:
        import os
        os.environ["MUSA_VISIBLE_DEVICES"] = devices
    @classmethod
    def unset_device_control_env_var(cls) -> None:
        import os
        os.environ.pop("MUSA_VISIBLE_DEVICES", None)
    @classmethod
    def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
        """Prefer native ops while compiling, otherwise prefer vLLM kernels."""
        from vllm.config.compilation import CompilationMode
        cc = vllm_config.compilation_config
        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
        default = ["native"] if using_inductor else ["vllm_c", "native"]
        return IrOpPriorityConfig.with_default(default)
@@ -206,12 +206,8 @@ class OmniGPUModelRunner(GPUModelRunner):
        cudagraph_mode = self.compilation_config.cudagraph_mode
        assert cudagraph_mode is not None
        has_separate_talker = getattr(self.model, "talker", None) is not None
-        talker_mtp_graph_safe = getattr(self.model, "talker_mtp_graph_safe", None)
+        talker_mtp_graph_safe = getattr(self.model, "talker_mtp_graph_safe", False)
-        # Preserve the legacy separate-talker default only when the model has
+        if cudagraph_mode.has_full_cudagraphs() and (has_separate_talker or talker_mtp_graph_safe):
        # not declared graph safety. An explicit False must disable wrapping.
        if talker_mtp_graph_safe is None:
            talker_mtp_graph_safe = has_separate_talker
        if cudagraph_mode.has_full_cudagraphs() and talker_mtp_graph_safe:
            graph_wrapper_cls = current_omni_platform.get_graph_wrapper_cls()
            self.talker_mtp = graph_wrapper_cls(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
        # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
Author	SHA1	Message	Date
Hongsheng Liu	b112c7fdee	docs: restore BitsAndBytes navigation Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-12 00:10:38 +08:00
Hongsheng Liu	e63ae881e7	docs: remove GGUF from active navigation Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 17:22:19 +08:00
Hongsheng Liu	2787f7f5f6	docs: clarify attention design coverage Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 17:07:49 +08:00
Hongsheng Liu	3cc05bf1c1	docs: address feature taxonomy review Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 16:55:39 +08:00
Hongsheng Liu	5a33440362	docs: align User Guide feature taxonomy Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 16:16:46 +08:00